Didelio našumo GPU atminties perkėlimas „AWS Sagemaker Hyperpod“

Moderni gilaus mokymo infrastruktūra dažnai reikalauja perkelti didelius duomenų kiekius tarp GPU skirtinguose įrenginiuose. „Perplexity“ susidūrėme su unikaliu techniniu iššūkiu: efektyviai perkelti neperšokančias GPU atminties sritis tarp įrenginių maksimaliu įmanomu greičiu. Mūsų tikslinė platforma, „AWS p5“ instancijos, siūlo įspūdingą 3200 Gbps tinklo pralaidumą per 32 tinklo plokštes. Šiame straipsnyje dalinamės savo kelione kuriant pasirinktinį didelio našumo tinklo sprendimą, kuris pasiekia 97,1 % šio teorinio pralaidumo.

Techninis iššūkis

Mūsų naudojimo atvejis iškėlė kelis pagrindinius techninius reikalavimus:

• Didelio pralaidumo perkėlimas tarp nuotolinių GPU iš neperšokančių atminties blokų
• Galimybė dinamiškai pridėti arba pašalinti mazgus iš „Kubernetes“ diegimų nesutrikdant vykdomų operacijų
• Palaikymas tarpusavio (peer-to-peer) komunikacijos šablonams

Nors „NVIDIA“ biblioteka NCCL yra de facto standartas paskirstytam gilaus mokymuisi, ji nebuvo idealus pasirinkimas mūsų atveju:

• NCCL puikiai tinka kolektyvinei komunikacijai, tačiau reikalauja sukurti statinį „pasaulį“, todėl koreguojant dalyvaujančius mazgus reikia iš naujo paleisti visą klasterį.
• NCCL sinchroninės komunikacijos modelis padidina mūsų asinchroninio darbo krūvio sudėtingumą
• Norėjome tiesioginio atminties perkėlimo šablonų valdymo optimizavimo tikslais
• Mūsų pačių sprendimo kūrimas suteikė vertingų mokymosi galimybių

Šiuolaikiniai didelio našumo tinklai

Norėdami suprasti mūsų sprendimą, pirmiausia išnagrinėkime, kuo šiuolaikiniai didelio našumo tinklai skiriasi nuo tradicinių tinklų.

Daugelis kasdien naudojamų tinklų remiasi TCP/IP protokolais, kuriuose programos bendrauja su tinklo plokšte per operacinės sistemos branduolį naudojant lizdus (sockets). Tačiau didelio našumo tinklai naudoja RDMA (nuotolinę tiesioginę atminties prieigą) – visiškai kitokią aparatinės ir programinės įrangos dėklą, leidžiančią tiesiogiai pasiekti atmintį tarp įrenginių be CPU įsikišimo.

AWS teikia „Elastic Fabric Adapter“ (EFA) – pasirinktinę tinklo sąsają, kuri įgyvendina „Amazon“ pasirinktinį protokolą, vadinamą „Scalable Reliable Datagram“ (SRD). Skirtingai nuo tradicinio TCP/IP tinklo, kur duomenys turi būti nukopijuoti kelis kartus tarp vartotojo erdvės, branduolio erdvės ir tinklo buferiui, EFA su RDMA leidžia tiesiogiai perkelti duomenis tarp GPU atminties ir tinklo plokštės, visiškai apeinant CPU.

Didelio našumo tinklo dizaino filosofija

Didelio našumo tinklo sistemų kūrimas reikalauja iš naujo apsvarstyti kelias pagrindines prielaidas:

