Transfer de memorie GPU de înaltă performanță pe AWS Sagemaker Hyperpod

Infrastructura modernă de deep learning necesită adesea transferul unor volume mari de date între GPU-uri, între mașini diferite. La Perplexity, ne-am confruntat cu o provocare tehnică unică: transferul eficient de regiuni de memorie GPU neconigue între mașini, la viteza maximă posibilă. Platforma noastră țintă, instanțele AWS p5, oferă o lățime de bandă de rețea impresionantă de 3200 Gbps prin 32 de plăci de rețea. Acest articol împărtășește parcursul nostru în construirea unei soluții de rețea personalizate de înaltă performanță, care atinge 97,1% din această lățime de bandă teoretică.

Provocarea tehnică

Cazul nostru de utilizare a prezentat mai multe cerințe tehnice cheie:

• Transfer de mare lățime de bandă între GPU-uri la distanță pentru fragmente de memorie necontigue
• Capacitatea de a adăuga sau elimina dinamic noduri din implementările Kubernetes fără a întrerupe operațiunile în curs
• Suport pentru modele de comunicare peer-to-peer

Deși biblioteca NCCL de la NVIDIA este standardul de facto pentru deep learning distribuit, aceasta nu a fost ideală pentru cazul nostru de utilizare:

• NCCL excelează în comunicarea colectivă, dar necesită stabilirea unei "lumi" statice, ceea ce implică repornirea întregului cluster la ajustarea nodurilor participante.
• Modelul de comunicare sincronă al NCCL adaugă complexitate pentru volumul nostru de lucru asincron
• Ne-am dorit control direct asupra modelelor noastre de transfer de memorie pentru optimizare
• Construirea propriei noastre soluții ne-a oferit oportunități valoroase de învățare

Rețele moderne de înaltă performanță

Pentru a înțelege soluția noastră, să explorăm mai întâi modul în care rețelele moderne de înaltă performanță diferă de cele tradiționale.

Majoritatea rețelelor pe care le folosim zilnic se bazează pe protocoale TCP/IP, unde aplicațiile comunică cu placa de rețea prin intermediul nucleului sistemului de operare, folosind socket-uri. Cu toate acestea, rețelele de înaltă performanță utilizează RDMA (Remote Direct Memory Access) - un stiv de hardware și software complet diferit, care permite accesul direct la memorie între mașini fără implicarea CPU-ului.

AWS oferă Elastic Fabric Adapter (EFA), o interfață de rețea personalizată care implementează protocolul personalizat Amazon numit Scalable Reliable Datagram (SRD). Spre deosebire de rețeaua tradițională TCP/IP, unde datele trebuie copiate de mai multe ori între spațiul utilizatorului, spațiul nucleului și bufferele de rețea, EFA cu RDMA permite transferul direct de date între memoria GPU și placa de rețea, ocolind complet CPU-ul.

Filosofia designului de rețea de înaltă performanță

Construirea unor sisteme de rețea de înaltă performanță necesită regândirea mai multor ipoteze fundamentale:

1. Proprietatea bufferului: Spre deosebire de socket-urile tradiționale, unde nucleul gestionează bufferele de rețea și necesită copierea între spațiul utilizatorului și cel al nucleului, RDMA cere ca aplicațiile să își gestioneze propriile buffere. Atunci când o aplicație inițiază o operațiune de rețea, aceasta transferă proprietatea bufferului către placa de rețea până la finalizarea operațiunii, eliminând necesitatea copierii datelor.
2. Înregistrarea memoriei: Aplicațiile trebuie să înregistreze regiunile de memorie în nucleul sistemului de operare. Nucleul setează mapări de adrese virtuale care permit CPU-ului, GPU-urilor și plăcilor de rețea să înțeleagă aceleași adrese virtuale. Această înregistrare este o operațiune unică ce permite transferuri ulterioare de date fără copiere (zero-copy).
3. Planul de control vs. Planul de date: Rețelele de înaltă performanță separă operațiunile în două categorii:
   1. Operațiunile planului de control (cum ar fi configurarea conexiunii și înregistrarea memoriei) trec prin nucleu pentru a asigura securitatea
   2. Operațiunile planului de date (transferul efectiv de date) ocolesc nucleul pentru performanță maximă
4. Recepție înainte de transmisie: Fără buffere gestionate de nucleu, aplicațiile trebuie să pre-publice operațiunile de recepție, specificând unde ar trebui plasate datele primite. Aceasta este o schimbare fundamentală față de modelul de socket, unde aplicațiile pot primi date în orice moment.
5. Finalizare bazată pe sondaj (Poll-based): În loc să aștepte notificări de la nucleu prin mecanisme precum epoll, aplicațiile interoghează direct cozile de finalizare hardware. Acest lucru elimină suprasolicitarea apelurilor de sistem și permite o reacție imediată la operațiunile finalizate.
6. Conștientizarea topologiei hardware: Înțelegerea și optimizarea pentru topologia hardware sunt esențiale pentru atingerea performanței maxime.

Înțelegerea topologiei hardware

Instanțele AWS p5 au o arhitectură internă sofisticată. Așa cum se arată mai jos, fiecare instanță conține două socket-uri CPU care formează două noduri NUMA, fiecare nod NUMA conectându-se la patru comutatoare PCIe:

Extinde

Sub fiecare comutator PCIe găsim:

