Høyytelses GPU-minneoverføring på AWS Sagemaker Hyperpod

Moderne infrastruktur for dyp læring krever ofte overføring av store datamengder mellom GPU-er på tvers av maskiner. Hos Perplexity møtte vi en unik teknisk utfordring: å effektivt overføre ikke-sammenhengende GPU-minneområder mellom maskiner med maksimal mulig hastighet. Vår målplattform, AWS p5-instanser, tilbyr en imponerende nettverksbåndbredde på 3200 Gbps gjennom 32 nettverkskort. Denne artikkelen deler vår reise med å bygge en tilpasset, høyytelses nettverksløsning som oppnår 97,1 % av denne teoretiske båndbredden.

Den tekniske utfordringen

Vårt bruksområde presenterte flere viktige tekniske krav:

• Høyhastighetsoverføring mellom eksterne GPU-er med ikke-sammenhengende minneblokker
• Mulighet til å dynamisk legge til eller fjerne noder fra Kubernetes-distribusjoner uten å avbryte pågående operasjoner
• Støtte for peker-til-peker-kommunikasjonsmønstre (peer-to-peer)

Selv om NVIDIAs NCCL-bibliotek er de facto-standarden for distribuert dyp læring, var det ikke ideelt for vårt bruksområde:

• NCCL utmerker seg på kollektiv kommunikasjon, men krever etablering av en statisk "verden", noe som krever omstart av hele klyngen når de deltakende nodene justeres.
• NCCLs synkrone kommunikasjonsmodell skaper ekstra kompleksitet for vår asynkrone arbeidsmengde
• Vi ønsket direkte kontroll over våre minneoverføringsmønstre for optimalisering
• Å bygge vår egen løsning ga verdifulle læringsmuligheter

Moderne høyytelsesnettverk

For å forstå løsningen vår, la oss først utforske hvordan moderne høyytelsesnettverk skil seg fra tradisjonell nettverksteknologi.

De fleste nettverk vi bruker daglig, er avhengige av TCP/IP-protokoller, der applikasjoner kommuniserer med nettverkskortet gjennom operativsystemkjernen ved hjelp av socket-er. Høyytelsesnettverk bruker imidlertid RDMA (Remote Direct Memory Access) – en helt annen maskinvare- og programvaretak som muliggjør direkte minnetilgang mellom maskiner uten å involvere CPU-en.

AWS tilbyr Elastic Fabric Adapter (EFA), et tilpasset nettverksgrensesnitt som implementerer Amazons egendefinerte protokoll kalt Scalable Reliable Datagram (SRD). I motsetning til tradisjonell TCP/IP-nettverksteknologi hvor data må kopieres flere ganger mellom brukerrom, kjernerom og nettverksbuffere, tillater EFA med RDMA direkte dataoverføring mellom GPU-minne og nettverkskortet, og omgår CPU-en fullstendig.

Filosofien bak design av høyytelsesnettverk

Å bygge nettverkssystemer med høy ytelse krever en gjenvisning av flere grunnleggende antakelser:

1. Buffer-eierskap: I motsetning til tradisjonelle socket-er hvor kjernen administrerer nettverksbuffere og krever kopiering mellom brukerrom og kjernerom, krever RDMA at applikasjoner administrerer sine egne buffere. Når en applikasjon initierer en nettverksoperasjon, overfører den buffer-eierskapet til nettverkskortet til operasjonen er fullført, noe som eliminerer behovet for datakopiering.
2. Minneregistrering: Applikasjoner må registrere minneområder hos operativsystemkjernen. Kjernen oppretter virtuelle adresse-mappinger som gjør at CPU-en, GPU-ene og nettverkskortene alle forstår de samme virtuelle adressene. Denne registreringen er en engangsoperasjon som muliggjør etterfølgende null-kopierings-dataoverføringer.
3. Kontrollplan vs. dataplan: Høyytelsesnettverk deler operasjoner inn i to kategorier:
   1. Kontrollplanoperasjoner (som tilkoblingsoppsett og minneregistrering) går gjennom kjernen for å sikre sikkerhet
   2. Dataplanoperasjoner (faktisk dataoverføring) omgår kjernen for maksimal ytelse
4. Mottak før overføring: Uten kjernestyrte buffere må applikasjoner forhåndsposte mottaoperasjoner og spesifisere hvor innkommende data skal plasseres. Dette er et fundamentalt skifte fra socket-modellen der applikasjoner kan motta data når som helst.
5. Avlesningsbasert fullføring: I stedet for å vente på kjernevarsler gjennom mekanismer som epoll, avleser applikasjoner maskinvarens fullføringskøer direkte. Dette eliminerer overhead fra systemkall og gir umiddelbar reaksjon på fullførte operasjoner.
6. Bevissthet om maskinvaretopologi: Å forstå og optimalisere for maskinvaretopologi er avgjørende for å oppnå maksimal ytelse.

Forstå maskinvaretopologi

AWS p5-instanser har en sofistikert intern arkitektur. Som vist nedenfor inneholder hver instans to CPU-sokkelen som danner to NUMA-noder, der hver NUMA-node kobles til fire PCIe-svitsjer:

Forstørr

