Activation de modèles à mille milliards de paramètres sur AWS EFA

Chez Perplexity, nous utilisons les meilleurs modèles pour notre produit, nos API et nos équipes de recherche. Les grands modèles open-source de type Mixture-of-Experts, tels que Kimi-K2, posent des défis particuliers, car les nœuds d'inférence les plus vastes équipés de 8 GPU NVIDIA H200 ne peuvent pas les prendre en charge efficacement, ce qui nécessite des déploiements multi-nœuds. Nous présentons un ensemble de noyaux pour le parallélisme d'experts qui atteignent des latences de pointe sur ConnectX-7, dépassant les performances de DeepEP. Ces mêmes noyaux sont également les premiers à atteindre des latences viables sur l'AWS Elastic Fabric Adapter (EFA), permettant le déploiement de modèles à mille milliards de paramètres.

Essayez nos noyaux sur GitHub et lisez l'article de recherche complet sur arXiv.

Introduction

L'architecture Mixture-of-Experts (MoE) s'est imposée comme la norme pour passer à l'échelle des centaines de milliards ou des milliers de milliards de paramètres tout en offrant une latence raisonnable lors de l'inférence. Les modèles MoE remplacent la couche dense d'un transformeur par un ensemble d'experts et une petite couche de routage linéaire qui décide avec quels experts multiplier un jeton. En regroupant les jetons associés au même expert, le problème devient épars et trivialement parallélisable, car différents experts peuvent désormais être hébergés sur différents GPU répartis sur plusieurs nœuds.

Contrairement à d'autres formes de parallélisme, telles que le parallélisme de tenseur (TP) et le parallélisme de données (DP) qui peuvent être facilement implémentés à l'aide de bibliothèques de communication collective comme torch.distributed et NCCL, le routage MoE implique une communication point à point éparse. Un noyau de répartition doit diviser l'ensemble des jetons entrants et les envoyer aux rangs hébergeant les experts, tandis que son problème dual, le noyau de combinaison, doit recevoir leurs homologues traités et calculer leur moyenne pondérée. Ces transferts ne correspondent pas à des collectifs et nécessitent des noyaux personnalisés pour atteindre une faible latence.

Chez Perplexity, nous avons développé des noyaux inter-nœuds portables et des noyaux intra-nœuds spécialisés, présentant nos travaux plus en détail dans un article précédent sur la communication MoE. Nous avons également montré dans un autre article de blog que le déploiement multi-nœuds peut réduire la latence et augmenter le débit pour les grands modèles MoE.

Dans cet article, nous analysons en profondeur l'implémentation de nos nouveaux noyaux inter-nœuds, qui améliorent considérablement les performances de nos travaux antérieurs basés sur NVSHMEM. En plus de présenter des résultats de pointe sur ConnectX-7, nous démontrons également qu'au sein de notre moteur d'exécution d'inférence, ces noyaux peuvent être utilisés pour servir efficacement les plus grands modèles open-source.

Inférence sur InfiniBand

Le parallélisme de tenseur et le parallélisme d'experts peuvent efficacement monter en charge au sein d'un même nœud, car NVLink offre un débit élevé (900 Go/s) avec une latence de l'ordre de la microseconde. Bien que les systèmes NVIDIA GB200 et GB300 soient dotés d'une mémoire suffisante pour contenir les poids et les caches KV des modèles à mille milliards de paramètres, les déploiements sur les instances AWS p5en avec jusqu'à 8 GPU H200 sont limités à 1120 Go de HBM, qui doivent être partagés entre les poids du modèle et les caches KV. Par conséquent, certains modèles nécessitent des déploiements inter-nœuds sur InfiniBand.

Bien qu'InfiniBand soit toujours extrêmement rapide, il n'offre qu'un débit de 400 Gbps (50 Go/s), ajoutant des dizaines à des centaines de microsecondes de latence pour le routage MoE. Heureusement, ces latences ajoutées peuvent être masquées par l'utilisation d'experts partagés, de micro-lots et du chevauchement calcul-communication pour offrir un débit viable à un coût compétitif. DeepSeek a été le pionnier de ces techniques via DeepSeek-R1, déployé sur des clusters de GPU H800 et d'adaptateurs InfiniBand NVIDIA ConnectX-7.

