Transferència de memòria de GPU d'alt rendiment a AWS Sagemaker Hyperpod

La infraestructura moderna d'aprenentatge profund sovint requereix transferir grans quantitats de dades entre GPU de diferents màquines. A Perplexity, ens vam trobar amb un repte tècnic únic: transferir de manera eficient regions de memòria de GPU no contígues entre màquines a la màxima velocitat possible. La nostra plataforma de destinació, les instàncies AWS p5, ofereix un ample de banda de xarxa impressionant de 3200 Gbps a través de 32 targetes de xarxa. Aquest article comparteix el nostre camí per construir una solució de xarxa personalitzada d'alt rendiment que aconsegueix el 97,1% d'aquest ample de banda teòric.

El repte tècnic

El nostre cas d'ús va presentar diversos requisits tècnics clau:

• Transferència d'alt ample de banda entre GPU remotes de fragments de memòria no contigus
• Capacitat d'afegir o eliminar dinàmicament nodes de desplejaments de Kubernetes sense interrompre les operacions en curs
• Suport per a patrons de comunicació d'igual a igual (peer-to-peer)

Tot i que la biblioteca NCCL de NVIDIA és l'estàndard de facto per a l'aprenentatge profund distribuït, no era ideal per al nostre cas d'ús:

• NCCL excel·leix en la comunicació col·lectiva, però requereix establir un "món" estàtic, la qual cosa obliga a reiniciar tot el clúster quan s'ajusten els nodes participants.
• El model de comunicació síncrona de NCCL afegeix complexitat per a la nostra càrrega de treball asíncrona
• Volíem un control directe sobre els nostres patrons de transferència de memòria per a l'optimització
• Construir la nostra pròpia solució va proporcionar oportunitats d'aprenentatge valuoses

Xarxes modernes d'alt rendiment

Per entendre la nostra solució, primer explorem com les xarxes modernes d'alt rendiment difereixen de les xarxes tradicionals.

La majoria de xarxes que fem servir diàriament depenen de protocols TCP/IP, on les aplicacions es comuniquen amb la targeta de xarxa a través del nucli del sistema operatiu utilitzant sòcols. No obstant això, les xarxes d'alt rendiment utilitzen RDMA (Remote Direct Memory Access), una pila de programari i maquinari completament diferent que permet l'accés directe a la memòria entre màquines sense implicar la CPU.

AWS proporciona Elastic Fabric Adapter (EFA), una interfície de xarxa personalitzada que implementa el protocol personalitzat d'Amazon anomenat Scalable Reliable Datagram (SRD). A diferència de les xarxes TCP/IP tradicionals, on les dades s'han de copiar diverses vegades entre l'espai d'usuari, l'espai del nucli i els búfers de xarxa, EFA amb RDMA permet la transferència directa de dades entre la memòria de la GPU i la targeta de xarxa, evitant completament la CPU.

Filosofia del disseny de xarxes d'alt rendiment

Construir sistemes de xarxa d'alt rendiment requereix replantejar diverses hipòtesis fonamentals:

1. Propietat del búfer: A diferència dels sòcols tradicionals on el nucli gestiona els búfers de xarxa i requereix la còpia entre l'espai d'usuari i l'espai del nucli, RDMA requereix que les aplicacions gestionin els seus propis búfers. Quan una aplicació inicia una operació de xarxa, transfereix la propietat del búfer a la targeta de xarxa fins que l'operació es completa, eliminant la necessitat de copiar dades.
2. Registre de memòria: Les aplicacions han de registrar regions de memòria amb el nucli del sistema operatiu. El nucli configura assignacions d'adreces virtuals que permeten a la CPU, les GPU i les targetes de xarxa entendre totes les mateixes adreces virtuals. Aquest registre és una operació d'un sol cop que permet transferències de dades posteriors sense còpies.
3. Pla de control enfront de pla de dades: Les xarxes d'alt rendiment separen les operacions en dues categories:
   1. Les operacions del pla de control (com la configuració de connexions i el registre de memòria) passen pel nucli per garantir la seguretat
   2. Les operacions del pla de dades (transferència real de dades) passen per alt el nucli per obtenir el màxim rendiment
4. Recepció abans de la transmissió: Sense búfers gestionats pel nucli, les aplicacions han de publicar prèviament operacions de recepció, especificant on s'han de col·locar les dades entrants. Aquest és un canvi fonamental respecte al model de sòcols on les aplicacions poden rebre dades en qualsevol moment.
5. Completat basat en sondeig: En lloc d'esperar notificacions del nucli a través de mecanismes com epoll, les aplicacions consulten directament les cues de completat del maquinari. Això elimina la sobrecàrrega de les crides al sistema i permet una reacció immediata a les operacions completades.
6. Consciència de la topologia del maquinari: Comprendre i optimitzar la topologia del maquinari és crucial per aconseguir el màxim rendiment.

Comprendre la topologia del maquinari

Les instàncies AWS p5 tenen una arquitectura interna sofisticada. Com es mostra a continuació, cada instància conté dos sòcols de CPU que formen dos nodes NUMA, i cada node NUMA es connecta a quatre commutadors PCIe:

