Transfert de mémoire GPU haute performance sur AWS Sagemaker Hyperpod

L'infrastructure moderne d'apprentissage profond nécessite souvent le transfert de grandes quantités de données entre des GPU situés sur des machines différentes. Chez Perplexity, nous avons été confrontés à un défi technique unique : transférer efficacement des régions de mémoire GPU non contiguës entre des machines à la vitesse maximale possible. Notre plateforme cible, les instances AWS p5, offre une bande passante réseau impressionnante de 3200 Gbps via 32 cartes réseau. Cet article partage notre parcours dans la création d'une solution réseau personnalisée haute performance qui atteint 97,1 % de cette bande passante théorique.

Le défi technique

Notre cas d'utilisation présentait plusieurs exigences techniques clés :

• Transfert à haut débit entre des GPU distants de segments de mémoire non contigus
• Capacité à ajouter ou supprimer dynamiquement des nœuds des déploiements Kubernetes sans interrompre les opérations en cours
• Prise en charge des modèles de communication pair à pair

Bien que la bibliothèque NCCL de NVIDIA soit la norme de fait pour l'apprentissage profond distribué, elle n'était pas idéale pour notre cas d'utilisation :

• NCCL excelle dans la communication collective mais nécessite l'établissement d'un "monde" statique, ce qui exige le redémarrage de l'ensemble du cluster lors de l'ajustement des nœuds participants.
• Le modèle de communication synchrone de NCCL ajoute de la complexité à notre charge de travail asynchrone
• Nous souhaitions un contrôle direct sur nos modèles de transfert de mémoire à des fins d'optimisation
• La construction de notre propre solution a offert de précieuses opportunités d'apprentissage

Réseaux modernes haute performance

Pour comprendre notre solution, explorons d'abord en quoi les réseaux modernes haute performance diffèrent du réseau traditionnel.

La plupart des réseaux que nous utilisons quotidiennement reposent sur des protocoles TCP/IP, où les applications communiquent avec la carte réseau via le noyau du système d'exploitation à l'aide de sockets. Cependant, les réseaux haute performance utilisent le RDMA (Remote Direct Memory Access) - une pile matérielle et logicielle complètement différente qui permet un accès direct à la mémoire entre les machines sans impliquer le processeur.

AWS fournit Elastic Fabric Adapter (EFA), une interface réseau personnalisée qui implémente le protocole personnalisé d'Amazon appelé Scalable Reliable Datagram (SRD). Contrairement au réseau TCP/IP traditionnel où les données doivent être copiées plusieurs fois entre l'espace utilisateur, l'espace noyau et les tampons réseau, l'EFA avec RDMA permet un transfert de données direct entre la mémoire GPU et la carte réseau, en contournant entièrement le processeur.

Philosophie de conception des réseaux haute performance

La construction de systèmes réseau haute performance nécessite de repenser plusieurs hypothèses fondamentales :

