Activer des modles dun billion de paramtres sur AWS EFA

Chez Perplexity, nous utilisons les meilleurs modles pour notre produit, nos API et nos quipes de recherche. Les grands modles de mlange dexperts open-source, tels que Kimi-K2, posent des difficults particulires, car les plus grands nuds dinfrence dots de 8 processeurs graphiques NVIDIA H200 ne peuvent pas les accueillir efficacement, ce qui ncessite des dploiements multi-nuds. Nous prsentons un ensemble de noyaux pour le paralllisme des experts qui atteignent des latences de pointe sur ConnectX-7, dpassant les performances de DeepEP. Ces mmes noyaux sont galement les premiers atteindre des latences viables sur lElastic Fabric Adapter (EFA) dAWS, permettant le dploiement de modles comptant un billion de paramtres.

Essayez nos noyaux sur GitHub et lisez larticle de recherche complet sur arXiv.

Le modle de mlange dexperts (MoE) sest impos comme larchitecture standard pour passer des centaines de milliards ou des billions de paramtres tout en offrant une latence raisonnable au moment de linfrence. Les modles MoE remplacent la couche dense dun transformateur par un ensemble dexperts et une petite couche de routage linaire qui dcide avec quel expert multiplier un jeton. En regroupant les jetons associs au mme expert, le problme devient clairsem et facilement paralllisable, car les diffrents experts peuvent dsormais tre hbergs sur des processeurs graphiques distincts rpartis sur plusieurs nuds.

Contrairement dautres formes de paralllisme, telles que le paralllisme de tenseurs (TP) et le paralllisme de donnes (DP) qui peuvent tre facilement implments laide de bibliothques de communication collective telles que torch.distributed et NCCL, le routage MoE implique une communication poste poste clairseme. Un noyau de rpartition doit diviser lensemble des jetons entrants et les acheminer vers les rangs hbergeant les experts, tandis que son problme dual, le noyau de combinaison, doit recevoir leurs homologues traits et calculer leur moyenne pondre. Ces transferts ne correspondent pas des oprations collectives et ncessitent des noyaux personnaliss pour atteindre une faible latence.

Chez Perplexity, nous avons dvelopp la fois des noyaux inter-nuds portables et des noyaux intra-nuds spcialiss, prsentant notre travail de manire plus approfondie dans un article prcdent sur la communication MoE. Nous avons galement montr dans un autre article de blog que le dploiement multi-nuds peut apporter une latence plus faible et un dbit plus lev pour les grands modles MoE.

Dans cet article, nous analysons en profondeur limplmentation de nos nouveaux noyaux inter-nuds qui amliorent considrablement les performances de nos travaux prcdents fonds sur NVSHMEM. En plus de prsenter des rsultats de pointe sur ConnectX-7, nous dmontrons galement quau sein de notre excution dinfrence, les noyaux peuvent tre utiliss pour servir efficacement les plus grands modles open-source.

Le paralllisme de tenseurs et le paralllisme dexperts peuvent tre mis lchelle efficacement au sein dun seul nud, car NVLink offre un dbit lev (900 Go/s) pour une latence de lordre de la microseconde. Bien que les systmes NVIDIA GB200 et GB300 soient dots dune mmoire suffisante pour contenir les poids et les caches KV des modles comptant un billion de paramtres, les dploiements sur linstance AWS p5en avec jusqu 8 processeurs graphiques H200 sont limts 1 120 Go de HBM, qui doivent tre partags entre les poids du modle et les caches KV. Par consquent, certains modles ncessitent des dploiements inter-nuds sur InfiniBand.

Bien quInfiniBand soit toujours trs rapide, il noffre quun dbit de 400 Gb/s (50 Go/s), ce qui ajoute des dizaines des centaines de microsecondes de latence pour le routage MoE. Heureusement, ces latences supplmentaires peuvent tre masques grce lutilisation dexperts partags, de micro-lots et du chevauchement calcul-communication pour offrir un dbit viable un cot comptitif. DeepSeek a ouvert la voie ces techniques via DeepSeek-R1, dploy sur des grappes de processeurs graphiques H800 et dadaptateurs NVIDIA ConnectX-7 InfiniBand.

