GPU-hukommelsesoverførsel med høj ydeevne på AWS Sagemaker Hyperpod

Moderne infrastruktur til deep learning kræver ofte overførsel af store datamængder mellem GPU'er på tværs af maskiner. Hos Perplexity stødte vi på en unik teknisk udfordring: effektiv overførsel af ikke-sammenhængende GPU-hukommelsesområder mellem maskiner med højest mulige hastighed. Vores målutstyr, AWS p5-instanser, leverer en imponerende netværksbåndbredde på 3200 Gbps via 32 netværkskort. Denne artikel deler vores rejse med at bygge en skræddersyet netværksløsning med høj ydeevne, der opnår 97,1 % af denne teoretiske båndbredde.

Den tekniske udfordring

Vores brugsscenarie præsenterede flere centrale tekniske krav:

• Overførsel med høj båndbredde mellem fjerntliggende GPU'er af ikke-sammenhængende hukommelsesbidder
• Mulighed for dynamisk at tilføje eller fjerne noder fra Kubernetes-udrulninger uden at afbryde igangværende operationer
• Understøttelse af peer-to-peer-kommunikationsmønstre

Selvom NVIDIAs NCCL-bibliotek er de facto-standarden for distribueret deep learning, var det ikke ideelt til vores brugsscenarie:

• NCCL udmærker sig ved kollektiv kommunikation, men kræver etablering af en statisk "verden", hvilket kræver genstart af hele klyngen ved justering af de deltagende noder.
• NCCL's synkrone kommunikationsmodel tilføjer kompleksitet til vores asynkrone arbejdsbelastning
• Vi ønskede direkte kontrol over vores hukommelsesoverførselsmønstre med henblik på optimering
• At bygge vores egen løsning gav værdifulde læringsmuligheder

Moderne netværk med høj ydeevne

For at forstå vores løsning skal vi først udforske, hvordan moderne netværk med høj ydeevne adskiller sig fra traditionelle netværk.

De fleste netværk, vi bruger dagligt, er baseret på TCP/IP-protokoller, hvor applikationer kommunikerer med netværkskortet gennem operativsystemets kerne ved hjælp af sockets. Netværk med høj ydeevne bruger dog RDMA (Remote Direct Memory Access) - en fuldstændig anden hardware- og softwarestak, der muliggør direkte hukommelsesadgang mellem maskiner uden at inddrage CPU'en.

AWS leverer Elastic Fabric Adapter (EFA), et skræddersyet netværksinterface, der implementerer Amazons egen protokol kaldet Scalable Reliable Datagram (SRD). I modsætning til traditionel TCP/IP-netværk, hvor data skal kopieres flere gange mellem brugerplads, kerneplads og netværksbuffer, tillader EFA med RDMA direkte dataoverførsel mellem GPU-hukommelse og netværkskortet, idet CPU'en omgås fuldstændigt.

Filosofien bag design af netværk med høj ydeevne

Opbygning af netværkssystemer med høj ydeevne kræver genovervejelse af flere grundlæggende antagelser:

1. Bufferejerskab: I modsætning til traditionelle sockets, hvor kernen administrerer netværksbuffere og kræver kopiering mellem brugerplads og kerneplads, kræver RDMA, at applikationer selv administrerer deres buffere. Når en applikation initierer en netværksoperation, overdrager den bufferejerskabet til netværkskortet, indtil operationen er fuldført, hvilket fjerner behovet for datakopiering.
2. Hukommelsesregistrering: Applikationer skal registrere hukommelsesområder hos operativsystemets kerne. Kernen opretter virtuelle adresse-mappings, der gør det muligt for CPU, GPU'er og netværkskort alle at forstå de samme virtuelle adresser. Denne registrering er en engangsoperation, der muliggør efterfølgende zero-copy-dataoverførsler.
3. Kontrolplan vs. dataplan: Netværk med høj ydeevne opdeler operationer i to kategorier:
   1. Kontrolplanoperationer (såsom opsætning af forbindelse og registrering af hukommelse) går gennem kernen for at sikre sikkerhed
   2. Dataplanoperationer (faktisk dataoverførsel) omgår kernen for at opnå maksimal ydeevne
4. Modtagelse før transmission: Uden kernestyrede buffere skal applikationer på forhånd indsende modtageoperationer og angive, hvor indkommende data skal placeres. Dette er et fundamentalt skift fra socket-modellen, hvor applikationer kan modtage data når som helst.
5. Polling-baseret fuldførelse: I stedet for at vente på kernemeddelelser gennem mekanismer som epoll poller applikationer direkte hardwarens fuldførelseskøer. Dette fjerner overhead fra systemkald og muliggør øjeblikkelig reaktion på fuldførte operationer.
6. Bevidsthed om hardwaretopologi: Forståelse af og optimering til hardwaretopologi er afgørende for at opnå maksimal ydeevne.

Forståelse af hardwaretopologi

AWS p5-instanser har en sofistikeret intern arkitektur. Som vist herunder indeholder hver instans to CPU-sockets, der danner to NUMA-noder, hvor hver NUMA-node forbinder til fire PCIe-switches:

