Högpresterande GPU-minnesöverföring på AWS Sagemaker Hyperpod

Modern infrastruktur för djupinlärning kräver ofta att stora mängder data överförs mellan GPU:er över maskiner. På Perplexity stötte vi på en unik teknisk utmaning: att effektivt överföra icke-sammanhängande GPU-minnesregioner mellan maskiner med högsta möjliga hastighet. Vår målplattform, AWS p5-instanser, erbjuder en imponerande nätverksbandbredd på 3200 Gbps genom 32 nätverkskort. Den här artikeln delar vår resa med att bygga en skräddarsydd, högpresterande nätverkslösning som uppnår 97,1 % av denna teoretiska bandbredd.

Den tekniska utmaningen

Vårt användningsfall presenterade flera viktiga tekniska krav:

• Överföring med hög bandbredd mellan fjärr-GPU:er för icke-sammanhängande minnessegment
• Möjlighet att dynamiskt lägga till eller ta bort noder från Kubernetes-distributioner utan att störa pågående åtgärder
• Stöd för kommunikationsmönster från peer-to-peer

Även om NVIDIAs NCCL-bibliotek är de facto-standarden för distribuerad djupinlärning, var det inte idealiskt för vårt användningsfall:

• NCCL utmärker sig på kollektiv kommunikation men kräver att man upprättar en statisk "värld", vilket kräver att hela klustret startas om när de deltagande noderna justeras.
• NCCL:s synkrona kommunikationsmodell tillför komplexitet för vår asynkrona arbetsbelastning
• Vi ville ha direkt kontroll över våra minnesöverföringsmönster för optimering
• Att bygga vår egen lösning gav värdefulla inlärningsmöjligheter

Moderna högpresterande nätverk

För att förstå vår lösning kan vi först utforska hur moderna högpresterande nätverk skiljer sig från traditionella nätverk.

De flesta nätverk vi använder dagligen förlitar sig på TCP/IP-protokoll, där applikationer kommunicerar med nätverkskortet genom operativsystemets kärna med hjälp av socketar. Högpresterande nätverk använder dock RDMA (Remote Direct Memory Access) – en helt annan hård- och mjukvarustack som möjliggör direkt minnesåtkomst mellan maskiner utan inblandning av CPU:n.

AWS tillhandahåller Elastic Fabric Adapter (EFA), ett anpassat nätverksgränssnitt som implementerar Amazons anpassade protokoll som kallas Scalable Reliable Datagram (SRD). Till skillnad från traditionell TCP/IP-nätverksteknik där data måste kopieras flera gånger mellan användarrymd, kärnrymd och nätverksbuffertar, tillåter EFA med RDMA direkt dataöverföring mellan GPU-minnet och nätverkskortet, vilket helt kringgår CPU:n.

Filosofin bakom design av högpresterande nätverk

Att bygga högpresterande nätverkssystem kräver att man omprövar flera grundläggande antaganden:

1. Buffertägarskap: Till skillnad från traditionella socketar där kärnan hanterar nätverksbuffertar och kräver kopiering mellan användarrymd och kärnrymd, kräver RDMA att applikationer hanterar sina egna buffertar. När en applikation initierar en nätverksåtgärd överför den buffertägarskapet till nätverkskortet tills åtgärden är slutförd, vilket eliminerar behovet av datakopiering.
2. Minnesregistrering: Applikationer måste registrera minnesregioner hos operativsystemets kärna. Kärnan sätter upp virtuella adressmappningar som gör att CPU:n, GPU:erna och nätverkskorten alla förstår samma virtuella adresser. Denna registrering är en engångsåtgärd som möjliggör efterföljande nollkopieringsdataöverföringar.
3. Styrplan vs dataplan: Högpresterande nätverk delar upp operationer i två kategorier:
   1. Styrplansoperationer (som anslutningsinställning och minnesregistrering) går igenom kärnan för att säkerställa säkerhet
   2. Dataplansoperationer (faktisk dataöverföring) kringgår kärnan för maximal prestanda
4. Mottagning före sändning: Utan kärnhanterade buffertar måste applikationer i förväg lägga upp mottagningsoperationer och ange var inkommande data ska placeras. Detta är ett grundläggande skifte från socket-modellen där applikationer kan ta emot data när som helst.
5. Avsökningsbaserat slutförande: Istället för att vänta på kärnnotifieringar genom mekanismer som epoll avsöker applikationer direkt hårdvarans slutförandeköer. Detta eliminerar omkostnader för systemanrop och möjliggör omedelbar reaktion på slutförda operationer.
6. Medvetenhet om hårdvarutopologi: Att förstå och optimera för hårdvarutopologi är avgörande för att uppnå maximal prestanda.

Att förstå hårdvarutopologi

AWS p5-instanser har en sofistikerad intern arkitektur. Som visas nedan innehåller varje instans två CPU-socklar som bildar två NUMA-noder, där varje NUMA-nod ansluter till fyra PCIe-switchar:

Förstora

