Trasferimento di memoria GPU ad alte prestazioni su AWS Sagemaker Hyperpod

Le moderne infrastrutture di deep learning richiedono spesso il trasferimento di grandi quantità di dati tra GPU in diverse macchine. In Perplexity, abbiamo riscontrato una sfida tecnica unica: trasferire in modo efficiente regioni di memoria GPU non contigue tra le macchine alla massima velocità possibile. La nostra piattaforma di destinazione, le istanze AWS p5, offre un'impressionante larghezza di banda di rete di 3200 Gbps attraverso 32 schede di rete. Questo articolo illustra il nostro percorso nello sviluppo di una soluzione di rete ad alte prestazioni personalizzata che raggiunge il 97,1% di questa larghezza di banda teorica.

La sfida tecnica

Il nostro caso d'uso presentava diversi requisiti tecnici chiave:

• Trasferimento ad alta larghezza di banda tra GPU remote di blocchi di memoria non contigui
• Capacità di aggiungere o rimuovere dinamicamente nodi dai deployment Kubernetes senza interrompere le operazioni in corso
• Supporto per modelli di comunicazione peer-to-peer

Sebbene la libreria NCCL di NVIDIA sia lo standard de facto per il deep learning distribuito, non era ideale per il nostro caso d'uso:

• NCCL eccelle nella comunicazione collettiva ma richiede la creazione di un "mondo" statico, il che necessita il riavvio dell'intero cluster quando si modificano i nodi partecipanti.
• Il modello di comunicazione sincrona di NCCL aggiunge complessità per il nostro carico di lavoro asincrono
• Volevamo il controllo diretto sui nostri modelli di trasferimento della memoria per l'ottimizzazione
• Costruire la nostra soluzione ha offerto preziose opportunità di apprendimento

Reti moderne ad alte prestazioni

Per comprendere la nostra soluzione, esploriamo innanzitutto come le moderne reti ad alte prestazioni differiscono dal networking tradizionale.

La maggior parte delle reti che utilizziamo quotidianamente si basa sui protocolli TCP/IP, dove le applicazioni comunicano con la scheda di rete attraverso il kernel del sistema operativo utilizzando socket. Tuttavia, le reti ad alte prestazioni utilizzano RDMA (Remote Direct Memory Access) - uno stack hardware e software completamente diverso che abilita l'accesso diretto alla memoria tra le macchine senza coinvolgere la CPU.

AWS fornisce Elastic Fabric Adapter (EFA), un'interfaccia di rete personalizzata che implementa il protocollo personalizzato di Amazon chiamato Scalable Reliable Datagram (SRD). A differenza del networking TCP/IP tradizionale, in cui i dati devono essere copiati più volte tra lo spazio utente, lo spazio kernel e i buffer di rete, EFA con RDMA consente il trasferimento diretto dei dati tra la memoria della GPU e la scheda di rete, ignorando completamente la CPU.

Filosofia della progettazione di reti ad alte prestazioni

Costruire sistemi di rete ad alte prestazioni richiede di ripensare diverse ipotesi fondamentali:

1. Proprietà del buffer: A differenza dei socket tradizionali in cui il kernel gestisce i buffer di rete e richiede la copia tra spazio utente e spazio kernel, RDMA richiede che le applicazioni gestiscano i propri buffer. Quando un'applicazione avvia un'operazione di rete, trasferisce la proprietà del buffer alla scheda di rete fino al completamento dell'operazione, eliminando la necessità di copiare i dati.
2. Registrazione della memoria: Le applicazioni devono registrare le regioni di memoria con il kernel del sistema operativo. Il kernel imposta le mappature degli indirizzi virtuali che consentono alla CPU, alle GPU e alle schede di rete di comprendere gli stessi indirizzi virtuali. Questa registrazione è un'operazione una tantum che abilita i successivi trasferimenti di dati zero-copy.
3. Piano di controllo vs Piano dati: Le reti ad alte prestazioni separano le operazioni in due categorie:
   1. Le operazioni del piano di controllo (come la configurazione della connessione e la registrazione della memoria) passano attraverso il kernel per garantire la sicurezza
   2. Le operazioni del piano dati (trasferimento effettivo dei dati) ignorano il kernel per le massime prestazioni
4. Ricezione prima della trasmissione: Senza buffer gestiti dal kernel, le applicazioni devono pre-inviare le operazioni di ricezione, specificando dove devono essere posizionati i dati in arrivo. Si tratta di un cambiamento fondamentale rispetto al modello socket in cui le applicazioni possono ricevere dati in qualsiasi momento.
5. Completamento basato su polling: Invece di attendere notifiche del kernel attraverso meccanismi come epoll, le applicazioni interrogano direttamente le code di completamento dell'hardware. Ciò elimina il sovraccarico delle chiamate di sistema e consente una reazione immediata alle operazioni completate.
6. Consapevolezza della topologia hardware: Comprendere e ottimizzare per la topologia hardware è fondamentale per ottenere le massime prestazioni.

Comprensione della topologia hardware

Le istanze AWS p5 hanno un'architettura interna sofisticata. Come mostrato di seguito, ogni istanza contiene due socket CPU che formano due nodi NUMA, con ogni nodo NUMA che si collega a quattro switch PCIe:

