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 processeurs graphiques (GPU) sur plusieurs machines. Chez Perplexity, nous avons rencontré 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 p5 d'AWS, offre une impressionnante largeur de bande réseau de 3 200 Gbps grâce à 32 cartes réseau. Cet article partage notre parcours dans la création d'une solution de mise en réseau personnalisée et haute performance qui atteint 97,1 % de cette largeur de bande théorique.

Le défi technique

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

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

Bien que la bibliothèque NCCL de NVIDIA soit le standard 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 oblige à redémarrer l'ensemble du cluster lors de l'ajustement des nœuds participants.
• Le modèle de communication synchrone de NCCL ajoute de la complexité pour notre charge de travail asynchrone
• Nous voulions 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 haute performance modernes

Pour comprendre notre solution, explorons d'abord en quoi les réseaux haute performance modernes diffèrent de la mise en réseau traditionnelle.

La plupart des réseaux que nous utilisons quotidiennement reposent sur les 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 protocole 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 des machines sans faire intervenir le processeur (CPU).

AWS fournit Elastic Fabric Adapter (EFA), une interface réseau personnalisée qui implémente le protocole propriétaire 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, EFA avec RDMA permet un transfert direct des données entre la mémoire du GPU et la carte réseau, en contournant complètement le CPU.

Philosophie de la conception de réseaux haute performance

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

1. Propriété des tampons : Contrairement aux sockets traditionnelles où le noyau gère les tampons réseau et exige des copies entre l’espace utilisateur et l’espace noyau, 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’à la fin de l’opération, éliminant ainsi le besoin de copier les 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 correspondances d'adresses virtuelles qui permettent au processeur, aux GPU et aux cartes réseau de comprendre tous les mêmes adresses virtuelles. Cet enregistrement est une opération ponctuelle qui permet les 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 des performances maximales
4. Réception avant transmission : En l'absence de 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. Complétion basée sur l'interrogation : Au lieu d'attendre les notifications du noyau par le biais de mécanismes comme epoll, les applications interrogent directement les files d'attente de complétion matérielles. Cela élimine la surcharge des appels système et permet une réaction immédiate aux opérations terminées.
6. Conscience 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 p5 d'AWS possèdent une architecture interne sophistiquée. Comme illustré ci-dessous, chaque instance contient deux sockets de processeur 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 du GPU au tampon d’application dans l’espace utilisateur
2. Du tampon de l’espace utilisateur au tampon du noyau
3. Du tampon du noyau vers la carte réseau
4. Le processus inverse se produit du côté de la réception

Chaque opération de copie consomme des cycles CPU et de la bande passante mémoire. L’application doit aussi effectuer une commutation de contexte entre l’espace utilisateur et l’espace noyau pour chaque opération réseau.

En revanche, le RDMA (côté droit) offre un véritable transfert de données sans copie :

1. La carte réseau lit directement à partir de la mémoire du 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 du GPU de destination
4. L’application n’a qu’à vérifier une file d’attente de complétion (CQ) dans l’espace utilisateur pour savoir quand le transfert est terminé

Avec le RDMA et un appariement matériel adéquat, le transfert de données entre deux GPU nécessite uniquement de traverser le commutateur PCIe local et le réseau :

Agrandir

En revanche, les transferts TCP/IP doivent copier les données plusieurs fois à travers la mémoire principale, ce qui entraîne une congestion importante du bus PCIe :

Agrandir

Construction avec libfabric

Nous avons utilisé libfabric, un cadre qui fournit une interface générique pour les services de fabric. Notre mise en œuvre 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 effectif des données, où chaque opération WRITE gère un bloc de mémoire contigu.

Parcours vers la performance de pointe

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

1. Mise en œuvre du transfert de message unidirectionnel de base à l'aide de SEND/RECV
2. Extension à la communication bidirectionnelle
3. Ajout de GPUDirect RDMA WRITE pour le transfert direct de GPU à GPU
4. Élargissement pour gérer de multiples 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 largeur de bande sur une seule carte réseau

Lors de la mise à l'échelle à 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 du réseau : L'établissement préalable des connexions a amélioré les performances de démarrage.
3. Multithreading : Des threads dédiés pour les opérations réseau de chaque GPU.
4. Épinglage des cœurs de processeur : Liaison des threads à des cœurs de processeur spécifiques pour éviter les effets NUMA et les défauts de cache.
5. Partage de l'état : Réduction de la contention entre les threads.
6. Mise en lots des opérations : Soumission simultanée de plusieurs opérations.
7. Publication paresseuse des opérations : Les opérations sont d'abord mises en file d'attente dans l'application. Après l'interrogation de la file d'attente de complétion, nous tentons de soumettre les opérations en attente à la carte réseau, garantissant une utilisation efficace des ressources réseau.
8. Allocation de ressources tenant compte du NUMA : Allouer les ressources libfabric telles que les files d'attente de complétion 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 largeur de bande maximale théorique.

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

Conclusion

La construction d'un système de mise en 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 proches du théorique en accordant une attention particulière à l'architecture système, à la topologie matérielle et à diverses techniques d'optimisation. La clé n'est pas seulement de comprendre les composants individuels, mais la façon dont ils interagissent pour former un système complet.

L'analyse technique approfondie de ce parcours, comprenant les détails de mise en œuvre et des exemples de code, est disponible dans notre répertoire 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}
}