Under varje PCIe-switch hittar vi:

• Fyra 100 Gbps EFA-nätverkskort
• En NVIDIA H100 GPU
• En NVMe SSD

Datavägarna för TCP/IP- och RDMA-överföringar visar på grundläggande skillnader i deras tillvägagångssätt:

Förstora

Med TCP/IP (vänster sida) måste data kopieras flera gånger:

1. Från GPU-minne till applikationsbuffert i användarrymd
2. Från buffert i användarrymd till kärnbuffert
3. Från kärnbuffert till nätverkskort
4. Den omvända processen sker på mottagarsidan

Varje kopieringsoperation förbrukar CPU-cykler och minnesbandbredd. Applikationen måste också kontextväxla mellan användarrymd och kärnrymd för varje nätverksoperation.

Däremot erbjuder RDMA (höger sida) äkta nollkopieringsdataöverföring:

1. Nätverkskortet läser direkt från GPU-minnet
2. Data färdas direkt till fjärrnätverkskortet
3. Fjärrnätverkskortet skriver direkt till målets GPU-minne
4. Applikationen behöver bara kontrollera en slutförandekö (CQ) i användarrymd för att veta när överföringen är klar

Med RDMA och korrekt hårdvaruparning kräver överföring av data mellan två GPU:er endast att man passerar den lokala PCIe-switchen och nätverket:

Förstora

Däremot måste TCP/IP-överföringar kopiera data flera gånger genom huvudminnet, vilket orsakar betydande överbelastning på PCIe-bussen:

Förstora

Bygga med libfabric

Vi använde libfabric – ett ramverk som tillhandahåller ett generiskt gränssnitt för tygtjänster. Vår implementation använder två typer av RDMA-operationer:

• Tvåsidig RDMA (SEND/RECV) för kontrollmeddelanden som bär metadata om minnesregioner (t.ex. förskjutningar och storlekar).
• Ensidig RDMA WRITE för faktisk dataöverföring, där varje WRITE-operation hanterar ett sammanhängande minnessegment.

Resan mot toppen av prestanda

Vår utveckling fortlöpte genom flera stadier:

1. Implementerade grundläggande enkelriktad meddelandeöverföring med SEND/RECV
2. Utökades till dubbelriktad kommunikation
3. Lade till GPUDirect RDMA WRITE för direkt GPU-till-GPU-överföring
4. Expanderades för att hantera flera samtidiga överföringar
5. Införde operationsköer för robusthet
6. Uppnådde 97,4 % bandbreddsutnyttjande på ett enda nätverkskort

När vi skalade upp till 32 nätverkskort implementerade vi flera avgörande optimeringar:

1. Operationsköer: Istället för att skicka in operationer direkt till nätverkskort underhåller vi en kö på applikationsnivå. Detta ger robusthet mot nätverksstockning och förenklar programmeringsmodellen.
2. Nätverksuppvärmning: Förhandsupprättande av anslutningar förbättrade startprestandan.
3. Fleretrådning: Dedikerade trådar för varje GPU:s nätverksoperationer.
4. CPU-kärnfastsättning: Att binda trådar till specifika CPU-kärnor för att undvika NUMA-effekter och cachemissar.
5. Tillståndssharding: Att minska konkurrensen mellan trådar.
6. Batchhantering av operationer: Att skicka in flera operationer tillsammans.
7. Lazy Operation Posting: Operationer köas först i applikationen. Efter avsökning av slutförandekön försöker vi skicka in väntande operationer till nätverkskortet, vilket säkerställer effektiv användning av nätverksresurser.
8. NUMA-medveten resursallokering: Allokera libfabric-resurser som slutförandeköer på rätt NUMA-nod för att minimera latensen för minnesåtkomst.

Genom dessa optimeringar uppnådde vi en slutlig prestanda på 3 108 Gbps – 97,1 % av den teoretiska maximala bandbredden.

Videon nedan visar vårt kommandoradsprogram i action. Det överför icke-sammanhängande bitar av GPU-minne till en fjärrnod och uppnår en överföringshastighet på 3108,283 Gbps – vilket demonstrerar en nästan teoretisk bandbreddsanvändning av nätverksinfrastrukturen:

Slutsats

Att bygga ett högpresterande nätverkssystem kräver förståelse för både hårdvaruarkitektur och systemdesignprinciper. Även om bibliotek som NCCL erbjuder utmärkningsvisa lösningar för vanliga mönster, är anpassade lösningar ibland nödvändiga för specifika krav.

Vår resa visar att det är möjligt att uppnå en nätverksprestanda som är nära den teoretiska med noggrann uppmärksamhet på systemarkitektur, hårdvarutopologi och olika optimeringstekniker. Nyckeln är inte bara att förstå enskilda komponenter, utan hur de samverkar för att bilda ett komplett system.

Den fullständiga tekniska genomgången av denna resa, inklusive implementationsdetaljer och kodexempel, finns tillgänglig i vårt öppen källkods-förråd.

Citering

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