Prenos GPU memorije visokih performansi na AWS Sagemaker Hyperpod

Moderna infrastruktura dubokog učenja često zahteva prenos velikih količina podataka između GPU-ova na različitim mašinama. U kompaniji Perplexity susreli smo se sa jedinstvenim tehničkim izazovom: efikasnim prenosom nekoniguitetnih regiona GPU memorije između mašina pri maksimalnoj mogućoj brzini. Naša ciljna platforma, AWS p5 instance, nudi impresivnih 3200 Gbps mrežne propusne moći kroz 32 mrežne kartice. Ovaj članak deli naš put izgradnje prilagođenog mrežnog rešenja visokih performansi koje postiže 97,1% ove teorijske propusne moći.

Tehnički izazov

Naš slučaj upotrebe predstavio je nekoliko ključnih tehničkih zahteva:

• Prenos visoke propusne moći između udaljenih GPU-ova nekoniguitetnih memorijskih blokova
• Mogućnost dinamičkog dodavanja ili uklanjanja čvorova iz Kubernetes deployment-a bez ometanja tekućih operacija
• Podrška za peer-to-peer obrasce komunikacije

Iako je NVIDIA-ina NCCL biblioteka de fakto standard za distribuirano duboko učenje, nije bila idealna za naš slučaj upotrebe:

• NCCL se ističe u kolektivnoj komunikaciji, ali zahteva uspostavljanje statičnog „sveta“, što zahteva restartovanje celog klastera pri podešavanju čvorova učesnika.
• NCCL-ov sinhroni model komunikacije dodaje složenost za naše asinhrono opterećenje
• Želeli smo direktnu kontrolu nad našim obrascima prenosa memorije radi optimizacije
• Izgradnja sopstvenog rešenja pružila je dragocene mogućnosti učenja

Moderne mreže visokih performansi

Da bismo razumeli naše rešenje, hajde prvo da istražimo kako se moderne mreže visokih performansi razlikuju od tradicionalnog umrežavanja.

Većina mreža koje svakodnevno koristimo oslanja se na TCP/IP protokol, gde aplikacije komuniciraju sa mrežnom karticom preko kernela operativnog sistema koristeći sokete. Međutim, mreže visokih performansi koriste RDMA (Remote Direct Memory Access) - potpuno drugačiji hardverski i softverski stak koji omogućava direktan pristup memoriji između mašina bez uključivanja CPU-a.

AWS obezbeđuje Elastic Fabric Adapter (EFA), prilagođeni mrežni interfejs koji implementira Amazonov prilagođeni protokol pod nazivom Scalable Reliable Datagram (SRD). Za razliku od tradicionalnog TCP/IP umrežavanja gde se podaci moraju kopirati više puta između korisničkog prostora, prostora kernela i mrežnih bafera, EFA sa RDMA-om omogućava direktan prenos podataka između GPU memorije i mrežne kartice, potpuno zaobilazeći CPU.

Filozofija dizajna mreža visokih performansi

Izgradnja mrežnih sistema visokih performansi zahteva preispitivanje nekoliko fundamentalnih pretpostavki:

1. Vlasništvo nad baferom: Za razliku od tradicionalnih soketa gde kernel upravlja mrežnim baferima i zahteva kopiranje između korisničkog prostora i prostora kernela, RDMA zahteva da aplikacije upravljaju sopstvenim baferima. Kada aplikacija pokrene mrežnu operaciju, ona prenosi vlasništvo nad baferom mrežnoj kartici dok se operacija ne završi, čime se eliminiše potreba za kopiranjem podataka.
2. Registracija memorije: Aplikacije moraju da registruju regione memorije sa kernelom operativnog sistema. Kernel postavlja mapiranja virtuelnih adresa koja omogućavaju CPU-u, GPU-ovima i mrežnim karticama da razumeju iste virtuelne adrese. Ova registracija je jednokratna operacija koja omogućava naknadne zero-copy prenose podataka.
3. Kontrolna ravan u odnosu na ravan podataka: Mreže visokih performansi dele operacije u dve kategorije:
   1. Operacije kontrolne ravni (poput podešavanja veze i registracije memorije) prolaze kroz kernel radi osiguranja bezbednosti
   2. Operacije ravni podataka (stvarni prenos podataka) zaobilaze kernel radi maksimalnih performansi
4. Prijem pre prenosa: Bez bafera kojima upravlja kernel, aplikacije moraju unapred objaviti operacije prijema, navodeći gde dolazni podaci treba da budu smešteni. Ovo je fundamentalna promena u odnosu na model soketa gde aplikacije mogu primati podatke u bilo kom trenutku.
5. Završetak zasnovan na anketiranju: Umesto čekanja na obaveštenja kernela kroz mehanizme poput epoll-a, aplikacije direktno anketiraju redove čekanja završetaka hardvera. Ovo eliminiše režijske troškove sistemskih poziva i omogućava trenutnu reakciju na završene operacije.
6. Svest o hardverskoj topologiji: Razumevanje i optimizacija za hardversku topologiju je ključna za postizanje maksimalnih performansi.

Razumevanje hardverske topologije

AWS p5 instance imaju sofisticiranu internu arhitekturu. Kao što je prikazano ispod, svaka instanca sadrži dva CPU soketa koji formiraju dva NUMA čvora, pri čemu se svaki NUMA čvor povezuje sa četiri PCIe sviča:

Uvećaj

Ispod svakog PCIe sviča nalazimo:

• Četiri EFA mrežne kartice od 100 Gbps
• Jedan NVIDIA H100 GPU
• Jedan NVMe SSD

