Krachtige GPU-geheugenoverdracht op AWS Sagemaker Hyperpod

Moderne infrastructuur voor deep learning vereist vaak het overdragen van grote hoeveelheden data tussen GPU's over verschillende machines heen. Bij Perplexity stuitten we op een unieke technische uitdaging: het efficiënt overdragen van niet-aaneengesloten GPU-geheugenregio's tussen machines met de hoogst mogelijke snelheid. Ons doelplatform, AWS p5-instanties, biedt een indrukwekkende 3200 Gbps aan netwerkbandbreedte via 32 netwerkkaarten. Dit artikel deelt onze reis bij het bouwen van een op maat gemaakte, krachtige netwerkoplossing die 97,1% van deze theoretische bandbreedte bereikt.

De technische uitdaging

Onze use case stelde verschillende belangrijke technische vereisten:

• Hoge bandbreedte-overdracht tussen externe GPU's van niet-aaneengesloten geheugenblokken
• Mogelijkheid om dynamisch nodes toe te voegen aan of te verwijderen uit Kubernetes-implementaties zonder lopende operaties te verstoren
• Ondersteuning voor peer-to-peer communicatiepatronen

Hoewel NVIDIA's NCCL-bibliotheek de de facto standaard is voor gedistribueerd deep learning, was deze niet ideaal voor onze use case:

• NCCL blinkt uit in collectieve communicatie, maar vereist het opzetten van een statische "wereld", wat het herstarten van het gehele cluster vereist bij het aanpassen van de deelnemende nodes.
• NCCL's synchrone communicatiemodel voegt complexiteit toe voor onze asynchrone werklast
• Wij wilden directe controle over onze geheugenoverdrachtpatronen voor optimalisatie
• Het bouwen van onze eigen oplossing bood waardevolle leermogelijkheden

Moderne krachtige netwerken

Om onze oplossing te begrijpen, onderzoeken we eerst hoe moderne krachtige netwerken verschillen van traditionele netwerken.

De meeste netwerken die we dagelijks gebruiken, vertrouwen op TCP/IP-protocollen, waarbij applicaties communiceren met de netwerkkaart via de kernel van het besturingssysteem met behulp van sockets. Krachtige netwerken gebruiken echter RDMA (Remote Direct Memory Access) - een totaal andere hardware- en softwarestack die directe geheugentoegang tussen machines mogelijk maakt zonder de CPU te betrekken.

AWS biedt de Elastic Fabric Adapter (EFA), een aangepaste netwerkinterface die het aangepaste protocol van Amazon genaamd Scalable Reliable Datagram (SRD) implementeert. In tegenstelling tot traditionele TCP/IP-netwerken waarbij data meerdere malen moet worden gekopieerd tussen gebruikersruimte, kernelruimte en netwerkbuffers, maakt EFA met RDMA directe gegevensoverdracht mogelijk tussen GPU-geheugen en de netwerkkaart, waarbij de CPU volledig wordt omzeild.

Filosofie van krachtig netwerkontwerp

Het bouwen van krachtige netwerksystemen vereist het heroverwegen van verschillende fundamentele aannames:

1. Buffer-eigenaarschap: In tegenstelling tot traditionele sockets waarbij de kernel netwerkbuffers beheert en kopiëren tussen gebruikersruimte en kernelruimte vereist, vereist RDMA dat applicaties hun eigen buffers beheren. Wanneer een applicatie een netwerkoperatie initieert, draagt het buffer-eigenaarschap over aan de netwerkkaart totdat de operatie is voltooid, waardoor gegevens kopiëren overbodig wordt.
2. Geheugenregistratie: Applicaties moeten geheugenregio's registreren bij de kernel van het besturingssysteem. De kernel stelt virtuele adresmappings in die de CPU, GPU's en netwerkkaarten in staat stellen allemaal dezelfde virtuele adressen te begrijpen. Deze registratie is een eenmalige operatie die daaropvolgende zero-copy gegevensoverdrachten mogelijk maakt.
3. Besturingsvlak vs. datavlak: Krachtige netwerken scheiden operaties in twee categorieën:
   1. Besturingsvlak-operaties (zoals verbindingsopzet en geheugenregistratie) verlopen via de kernel om veiligheid te waarborgen
   2. Datavlak-operaties (werkelijke gegevensoverdracht) omzeilen de kernel voor maximale prestaties
4. Ontvangst vóór verzending: Zonder door de kernel beheerde buffers moeten applicaties ontvangstoperaties vooraf plaatsen, waarbij wordt gespecificeerd waar inkomende data moet worden geplaatst. Dit is een fundamentele verschuiving ten opzichte van het socketmodel waarbij applicaties op elk moment data kunnen ontvangen.
5. Poll-gebaseerde voltooiing: In plaats van te wachten op kernelmeldingen via mechanismen zoals epoll, pollen applicaties direct hardware-voltooingswachtrijen. Dit elimineert overhead van systeemaanroepen en maakt directe reactie op voltooide operaties mogelijk.
6. Hardware-topologiebewustzijn: Begrip van en optimalisatie voor hardware-topologie is cruciaal voor het bereiken van maximale prestaties.

Inzicht in hardware-topologie

AWS p5-instanties hebben een geavanceerde interne architectuur. Zoals hieronder getoond, bevat elke instantie twee CPU-sockets die twee NUMA-nodes vormen, waarbij elke NUMA-node verbinding maakt met vier PCIe-switches:

Vergroten