Sur les grappes AWS, la situation est diffrente en raison de lutilisation de leurs adaptateurs personnaliss Elastic Fabric Adapters (EFA). Bien que ceux-ci offrent un dbit de pointe de 400 Gb/s sur les communications collectives gnralement utilises dans les charges de travail dapprentissage, ils sont lgrement infrieurs ConnectX-7 en ce qui concerne la taille des messages changs lors de la rpartition et de la combinaison MoE. De plus, EFA ne prend pas en charge GPUDirect Async, ce qui ncessite un fil dexcution proxy de UC faisant le pont entre le processeur graphique et la carte rseau afin dinitier les transferts. Outre la complexit accrue, les transactions PCIe supplmentaires ajoutent des microsecondes de surcharge tous les transferts. Des frameworks tels que NVSHMEM permettent aux auteurs de noyaux de fournir des implmentations indpendantes du priphrique, mais les API gnriques engendrent des surcharges qui rendent les dploiements inefficaces en termes de cot et de latence.

Nos noyaux prcdents ont abstrait limplmentation du transport sous-jacent grce NVSHMEM. Ils reposaient sur nvshmemx_putmem_signal_nbi_warp pour transfrer les jetons un par un, en utilisant des compteurs atomiques pour synchroniser les pairs. Bien que cela offrait des performances raisonnables sur les cartes rseau ConnectX-7 avec IBGDA, limplmentation ConnectX-7 IBRC base sur un proxy tait nettement plus lente, tandis que limplmentation EFA na pas parvenu offrir des latences infrieures une milliseconde.

Nous avons utilis les noyaux comme banc dessai pour le matriel sous-jacent, tout en dterminant linterface la plus approprie exposer nos modles. Nous avons mis lhypothse que les transferts frquents et les interactions entre les noyaux et les cartes rseau ou les fils dexcution proxy entravaient les performances. De plus, notre implmentation ne reposait sur aucune garantie de la part du matriel autre que la fiabilit, ce qui suggrait que nous pouvions acclrer les choses en rduisant la surcharge des primitives de synchronisation.

la suite de ces observations, nous avons conu nos noyaux bass sur EFA exposant la mme interface, optimisant les interactions entre le noyau du priphrique et le proxy de lhte sous le capot pour des performances maximales.

Communication hybride RDMA et NVLink avec un proxy hôte, montrant les rangs d’initiation et de réception, les workers du proxy et les NIC.

Les nouveaux noyaux suivent une architecture hybride UC-processeur graphique, o les noyaux de processeur graphique interagissent avec le modle sur le priphrique et un fil dexcution proxy sexcutant sur le processeur hte gre les interactions avec la carte rseau. Les nouveaux noyaux rutilisent le mme TransferEngine que nous avons bti pour les transferts de cache KV.

Le routage MoE est effectu par un noyau de rpartition, qui reoit du modle les jetons et les experts par jeton afin de les acheminer, en les expdiant sur le rseau vers les rangs qui hbergent les experts. Une fois que les noyaux GEMM groups multiplient les jetons avec les poids des experts, le noyau de combinaison transfre les jetons vers leurs rangs source et calcule leur moyenne pondre avant de les renvoyer au modle.

Les noyaux de rpartition et de combinaison sont diviss en une moiti dmission et une moiti de rception pour faciliter le chevauchement et le traitement par micro-lots. Lmetteur organise les jetons dans des tampons dmission afin quils puissent tre crits aux pairs par le biais dune seule opration dcriture RDMA. Aprs lexcution de lmetteur, le processeur graphique peut effectuer ds autres tches (telles que la multiplication dexperts partags) pendant que les transferts sur le rseau sont en attente en arrire-plan. Les noyaux de rception se bloquent jusqu ce que tous les transferts aient t reus. En arrire-plan, un fil dexcution de travail ct hte gre les transferts : il est signal par le processeur graphique lorsque les tampons dmission ont t remplis, aprs quoi il rpartit les critures RDMA. Une fois que la carte rseau signale lachvement, elle dbloque le rcepteur, qui renvoie les jetons au modle.

Les noyaux sont troitement coupls lhte, communiquant entre eux par lintermdiaire de la mmoire unifie et de GDRCopy. La mmoire unifie CUDA permet de mapper la mmoire du priphrique ou de lhte dans les Espace dadressage des deux, grant de manire transparente la communication via PCIe.

Bien que la mmoire unifie convienne aux transferts en bloc, linterrogation faible latence est effectue par GDRCopy. GDRCopy utilise GPUDirect RDMA pour lire et crire dans la mmoire du processeur graphique avec une latence de lordre de la microseconde.

En plus de RDMA, au sein dun nud, les noyaux tirent galement parti de NVLink pour changer des jetons. Cela permet de dlester une fraction importante du trafic, jusqu un quart pour EP32 et jusqu un huitime pour EP64, permettant datteindre des latences globales plus faibles.

