Prijenos GPU memorije visokih performansi na AWS Sagemaker Hyperpod

Moderna infrastruktura dubokog učenja često zahtijeva prijenos velikih količina podataka između GPU-ova na različitim strojevima. U tvrtki Perplexity naišli smo na jedinstven tehnički izazov: učinkovit prijenos nesusjednih regija GPU memorije između strojeva uz najveću moguću brzinu. Naša ciljna platforma, AWS p5 instance, nudi impresivnu mrežnu propusnost od 3200 Gbps putem 32 mrežne kartice. Ovaj članak dijeli naše putovanje u izgradnji prilagođenog mrežnog rješenja visokih performansi koje postiže 97,1% ove teorijske propusnosti.

Tehnički izazov

Naš slučaj korištenja predstavio je nekoliko ključnih tehničkih zahtjeva:

• Prijenos visoke propusnosti između udaljenih GPU-ova nesusjednih memorijskih dijelova
• Mogućnost dinamičkog dodavanja ili uklanjanja čvorova iz Kubernetes implementacija bez ometanja tekućih operacija
• Podrška za uzajamnu komunikaciju (peer-to-peer)

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

• NCCL se ističe u kolektivnoj komunikaciji, ali zahtijeva uspostavu statičnog „svijeta“, što zahtijeva ponovno pokretanje cijelog klastera prilikom prilagodbe čvorova koji sudjeluju.
• NCCL-ov sinkroni komunikacijski model dodaje složenost našem asinkronom radnom opterećenju
• Željeli smo izravnu kontrolu nad našim uzorcima prijenosa memorije radi optimizacije
• Izgradnja vlastitog rješenja pružila je dragocjene prilike za učenje

Moderne mreže visokih performansi

Kako bismo razumjeli naše rješenje, najprije istražimo kako se moderne mreže visokih performansi razlikuju od tradicionalnog umrežavanja.

Većina mreža koje koristimo svakodnevno oslanja se na TCP/IP protokole, pri čemu aplikacije komuniciraju s mrežnom karticom kroz jezgru operativnog sustava pomoću utičnica (socketa). Međutim, mreže visokih performansi koriste RDMA (Remote Direct Memory Access) - potpuno drugačiji sklop sklopovlja i programske podrške koji omogućuje izravan pristup memoriji između strojeva bez uključivanja CPU-a.

AWS pruža Elastic Fabric Adapter (EFA), prilagođeno mrežno sučelje koje implementira Amazonov prilagođeni protokol nazvan Scalable Reliable Datagram (SRD). Za razliku od tradicionalnog TCP/IP umrežavanja gdje se podaci moraju kopirati više puta između korisničkog prostora, prostora jezgre i mrežnih međuspremnika, EFA s RDMA-om omogućuje izravan prijenos podataka između GPU memorije i mrežne kartice, u potpunosti zaobilazeći CPU.

Filozofija dizajna mreža visokih performansi

Izgradnja mrežnih sustava visokih performansi zahtijeva preispitivanje nekoliko temeljnih pretpostavki:

1. Vlasništvo nad međuspremnikom: Za razliku od tradicionalnih utičnica gdje jezgra upravlja mrežnim međuspremnicima i zahtijeva kopiranje između korisničkog prostora i prostora jezgre, RDMA zahtijeva da aplikacije upravljaju vlastitim međuspremnicima. Kada aplikacija pokrene mrežnu operaciju, ona prenosi vlasništvo nad međuspremnikom na mrežnu karticu dok operacija ne završi, čime se eliminira potreba za kopiranjem podataka.
2. Registracija memorije: Aplikacije moraju registrirati memorijske regije s jezgrom operativnog sustava. Jezgra postavlja mapiranja virtualnih adresa koja omogućuju CPU-u, GPU-ovima i mrežnim karticama da svi razumiju iste virtualne adrese. Ova registracija je jednokratna operacija koja omogućuje naknadne prijenose podataka bez kopiranja.
3. Kontrolna ravnina vs ravnina podataka: Mreže visokih performansi dijele operacije u dvije kategorije:
   1. Operacije kontrolne ravnine (poput uspostave veze i registracije memorije) prolaze kroz jezgru radi osiguranja sigurnosti
   2. Operacije ravnine podataka (stvarni prijenos podataka) zaobilaze jezgru radi maksimalnih performansi
4. Prijem prije prijenosa: Bez međuspremnika kojima upravlja jezgra, aplikacije moraju unaprijed objaviti operacije primanja, navodeći gdje dolazni podaci trebaju biti smješteni. Ovo je temeljna promjena u odnosu na model utičnica gdje aplikacije mogu primati podatke u bilo kojem trenutku.
5. Dovršetak temeljen na anketiranju: Umjesto čekanja obavijesti jezgre putem mehanizama kao što je epoll, aplikacije izravno anketiraju redove čekanja dovršetaka sklopovlja. To eliminira trošak sistemskih poziva i omogućuje trenutnu reakciju na dovršene operacije.
6. Svijest o sklopovskoj topologiji: Razumijevanje sklopovske topologije i optimizacija za nju ključni su za postizanje maksimalnih performansi.

Razumijevanje sklopovske topologije

AWS p5 instance imaju sofisticiranu unutarnju arhitekturu. Kao što je prikazano u nastavku, svaka instanca sadrži dva CPU utičnica koji tvore dva NUMA čvora, pri čemu se svaki NUMA čvor spaja na četiri PCIe preklopnika:

Povećaj

Pod svakim PCIe preklopnikom nalazimo:

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

Putanje podataka za TCP/IP i RDMA prijenose pokazuju temeljne razlike u njihovim pristupima:

Povećaj

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

1. Iz GPU memorije u međuspremnik aplikacije u korisničkom prostoru
2. Iz međuspremnika korisničkog prostora u međuspremnik jezgre
3. Iz međuspremnika jezgre u mrežnu karticu
4. Obrnuti proces događa se na strani primanja

Svaka operacija kopiranja troši CPU cikluse i propusnost memorije. Aplikacija također mora prebacivati kontekst između korisničkog prostora i prostora jezgre za svaku mrežnu operaciju.

Nasuprot tome, RDMA (desna strana) omogućuje pravi prijenos podataka bez kopiranja (zero-copy):

1. Mrežna kartica čita izravno iz GPU memorije
2. Podaci putuju izravno do udaljene mrežne kartice
3. Udaljena mrežna kartica piše izravno u odredišnu GPU memoriju
4. Aplikacija treba samo provjeriti red čekanja dovršetaka (CQ) u korisničkom prostoru kako bi saznala kada je prijenos gotov

S RDMA-om i odgovarajućim uparivanjem sklopovlja, prijenos podataka između dva GPU-a zahtijeva samo prolazak kroz lokalni PCIe preklopnik i mrežu:

Povećaj

Nasuprot tome, TCP/IP prijenosi moraju kopirati podatke više puta kroz glavnu memoriju, uzrokujući značajnu gužvu na PCIe sabirnici:

Povećaj

Izgradnja pomoću libfabrica

Koristili smo libfabric - okvire koji pružaju generičko sučelje za usluge tkanja (fabric). Naša implementacija koristi dvije vrste RDMA operacija:

• Dvostrane RDMA (SEND/RECV) poruke za kontrolne poruke koje nose metapodatke o memorijskim regijama (npr. pomaci i veličine).
• Jednostrani RDMA WRITE za stvarni prijenos podataka, pri čemu svaka WRITE operacija upravlja jednim susjednim memorijskim dijelom.

Put do vršnih performansi

Naš razvoj napredovao je kroz nekoliko faza:

1. Implementiran osnovni jednosmjerni prijenos poruka pomoću SEND/RECV
2. Prošireno na dvosmjernu komunikaciju
3. Dodan GPUDirect RDMA WRITE za izravan GPU-GPU prijenos
4. Prošireno za rukovanje višestrukim istovremenim prijenosima
5. Uvedeno stavljanje operacija u red čekanja radi robusnosti
6. Postignuto 97,4% iskorištenja propusnosti na jednoj mrežnoj kartici

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

1. Stavljanje operacija u red čekanja: Umjesto izravnog slanja operacija mrežnim karticama, održavamo red čekanja na razini aplikacije. To pruža robusnost protiv mrežne gužve i pojednostavljuje programski model.
2. Mrežno zagrijavanje: Unaprijed uspostavljene veze poboljšale su performanse pokretanja.
3. Višenitnost: Namjenske niti za mrežne operacije svakog GPU-a.
4. Prikucavanje procesorskih jezgri: Povezivanje niti s određenim procesorskim jezgrama kako bi se izbjegli NUMA efekti i promašaji predmemorije.
5. Dijeljenje stanja (sharding): Smanjenje sukoba između niti.
6. Grupiranje operacija: Slanje više operacija zajedno.
7. Lijeno objavljivanje operacija: Operacije se prvo stavljaju u red čekanja u aplikaciji. Nakon anketiranja (polling) reda čekanja dovršetaka, pokušavamo poslati operacije na čekanju mrežnoj kartici, čime se osigurava učinkovita uporaba mrežnih resursa.
8. Alokacija resursa svjesna NUMA-e: Alocirajte libfabric resurse kao što su redovi čekanja dovršetaka na ispravnom NUMA čvoru kako biste smanjili latenciju pristupa memoriji.

Kroz ove optimizacije postigli smo konačne performanse od 3.108 Gbps - 97,1% teoretske maksimalne propusnosti.

Videozapis u nastavku prikazuje naš program naredbenog retka u akciji. On prenosi nesusjedne dijelove GPU memorije na udaljeni čvor, postižući brzinu prijenosa od 3108,283 Gbps - što pokazuje gotovo teoretsko iskorištenje propusnosti mrežne infrastrukture:

Zaključak

Izgradnja mrežnog sustava visokih performansi zahtijeva razumijevanje i sklopovske arhitekture i načela dizajna sustava. Iako biblioteke poput NCCL-a pružaju izvrsna rješenja za uobičajene uzorke, ponekad su prilagođena rješenja neophodna za specifične zahtjeve.

Naše putovanje pokazuje da je postizanje gotovo teoretskih mrežnih performansi moguće uz pažljivu pozornost na arhitekturu sustava, sklopovsku topologiju i razne tehnike optimizacije. Ključ nije samo u razumijevanju pojedinačnih komponenti, već u tome kako one međusobno djeluju kako bi činile cjelovit sustav.

Potpuna tehnička analiza ovog putovanja, uključujući detalje implementacije i primjere koda, dostupna je u našem spremištu otvorenog koda.

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}
}