Vysokovýkonný prenos pamäte GPU na AWS Sagemaker Hyperpod

Moderná infraštruktúra hlbokého učenia často vyžaduje prenos veľkého množstva dát medzi GPU naprieč strojmi. V spoločnosti Perplexity sme narazili na jedinečnú technickú výzvu: efektívny prenos nesúvislých oblastí pamäte GPU medzi strojmi pri maximálnej možnej rýchlosti. Naša cieľová platforma, inštancie AWS p5, ponúka pôsobivú sieťovú šírku pásma 3200 Gbps prostredníctvom 32 sieťových kariet. Tento článok zdieľa našu cestu budovania vlastného vysoko výkonného sieťového riešenia, ktoré dosahuje 97,1 % tejto teoretickej šírky pásma.

Technická výzva

Náš prípad použitia priniesol niekoľko kľúčových technických požiadaviek:

• Vysokorýchlostný prenos medzi vzdialenými GPU nesúvislých pamäťových blokov
• Schopnosť dynamicky pridávať alebo odoberať uzly z nasadení Kubernetes bez narušenia prebiehajúcich operácií
• Podpora pre vzorce komunikácie typu peer-to-peer

Zatiaľ čo knižnica NCCL od spoločnosti NVIDIA je de facto štandardom pre distribuované hlboké učenie, pre náš prípad použitia nebola ideálna:

• NCCL vyniká v kolektívnej komunikácii, ale vyžaduje vytvorenie statického "sveta", čo vyžaduje reštartovanie celého klastra pri úprave zúčastnených uzlov.
• Synchrónny komunikačný model NCCL pridáva komplexnosť pre našu asynchrónnu pracovnú záťaž
• Chceli sme priamu kontrolu nad vzorcami prenosu pamäte na účely optimalizácie
• Budovanie vlastného riešenia poskytlo cenné príležitosti na učenie

Moderné vysoko výkonné siete

Aby sme porozumeli nášmu riešeniu, poďme najprv preskúmať, ako sa moderné vysoko výkonné siete líšia od tradičných sietí.

Väčšina sietí, ktoré denne používame, sa spolieha na protokoly TCP/IP, kde aplikácie komunikujú so sieťovou kartou cez jadro operačného systému pomocou soketov. Vysokorýchlostné siete však využívajú RDMA (Remote Direct Memory Access) – úplne iný hardvérový a softvérový zásobník, ktorý umožňuje priamy prístup do pamäte medzi strojmi bez zapojenia CPU.

AWS poskytuje Elastic Fabric Adapter (EFA), vlastné sieťové rozhranie, ktoré implementuje vlastný protokol od Amazonu s názvom Scalable Reliable Datagram (SRD). Na rozdiel od tradičných sietí TCP/IP, kde sa dáta musia viackrát kopírovať medzi používateľským priestorom, priestorom jadra a sieťovými vyrovnávacími pamäťami, EFA s RDMA umožňuje priamy prenos dát medzi pamäťou GPU a sieťovou kartou, pričom úplne obchádza CPU.

Filozofia návrhu vysoko výkonných sietí

Budovanie vysoko výkonných sieťových systémov si vyžaduje prehodnotenie niekoľkých základných predpokladov:

1. Vlastníctvo vyrovnávacej pamäte: Na rozdiel od tradičných soketov, kde jadro spravuje sieťové vyrovnávacie pamäte a vyžaduje kopírovanie medzi používateľským priestorom a priestorom jadra, RDMA vyžaduje, aby si aplikácie spravovali vlastné vyrovnávacie pamäte. Keď aplikácia iniciuje sieťovú operáciu, prenesie vlastníctvo vyrovnávacej pamäte na sieťovú kartu, kým sa operácia dokončí, čím sa eliminuje potreba kopírovania dát.
2. Registrácia pamäte: Aplikácie musia zaregistrovať oblasti pamäte v jadre operačného systému. Jadro nastaví mapovania virtuálnych adries, ktoré umožňujú CPU, GPU a sieťovým kartám porozumieť rovnakým virtuálnym adresám. Táto registrácia je jednorazová operácia, ktorá umožňuje následné prenosy dát bez kopírovania.
3. Riadiaca rovina vs. Dátová rovina: Vysoko výkonné siete delia operácie do dvoch kategórií:
   1. Operácie riadiacej roviny (ako nastavenie pripojenia a registrácia pamäte) prechádzajú cez jadro, aby sa zaistila bezpečnosť
   2. Operácie dátovej roviny (skutočný prenos dát) obchádzajú jadro pre maximálny výkon
4. Príjem pred prenosom: Bez vyrovnávacích pamätí spravovaných jadrom musia aplikácie vopred zverejniť operácie príjmu a špecifikovať, kde sa majú prichádzajúce dáta umiestniť. Ide o zásadný posun oproti modelu soketov, kde aplikácie môžu prijímať dáta kedykoľvek.
5. Dokončenie založené na dotazovaní: Namiesto čakania na upozornenia jadra prostredníctvom mechanizmov ako epoll aplikácie priamo dotazujú fronty dokončení hardvéru. To eliminuje réžiu systémových volaní a umožňuje okamžitú reakciu na dokončené operácie.
6. Povedomie o hardvérovej topológii: Pochopenie a optimalizácia hardvérovej topológie sú kľúčové pre dosiahnutie maximálneho výkonu.

Pochopenie hardvérovej topológie

Inštancie AWS p5 majú prepracovanú vnútornú architektúru. Ako je ukázané nižšie, každá inštancia obsahuje dva sokety CPU tvoriace dva uzly NUMA, pričom každý uzol NUMA sa pripája k štyrom prepínačom PCIe:

Zväčšiť

Pod každým prepínačom PCIe nájdeme:

• Štyri 100 Gbps sieťové karty EFA
• Jednu NVIDIA H100 GPU
• Jeden NVMe SSD

Dátové cesty pre prenosy TCP/IP a RDMA demonštrujú zásadné rozdiely v ich prístupoch:

Zväčšiť

S TCP/IP (ľavá strana) sa dáta musia kopírovať viackrát:

1. Z pamäte GPU do vyrovnávacej pamäte aplikácie v používateľskom priestore
2. Z vyrovnávacej pamäte používateľského priestoru do vyrovnávacej pamäte jadra
3. Z vyrovnávacej pamäte jadra do sieťovej karty
4. Obrátený proces prebieha na strane príjemcu

Každá operácia kopírovania spotrebúva cykly CPU a šírku pásma pamäte. Aplikácia musí tiež prepínať kontext medzi používateľským priestorom a priestorom jadra pre každú sieťovú operáciu.

Naproti tomu RDMA (pravá strana) poskytuje skutočný prenos dát bez kopírovania (zero-copy):

1. Sieťová karta číta priamo z pamäte GPU
2. Dáta putujú priamo do vzdialenej sieťovej karty
3. Vzdialená sieťová karta zapisuje priamo do cieľovej pamäte GPU
4. Aplikácia potrebuje iba skontrolovať front dokončení (CQ) v používateľskom priestore, aby zistila, kedy je prenos dokončený

S RDMA a správnym spárovaním hardvéru si prenos dát medzi dvoma GPU vyžaduje iba prechod cez lokálny prepínač PCIe a sieť:

Zväčšiť

Na rozdiel od toho prenosy TCP/IP musia kopírovať dáta viackrát cez hlavnú pamäť, čo spôsobovať výrazné preťaženie zbernice PCIe:

Zväčšiť

Budovanie pomocou libfabric

Použili sme libfabric – framework, ktorý poskytuje generické rozhranie pre služby fabric. Naša implementácia využíva dva typy operácií RDMA:

• Obojsmerné RDMA (SEND/RECV) pre riadiace správy, ktoré nesú metadáta o oblastiach pamäte (napr. posuny a veľkosti).
• Jednostranné RDMA WRITE pre samotný prenos dát, kde každá operácia WRITE spracováva jeden súvislý pamäťový blok.

Cesta k maximálnemu výkonu

Náš vývoj prebiehal v niekoľkých fázach:

1. Implementovaný základný jednosmerný prenos správ pomocou SEND/RECV
2. Rozšírené na obojsmernú komunikáciu
3. Pridané GPUDirect RDMA WRITE pre priamy prenos medzi GPU
4. Rozšírené na zvládnutie viacerých súbežných prenosov
5. Zavedené radenie operácií do frontu kvôli robustnosti
6. Dosiahnuté 97,4% využitie šírky pásma na jednej sieťovej karte

Pri škálovaní na 32 sieťových kariet sme implementovali niekoľko kľúčových optimalizácií:

1. Radenie operácií do frontu: Namiesto priameho odosielania operácií do sieťových kariet udržiavame front na úrovni aplikácie. To poskytuje robustnosť voči preťaženiu siete a zjednodušuje model programovania.
2. Zahriatie siete: Predbežné nadviazanie spojení zlepšilo výkon pri spúšťaní.
3. Viacvláknovosť: Vyhradené vlákna pre sieťové operácie každej GPU.
4. Pripnutie jadier CPU: Naviazanie vlákien na konkrétne jadrá CPU, aby sa predišlo účinkom NUMA a výpadkom vyrovnávacej pamäte.
5. Sharding stavu: Zníženie konkurencie medzi vláknami.
6. Dávkovanie operácií: Odosielanie viacerých operácií naraz.
7. Lajdácke zverejňovanie operácií: Operácie sa najprv zaraďujú do frontu v aplikácii. Po dotazovaní frontu dokončení sa pokúsime odoslať čakajúce operácie do sieťovej karty, čím sa zabezpečí efektívne využitie sieťových zdrojov.
8. Pridelenie zdrojov s ohľadom na NUMA: Priradenie zdrojov libfabric, ako sú fronty dokončení, na správny uzol NUMA, aby sa minimalizovala latencia prístupu do pamäte.

Prostredníctvom týchto optimalizácií sme dosiahli konečný výkon 3 108 Gbps – 97,1 % teoretického maxima šírky pásma.

Video nižšie ukazuje náš program príkazového riadka v akcii. Prenáša nesúvislé bloky pamäte GPU do vzdialeného uzla, čím dosahuje prenosovú rýchlosť 3108,283 Gbps – čo demonštruje takmer teoretické využitie šírky pásma sieťovej infraštruktúry:

Záver

Budovanie vysoko výkonného sieťového systému si vyžaduje pochopenie hardvérovej architektúry aj princípov návrhu systému. Zatiaľ čo knižnice ako NCCL poskytujú vynikajúce riešenia pre bežné vzorce, niekedy sú pre špecifické požiadavky nutné vlastné riešenia.

Naša cesta dokazuje, že dosiahnutie takmer teoretického sieťového výkonu je možné s dôkladnou pozornosťou venovanou systémovej architektúre, hardvérovej topológii a rôznym optimalizačným technikám. Kľúčom nie je len pochopenie jednotlivých komponentov, ale aj to, ako vzájomne interagujú a vytvárajú kompletný systém.

Úplný technický rozbor tejto cesty vrátane podrobností implementácie a príkladov kódu je k dispozícii v našom úložisku open source.

Citácia

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