1. Buferio nuosavybė: Skirtingai nuo tradicinių lizdų, kur branduolys valdo tinklo buferius ir reikalauja kopijavimo tarp vartotojo erdvės ir branduolio erdvės, RDMA reikalauja, kad programos pačios valdytų savo buferius. Kai programa inicijuoja tinklo operaciją, ji perkelia buferio nuosavybę tinklo plokštei, kol operacija baigiasi, taip pašalindama duomenų kopijavimo poreikį.
2. Atminties registracija: Programos turi užregistruoti atminties sritis operacinės sistemos branduolyje. Branduolys nustato virtualaus adreso susiejimus, leidžiančius CPU, GPU ir tinklo plokštėms suprasti tuos pačius virtualius adresus. Ši registracija yra vienkartinė operacija, įgalinanti vėlesnius nulinio kopijavimo duomenų perkėlimus.
3. Valdymo plokštė ir duomenų plokštė: Didelio našumo tinklai skirsto operacijas į dvi kategorijas:
   1. Valdymo plokštės operacijos (tokios kaip ryšio sąranka ir atminties registracija) eina per branduolį, siekiant užtikrinti saugumą
   2. Duomenų plokštės operacijos (faktinis duomenų perkėlimas) apeina branduolį, siekiant maksimalaus našumo
4. Guvimas prieš perdavimą: Be branduolio valdomų buferiui, programos turi iš anksto paskelbti gavimo operacijas, nurodydamos, kur turėtų būti padėti gaunami duomenys. Tai esminis pokytis nuo lizdo modelio, kur programos gali gauti duomenis bet kuriuo metu.
5. Užbaigimas apklausos pagrindu: Vietoj laukimo branduolio pranešimų per tokius mechanizmus kaip epoll, programos tiesiogiai tikrina aparatinės įrangos užbaigimo eiles. Tai pašalina sisteminių iškvietimų pridėtinę išlaidą ir leidžia nedelsiant reaguoti į užbaigtas operacijas.
6. Aparatinės įrangos topologijos supratimas: Aparatinės įrangos topologijos supratimas ir optimizavimas yra gyvybiškai svarbūs norint pasiekti maksimalų našumą.

Aparatinės įrangos topologijos supratimas

„AWS p5“ instancijos turi sudėtingą vidinę architektūrą. Kaip parodyta žemiau, kiekvienoje instancijoje yra du CPU lizdai, sudarantys du NUMA mazgus, o kiekvienas NUMA mazgas jungiasi prie keturių PCIe jungiklių:

Padidinti

Po kiekvienu PCIe jungikliu randame:

• Keturias 100 Gbps EFA tinklo plokštes
• Vieną „NVIDIA H100 GPU“
• Vieną NVMe SSD

TCP/IP ir RDMA perkėlimų duomenų keliai demonstruoja esminius jų metodų skirtumus:

Padidinti

Naudojant TCP/IP (kairėje pusėje), duomenys turi būti nukopijuoti kelis kartus:

1. Iš GPU atminties į programos buferį vartotojo erdvėje
2. Iš vartotojo erdvės buferio į branduolio buferį
3. Iš branduolio buferio į tinklo plokštę
4. Atvirkštinis procesas vyksta gavimo pusėje

Kiekviena kopijavimo operacija eikvoja CPU ciklus ir atminties pralaidumą. Programa taip pat turi perjungti kontekstą tarp vartotojo erdvės ir branduolio erdvės kiekvienai tinklo operacijai.

Priešingai, RDMA (dešinėje pusėje) užtikrina tikrą nulinio kopijavimo (zero-copy) duomenų perkėlimą:

1. Tinklo plokštė skaito tiesiogiai iš GPU atminties
2. Duomenys keliauja tiesiogiai į nuotolinę tinklo plokštę
3. Nuotolinė tinklo plokštė rašo tiesiogiai į paskirties GPU atmintį
4. Programai tereikia patikrinti užbaigimo eilę (CQ) vartotojo erdvėje, kad sužinotų, jog perkėlimas baigtas

Su RDMA ir tinkamu aparatinės įrangos suporavimu, duomenų perkėlimas tarp dviejų GPU reikalauja kirsti tik vietinį PCIe jungiklį ir tinklą:

Padidinti

Priešingai, TCP/IP perkėlimai turi nukopijuoti duomenis kelis kartus per pagrindinę atmintį, sukeldami didelę PCIe magistralės spūstį:

Padidinti

Kūrimas naudojant „libfabric“