La conception des tampons dmission pour le noyau de rpartition est cruciale pour atteindre des latences viables impliquant le proxy de lhte. Sil y a N rangs rpartissant chacun T jetons, chacun achemin vers R experts sur E, un rang peut recevoir jusqu N * T * max(R, E / N) jetons dans le pire des cas, si tous y sont achemins. De plus, cette limite minimale ne peut tre atteinte que lorsque les metteurs disposent leurs critures de manire contigue, car des tampons de rception privs consommeraient trop de mmoire. Pour que les metteurs puissent disposer leurs critures conscutivement sans conflit, une connaissance approfondie des informations de routage doit tre distribue chaque rang. Par consquent, nos noyaux de rpartition changent dabord les dcomptes de jetons par expert, aprs quoi chacun peut dterminer o vont ses critures et partir de quel dcalage il reoit des jetons de chaque rang.

Le simple change dinformations de routage entranerait des latences inopinment leves, car chaque rang devrait attendre larrive des critures avant dexpdier ses propres critures. Au lieu de cela, nous rservons une petite quantit dEspace priv par metteur sur chaque rang de rception et expdions certains jetons en mme temps que les critures des informations de routage. Cela garantit une utilisation maximale de la bande passante tout au long de la majeure partie du processus de rpartition. Une fois les routes changes, les jetons restants sont transfrs via une seule criture par pair distant. Lchange dinformations de routage et des tampons de taille fixe, ainsi que les transferts les alimentant, sont illustrs dans la figure ci-dessous.

Disposition des tampons pour la répartition MoE, montrant comment les données de jetons et les métadonnées de routage sont organisées et échangées entre deux rangs.

Les noyaux de combinaison sont sensiblement plus simples, car les informations de routage leur sont dj disponibles partir de ltape de rpartition. Ils peuvent copier les jetons en une seule criture par pair, aprs quoi ils lisent le tampon de rception pour calculer la moyenne pondre des jetons dans les tenseurs de sortie.

Pour NVLink, les tampons dmission et de rception au sein dun nud sont mapps dans les pairs pour leur permettre de pousser ou de tirer directement des jetons des rangs voisins.

Les figures ci-dessous dtaillent le calendrier et les interactions entre le processeur graphique, le processeur et la carte rseau, montrant une itration de rpartition et de combinaison entrelace avec les noyaux du modle.

Le travail commence dans lmission de rpartition, le noyau agrgeant les informations de routage et calculant les dcomptes de jetons par expert par le biais dune accumulation atomique dans la mmoire partage au sein dun seul bloc. Une fois que les dcomptes de jetons sont disponibles, ils sont copis dans la mmoire de lhte et le proxy est immdiatement averti de commencer rpartir les informations de routage vers tous les autres rangs de pairs. En parallle, la rpartition continue de calculer la somme cumulative des dcomptes de jetons, identifiant les dcalages ncessaires pour emballer les jetons dans des tampons dmission contigus par pair. Une fois que les dcalages sont disponibles, tous les jetons envoyer via EFA sont emballs dans les tampons dmission. Le proxy est invit publier les critures dans les tampons dmission des pairs inter-nuds. Au niveau intra-nud, les noyaux poussent galement immdiatement des jetons dans des tampons privs au sein du mme nud via NVLink.

Synchronisation et coordination de la répartition MoE, montrant les opérations GPU, CPU et NIC par étapes ainsi que leur synchronisation.

Synchronisation et coordination du regroupement MoE, montrant les opérations GPU, CPU et NIC ainsi que leur synchronisation.

Une fois quun rang reoit des informations de routage de tous les pairs, il les centralise et calcule divers dcalages ncessaires pour identifier les jetons destins la rpartition et la combinaison. Ce travail est effectu sur le processeur alors que des transferts sont dj en attente, il nest donc pas critique en termes de latence. Par la suite, une fois que tous les dcalages sont connus, les critures pour le reste des jetons, qui dpassent la taille du tampon priv, sont publies et la phase dmission de la rpartition est termine.

Du ct du rcepteur, le noyau tire dabord du mme nud des jetons qui dpassent la capacit des tampons privs via NVLink. Ensuite, il se bloque jusqu ce que le proxy de lhte signale lachvement de tous les transferts. Une fois le signal reu, les jetons sont mlangs des tampons de rception vers le tenseur de sortie transmis au groupe GEMM. tant donn que les tampons sont rutiliss entre la rpartition et la combinaison, les barres NVLink et RDMA garantissent que les tampons de rception ont t lus avant de permettre au noyau suivant de les craser.

