Nagy teljesítményű GPU-memóriaátvitel AWS Sagemaker Hyperpodon

A modern deep learning infrastruktúra gyakran igényli nagy mennyiségű adat továbbítását a gépek közötti GPU-k között. A Perplexitynél egy egyedi technikai kihívással találkoztunk: a nem folytonos GPU-memóriaterületek gépek közötti, maximális lehetséges sebességgel történő, hatékony átvitelével. Célplatformunk, az AWS p5 példányok, lenyűgöző 3200 Gbps hálózati sávszélességet biztosítanak 32 hálózati kártyán keresztül. Ez a cikk bemutatja útját egy olyan egyedi, nagy teljesítményű hálózati megoldás felépítésének, amely eléri e elméleti sávszélesség 97,1%-át.

A technikai kihívás

Használati esetünk számos kulcsfontosságú technikai követelményt támasztott:

• Nagy sávszélességű adatátvitel távoli GPU-k nem folytonos memóriadarabjai között
• Lehetőség csomópontok dinamikus hozzáadására vagy eltávolítására a Kubernetes-telepítésekből a folyamatos működés megszakítása nélkül
• Peer-to-peer kommunikációs minták támogatása

Bár az NVIDIA NCCL könyvtára a de facto szabvány elosztott deep learninghez, nem volt ideális a mi használati esetünkhöz:

• Az NCCL kiválóan alkalmas kollektív kommunikációra, de statikus "világ" létrehozását igényli, ami a részt vevő csomópontok módosításakor a teljes fürt újraindítását követeli meg.
• Az NCCL szinkron kommunikációs modellje bonyolultságot ad az aszinkron munkaterhelésünkhöz
• Közvetlen irányítást akartunk gyakorolni memóriaátviteli mintáink felett az optimalizálás érdekében
• Saját megoldásunk megépítése értékes tanulási lehetőségeket biztosított

Modern nagy teljesítményű hálózatok

Megoldásunk megértéséhez először vizsgáljuk meg, hogy a modern nagy teljesítményű hálózatok miben különböznek a hagyományos hálózatoktól.

A mindennap használt hálózatok többsége TCP/IP protokollokra támaszkodik, ahol az alkalmazások az operációs rendszer kernelén keresztül, socketek használatával kommunikálnak a hálózati kártyával. A nagy teljesítményű hálózatok azonban RDMA-t (Remote Direct Memory Access) használnak – egy teljesen eltérő hardver- és szoftververmet, amely közvetlen memóriahozzáférést tesz lehetővé a gépek között a CPU bevonása nélkül.

Az AWS az Elastic Fabric Adapter (EFA) szolgáltatást nyújtja, amely egy egyedi hálózati interfész, és az Amazon egyedi, Scalable Reliable Datagram (SRD) nevű protokollját valósítja meg. A hagyományos TCP/IP hálózatokkal ellentétben, ahol az adatokat többször is másolni kell a felhasználói tér, a kernel tér és a hálózati pufferek között, az EFA az RDMA-val lehetővé teszi a közvetlen adatátvitelt a GPU-memória és a hálózati kártya között, teljesen megkerülve a CPU-t.

A nagy teljesítményű hálózati tervezés filozófiája

A nagy teljesítményű hálózati rendszerek építése számos alapvető feltételezés újragondolását igényli:

1. Puffer-tulajdonlás: Ellentétben a hagyományos socketekkel, ahol a kernel kezeli a hálózati puffereket, és megköveteli a másolást a felhasználói tér (user space) és a kernel tér között, az RDMA megköveteli, hogy az alkalmazások saját maguk kezeljék a puffereiket. Amikor egy alkalmazás hálózati műveletet indít, átruházza a puffer tulajdonjogát a hálózati kártyára a művelet befejeződéséig, ezzel kiküszöbölve az adatmásolás szükségességét.
2. Memóriaregisztráció: Az alkalmazásoknak regisztrálniuk kell a memóriaterületeket az operációs rendszer kernelénél. A kernel olyan virtuális címleképezéseket hoz létre, amelyek lehetővé teszik, hogy a CPU, a GPU-k és a hálózati kártyák mind ugyanazokat a virtuális címeket értsék meg. Ez a regisztráció egy egyszeri művelet, amely lehetővé teszi a későbbi zéró másolású (zero-copy) adatátviteleket.
3. Vezérlési sík vs. Adatsík (Control Plane vs Data Plane): A nagy teljesítményű hálózatok két kategóriára osztják a műveleteket:
   1. A vezérlési sík (control plane) műveletei (például a kapcsolatbeállítás és a memóriaregisztráció) a kernelen keresztül haladnak a biztonság garantálása érdekében
   2. Az adatsík (data plane) műveletei (a tényleges adatátvitel) megkerülik a kernelt a maximális teljesítmény érdekében