Onder elke PCIe-switch vinden we:

• Vier 100 Gbps EFA-netwerkkaarten
• Eén NVIDIA H100 GPU
• Eén NVMe SSD

De datapaden voor TCP/IP- en RDMA-overdrachten tonen fundamentele verschillen in hun benadering:

Vergroten

Bij TCP/IP (linkerkant) moet data meerdere malen worden gekopieerd:

1. Van GPU-geheugen naar applicatiebuffer in gebruikersruimte
2. Van gebruikersruimtebuffer naar kernelbuffer
3. Van kernelbuffer naar netwerkkaart
4. Het omgekeerde proces vindt plaats aan de ontvangende kant

Elke kopieeroperatie verbruikt CPU-cycli en geheugenbandbreedte. De applicatie moet ook context schakelen tussen gebruikersruimte en kernelruimte voor elke netwerkoperatie.

Daarentegen biedt RDMA (rechterkant) echte zero-copy gegevensoverdracht:

1. De netwerkkaart leest direct uit het GPU-geheugen
2. Data reist rechtstreeks naar de externe netwerkkaart
3. De externe netwerkkaart schrijft direct naar het doel-GPU-geheugen
4. De applicatie hoeft alleen een voltooingswachtrij (CQ) in de gebruikersruimte te controleren om te weten wanneer de overdracht is voltooid

Met RDMA en de juiste hardwarekoppeling vereist het overdragen van data tussen twee GPU's alleen het doorkruisen van de lokale PCIe-switch en het netwerk:

Vergroten

TCP/IP-overdrachten moeten daarentegen data meerdere keren kopiëren via het hoofdgeheugen, wat aanzienlijke PCIe-buscongestie veroorzaakt:

Vergroten

Bouwen met libfabric

We gebruikten libfabric - een framework dat een algemene interface biedt voor fabric-services. Onze implementatie gebruikt twee soorten RDMA-operaties:

• Tweezijdige RDMA (SEND/RECV) voor controleberichten die metadata over geheugenregio's bevatten (bijv. offsets en groottes).
• Eenzijdige RDMA WRITE voor werkelijke gegevensoverdracht, waarbij elke WRITE-operatie één aaneengesloten geheugenblok afhandelt.

Reis naar topprestaties

Onze ontwikkeling verliep in verschillende fasen:

1. Basis unidirectionele berichtoverdracht geïmplementeerd met behulp van SEND/RECV
2. Uitgebreid naar bidirectionele communicatie
3. GPUDirect RDMA WRITE toegevoegd voor directe GPU-GPU overdracht
4. Uitgebreid om meerdere gelijktijdige overdrachten te verwerken
5. Wachtrij voor operaties geïntroduceerd voor robuustheid
6. 97,4% bandbreedtebenutting bereikt op een enkele netwerkkaart

Bij het schalen naar 32 netwerkkaarten hebben we verschillende cruciale optimalisaties doorgevoerd:

1. Operatiewachtrij: In plaats van operaties direct naar netwerkkaarten te sturen, onderhouden we een wachtrij op applicatieniveau. Dit biedt robuustheid tegen netwerkcongestie en vereenvoudigt het programmeermodel.
2. Netwerkopwarming: Het vooraf tot stand brengen van verbindingen verbeterde de opstartprestaties.
3. Multi-threading: Toegewijde draden voor de netwerkoperaties van elke GPU.
4. CPU Core Pinning: Het binden van draden aan specifieke CPU-cores om NUMA-effecten en cachemissers te voorkomen.
5. State Sharding: Het verminderen van strijd tussen draden.
6. Operatie-batching: Meerdere operaties samen indienen.
7. Lazy Operation Posting: Operaties worden eerst in de applicatie in de wachtrij geplaatst. Na het pollen van de voltooingswachtrij proberen we in behandeling zijnde operaties in te dienen bij de netwerkkaart, wat zorgt voor efficiënt gebruik van netwerkbronnen.
8. NUMA-bewuste toewijzing van bronnen: Wijs libfabric-bronnen zoals voltooingswachtrijen toe op de juiste NUMA-node om de latentie van geheugentoegang te minimaliseren.

Door deze optimalisaties bereikten we een uiteindelijke prestatie van 3.108 Gbps - 97,1% van de theoretische maximale bandbreedte.

De onderstaande video toont ons opdrachtregelprogramma in actie. Het draagt niet-aaneengesloten stukken GPU-geheugen over naar een externe node, waarbij een transmissiesnelheid van 3108,283 Gbps wordt bereikt - wat bijna theoretische bandbreedtebenutting van de netwerkinfrastructuur aantoont:

Conclusie

Het bouwen van een krachtig netwerksysteem vereist inzicht in zowel hardware-architectuur als systeemontwerpprincipes. Hoewel bibliotheken zoals NCCL uitstekende oplossingen bieden voor veelvoorkomende patronen, zijn soms aangepaste oplossingen noodzakelijk voor specifieke vereisten.

Onze reis toont aan dat het bereiken van bijna theoretische netwerkprestaties mogelijk is met zorgvuldige aandacht voor systeemarchitectuur, hardware-topologie en diverse optimalisatietechnieken. De sleutel is niet alleen het begrijpen van individuele componenten, maar ook hoe ze interageren om een compleet systeem te vormen.

De volledige technische diepe duik in deze reis, inclusief implementatiedetails en codevoorbeelden, is beschikbaar in onze open source repository.

Citatie

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