1. Propriété des tampons : Contrairement aux sockets traditionnels où le noyau gère les tampons réseau et nécessite des copies entre l'espace utilisateur et l'espace noyau, le RDMA exige que les applications gèrent leurs propres tampons. Lorsqu'une application lance une opération réseau, elle transfère la propriété du tampon à la carte réseau jusqu'à ce que l'opération soit terminée, éliminant ainsi le besoin de copie de données.
2. Enregistrement de la mémoire : Les applications doivent enregistrer les régions de mémoire auprès du noyau du système d'exploitation. Le noyau établit des mappages d'adresses virtuelles qui permettent au CPU, aux GPU et aux cartes réseau de comprendre les mêmes adresses virtuelles. Cet enregistrement est une opération ponctuelle qui permet des transferts de données ultérieurs sans copie.
3. Plan de contrôle vs Plan de données : Les réseaux haute performance séparent les opérations en deux catégories :
   1. Les opérations du plan de contrôle (comme la configuration de la connexion et l'enregistrement de la mémoire) passent par le noyau pour garantir la sécurité
   2. Les opérations du plan de données (transfert de données réel) contournent le noyau pour une performance maximale
4. Réception avant transmission : Sans tampons gérés par le noyau, les applications doivent pré-publier les opérations de réception, en spécifiant où les données entrantes doivent être placées. Il s'agit d'un changement fondamental par rapport au modèle de socket où les applications peuvent recevoir des données à tout moment.
5. Achèvement basé sur le sondage : Au lieu d'attendre les notifications du noyau via des mécanismes comme epoll, les applications sondent directement les files d'attente d'achèvement matérielles. Cela élimine la surcharge des appels système et permet une réaction immédiate aux opérations terminées.
6. Connaissance de la topologie matérielle : Comprendre et optimiser la topologie matérielle est crucial pour atteindre des performances maximales.

Comprendre la topologie matérielle

Les instances AWS p5 ont une architecture interne sophistiquée. Comme illustré ci-dessous, chaque instance contient deux sockets CPU formant deux nœuds NUMA, chaque nœud NUMA se connectant à quatre commutateurs PCIe :

Agrandir

Sous chaque commutateur PCIe, nous trouvons :

• Quatre cartes réseau EFA de 100 Gbps
• Un GPU NVIDIA H100
• Un SSD NVMe

Les chemins de données pour les transferts TCP/IP et RDMA démontrent des différences fondamentales dans leurs approches :

Agrandir

Avec TCP/IP (côté gauche), les données doivent être copiées plusieurs fois :

1. De la mémoire GPU vers le tampon d'application dans l'espace utilisateur
2. Du tampon de l'espace utilisateur vers le tampon du noyau
3. Du tampon du noyau vers la carte réseau
4. Le processus inverse se produit côté réception

Chaque opération de copie consomme des cycles CPU et de la bande passante mémoire. L'application doit également changer de contexte entre l'espace utilisateur et l'espace noyau pour chaque opération réseau.

En revanche, le RDMA (côté droit) fournit un transfert de données sans copie (zéro-copie) :

1. La carte réseau lit directement depuis la mémoire GPU
2. Les données voyagent directement vers la carte réseau distante
3. La carte réseau distante écrit directement dans la mémoire GPU de destination
4. L'application doit seulement vérifier une file d'attente d'achèvement (CQ) dans l'espace utilisateur pour savoir quand le transfert est terminé

Avec le RDMA et un couplage matériel approprié, le transfert de données entre deux GPU ne nécessite que de traverser le commutateur PCIe local et le réseau :

Agrandir

À l'inverse, les transferts TCP/IP doivent copier les données plusieurs fois à travers la mémoire principale, provoquant une congestion importante du bus PCIe :

Agrandir

Construire avec libfabric

Nous avons utilisé libfabric - un framework qui fournit une interface générique pour les services fabric. Notre implémentation utilise deux types d'opérations RDMA :

• RDMA bidirectionnel (SEND/RECV) pour les messages de contrôle qui transportent des métadonnées sur les régions de mémoire (par exemple, les décalages et les tailles).
• RDMA WRITE unidirectionnel pour le transfert de données réel, où chaque opération WRITE gère un segment de mémoire contigu.

Vers une performance maximale

Notre développement a progressé à travers plusieurs étapes :

1. Implémentation du transfert de message unidirectionnel de base en utilisant SEND/RECV
2. Extension à la communication bidirectionnelle
3. Ajout du GPUDirect RDMA WRITE pour un transfert direct GPU-GPU
4. Extension pour gérer plusieurs transferts simultanés
5. Introduction de la mise en file d'attente des opérations pour plus de robustesse
6. Atteinte de 97,4 % d'utilisation de la bande passante sur une seule carte réseau

Lors du passage à 32 cartes réseau, nous avons mis en œuvre plusieurs optimisations cruciales :

1. Mise en file d'attente des opérations : Plutôt que de soumettre directement les opérations aux cartes réseau, nous maintenons une file d'attente au niveau de l'application. Cela offre une robustesse contre la congestion du réseau et simplifie le modèle de programmation.
2. Préchauffage réseau : L'établissement préalable des connexions a amélioré les performances au démarrage.
3. Multi-threading : Des fils dédiés pour les opérations réseau de chaque GPU.
4. Affinité de cœur CPU : Liaison des fils à des cœurs CPU spécifiques pour éviter les effets NUMA et les défauts de cache.
5. Sharding d'état : Réduction de la contention entre les fils.
6. Traitement par lots des opérations : Soumission de plusieurs opérations ensemble.
7. Envoi différé des opérations : Les opérations sont d'abord placées dans la file d'attente de l'application. Après avoir sondé la file d'attente d'achèvement, nous tentons de soumettre les opérations en attente à la carte réseau, assurant une utilisation efficace des ressources réseau.
8. Allocation de ressources tenant compte de NUMA : Allouer les ressources libfabric comme les files d'attente d'achèvement sur le bon nœud NUMA pour minimiser la latence d'accès à la mémoire.

Grâce à ces optimisations, nous avons atteint une performance finale de 3 108 Gbps - soit 97,1 % de la bande passante théorique maximale.

La vidéo ci-dessous montre notre programme en ligne de commande en action. Il transfère des segments non contigus de mémoire GPU vers un nœud distant, atteignant une vitesse de transmission de 3108,283 Gbps - démontrant une utilisation quasi théorique de la bande passante de l'infrastructure réseau :

Conclusion

La construction d'un système réseau haute performance nécessite de comprendre à la fois l'architecture matérielle et les principes de conception des systèmes. Bien que des bibliothèques comme NCCL fournissent d'excellentes solutions pour les modèles courants, des solutions personnalisées sont parfois nécessaires pour des exigences spécifiques.

Notre parcours démontre qu'il est possible d'atteindre des performances réseau quasi théoriques avec une attention particulière portée à l'architecture du système, à la topologie matérielle et à diverses techniques d'optimisation. La clé n'est pas seulement de comprendre les composants individuels, mais comment ils interagissent pour former un système complet.

L'analyse technique approfondie de ce parcours, y compris les détails de mise en œuvre et des exemples de code, est disponible dans notre dépôt open source.

Citation

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