4. Fogadás az átvitel előtt: Kernel által nem kezelt pufferek hiányában az alkalmazásoknak elő kell készíteniük és közzé kell tenniük a fogadási (receive) műveleteket, megadva, hogy hová kerüljenek a bejövő adatok. Ez alapvető elmozdulás a socket modelltől, ahol az alkalmazások bármikor képesek adatokat fogadni.
5. Lekérdezésen alapuló befejezés (Poll-based Completion): Ahelyett, hogy olyan mechanizmusokon keresztül várnának a kernel értesítéseire, mint az epoll, az alkalmazások közvetlenül lekérdezik a hardveres befejezési várólistákat. Ez kiküszöböli a rendszerhívások többletterhét, és azonnali reakciót tesz lehetővé a befejeződött műveletekre.
6. Hardver topológia tudatosság: A hardver topológia megértése és az arra való optimalizálás elengedhetetlen a maximális teljesítmény eléréséhez.

A hardver topológia megértése

Az AWS p5 példányok kifinomult belső architektúrával rendelkeznek. Ahogy az alábbiakban látható, minden példány két CPU-foglalatot tartalmaz, amelyek két NUMA-csomópontot alkotnak, és minden NUMA-csomópont négy PCIe-kapcsolóhoz csatlakozik:

Nagyítás

Minden PCIe-kapcsoló alatt az alábbiak találhatók:

• Négy 100 Gbps sebességű EFA hálózati kártya
• Egy NVIDIA H100 GPU
• Egy NVMe SSD

A TCP/IP és az RDMA átvitelek adatútvonalai szemléltetik a megközelítéseik közötti alapvető különbségeket:

Nagyítás

TCP/IP esetén (bal oldalon) az adatokat többször át kell másolni:

1. A GPU-memóriából a felhasználói térben (user space) lévő alkalmazás pufferbe
2. A felhasználói térben (user space) lévő pufferből a kernel pufferbe
3. A kernel pufferből a hálózati kártyára
4. A fordított folyamat megy végbe a fogadó oldalon

Minden másolási művelet CPU-ciklusokat és memóriasávszélességet emészt fel. Az alkalmazásnak ráadásul kontextust kell váltania a felhasználói tér (user space) és a kernel tér között minden egyes hálózati művelethez.

Ezzel szemben az RDMA (jobb oldalon) valódi zéró másolású (zero-copy) adatátvitelt biztosít:

1. A hálózati kártya közvetlenül olvassa a GPU-memóriát
2. Az adatok közvetlenül a távoli hálózati kártyára utaznak
3. A távoli hálózati kártya közvetlenül a cél GPU-memóriába ír
4. Az alkalmazásnak csak egy felhasználói térben (user space) lévő befejezési várólistát (completion queue - CQ) kell ellenőriznie, hogy megtudja, mikor fejeződött be az átvitel

RDMA-val és megfelelő hardverpárosítással a két GPU közötti adatátvitelhez csupán a helyi PCIe-kapcsolón és a hálózaton kell áthaladni:

Nagyítás

Ezzel szemben a TCP/IP-átvitelek során az adatokat többször át kell másolni a főmemórián keresztül, ami jelentős PCIe-busztorlódást okoz:

Nagyítás

Fejlesztés a libfabric használatával

