Communication RDMA point à point pour les systèmes LLM
Outil élégant pour répondre aux modèles de communication LLM émergents
Les systèmes de LLM reposent traditionnellement sur des communications collectives pour les modèles parallèles structurés lors de l'entraînement et de l'inférence. Cependant, les charges de travail émergentes, telles que l'inférence désagrégée, le routage par mélange d'experts (MoE) et le réglage fin par apprentissage par renforcement asynchrone, exigent la flexibilité d'une communication point à point.
Notre dernier article de recherche présente TransferEngine, une bibliothèque de communication portable basée sur RDMA qui fonctionne sur différents contrôleurs d'interface réseau (NIC), notamment NVIDIA ConnectX-7 et AWS Elastic Fabric Adapter (EFA). Nous démontrons son efficacité dans trois systèmes de production chez Perplexity :
Préréglage et décodage désagrégés : Transfert de KvCache couche par couche utilisant des écritures paginées et des observateurs UVM, permettant aux instances de pré-remplissage et de décodage de monter en charge indépendamment.
Transfert de poids pour le post-entraînement par RL en moins de 2 secondes : Synchronisation ultra-rapide des poids entre les nœuds d'entraînement et d'inférence.
Activation de modèles à mille milliards de paramètres sur AWS EFA : Opérations de dispatch et de combinaison MoE à faible latence utilisant des primitives de diffusion et de barrière, atteignant une latence de pointe sur ConnectX-7 et la première implémentation viable compatible avec AWS EFA.
Cet article se concentre sur TransferEngine lui-même, la bibliothèque de communication qui alimente ces systèmes. Lisez l'article de recherche complet sur arXiv et explorez le code sur GitHub.
Motivation
Les communications collectives excellent dans les modèles parallèles structurés tels que le parallélisme de tenseurs et le parallélisme de données, mais elles imposent des contraintes qui limitent leur applicabilité aux charges de travail émergentes :
Appartenance fixe : Les collectifs nécessitent l'établissement d'un monde de communication global lors de l'initialisation, avec une configuration synchronisée entre tous les participants. Pour l'inférence désagrégée, cela est problématique car le nombre d'instances de pré-remplissage et de décodage doit évoluer dynamiquement en fonction des modèles de trafic.
Formes de tenseurs uniformes : Les collectifs supposent que tous les participants échangent des tenseurs avec des formes et des types de données identiques. Dans le routage MoE, le nombre de jetons envoyés à chaque expert dépend des données et est découvert au moment de l'exécution, forçant les collectifs denses tout-à-tout à remplir jusqu'à des tailles de tampon les plus défavorables.
Ordre des opérations : Les API collectives imposent des sémantiques synchronisées, même lorsque les applications ont seulement besoin de savoir qu'un ensemble d'opérations est terminé. Par exemple, dans le transfert KvCache, toutes les pages de l'instance de pré-remplissage vers l'instance de décodage peuvent arriver dans n'importe quel ordre. De même, dans les mises à jour de poids RL, tous les tenseurs de poids peuvent être transférés indépendamment sans contraintes d'ordre.
La communication point à point basée sur RDMA résout naturellement ces limitations, mais trois obstacles ont empêché son adoption :
Limitations de la bibliothèque : Bien que torch.distributed et NCCL fournissent des API point à point telles que Send/Recv, ces opérations héritent toujours des sémantiques et des limitations collectives.
Diversité matérielle : Les cartes réseau NVIDIA ConnectX avec transport Reliable Connection (RC) sont la configuration RDMA la plus courante, garantissant une livraison dans l'ordre. Les fournisseurs de cloud comme AWS utilisent un matériel différent, l'Elastic Fabric Adapter (EFA) avec le transport Scalable Reliable Datagram (SRD) délivre les messages dans le désordre. L'EFA nécessite également l'agrégation de quatre cartes réseau de 100 Gbps (ou deux cartes de 200 Gbps) pour atteindre une bande passante de 400 Gbps.
Portabilité du logiciel : Les bibliothèques RDMA existantes souffrent d'un verrouillage par le fournisseur. DeepEP nécessite le RDMA initié par GPU (IBGDA) disponible uniquement sur ConnectX. NVSHMEM présente une dégradation sévère des performances sur EFA. Des bibliothèques comme Mooncake et NIXL manquent de support EFA ou restent au début de leur développement. Par conséquent, il n'existe aucune solution viable inter-fournisseurs pour l'inférence LLM.
TransferEngine
TransferEngine répond à ces obstacles grâce à quatre capacités fondamentales :
Hypothèses de transport communes : Plutôt que de cibler une carte réseau spécifique, TransferEngine suppose uniquement que le transport sous-jacent est fiable mais non ordonné, c'est-à-dire l'intersection des garanties fournies par les transports RC et SRD.
Synchronisation ImmCounter : Une primitive innovante pour la notification d'achèvement qui fonctionne sans dépendre de l'ordonnancement des messages, permettant au même code de s'exécuter sur des transports ordonnés (RC) et non ordonnés (SRD).
Support multi-NIC : Gestion transparente de plusieurs cartes réseau par GPU, fragmentant automatiquement les transferts pour maximiser la bande passante.
Surface API minimale : Un petit ensemble d'opérations suffisant pour construire des modèles de communication complexes.
TransferEngine génère un fil de travail par GPU. Chaque travailleur gère un DomainGroup qui coordonne 1 à 4 cartes réseau RDMA, en fonction de la plateforme matérielle (un seul ConnectX-7 fournit 400 Gbps, tandis que l'EFA nécessite quatre adaptateurs de 100 Gbps ou deux adaptateurs de 200 Gbps). Au sein d'un DomainGroup, chaque Domain est spécifique au matériel et gère une seule carte réseau, traitant la gestion des paires de files d'attente, la soumission du travail et l'interrogation de l'achèvement.
Chaque instance TransferEngine expose une seule adresse principale pour l'identification et la découverte des pairs. Les pairs distants échangent des structures NetAddr qui encodent les adresses réseau et peuvent être sérialisées pour une communication hors bande. Tous les pairs d'un système doivent utiliser le même nombre de cartes réseau par GPU. Cette configuration uniforme permet à TransferEngine de fragmenter intelligemment les transferts : lors de l'initiation d'un transfert, la source et la destination savent exactement quelles cartes réseau sont disponibles, permettant un équilibrage de charge automatique sur tous les adaptateurs.
Conception de l'API
TransferEngine expose une API minimale qui abstrait la complexité des interfaces RDMA sous-jacentes.
#[serde] struct NetAddr(Bytes);
#[serde] struct MrDesc{ ptr: u64, rkeys: Vec<(NetAddr, u64)> }
struct MrHandle(NonNull<c_void>); type Offset = u64;
struct Pages{ indices: Vec<u32>, stride: u64, offset: Offset }
struct PeerGroupHandle(u64);
struct ScatterDst{ len: u64, src: Offset, dst: (MrDesc,Offset)}
enum OnDone { Callback(fn () -> ()), Flag(Atomic<bool>) }
trait TransferEngine {
fn main_address() -> NetAddr;
// Memory Region Management
fn reg_mr(ptr, len, device) -> (MrHandle, MrDesc);
// Two-sided Send/Recv
fn submit_send(addr: NetAddr, msg: &[u8], cb: fn () -> ());
fn submit_recvs(len: u64, cnt: u64, cb: fn (&[u8]) -> ());
// One-sided Write
fn expect_imm_count(imm: u32, count: u32, cb: fn () -> ());
fn submit_single_write(len: u64, imm: Option<u32>
src: (MrHandle, Offset), dst: (MrDesc, Offset), OnDone);
fn submit_paged_writes(page_len: u64, imm: Option<u32>,
src: (MrHandle, Pages), dst: (MrDesc, Pages), OnDone);
// One-sided Write to a group of peers
fn add_peer_group(addrs: Vec<NetAddr>) -> PeerGroupHandle;
fn submit_scatter(h: Option<PeerGroupHandle>, OnDone,
imm: Option<u32>, src: MrHandle, dst: Vec<ScatterDst>);
fn submit_barrier(h: Option<PeerGroupHandle>, OnDone,
imm: u32, dst: Vec<MrDesc>);
// Watcher for CPU-GPU synchronization
fn alloc_uvm_watcher(cb: fn(u64,u64) -> ()) -> NonNull<u64>;
}L'API se compose de plusieurs composants clés :
Enregistrement de la mémoire
Les opérations RDMA nécessitent que les régions de mémoire soient enregistrées auprès de la carte réseau avant le transfert de données. La fonction reg_mr prend un pointeur, une longueur et un périphérique, et renvoie deux handles :
MrHandle : Un handle local utilisé comme source pour les transferts
MrDesc : Un descripteur sérialisable qui peut être échangé avec des pairs distants pour leur permettre d'écrire dans la mémoire enregistrée
La structure MrDesc contient une liste d'adresses réseau distantes et de clés distantes (une par carte réseau) qui accordent un accès distant à la région mémoire. Cette conception gère naturellement les configurations multi-NIC.
Envoi/Réception bilatéral
Pour les petits messages et les messages de plan de contrôle de style RPC, TransferEngine fournit submit_send et submit_recvs :
submit_send : Envoie le message à un pair distant identifié par NetAddr. Le message est copié en interne, de sorte que l'appelant peut immédiatement réutiliser ou libérer le tampon.
submit_recvs : Publie un pool de tampons de réception. Lorsqu'un message arrive, un rappel est invoqué avec les données reçues. Le tampon est automatiquement retourné au pool après la fin du rappel.
Écriture unilatérale
Pour les transferts de données en masse, TransferEngine fournit des opérations d'écriture unilatérales sans copie :
submit_single_write : Transfère une région contiguë de la source à la destination
submit_paged_writes : Transfère des pages non contiguës spécifiées par des indices, des foulées et des décalages
Les deux opérations acceptent une valeur immédiate optionnelle de 32 bits qui est transmise au récepteur de manière atomique lors de l'achèvement.
Le moteur fragmente automatiquement les transferts importants sur plusieurs cartes réseau pour maximiser la bande passante.
Groupes de pairs et opérations en masse
Pour les opérations ciblant plusieurs pairs (courant dans le routage MoE), TransferEngine fournit des primitives optimisées :
add_peer_group : Pré-enregistre un ensemble d'adresses de pairs pour une recherche rapide
submit_scatter : Envoie différentes tranches de données à chaque pair du groupe
submit_barrier : Envoie une notification immédiate uniquement à tous les pairs (transfert zéro octet)
Ces opérations sont optimisées pour une faible latence en mettant en cache les structures de données internes et en créant des modèles de demandes de travail RDMA, évitant ainsi les frais généraux de configuration répétitifs.
Observateur UVM : Signalisation GPU vers CPU
Un défi clé dans les systèmes LLM consiste à coordonner les transferts RDMA avec la progression du GPU. Idéalement, le transfert de données devrait commencer dès que le GPU produit une sortie, mais l'interrogation du GPU à partir du CPU est inefficace.
TransferEngine répond à cela avec des observateurs UVM : des emplacements de mémoire virtuelle unifiée qui peuvent être mis à jour par les noyaux GPU (même au sein des graphiques CUDA) et interrogés efficacement par le CPU en utilisant GDRCopy. Lorsqu'un observateur UVM est alloué, un rappel est enregistré qui est invoqué chaque fois que le GPU met à jour la valeur.
Cela permet des modèles tels que le transfert KvCache couche par couche : après la fin de chaque couche de transformateur, le GPU incrémente l'observateur, déclenchant un transfert RDMA du KvCache de cette couche pendant que les couches suivantes continuent de calculer.
Notification d'achèvement
TransferEngine fournit des mécanismes de notification d'achèvement pour le côté initiateur (l'appelant de submit_write) et le côté récepteur (la cible des opérations d'écriture).
Notification côté initiateur : Lorsqu'une opération d'écriture est soumise, l'appelant peut choisir d'être notifié de l'achèvement par une fonction de rappel ou un indicateur atomique. Cette notification confirme que les données ont été transférées avec succès depuis le tampon source.
Notification côté récepteur : L'innovation clé est l' ImmCounter, un composant qui suit les compteurs indexés par des nombres immédiats de 32 bits. Lorsqu'une opération d'écriture porte une valeur immédiate, l'ImmCounter du récepteur est incrémenté atomiquement lors de la livraison des données. La fonction expect_imm_count permet au récepteur d'enregistrer un rappel qui se déclenche lorsqu'une valeur immédiate spécifique atteint un nombre cible. Les compteurs peuvent être synchronisés de manière transparente avec le GPU via GDRCopy, interrogés directement ou surveillés par des rappels.
Ce mécanisme côté récepteur fournit une synchronisation fiable sans nécessiter d'ordonnancement des messages. Il fonctionne tout aussi bien sur les cartes réseau ConnectX (qui garantissent une livraison dans l'ordre sur le transport RC) et les cartes réseau EFA (qui livrent les messages dans le désordre sur le transport SRD).
Implémentation
TransferEngine est implémenté en Rust, optimisant soigneusement les allocations, le filetage et la synchronisation pour minimiser la latence et atteindre un débit élevé. Le moteur génère un fil de travail par DomainGroup, chacun épinglé à un cœur de CPU sur le nœud NUMA auquel les périphériques sont attachés, minimisant à la fois la planification et la latence d'accès à la mémoire. Les structures de données spécifiques à un domaine sont allouées après l'épinglage pour garantir que la mémoire est réservée dans le bon nœud NUMA. Un fil de travail gère jusqu'à 4 Domains (chacun gérant une seule carte réseau), tandis qu'un autre fil dédié interroge le GPU pour mettre à jour les observateurs UVM. La communication inter-fils se fait via des files d'attente sans verrouillage.
L'API transmet les demandes au DomainGroup desservant le périphérique approprié, le premier servant également l'hôte. Dans une boucle serrée, le travailleur du domaine interroge le nouveau travail, en donnant la priorité à la soumission de nouvelles demandes. Selon le matériel et la configuration, les demandes sont fragmentées et équilibrées sur les Domains disponibles. La première écriture d'une demande composite est immédiatement publiée dans la file d'attente d'envoi de la carte réseau. Une fois que les nouvelles demandes sont épuisées, le travailleur procède à la progression des opérations en attente, en publiant des écritures pour remplir le pipeline matériel. Par la suite, les files d'attente d'achèvement sont interrogées pour demander les transferts terminés et les incréments de compteur immédiats. Les événements sont agrégés pour fournir des notifications par transfert, confiant le transfert à un fil de rappel dédié partagé par tous les groupes.
La fragmentation au sein d'un DomainGroup est flexible. Les transferts peuvent cibler des cartes réseau spécifiques par index. Une seule écriture peut être divisée. Les transferts paginés, la diffusion et les opérations de barrière, qui se traduisent tous par plusieurs écritures, peuvent être fragmentés sur toutes les cartes réseau.
Optimisations spécifiques au matériel
Les Domains au sein de TransferEngine sont spécialisés et optimisés pour le matériel sous leur contrôle :
AWS EFA : L'implémentation EFA utilise libfabric, gérant un domaine de structure par carte réseau au sein du DomainGroup. Les décisions d'implémentation clés et les optimisations incluent :
Écritures de taille nulle : L'EFA diverge de la spécification RDMA en exigeant des descripteurs cibles valides même pour les écritures de taille nulle. Par conséquent, TransferEngine impose des descripteurs valides pour tous les transferts.
Modélisation des demandes de travail : Pré-remplit et conserve les champs communs des structures libfabric avant la publication, réduisant les frais généraux par opération pour les transferts en masse et les opérations de groupe de pairs
NVIDIA ConnectX-7 : L'implémentation ConnectX-7 utilise libibverbs avec transport RC. Les décisions d'implémentation clés et les optimisations incluent :
Paires de files d'attente UD pour les poignées de main : Chaque Domain utilise une paire de files d'attente de datagramme non fiable (UD) pour échanger des poignées de main de connexion RC.
Paires de files d'attente RC séparées : Crée 2 paires de files d'attente RC par pair, une pour les opérations d'envoi/réception bilatérales et une autre pour les opérations d'écriture unilatérales et WriteImm. Cette séparation est nécessaire car les achèvements de réception et de WriteImm consomment des demandes de travail dans l'ordre de publication, permettant des sémantiques de réception de haut niveau tout en prenant en charge WriteImm sans interférence.
Modélisation des demandes de travail : Pré-remplit les champs communs des structures de demande de travail pour les transferts en masse et les groupes de pairs.
Chaînage des demandes de travail : Lie jusqu'à 4 demandes de travail via le pointeur next de ibv_send_wr, réduisant le nombre de sonneries de porte à la carte réseau.
Ordonnancement relaxé : Active IBV_ACCESS_RELAXED_ORDERING pour permettre des transactions PCIe dans le désordre entre la carte réseau et la mémoire GPU.
Conclusion
En ciblant l'intersection des garanties sur différents transports RDMA et en introduisant une synchronisation indépendante de l'ordre via ImmCounter, TransferEngine fournit une base portable pour la communication point à point dans les systèmes LLM. La bibliothèque est open source et disponible sur GitHub, et nous encourageons la communauté à explorer et à construire dessus pour les modèles de communication émergents dans les systèmes d'apprentissage automatique à grande échelle.
Citation
@online{perplexity2025RdmaPointToPointCommunicationForLlmSystems,
author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
title = {{RDMA point-to-point communication for LLM systems}},
date = {2025-11-05},
url = {https://www.perplexity.ai/hub/blog/rdma-point-to-point-communication-for-llm-systems}
}