Under hver PCIe-svitsj finner vi:

• Fire 100 Gbps EFA-nettverkskort
• Én NVIDIA H100 GPU
• Én NVMe SSD

Databanene for TCP/IP- og RDMA-overføringer viser grunnleggende forskjeller i tilnærmingene deres:

Forstørr

Med TCP/IP (venstre side) må data kopieres flere ganger:

1. Fra GPU-minne til applikasjonsbuffer i brukerrom
2. Fra brukerromsbuffer til kjerneromsbuffer
3. Fra kjerneromsbuffer til nettverkskort
4. Den omvendte prosessen skjer på mottakersiden

Hver kopieringsoperasjon forbruker CPU-sykluser og minnebåndbredde. Applikasjonen må også foreta kontekstbytter mellom brukerrom og kjernerom for hver nettverksoperasjon.

I kontrast gir RDMA (høyre side) ekte null-kopierings-dataoverføring (zero-copy):

1. Nettverkskortet leser direkte fra GPU-minnet
2. Data reiser direkte til det eksterne nettverkskortet
3. Det eksterne nettverkskortet skriver direkte til destinasjonens GPU-minne
4. Applikasjonen trenger bare å sjekke en fullføringskø (CQ) i brukerrommet for å vite når overføringen er ferdig

Med RDMA og riktig maskinvareparing krever dataoverføring mellom to GPU-er bare å krysse den lokale PCIe-svitsjen og nettverket:

Forstørr

I kontrast må TCP/IP-overføringer kopiere data flere ganger gjennom hovedminnet, noe som forårsaker betydelig køing på PCIe-bussen:

Forstørr

Bygging med libfabric

Vi brukte libfabric – et rammeverk som gir et generisk grensesnitt for fabriktjenester. Vår implementering bruker to typer RDMA-operasjoner:

• Tosidig RDMA (SEND/RECV) for kontrollmeldinger som bærer metadata om minneområder (f.eks. forskyvninger og størrelser).
• Ensidig RDMA WRITE for faktisk dataoverføring, der hver WRITE-operasjon håndterer én sammenhengende minneblokk.

Reisen mot toppytelse

Utviklingen vår gikk gjennom flere stadier:

1. Implementerte grunnleggende ensidig meldingeroverføring ved hjelp av SEND/RECV
2. Utvidet til tosidig kommunikasjon
3. La til GPUDirect RDMA WRITE for direkte GPU-til-GPU-overføring
4. Utvidet til å håndtere flere samtidige overføringer
5. Introduserte operasjonskøing for robusthet
6. Oppnådde 97,4 % båndbreddeutnyttelse på et enkelt nettverkskort

Da vi skalerte til 32 nettverkskort, implementerte vi flere avgjørende optimaliseringer:

1. Operasjonskøing: I stedet for å sende operasjoner direkte til nettverkskort, opprettholder vi en kø på applikasjonsnivå. Dette gir robusthet mot nettverkstetthet og forenkler programmeringsmodellen.
2. Nettverksoppvarming: Forhåndsetablering av tilkoblinger forbedret oppstartsytelsen.
3. Multitråding: Dedikerte tråder for hver GPUs nettverksoperasjoner.
4. CPU-kjernepinning: Binde tråder til spesifikke CPU-kjerner for å unngå NUMA-effekter og cache-misser.
5. Tilstandsharding: Redusere konkurranse mellom tråder.
6. Operasjonsgruppering: Innsending av flere operasjoner samlet.
7. Lat operasjonsposting: Operasjoner settes først i kø i applikasjonen. Etter avlesning av fullføringskøen forsøker vi å sende inn ventende operasjoner til nettverkskortet, noe som sikrer effektiv bruk av nettverksressurser.
8. NUMA-bevisst ressursallokering: Alloker libfabric-ressurser som fullføringskøer på riktig NUMA-node for å minimere ventetid for minnetilgang.

Gjennom disse optimaliseringene oppnådde vi en endelig ytelse på 3108 Gbps – 97,1 % av den teoretiske maksimale båndbredden.

Videoen nedenfor viser kommandolinjeprogrammet vårt i aksjon. Det overfører ikke-sammenhengende biter av GPU-minne til en ekstern node, og oppnår en overføringshastighet på 3108,283 Gbps – noe som demonstrerer en tilnærmet teoretisk båndbreddeutnyttelse av nettverksinfrastrukturen:

Konklusjon

Å bygge et nettverkssystem med høy ytelse krever forståelse for både maskinvarearkitektur og prinsipper for systemdesign. Selv om biblioteker som NCCL tilbyr utmerkede løsninger for vanlige mønstre, er tilpassede løsninger noen ganger nødvendige for spesifikke krav.

Reisen vår viser at det er mulig å oppnå tilnærmet teoretisk nettverksytelse med grundig oppmerksomhet på systemarkitektur, maskinvaretopologi og ulike optimaliseringsteknikker. Nøkkelen er ikke bare å forstå individuelle komponenter, men hvordan de samhandler for å danne et komplett system.

Hele den tekniske dypdykken i denne reisen, inkludert implementeringsdetaljer og kodeeksempler, er tilgjengelig i vårt åpen kildekode-lager.

Sitat

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