Sur les clusters AWS, la situation est différente en raison de l'utilisation de leurs adaptateurs personnalisés Elastic Fabric Adapters (EFA). Bien qu'ils atteignent un débit de pointe de 400 Gbps sur les collectifs généralement utilisés dans les charges de travail d'entraînement, ils sont légèrement inférieurs à ConnectX-7 en ce qui concerne les tailles de messages échangés lors de la répartition et de la combinaison MoE. De plus, l'EFA ne prend pas en charge GPUDirect Async, ce qui nécessite un fil de proxy CPU faisant le pont entre le GPU et la carte réseau afin d'initier les transferts. Outre la complexité accrue, les transactions PCIe supplémentaires ajoutent des microsecondes de surcoût à tous les transferts. Des frameworks tels que NVSHMEM permettent aux auteurs de noyaux de fournir des implémentations indépendantes du périphérique, cependant les API génériques entraînent des surcoûts qui rendent les déploiements inefficaces en termes de coût et de latence.

pplx-kernels basés sur NVSHMEM

Nos noyaux précédents abstraient l'implémentation de transport sous-jacente via NVSHMEM. Ils reposaient sur nvshmemx_putmem_signal_nbi_warp pour transférer les jetons un par un, en utilisant des compteurs atomiques pour synchroniser les pairs. Bien que cela offrait des performances raisonnables sur les cartes réseau ConnectX-7 avec IBGDA, l'implémentation IBRC basée sur un proxy pour ConnectX-7 était nettement plus lente, tandis que l'implémentation EFA n'a pas réussi à fournir des latences inférieures à une milliseconde.

Nous avons utilisé les noyaux comme banc d'essai pour le matériel sous-jacent, tout en déterminant l'interface la plus appropriée à exposer à nos modèles. Nous avons émis l'hypothèse que les transferts fréquents et les interactions entre les noyaux et les cartes réseau ou les fils de proxy entravaient les performances. De plus, notre implémentation ne reposait sur aucune garantie matérielle autre que la fiabilité, ce qui suggère que nous pourrions accélérer les choses en réduisant le surcoût des primitives de synchronisation.

Suite à ces observations, nous avons conçu nos noyaux basés sur EFA exposant la même interface, optimisant les interactions noyau de périphérique - proxy hôte sous le capot pour des performances maximales.

MoE all-to-all sur EFA

Agrandir

Les nouveaux noyaux suivent une architecture hybride CPU-GPU, où les noyaux GPU interagissent avec le modèle sur le périphérique et un fil de proxy s'exécutant sur le CPU hôte gère les interactions avec la carte réseau. Les nouveaux noyaux réutilisent le même TransferEngine que nous avons construit pour les transferts de cache KV.

Le routage MoE est effectué par un noyau de répartition, qui reçoit les jetons et les experts par jeton auxquels les router à partir du modèle, en les expédiant sur le réseau vers les rangs qui hébergent les experts. Une fois que les noyaux GEMM groupés ont multiplié les jetons avec les poids des experts, le noyau de combinaison transfère les jetons vers leurs rangs source et calcule leur moyenne pondérée avant de les renvoyer au modèle.

Les noyaux de répartition et de combinaison sont tous deux divisés en une moitié expéditeur et une moitié récepteur pour faciliter le chevauchement et le micro-lotissement. L'expéditeur mélange les jetons dans des tampons d'envoi afin qu'ils puissent être écrits chez les pairs via une seule opération d'écriture RDMA. Après l'exécution de l'expéditeur, le GPU peut effectuer d'autres tâches (telles que la multiplication par un expert partagé) pendant que les transferts sur le réseau sont en attente en arrière-plan. Les noyaux récepteurs se bloquent jusqu'à ce que tous les transferts aient été reçus. En arrière-plan, un fil de travail côté hôte gère les transferts : il est averti par le GPU lorsque les tampons d'envoi ont été remplis, après quoi il déclenche les écritures RDMA. Une fois que la carte réseau signale l'achèvement, il débloque le récepteur, qui renvoie les jetons au modèle.

Les noyaux sont étroitement couplés à l'hôte, communiquant entre eux via la mémoire unifiée et GDRCopy. La mémoire unifiée CUDA permet à la mémoire du périphérique ou à la mémoire de l'hôte d'être mappée dans les espaces d'adressage des deux, gérant de manière transparente la communication via PCIe.

Bien que la mémoire unifiée soit adaptée aux transferts en masse, l'interrogation à faible latence s'effectue via GDRCopy. GDRCopy utilise GPUDirect RDMA pour lire et écrire dans la mémoire GPU avec une latence de l'ordre de la microseconde.

En plus du RDMA, au sein d'un nœud, les noyaux tirent également parti de NVLink pour échanger des jetons. Cela peut décharger une fraction significative du trafic, jusqu'à un quart pour EP32 et jusqu'à un huitième pour EP64, obtenant ainsi des latences globales plus faibles.