Les noyaux de combinaison suivent un schma sensiblement similaire, mlangeant des jetons, signalant le proxy et attendant lachvement. Cependant, les mtadonnes permettant de publier leurs critures sont immdiatement disponibles et tous les dcalages distants sont connus, de sorte que tous les jetons transfrs en intra-nud sont pousss via NVLink vers les rcepteurs et les jetons inter-nuds sont transfrs via une seule criture. Avant de renvoyer les rsultats, le noyau accumule les jetons provenant des tampons de rception, en les crivant dans les tenseurs de sortie.

Pour prototyper notre implmentation, nous avons initialement rutilis le moteur de transfert (Transfer Engine) bti pour le pr-remplissage et le dcodage dsagrgs afin dmettre et de suivre les critures RDMA. Cependant, contrairement au dcodage dsagrg o chaque nud publie un grand nombre dcritures vers le mme pair, la rpartition MoE publie de petits transferts vers tous les pairs en une seule opration. Prt lemploi, les performances ntaient pas optimales et notre implmentation acclrait un retard denviron 20 s par rapport DeepEP sur ConnectX-7. Par consquent, nous avons spcialis le TransferEngine pour prendre en charge deux oprations, la dispersion (scatter) et la barrire, afin doptimiser les interactions avec un grand nombre de pairs.

Les oprations de dispersion et de barrire exposent des groupes de pairs qui peuvent tre pr-enregistrs pour pr-remplir les structures de donnes utilises par limplmentation du transport sous-jacent (modlisation de requtes de travail [WR]). La dispersion copie diffrentes tranches du mme tampon source dans diffrents dcalages sur tous les membres dun groupe de pairs. La barrire signale tous les pairs par le biais dune immdiate. Mettre en file dattente et grer ces oprations en bloc permet dconomiser de prcieuses microsecondes sur le chemin critique. Sur EFA, o deux cartes rseau de 200 Gb/s sagrgment pour fournir une bande passante de 400 Gb/s, nous fractionnons les pairs du groupe entre les deux cartes rseau au lieu de fractionner les octets de chaque transfert.

Bien quinitialement nous ayons dvelopp le TransferEngine uniquement sur EFA en utilisant libfabric, les noyaux MoE fonctionnent sur ConnectX-7 sans modifier une seule ligne de code aprs lajout de la prise en charge de ConnectX-7 dans le TransferEngine en utilisant libibverbs. Par rapport lEFA fond sur des datagrammes fiables et lchelle (SRD), ConnectX-7 ncessite un travail supplmentaire pour la configuration de la connexion et la gestion des pairs. Nous utilisons une paire de files dattente de datagrammes non fiables (UD) pour amorcer la connexion fiable (RC). Nous crons deux paires de files dattente pour chaque pair distant, lune pour SEND/RECV bilatral et lautre pour WRITE_IMM unilatral, car RECV et WRITE_IMM consomment tous deux lopration RECV dans lordre tandis que les deux oprations ncessitent des tailles de tampon diffrentes. En outre, nous assouplissons lordonnancement des transactions PCIe entre ConnectX-7 et la mmoire du processeur graphique. De manire similaire lEFA, nous implmentons loptimisation de la modlisation des requtes de travail (WR). De plus, lors de lenvoi de plusieurs petits messages au mme pair, nous utilisons le chanage de requtes de travail pour amliorer le dbit des messages.

tant donn que le travail du ct du processeur graphique est identique sur les deux adaptateurs, nous avons pouss limplmentation du transport sous-jacent ses limites pour dpasser les performances de DeepEP afin de souligner que lcart entre le routage MoE EFA et ConnectX-7 est dtermin par les capacits du matriel sous-jacent.

Nous avons valu les latences de rpartition et de combinaison sur des systmes H200 quips dadaptateurs ConnectX-7 et EFA sur la configuration standard DeepSeek-V3 et Kimi-K2 avec une dimension cache de 7168, une rpartition fp8 mise lchelle par blocs et une combinaison bf16. Nous considrons une taille de lot de 128 pour la rpartition et de 4096 pour la combinaison. Sur ConnectX-7, nous comparons DeepEP et nos noyaux pplx prcdents fonds sur NVSHMEM, en utilisant la fois GPUDirect Async (IBGDA) et un proxy hte (IBRC). Nous obtenons des latences combines de 459 s, 582 s et 692 s sur EP16, EP32 et EP64, respectivement. Nous surpassons ds autres noyaux fonds sur EFA reposant sur des proxys gnriques, tels que UCCL-EP, qui atteignent 519 s, 966 s et 1159 s.

Latence de répartition et de regroupement pour le décodage MoE EP64, comparant les implémentations et les chemins de communication illustrés.