A libfabricet használtuk – egy olyan keretrendszert, amely általános interfészt biztosít fabric szolgáltatásokhoz. Implementációnk kétféle RDMA-műveletet használ:

• Kétirányú (two-sided) RDMA (SEND/RECV) vezérlőüzenetekhez, amelyek memóriaterületekről szóló metaadatokat (pl. eltolásokat és méreteket) hordoznak.
• Egyirányú (one-sided) RDMA WRITE a tényleges adatátvitelhez, ahol minden egyes WRITE művelet egy folytonos memóriadarabot kezel.

Út a csúcsteljesítményhez

Fejlesztésünk több szakaszon ment keresztül:

1. Alapvető egyirányú üzenetátvitel implementálása SEND/RECV használatával
2. Kiterjesztés kétirányú kommunikációra
3. GPUDirect RDMA WRITE hozzáadása a közvetlen GPU-GPU átvitelhez
4. Kiterjesztés több párhuzamos átvitel kezelésére
5. Műveleti várakozási sor bevezetése a robusztusság érdekében
6. 97,4%-os sávszélesség-kihasználtság elérése egyetlen hálózati kártyán

A 32 hálózati kártyára történő skálázáskor számos kulcsfontosságú optimalizációt hajtottunk végre:

1. Műveleti sorba állítás: Ahelyett, hogy közvetlenül küldenénk be műveleteket a hálózati kártyáknak, alkalmazásszintű sort tartunk fenn. Ez robusztusságot biztosít a hálózati torlódásokkal szemben, és egyszerűsíti a programozási modellt.
2. Hálózat bemelegítése: A kapcsolatok előzetes létrehozása javította az indítási teljesítményt.
3. Többszálúság (Multi-threading): Külön szálak az egyes GPU-k hálózati műveleteihez.
4. CPU-magokhoz rögzítés (Core Pinning): A szálak hozzárendelése specifikus CPU-magokhoz a NUMA-hatások és a gyorsítótár-tévesztések elkerülése érdekében.
5. Állapot-sharding: A szálak közötti versengés csökkentése.
6. Műveletek kötegelése (Batching): Több művelet együttes beküldése.
7. Lusta műveletközzététel (Lazy Operation Posting): A műveletek először az alkalmazásban kerülnek sorbaállításra. A befejezési várólista (completion queue) lekérdezése után megkíséreljük beküldeni a függőben lévő műveleteket a hálózati kártyának, biztosítva a hálózati erőforrások hatékony felhasználását.
8. NUMA-tudatos erőforrás-foglalás: Libfabric erőforrások, például befejezési várólisták lefoglalása a helyes NUMA-csomóponton a memóriahozzáférési késleltetés minimalizálása érdekében.

Ezen optimalizációk révén 3108 Gbps végső teljesítményt értünk el – az elméleti maximális sávszélesség 97,1%-át.

Az alábbi videó a parancssori programunkat mutatja működés közben. A GPU-memória nem folytonos darabjait továbbítja egy távoli csomópontra, 3108,283 Gbps átviteli sebességet érve el – szemléltetve a hálózati infrastruktúra elméletit megközelítő sávszélesség-kihasználtságát:

Összegzés

A nagy teljesítményű hálózati rendszerek építése megköveteli mind a hardverarchitektúra, mind a rendszertervezési alapelvek megértését. Míg az olyan könyvtáraink, mint az NCCL, kiváló megoldásokat nyújtanak a gyakori mintákhoz, bizonyos egyedi követelményekhez egyedi megoldásokra van szükség.

Utunk bemutatja, hogy az elméletit megközelítő hálózati teljesítmény elérhető a rendszerarchitektúrára, a hardver topológiára és a különféle optimalizációs technikákra fordított aprólékos figyelemmel. A kulcs nem csupán az egyes komponensek megértése, hanem az is, hogy azok hogyan lépnek kapcsolatba egymással egy teljes rendszert alkotva.

Ennek az útnak a teljes technikai részletezése, beleértve az implementációs részleteket és a kódpéldákat, elérhető az nyílt forráskódú tárhelyünkön (repository).

Idézet

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