Tampons d'envoi et de réception

La conception des tampons d'envoi pour le noyau de répartition est cruciale pour atteindre des latences viables impliquant le proxy hôte. S'il y a N rangs expédiant chacun T jetons, chacun routé vers R experts parmi E, un rang peut recevoir jusqu'à N * T * max(R, E / N) jetons dans le pire des cas, si tous lui sont routés. De plus, cette borne minimale ne peut être atteinte que lorsque les expéditeurs disposent leurs écritures de manière contiguë, car des tampons de réception privés consommeraient trop de mémoire. Pour que les expéditeurs puissent disposer leurs écritures consécutivement sans conflit, une connaissance complète des informations de routage doit être distribuée à chaque rang. Par conséquent, nos noyaux de répartition échangent d'abord le nombre de jetons par expert, après quoi ils peuvent chacun déterminer où vont leurs écritures et à partir de quel décalage ils reçoivent des jetons de chaque rang.

Le simple fait d'échanger des informations de routage entraînerait des latences inutilement élevées, car chaque rang devrait attendre que les écritures arrivent avant d'expédier ses propres écritures. Au lieu de cela, nous réservons une petite quantité d'espace privé par expéditeur sur chaque rang récepteur et expédions certains des jetons parallèlement aux écritures des informations de routage. Cela garantit une utilisation complète de la bande passante tout au long du processus de répartition. Une fois les routes échangées, les jetons restants sont transférés via une seule écriture par pair distant. L'échange d'informations de routage et les tampons de taille fixe, ainsi que les transferts les remplissant, sont illustrés dans la figure ci-dessous.

Agrandir

Les noyaux de combinaison sont nettement plus simples, car les informations de routage sont déjà disponibles pour eux à partir de l'étape de répartition. Ils peuvent copier les jetons en une seule écriture par pair, après quoi ils lisent à partir du tampon de réception pour calculer la moyenne pondérée des jetons dans les tenseurs de sortie.

Pour NVLink, les tampons d'envoi et de réception au sein d'un nœud sont mappés dans les pairs afin qu'ils puissent pousser ou tirer directement des jetons des rangs voisins.

Répartition et combinaison

Les figures ci-dessous détaillent le timing et les interactions entre le GPU, le CPU et la carte réseau, montrant une itération de répartition et de combinaison entrelacée avec les noyaux du modèle.

Le travail commence dans l'envoi de répartition, avec le noyau agrégeant les informations de routage et calculant le nombre de jetons par expert via une accumulation atomique en mémoire partagée au sein d'un seul bloc. Une fois que le nombre de jetons est disponible, ils sont copiés dans la mémoire hôte et le proxy est immédiatement averti de commencer à répartir les informations de routage vers tous les autres rangs pairs. En parallèle, la répartition continue de calculer la somme cumulative du nombre de jetons, identifiant les décalages nécessaires pour emballer les jetons dans des tampons d'envoi contigus par pair. Une fois les décalages disponibles, tous les jetons à envoyer via EFA sont emballés dans les tampons d'envoi. Le proxy est averti d'inscrire les écritures dans les tampons d'envoi des pairs inter-nœuds. Au sein du nœud, les noyaux poussent également immédiatement les jetons dans des tampons privés au sein du même nœud via NVLink.

Agrandir

Agrandir

Une fois qu'un rang reçoit des informations de routage de tous les pairs, il les centralise et calcule divers décalages nécessaires pour identifier les jetons pour la répartition et la combinaison. Ce travail est effectué sur le CPU alors que des transferts sont déjà en attente, il n'est donc pas critique en termes de latence. Par la suite, une fois que tous les décalages sont connus, les écritures pour le reste des jetons, qui dépassent la taille du tampon privé, sont publiées et la phase d'envoi de la répartition est terminée.

Côté récepteur, le noyau tire d'abord les jetons du même nœud qui dépassent la capacité des tampons privés via NVLink. Ensuite, il se bloque jusqu'à ce que le proxy hôte signale l'achèvement de tous les transferts. Une fois le signal reçu, les jetons sont mélangés des tampons de réception dans le tenseur de sortie transmis au GEMM de groupe. Étant donné que les tampons sont réutilisés entre la répartition et la combinaison, les barrières NVLink et RDMA garantissent que les tampons de réception ont été lus avant de permettre au noyau suivant de les écraser.