Sur le dcodage, bien que nos noyaux soient quelques microsecondes plus lents que DeepEP, nous obtenons des latences globales plus faibles car nos noyaux de combinaison sont sensiblement plus rapides. Bien que nous utilisions un proxy hte, en raison de nos transferts en bloc, nous dpassons les performances de toutes les implmentations sans proxy fondes sur IBGDA. Nous montrons galement que limplmentation du proxy importe, car notre proxy hautement spcialis est un ordre de grandeur plus rapide que limplmentation NVSHMEM gnrique. Malgr le fait que lEFA 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 change par le routage MoE, lcart entre les deux adaptateurs est globalement moins important.

Latence de répartition et de regroupement pour le préremplissage MoE EP64, comparant les chemins de communication de pplx et de DeepEP.

Mme si nos noyaux optimiss pour le dcodage peuvent passer lchelle des tailles de lot de pr-remplissage sans aucun rglage supplmentaire, les performances accusent un retard par rapport DeepEP. Cela est principalement d au fait que nous accumulons les jetons uniquement du ct du rcepteur de la combinaison, alors que DeepEP accumule dans lmetteur pour rduire la quantit de donnes transfres. Cependant, cette optimisation peut avoir des rpercussions sur la prcision. De plus, comme les noyaux transfrent les donnes en une seule criture, nous utilisons plus de mmoire pour les tampons dmission et de rception.

Pour prouver la viabilit de nos noyaux, nous effectuons des tests de performance de bout en bout de notre moteur dinfrence interne desservant DeepSeek-V3 (671 milliards de paramtres) et Kimi-K2 (1 billion de paramtres) sur de multiples instances p5en quipes dadaptateurs Elastic Fabric Adapters (EFA) et de processeurs graphiques H200. Nous excutons le modle en mode de vrification de brouillon et signalons la latence en supposant un taux dacceptation de 80 % en utilisant une seule couche MTP.

Sur DeepSeek-V3, nous comparons par rapport une ligne de base intra-nud EP=8 excutant des noyaux de rpartition et de combinaison efficaces fonds uniquement sur NVLink. Les dploiements inter-nuds galent ou dpassent globalement les performances du dploiement sur un seul nud, en particulier sur les lots de taille moyenne, dmontrant que la mise lchelle entre les nuds apporte des amliorations au dbit global.

Sur Kimi-K2, les noyaux permettent le service. En raison de sa taille, le modle ne peut pas tenir dans le mme nud en utilisant le paralllisme de donnes (DP) sur un nud H200, un dploiement inter-nud est donc ncessaire pour servir le modle. Comme le modle a moins de ttes dattention, nous obtenons de meilleures latences que DeepSeek-V3 et nous pouvons servir le modle sur EFA une latence viable.

Débit par GPU de DeepSeek-V3 par rapport à la vitesse de génération pour les configurations EP32, EP16 et EP8.

Débit par GPU de Kimi-K2 par rapport à la vitesse de génération pour les configurations EP32 et EP16.

Nous avons galement examin le traitement par micro-lots, bien quil nait amlior les latences qu des tailles de lot trs leves ou lors de lutilisation de plus dun jeton de brouillon. Pour bnficier de la division des problmes GEMM groups en deux, les tailles de lot de jetons doivent tre suffisamment denses. Le plus clair du temps est toujours consacr la rpartition et la combinaison MoE, ce qui indique que toute amlioration supplmentaire apporte aux noyaux amliorera par consquent les latences et facilitera le traitement par micro-lots.

Nous travaillons en collaboration troite avec les ingnieurs dAWS pour amliorer les performances sur EFA. De notre ct, nous suivrons les mises jour de libfabric qui rduiront la surcharge du plan de donnes. En outre, nous exprimenterons avec efa-direct pour rduire davantage les surcharges de lespace utilisateur et interagir directement avec le matriel.

Rapport technique de DeepSeek R1

Rapport technique de Kimi K2

DeepEP

UCCL-EP

Amlioration des performances rseau des systmes HPC laide de NVIDIA Magnum IO NVSHMEM et GPUDirect Async

BibTeX

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

Introduction

Infrence sur InfiniBand

Noyaux pplx fonds sur NVSHMEM

Routage global MoE sur EFA

Tampons dmission et de rception

Rpartition et combinaison

TransferEngine pour le routage MoE

valuation

Latence de rpartition/combinaison

Dbit Kimi-K2 / DeepSeek-V3

Travaux futurs

Rfrences

Citation

Infrence sur InfiniBand

Noyaux pplx fonds sur NVSHMEM

Tampons dmission et de rception

Rpartition et combinaison

valuation

Latence de rpartition/combinaison

Dbit Kimi-K2 / DeepSeek-V3

Rfrences