• Patru plăci de rețea EFA de 100 Gbps
• Un GPU NVIDIA H100
• Un SSD NVMe

Căile de date pentru transferurile TCP/IP și RDMA demonstrează diferențe fundamentale în abordările lor:

Extinde

Cu TCP/IP (partea stângă), datele trebuie copiate de mai multe ori:

1. Din memoria GPU în bufferul aplicației din spațiul utilizatorului
2. Din bufferul din spațiul utilizatorului în bufferul nucleului
3. Din bufferul nucleului în placa de rețea
4. Procesul invers are loc în partea de recepție

Fiecare operațiune de copiere consumă cicluri de CPU și lățime de bandă a memoriei. De asemenea, aplicația trebuie să facă schimbare de context (context switch) între spațiul utilizatorului și cel al nucleului pentru fiecare operațiune de rețea.

În schimb, RDMA (partea dreaptă) oferă un transfer de date fără nicio copiere (zero-copy):

1. Placa de rețea citește direct din memoria GPU
2. Datele călătoresc direct către placa de rețea la distanță
3. Placa de rețea de la distanță scrie direct în memoria GPU de destinație
4. Aplicația trebuie doar să verifice o coadă de finalizare (CQ) în spațiul utilizatorului pentru a afla când s-a încheiat transferul

Cu RDMA și asocierea hardware adecvată, transferul de date între două GPU-uri necesită doar traversarea comutatorului PCIe local și a rețelei:

Extinde

În schimb, transferurile TCP/IP trebuie să copieze datele de mai multe ori prin memoria principală, cauzând o congestie semnificativă a magistralei PCIe:

Extinde

Construirea cu libfabric

Am folosit libfabric - un cadru care oferă o interfață generică pentru servicii de fabric. Implementarea noastră folosește două tipuri de operațiuni RDMA:

• RDMA bidirecțional (SEND/RECV) pentru mesaje de control care transportă metadate despre regiunile de memorie (de exemplu, deplasamente și dimensiuni).
• RDMA WRITE unidirecțional pentru transferul efectiv de date, unde fiecare operațiune WRITE gestionează un singur fragment de memorie contiguu.

Parcursul către performanța de vârf

Dezvoltarea noastră a progresat în mai multe etape:

1. Implementarea transferului de mesaje unidirecțional de bază folosind SEND/RECV
2. Extinderea la comunicarea bidirecțională
3. Adăugarea GPUDirect RDMA WRITE pentru transfer direct GPU-GPU
4. Extinderea pentru a gestiona mai multe transferuri concurente
5. Introducerea cozilor de operațiuni pentru robustețe
6. Atingerea unei utilizări a lățimii de bandă de 97,4% pe o singură placă de rețea

La scalarea la 32 de plăci de rețea, am implementat mai multe optimizări cruciale:

1. Coada de operațiuni: În loc să trimitem operațiunile direct către plăcile de rețea, menținem o coadă la nivelul aplicației. Acest lucru oferă robustețe împotriva aglomerației rețelei și simplifică modelul de programare.
2. Încălzirea rețelei: Pre-stabilirea conexiunilor a îmbunătățit performanța la pornire.
3. Multithreading: Fire de execuție dedicate pentru operațiunile de rețea ale fiecărui GPU.
4. Fixarea core-urilor de CPU (CPU Core Pinning): Asocierea firelor de execuție cu anumite nuclee CPU pentru a evita efectele NUMA și cache misses.
5. Partajarea stării (State Sharding): Reducerea contingenței între firele de execuție.
6. Gruparea operațiunilor (Batching): Trimiterea mai multor operațiuni împreună.
7. Publicarea leneșă a operațiunilor (Lazy Operation Posting): Operațiunile sunt mai întâi puse în coadă în aplicație. După interogarea cozii de finalizare, încercăm să trimitem operațiunile în așteptare către placa de rețea, asigurând o utilizare eficientă a resurselor de rețea.
8. Alocarea resurselor conștientă de NUMA: Alocarea resurselor libfabric, cum ar fi cozile de finalizare, pe nodul NUMA corect pentru a minimiza latența de acces la memorie.

Prin aceste optimizări, am obținut o performanță finală de 3.108 Gbps - 97,1% din maximul teoretic al lățimii de bandă.

Videoclipul de mai jos prezintă programul nostru în linie de comandă în acțiune. Acesta transferă fragmente necontigue de memorie GPU către un nod la distanță, atingând o viteză de transmisie de 3108,283 Gbps - demonstrând o utilizare aproape teoretică a lățimii de bandă a infrastructurii de rețea:

Concluzie

Construirea unui sistem de rețea de înaltă performanță necesită înțelegerea atât a arhitecturii hardware, cât și a principiilor de proiectare a sistemului. Deși biblioteci precum NCCL oferă soluții excelente pentru modele comune, uneori sunt necesare soluții personalizate pentru cerințe specifice.

Parcursul nostru demonstrează că atingerea unei performanțe de rețea aproape teoretice este posibilă acordând o atenție deosebită arhitecturii sistemului, topologiei hardware și diferitelor tehnici de optimizare. Cheia nu este doar înțelegerea componentelor individuale, ci a modului în care acestea interacționează pentru a forma un sistem complet.

Analiza tehnică detaliată a acestui parcurs, inclusiv detaliile de implementare și exemplele de cod, este disponibilă în repositoriul nostru open-source.

Citare

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