Putanja podataka za TCP/IP i RDMA prenose demonstriraju fundamentalne razlike u njihovim pristupu:

Uvećaj

Sa TCP/IP-om (leva strana), podaci se moraju kopirati više puta:

1. Iz GPU memorije u bafer aplikacije u korisničkom prostoru
2. Iz bafera korisničkog prostora u bafer kernela
3. Iz bafera kernela u mrežnu karticu
4. Obrnuti proces se dešava na strani prijema

Svaka operacija kopiranja troši CPU cikluse i memorijsku propusnu moć. Aplikacija takođe mora da vrši kontekstnu promenu (context switch) između korisničkog prostora i prostora kernela za svaku mrežnu operaciju.

Nasuprot tome, RDMA (desna strana) obezbeđuje pravi zero-copy prenos podataka:

1. Mrežna kartica čita direktno iz GPU memorije
2. Podaci putuju direktno do udaljene mrežne kartice
3. Udaljena mrežna kartica piše direktno u odredišnu GPU memoriju
4. Aplikacija samo treba da proveri red čekanja završetaka (completion queue - CQ) u korisničkom prostoru kako bi saznala kada je prenos završen

Sa RDMA-om i pravilnim uparivanjem hardvera, prenos podataka između dva GPU-a zahteva samo prelaženje preko lokalnog PCIe sviča i mreže:

Uvećaj

Nasuprot tome, TCP/IP prenosi moraju da kopiraju podatke više puta kroz glavnu memoriju, uzrokujući značajnu gužvu na PCIe magistrali:

Uvećaj

Izgradnja pomoću libfabric-a

Koristili smo libfabric - okruženje koje obezbeđuje generički interfejs za usluge tkanine. Naša implementacija koristi dve vrste RDMA operacija:

• Dvosmerni RDMA (SEND/RECV) za kontrolne poruke koje nose metatipove o memorijskim regionima (npr. pomake i veličine).
• Jednosmerni RDMA WRITE za stvarni prenos podataka, gde svaka WRITE operacija upravlja jednim neprekidnim memorijskim blokom.

Put do vršnih performansi

Naš razvoj je napredovao kroz nekoliko faza:

1. Implementiran osnovni jednosmerni prenos poruka pomoću SEND/RECV
2. Prošireno na dvosmernu komunikaciju
3. Dodat GPUDirect RDMA WRITE za direktan GPU-GPU prenos
4. Prošireno za rukovanje višestrukim istovremenim prenosima
5. Uvedeno stavljanje operacija u red čekanja radi robusnosti
6. Postignuto iskorišćenje propusne moći od 97,4% na jednoj mrežnoj kartici

Prilikom skaliranja na 32 mrežne kartice, implementirali smo nekoliko ključnih optimizacija:

1. Stavljanje operacija u red čekanja: Umesto direktnog prosleđivanja operacija mrežnim karticama, održavamo red čekanja na nivou aplikacije. Ovo obezbeđuje robusnost protiv mrežnih zagušenja i pojednostavljuje programski model.
2. Zagrevanje mreže: Unapred uspostavljene veze poboljšale su performanse pokretanja.
3. Višenitnost: Namenske niti za mrežne operacije svakog GPU-a.
4. Pričvršćivanje CPU jezgara: Vezivanje niti za specifična CPU jezgra radi izbegavanja NUMA efekata i promašaja keša.
5. Deljenje stanja (Sharding): Smanjenje konflikta između niti.
6. Grupisanje operacija (Batching): Slanje više operacija zajedno.
7. Lenje objavljivanje operacija: Operacije se prvo stavljaju u red čekanja u aplikaciji. Nakon anketiranja reda čekanja završetaka, pokušavamo da prosledimo operacije na čekanju mrežnoj kartici, čime se obezbeđuje efikasno korišćenje mrežnih resursa.
8. Alokacija resursa svesna NUMA-e: Alociranje libfabric resursa kao što su redovi čekanja završetaka na tačnom NUMA čvoru radi minimiziranja kašnjenja pristupa memoriji.

Kroz ove optimizacije postigli smo konačne performanse od 3.108 Gbps - 97,1% teorijskog maksimalnog opsega.

Video ispod prikazuje naš komandni program u akciji. On prenosi nekoniguitetne blokove GPU memorije na udaljeni čvor, postižući brzinu prenosa od 3108,283 Gbps - demonstrirajući gotovo teorijsko iskorišćenje propusne moći mrežne infrastrukture:

Zaključak

Izgradnja mrežnog sistema visokih performansi zahteva razumevanje i hardverske arhitekture i principa dizajna sistema. Iako biblioteke poput NCCL-a pružaju odlična rešenja za uobičajene obrasce, ponekad su prilagođena rešenja neophodna za specifične zahteve.

Naš put pokazuje da je postizanje gotovo teorijskih mrežnih performansi moguće uz pažljivu posvećenost arhitekturi sistema, hardverskoj topologiji i različitim tehnikama optimizacije. Ključ nije samo u razumevanju pojedinačnih komponenti, već u načinu na koji one međusobno deluju kako bi činile kompletan sistem.

Detaljna tehnička analiza ovog puta, uključujući detalje implementacije i primere koda, dostupna je u našem open source skladištu.

Citiranje

@online{perplexity2025HighPerformanceGpuMemoryTransferOnAwsSagemakerHyperpod,
  author = {{Perplexity} and Lequn Chen and Phi Hoang},
  title = {{High-performance GPU memory transfer on AWS Sagemaker Hyperpod}},
  date = {2025-02-10},
  url = {https://www.perplexity.ai/hub/blog/high-performance-gpu-memory-transfer-on-aws-sagemaker-hyperpod}
}