Amplia

Sota cada commutador PCIe, trobem:

• Quatre targetes de xarxa EFA de 100 Gbps
• Una GPU NVIDIA H100
• Un SSD NVMe

Els camins de dades per a les transferències TCP/IP i RDMA demostren diferències fonamentals en els seus enfocaments:

Amplia

Amb TCP/IP (costat esquerre), les dades s'han de copiar diverses vegades:

1. De la memòria de la GPU al búfer de l'aplicació a l'espai d'usuari
2. Del búfer de l'espai d'usuari al búfer del nucli
3. Del búfer del nucli a la targeta de xarxa
4. El procés invers té lloc al costat receptor

Cada operació de còpia consumeix cicles de CPU i ample de banda de memòria. L'aplicació també ha de canviar de context entre l'espai d'usuari i l'espai del nucli per a cada operació de xarxa.

En contrast, RDMA (costat dret) proporciona una transferència de dades real sense còpies (zero-copy):

1. La targeta de xarxa llegeix directament de la memòria de la GPU
2. Les dades viatgen directament a la targeta de xarxa remota
3. La targeta de xarxa remota escriu directament a la memòria de la GPU de destinació
4. L'aplicació només necessita comprovar una cua de completat (CQ) a l'espai d'usuari per saber quan s'ha completat la transferència

Amb RDMA i l'aparellament de maquinari adequat, transferir dades entre dues GPU només requereix travessar el commutador PCIe local i la xarxa:

Amplia

En canvi, les transferències TCP/IP han de copiar dades diverses vegades a través de la memòria principal, provocant una congestió significativa del bus PCIe:

Amplia

Construcció amb libfabric

Hem utilitzat libfabric, un marc que proporciona una interfície genèrica per a serveis de fabric. La nostra implementació utilitza dos tipus d'operacions RDMA:

• RDMA bidireccional (SEND/RECV) per a missatges de control que transporten metadades sobre regions de memòria (p. ex., desplaçaments i mides).
• RDMA WRITE unidireccional per a la transferència real de dades, on cada operació WRITE gestiona un fragment de memòria contigu.

Camí cap al rendiment màxim

El nostre desenvolupament ha progressat a través de diverses etapes:

1. S'ha implementat la transferència de missatges unidireccional bàsica utilitzant SEND/RECV
2. S'ha ampliat a la comunicació bidireccional
3. S'ha afegit GPUDirect RDMA WRITE per a la transferència directa entre GPU
4. S'ha ampliat per gestionar múltiples transferències concurrents
5. S'ha introduït l'encua d'operacions per a la robustesa
6. S'ha aconseguit una utilització de l'ample de banda del 97,4% en una sola targeta de xarxa

En escalar a 32 targetes de xarxa, hem implementat diverses optimitzacions crucials:

1. Encua d'operacions: En lloc de trametre operacions directament a les targetes de xarxa, mantenim una cua a nivell d'aplicació. Això proporciona robustesa contra la congestió de la xarxa i simplifica el model de programació.
2. Escalfament de la xarxa: Preestablir connexions va millorar el rendiment d'inici.
3. Multifil: Fils dedicats per a les operacions de xarxa de cada GPU.
4. Fixació de nuclis de CPU: Vincular fils a nuclis de CPU específics per evitar efectes NUMA i fallades de memòria cau.
5. Fragmentació d'estat: Reducció de la contenció entre fils.
6. Lot d'operacions: Tramesa conjunta de múltiples operacions.
7. Publicació ajornada d'operacions: Les operacions s'encuen primer a l'aplicació. Després de consultar la cua de completat, intentem trametre les operacions pendents a la targeta de xarxa, garantint un ús eficient dels recursos de xarxa.
8. Assignació de recursos compatible amb NUMA: Assignar recursos de libfabric com ara cues de completat al node NUMA correcte per minimitzar la latència d'accés a la memòria.

Mitjançant aquestes optimitzacions, hem aconseguit un rendiment final de 3.108 Gbps, el 97,1% de l'ample de banda màxim teòric.

El vídeo següent mostra el nostre programa de línia d'ordres en acció. Transfereix fragments no contigus de memòria de GPU a un node remot, aconseguint una velocitat de transmissió de 3108,283 Gbps, demostrant una utilització de l'ample de banda gairebé teòrica de la infraestructura de xarxa:

Conclusió

Construir un sistema de xarxa d'alt rendiment requereix comprendre tant l'arquitectura del maquinari com els principis de disseny de sistemes. Tot i que biblioteques com NCCL proporcionen excel·lents solucions per a patrons comuns, de vegades calen solucions personalitzades per a requisits específics.

El nostre viatge demostra que aconseguir un rendiment de xarxa gairebé teòric és possible amb una atenció acurada a l'arquitectura del sistema, la topologia del maquinari i diverses tècniques d'optimització. La clau no és només entendre els components individuals, com interactuen per formar un sistema complet.

L'anàlisi tècnica completa d'aquest viatge, inclosos els detalls d'implementació i exemples de codi, està disponible al nostre repositori de codi obert.

Citació

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