Les noyaux de combinaison suivent un modèle globalement similaire, mélangeant les jetons, signalant le proxy et attendant l'achèvement. Cependant, les métadonnées pour publier leurs écritures sont facilement disponibles et tous les décalages distants sont connus, ainsi tous les jetons transférés au sein du nœud sont poussés via NVLink vers les récepteurs et les jetons inter-nœuds sont transférés via une seule écriture. Avant de renvoyer les résultats, le noyau accumule les jetons des tampons de réception, en les écrivant dans les tenseurs de sortie.

TransferEngine pour le routage MoE

Pour prototyper notre implémentation, nous avons initialement réutilisé le moteur de transfert construit pour le pré-remplissage et le décodage désagrégés afin d'émettre et de suivre les écritures RDMA. Cependant, contrairement au décodage désagrégé où chaque nœud publie un grand nombre d'écritures vers le même pair, la répartition MoE publie de petits transferts vers tous les pairs en une seule opération. D'emblée, les performances n'étaient pas optimales et notre implémentation accusait un retard de 20 us sur DeepEP sur ConnectX-7. Par conséquent, nous avons spécialisé le TransferEngine pour prendre en charge deux opérations, scatter (dispersion) et barrier (barrière), afin d'optimiser les interactions avec un grand nombre de pairs.

Les opérations de dispersion et de barrière exposent toutes deux des groupes de pairs qui peuvent être pré-enregistrés pour pré-remplir les structures de données utilisées par l'implémentation de transport sous-jacente (templating de demande de travail (WR)). La dispersion copie différentes tranches du même tampon source dans différents décalages sur tous les membres d'un groupe de pairs. La barrière signale tous les pairs via un immédiat. La mise en file d'attente et le traitement de ces opérations en masse permettent d'économiser des microsecondes précieuses sur le chemin critique. Sur EFA, où deux cartes réseau de 200 Gbps s'agrègent pour fournir une bande passante de 400 Gbps, nous fragmentons les pairs du groupe sur les deux cartes réseau au lieu de fragmenter les octets de chaque transfert.

Bien qu'initialement nous ayons développé le TransferEngine uniquement sur EFA en utilisant libfabric, les noyaux MoE fonctionnent sur ConnectX-7 sans changer une seule ligne de code après avoir ajouté la prise en charge de ConnectX-7 dans le TransferEngine en utilisant libibverbs. Par rapport à l'EFA basée sur le SRD (Scalable Reliable Datagram), ConnectX-7 nécessite un travail supplémentaire pour la configuration de la connexion et la gestion des pairs. Nous utilisons une paire de files d'attente de datagrammes non fiables (UD) pour démarrer la connexion fiable (RC). Nous créons deux paires de files d'attente pour chaque pair distant, une pour l'ENVOI/RÉCEPTION bilatérale et une pour l'ÉCRITURE_IMM unilatérale, car la RÉCEPTION et l'ÉCRITURE_IMM consomment toutes deux l'opération de RÉCEPTION dans l'ordre alors que les deux opérations nécessitent des tailles de tampon différentes. En outre, nous assouplissons l'ordonnancement des transactions PCIe entre ConnectX-7 et la mémoire GPU. Comme pour l'EFA, nous implémentons l'optimisation du templating WR. De plus, lors de l'envoi de plusieurs petits messages au même pair, nous utilisons le chaînage WR pour améliorer le taux de messages.

Comme le travail côté GPU est identique sur les deux adaptateurs, nous avons poussé l'implémentation de transport sous-jacente à ses limites pour dépasser les performances de DeepEP afin de souligner que l'écart entre le routage MoE EFA et ConnectX-7 est déterminé par les capacités du matériel sous-jacent.

Évaluation

Latence de répartition/combinaison

Nous avons évalué les latences de répartition et de combinaison sur des systèmes H200 équipés d'adaptateurs ConnectX-7 et EFA sur la configuration standard DeepSeek-V3 et Kimi-K2 avec une dimension cachée de 7168, une répartition fp8 mise à l'échelle par bloc et une combinaison bf16. Nous considérons une taille de lot de 128 pour la répartition et 4096 pour la combinaison. Sur ConnectX-7, nous comparons à DeepEP et à nos précédents pplx-kernels basés sur NVSHMEM, en utilisant à la fois GPUDirect Async (IBGDA) et un proxy hôte (IBRC). Nous atteignons des latences combinées de 459 us, 582 us et 692 us sur EP16, EP32 et EP64, respectivement. Nous surpassons d'autres noyaux basés sur EFA reposant sur des proxies génériques, tels que UCCL-EP, qui atteignent 519 us, 966 us et 1159 us.