Espandi

Sotto ogni switch PCIe, troviamo:

• Quattro schede di rete EFA da 100 Gbps
• Una GPU NVIDIA H100
• Un SSD NVMe

I percorsi dati per i trasferimenti TCP/IP e RDMA dimostrano differenze fondamentali nei loro approcci:

Espandi

Con TCP/IP (lato sinistro), i dati devono essere copiati più volte:

1. Dalla memoria GPU al buffer dell'applicazione nello spazio utente
2. Dal buffer dello spazio utente al buffer del kernel
3. Dal buffer del kernel alla scheda di rete
4. Il processo inverso avviene sul lato ricevente

Ogni operazione di copia consuma cicli CPU e larghezza di banda della memoria. L'applicazione deve inoltre effettuare il context switch tra spazio utente e spazio kernel per ogni operazione di rete.

Al contrario, RDMA (lato destro) fornisce un vero trasferimento dati zero-copy:

1. La scheda di rete legge direttamente dalla memoria GPU
2. I dati viaggiano direttamente verso la scheda di rete remota
3. La scheda di rete remota scrive direttamente nella memoria GPU di destinazione
4. L'applicazione deve solo controllare una coda di completamento (CQ) nello spazio utente per sapere quando il trasferimento è terminato

Con RDMA e un corretto accoppiamento hardware, il trasferimento di dati tra due GPU richiede solo l'attraversamento dello switch PCIe locale e della rete:

Espandi

Al contrario, i trasferimenti TCP/IP devono copiare i dati più volte attraverso la memoria principale, causando una significativa congestione del bus PCIe:

Espandi

Costruire con libfabric

Abbiamo utilizzato libfabric, un framework che fornisce un'interfaccia generica per i servizi fabric. La nostra implementazione utilizza due tipi di operazioni RDMA:

• RDMA a due vie (SEND/RECV) per messaggi di controllo che trasportano metadati sulle regioni di memoria (es. offset e dimensioni).
• RDMA WRITE a una via per il trasferimento effettivo dei dati, dove ogni operazione WRITE gestisce un blocco di memoria contiguo.

Il percorso verso le massime prestazioni

Il nostro sviluppo è progredito attraverso diverse fasi:

1. Implementato il trasferimento unidirezionale di base dei messaggi utilizzando SEND/RECV
2. Esteso alla comunicazione bidirezionale
3. Aggiunta di GPUDirect RDMA WRITE per il trasferimento diretto GPU-GPU
4. Espanso per gestire molteplici trasferimenti simultanei
5. Introdotta la coda delle operazioni per la robustezza
6. Raggiunto il 97,4% di utilizzo della larghezza di banda su una singola scheda di rete

Durante il ridimensionamento a 32 schede di rete, abbiamo implementato diverse ottimizzazioni cruciali:

1. Coda delle operazioni: Invece di inviare direttamente le operazioni alle schede di rete, manteniamo una coda a livello di applicazione. Ciò fornisce robustezza contro la congestione della rete e semplifica il modello di programmazione.
2. Riscaldamento della rete: La pre-stabilizzazione delle connessioni ha migliorato le prestazioni all'avvio.
3. Multi-threading: Thread dedicati per le operazioni di rete di ogni GPU.
4. CPU Core Pinning: Binding dei thread a specifici core CPU per evitare effetti NUMA e cache miss.
5. State Sharding: Riduzione della contesa tra i thread.
6. Batch delle operazioni: Invio di più operazioni insieme.
7. Lazy Operation Posting: Le operazioni vengono prima accodate nell'applicazione. Dopo aver interrogato la coda di completamento, tentiamo di inviare le operazioni in sospeso alla scheda di rete, garantendo un uso efficiente delle risorse di rete.
8. Allocazione delle risorse compatibile con NUMA: Allocare risorse libfabric come le code di completamento sul nodo NUMA corretto per ridurre al minimo la latenza di accesso alla memoria.

Grazie a queste ottimizzazioni, abbiamo ottenuto una performance finale di 3.108 Gbps, il 97,1% della larghezza di banda massima teorica.

Il video qui sotto mostra il nostro programma da riga di comando in azione. Trasferisce blocchi non contigui di memoria GPU a un nodo remoto, raggiungendo una velocità di trasmissione di 3108.283 Gbps, dimostrando un utilizzo quasi teorico della larghezza di banda dell'infrastruttura di rete:

Conclusione

Costruire un sistema di rete ad alte prestazioni richiede la comprensione sia dell'architettura hardware che dei principi di progettazione del sistema. Sebbene librerie come NCCL forniscano soluzioni eccellenti per modelli comuni, a volte le soluzioni personalizzate sono necessarie per requisiti specifici.

Il nostro percorso dimostra che raggiungere prestazioni di rete quasi teoriche è possibile con un'attenta attenzione all'architettura del sistema, alla topologia hardware e a varie tecniche di ottimizzazione. La chiave non è solo comprendere i singoli componenti, ma come interagiscono per formare un sistema completo.

Il deep-dive tecnico completo di questo viaggio, inclusi i dettagli di implementazione ed esempi di codice, è disponibile nel nostro repository open source.

Citazione

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