Naudojome „libfabric“ – sistemą, teikiančią bendrą sąsają audinio paslaugoms. Mūsų diegimas naudoja dviejų tipų RDMA operacijas:

• Dviejų pusių RDMA (SEND/RECV) valdymo pranešimams, kurie neša metaduomenis apie atminties sritis (pvz., poslinkius ir dydžius).
• Vienos pusės RDMA WRITE faktiniam duomenų perkėlimui, kur kiekviena WRITE operacija apdoroja vieną gretimą atminties bloką.

Kelionė link maksimalaus našumo

Mūsų kūrimas perėjo kelis etapus:

1. Įdiegtas bazinis vienpusis pranešimų perkėlimas naudojant SEND/RECV
2. Išplėsta iki dvipusės komunikacijos
3. Pridėtas „GPUDirect RDMA WRITE“ tiesioginiam GPU-GPU perkėlimui
4. Išplėsta, kad būtų galima apdoroti kelis vienu metu vykstančius perkėlimus
5. Įvestas operacijų eilės sudarymas tvirtumui užtikrinti
6. Pasiektas 97,4 % pralaidumo išnaudojimas vienoje tinklo plokštėje

Paskirstant iki 32 tinklo plokščių, įdiegėme kelis esminius optimizavimus:

1. Operacijų eilės sudarymas: Vietoj tiesioginio operacijų pateikimo tinklo plokštėms, palaikome programos lygio eilę. Tai užtikrina tvirtumą prieš tinklo spūstis ir supaprastina programavimo modelį.
2. Tinklo apšilimas: Išankstinis ryšių užmezgimas pagerino paleidimo našumą.
3. Daugiagijis apdorojimas: Skiriamosios gijos kiekvieno GPU tinklo operacijoms.
4. CPU branduolių pririšimas: Gijų susiejimas su konkrečiais CPU branduoliais, siekiant išvengti NUMA efektų ir talpyklos praleidimų.
5. Būsenos šardinimas: Ginčų tarp gijų sumažinimas.
6. Operacijų pakavimas: Kelių operacijų pateikimas kartu.
7. Pasyvus operacijų skelbimas: Operacijos pirmiausia įtraukiamos į eilę programoje. Apdoroję užbaigimo eilę, bandome pateikti laukiančias operacijas tinklo plokštei, užtikrindami efektyvų tinklo išteklių naudojimą.
8. NUMA palaikantis išteklių paskirstymas: Paskirstyti „libfabric“ išteklius, tokius kaip užbaigimo eilės, teisingame NUMA mazge, siekiant sumažinti atminties prieigos vėlavimą.

Per šiuos optimizavimus pasiekėme galutinį 3 108 Gbps našumą – 97,1 % teorinio maksimalaus pralaidumo.

Žemiau pateiktame vaizdo įraše rodoma mūsų komandinės eilutės programa veiksme. Ji perkelia neperšokančius GPU atminties blokus į nuotolinį mazgą, pasiekdama 3108,283 Gbps perdavimo greitą – tai demonstruoja beveik teorinį tinklo infrastruktūros pralaidumo išnaudojimą:

Išvada

Didelio našumo tinklo sistemos kūrimas reikalauja suprasti tiek aparatinės įrangos architektūrą, tiek sistemos dizaino principus. Nors tokios bibliotekos kaip NCCL suteikia puikius sprendimus bendriems šablonams, kartais pasirinktiniai sprendimai yra būtini specifiniams reikalavimams.

Mūsų kelionė rodo, kad pasiekti beveik teorinį tinklo našumą įmanoma atidžiai atsižvelgiant į sistemos architektūrą, aparatinės įrangos topologiją ir įvairius optimizavimo metodus. Svarbiausia suprasti ne tik atskirus komponentus, bet ir tai, kaip jie sąveikauja, sudarydami pilną sistemą.

Pilną šios kelionės techninę apžvalgą, įskaitant diegimo detales ir kodo pavyzdžius, rasite mūsų atvirojo kodo saugykloje.

Citata

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