Agrandir

Sur le décodage, bien que nos noyaux soient quelques microsecondes plus lents que DeepEP, nous atteignons des latences globales plus faibles car nos noyaux de combinaison sont nettement plus rapides. Bien que nous utilisions un proxy hôte, grâce à nos transferts en masse, nous dépassons les performances de toutes les implémentations sans proxy basées sur IBGDA. Nous montrons également que l'implémentation du proxy compte, car notre proxy hautement spécialisé est un ordre de grandeur plus rapide que l'implémentation NVSHMEM générique. Malgré le fait que l'EFA soit deux fois plus lent que ConnectX-7 lors du transfert de paquets de 256 Ko, qui sont proches de la taille de paquet typique échangée par le routage MoE, l'écart entre les deux adaptateurs est globalement moins substantiel.

Agrandir

Même si nos noyaux optimisés pour le décodage peuvent s'adapter aux tailles de lots de pré-remplissage sans aucun ajustement supplémentaire, les performances sont à la traîne par rapport à DeepEP. Cela est principalement dû au fait que nous accumulons les jetons uniquement du côté récepteur de la combinaison, alors que DeepEP accumule dans l'expéditeur pour réduire la quantité de données transférées. Cependant, cette optimisation peut avoir des implications sur la précision. De plus, étant donné que les noyaux transfèrent les données en une seule écriture, nous utilisons plus de mémoire pour les tampons d'envoi et de réception.

Débit Kimi-K2 / DeepSeek-V3

Pour prouver la viabilité de nos noyaux, nous analysons les performances de bout en bout de notre moteur d'inférence interne servant DeepSeek-V3 (671B paramètres) et Kimi-K2 (1T paramètres) sur plusieurs instances p5en équipées d'Elastic Fabric Adapters (EFA) et de GPU H200. Nous exécutons le modèle en mode de vérification de brouillon et rapportons la latence en supposant un taux d'acceptation de 80 % en utilisant une seule couche MTP.

Sur DeepSeek-V3, nous comparons à une ligne de base intra-nœud EP=8 exécutant des noyaux de répartition et de combinaison efficaces basés uniquement sur NVLink. Les déploiements inter-nœuds égalent ou dépassent globalement les performances du déploiement mono-nœud, en particulier sur des lots de taille moyenne, démontrant que la mise à l'échelle sur plusieurs nœuds apporte des améliorations au débit global.

Sur Kimi-K2, les noyaux permettent le service. En raison de sa taille, le modèle ne peut pas tenir dans le même nœud en utilisant le parallélisme de données (DP) sur un nœud H200, un déploiement inter-nœud est donc nécessaire pour servir le modèle. Comme le modèle possède moins de têtes d'attention, nous obtenons de meilleures latences que DeepSeek-V3 et nous pouvons servir le modèle sur EFA à une latence viable.

Agrandir

Agrandir

Nous avons également envisagé le micro-lotissement, bien qu'il n'ait amélioré les latences qu'à des tailles de lots très élevées ou lors de l'utilisation de plus d'un jeton de brouillon. Pour bénéficier de la division des problèmes GEMM groupés en deux, les tailles de lots de jetons doivent être suffisamment denses. La majeure partie du temps est encore consacrée à la répartition et à la combinaison MoE, ce qui indique que toute amélioration supplémentaire des noyaux améliorera par conséquent les latences et facilitera le micro-lotissement.

Travaux futurs

Nous travaillons en étroite collaboration avec les ingénieurs d'AWS pour améliorer les performances sur EFA. De notre côté, nous suivrons les mises à jour de libfabric qui réduiront la surcharge du plan de données. De plus, nous expérimenterons avec efa-direct pour réduire davantage les surcharges de l'espace utilisateur et interfacer directement avec le matériel.

Références

Rapport technique DeepSeek R1

Rapport technique Kimi K2

DeepEP

UCCL-EP

Amélioration des performances réseau des systèmes HPC utilisant NVIDIA Magnum IO NVSHMEM et GPUDirect Async

Citation

@online{perplexity2025EnablingTrillionParameterModelsOnAwsEfa,
  author = {{Perplexity} and Nandor Licker and Kevin Hu and Vladimir Zaytsev and Lequn Chen},
  title = {{Enabling trillion-parameter models on AWS EFA}},
  date = {2025-11-04},
  url = {https://www.perplexity.ai/hub/blog/enabling-trillion-parameter-models-on-aws-efa}
}