Udvid

Under hver PCIe-switch finder vi:

• Fire 100 Gbps EFA-netværkskort
• Én NVIDIA H100 GPU
• Én NVMe SSD

Datastierne for TCP/IP- og RDMA-overførsler viser fundamentale forskelle i deres tilgange:

Udvid

Med TCP/IP (venstre side) skal data kopieres flere gange:

1. Fra GPU-hukommelse til applikationsbuffer i brugerplads
2. Fra brugerpladsbuffer til kernebuffer
3. Fra kernebuffer til netværkskort
4. Den modsatte proces foregår på modtagesiden

Hver kopieringsoperation forbruger CPU-cyklusser og hukommelsesbåndbredde. Applikationen skal også udføre et kontekstskifte mellem brugerplads og kerneplads for hver netværksoperation.

Til gengæld leverer RDMA (højre side) ægte zero-copy-dataoverførsel:

1. Netværkskortet læser direkte fra GPU-hukommelsen
2. Data rejser direkte til det fjerntliggende netværkskort
3. Det fjerntliggende netværkskort skriver direkte til destinations-GPU-hukommelsen
4. Applikationen behøver kun at tjekke en fuldførelseskø (CQ) i brugerplads for at vide, hvornår overførslen er færdig

Med RDMA og korrekt hardwareparring kræver overførsel af data mellem to GPU'er kun passage gennem den lokale PCIe-switch og netværket:

Udvid

I modsætning hertil skal TCP/IP-overførsler kopiere data flere gange gennem hovedhukommelsen, hvilket forårsager betydelig overbelastning af PCIe-bussen:

Udvid

Opbygning med libfabric

Vi brugte libfabric - et framework, der leverer et generisk interface til fabric-tjenester. Vores implementering bruger to typer RDMA-operationer:

• Tosidet RDMA (SEND/RECV) til kontrolbeskeder, der indeholder metadata om hukommelsesområder (f.eks. offsets og størrelser).
• Ensidet RDMA WRITE til den faktiske dataoverførsel, hvor hver WRITE-operation håndterer én sammenhængende hukommelsesbid.

Rejsen mod top-ydeevne

Vores udvikling forløb gennem flere stadier:

1. Implementerede grundlæggende ensrettet beskedoverførsel ved hjælp af SEND/RECV
2. Udskiftede til tovejskommunikation
3. Tilføjede GPUDirect RDMA WRITE til direkte GPU-til-GPU-overførsel
4. Udvidet til at håndtere flere samtidige overførsler
5. Introducerede operationskøer for robusthed
6. Opnåede 97,4 % båndbreddeudnyttelse på et enkelt netværkskort

Da vi skalerede til 32 netværkskort, implementerede vi flere afgørende optimeringer:

1. Operationskøer: I stedet for at indsende operationer direkte til netværkskort opretholder vi en kø på applikationsniveau. Dette giver robusthed mod netværksbelastning og forenkler programmeringsmodellen.
2. Netværksopvarmning: Forudgående etablering af forbindelser forbedrede opstartsydeevnen.
3. Multithreading: Dedikerede tråde til hver GPU's netværksoperationer.
4. CPU-kernebinding: Binding af tråde til specifikke CPU-kerner for at undgå NUMA-effekter og cache-misses.
5. Tilstandsharding: Reduktion af strid mellem tråde.
6. Batch-behandling af operationer: Indsendelse af flere operationer samlet.
7. Laizy operation-indsendelse: Operationer sættes først i kø i applikationen. Efter afstemning af fuldførelseskøen forsøger vi at indsende udestående operationer til netværkskortet, hvilket sikrer effektiv udnyttelse af netværksressourcer.
8. NUMA-bevidst ressourcetildeling: Tildel libfabric-ressourcer som fuldførelseskøer på den korrekte NUMA-node for at minimere latenstiden for hukommelsesadgang.

Gennem disse optimeringer opnåede vi en endelig ydeevne på 3.108 Gbps - 97,1 % af den teoretiske maksimale båndbredde.

Videoen herunder viser vores kommandolinjeprogram i aktion. Det overfører ikke-sammenhængende bidder af GPU-hukommelse til en fjerntliggende node og opnår en transmissionshastighed på 3108,283 Gbps - hvilket demonstrerer en næsten teoretisk båndbreddeudnyttelse af netværksinfrastrukturen:

Konklusion

Opbygning af et netværkssystem med høj ydeevne kræver forståelse af både hardwarearkitektur og systemdesignprincipper. Selvom biblioteker som NCCL leverer fremragende løsninger til almindelige mønstre, er skræddersyede løsninger undertiden nødvendige for specifikke krav.

Vores rejse viser, at det er muligt at opnå en næsten teoretisk netværksydelse med omhyggelig opmærksomhed på systemarkitektur, hardwaretopologi og forskellige optimeringsteknikker. Nøglen er ikke kun at forstå individuelle komponenter, men hvordan de interagerer for at danne et komplet system.

Den fulde tekniske gennemgang af denne rejse, herunder implementeringsdetaljer og kodeeksempler, er tilgængelig i vores open source-lager.

Citation

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