Communication efficace et portable pour les modèles à mélange d'experts

Nous présentons une bibliothèque open source, portable et haute performance pour la communication des modèles à mélange d'experts (MoE) qui offre des performances 10 fois plus rapides que les primitives de communication All-to-All standard.

Notre implémentation présente plusieurs innovations techniques clés qui garantissent une efficacité supérieure de la communication MoE :

Communication initiée par le GPU (IBGDA) : prend en charge la communication directe de GPU à carte réseau (NIC), ce en réduisant considérablement la latence en contournant l'implication du processeur

Chevauchement de la communication et du calcul : une architecture de noyau divisée avec des étapes d'envoi et de réception distinctes permet au calcul de se dérouler simultanément avec les transferts réseau

Performances les plus rapides sur un seul nœud : une latence 2,5 fois inférieure à celle de l'implémentation la plus rapide précédente sur des configurations à un seul nœud

Performances multi-nœuds efficaces et portables : notre implémentation atteint des vitesses jusqu'à 10 fois plus rapides que la communication all-to-all standard. Bien qu'environ 2 fois plus lente que les implémentations hautement spécialisées, notre approche offre une meilleure portabilité entre les versions de NVSHMEM et les environnements réseau (NVLink, CX-7 et EFA).

La bibliothèque est entièrement open source et disponible sur https://github.com/ppl-ai/pplx-kernels.

Dans cet article, nous explorons les défis du parallélisme des experts dans les grands modèles MoE et décrivons notre approche pour une distribution et une combinaison efficaces des jetons dans des environnements GPU distribués.

Introduction

Les modèles à mélange d'experts (MoE), tels que DeepSeek R1 et Mixtral 8x7B, améliorent les modèles denses en limitant le nombre de poids activés pour chaque token. Au lieu d'obtenir un nombre élevé de paramètres en augmentant le nombre de couches ou la taille des couches linéaires dans le perceptron multicouche (MLP) de chaque couche de décodeur, les modèles MoE remplacent le MLP traditionnel par plusieurs experts et un routeur. Par exemple, sur les 671 milliards de paramètres de DeepSeek R1, seuls 37 milliards sont multipliés par un token donné lors de l'inférence. Cela réduit considérablement la quantité de calculs requis pour décoder un token et améliore la latence par rapport à un modèle dense.

Les modèles MoE présentent des défis supplémentaires pour l'inférence par rapport aux modèles denses. Alors que les experts eux-mêmes sont de petites couches MLP, chaque couche de décodeur comprend un routeur qui décide vers quels experts un token est acheminé, chaque token étant acheminé vers plusieurs experts. Le routeur est généralement une petite couche linéaire produisant une distribution de probabilité. Habituellement, les experts ayant les scores top-K les plus élevés sont sélectionnés et l'activation finale est calculée comme une moyenne pondérée, en additionnant les sorties des experts multipliées par un poids dérivé de la distribution de probabilité du routeur.

Pour minimiser la latence dans les systèmes MoE distribués, le parallélisme peut être exploité sur plusieurs appareils, mais cela introduit des défis de communication. Des modèles tels que Mixtral 8x7B ou Llama-70B tiennent dans 8 appareils sur un seul nœud, bénéficiant d'interconnexions NVLink rapides et à faible latence (jusqu'à 900 Gbps). Cependant, les modèles plus volumineux nécessitent le partitionnement (« sharding ») des experts sur plusieurs nœuds à l'aide d'InfiniBand (culminant à 400 Gbps), ce qui introduit des défis de latence supplémentaires.

Agrandir

Dans cet article, nous explorons le problème du parallélisme des experts, en décrivant notre implémentation de noyaux (« kernels ») GPU qui acheminent efficacement les tokens vers les appareils hébergeant les experts assignés et les récupèrent en calculnant la somme pondérée des sorties des experts. Bien qu'il existe des implémentations fortement optimisées ciblant un matériel réseau spécifique, nous présentons une implémentation hautement portable qui repose sur un ensemble minimal de primitives NVSHMEM tout en offrant des performances 10 fois plus rapides que la communication all-to-all standard. Notre implémentation atteint des performances de pointe sur les configurations à un seul nœud tout en maintenant une excellente portabilité dans divers environnements réseau.

Partitionnement (« Sharding ») et parallélisme

Pour une inférence efficace, les poids d'un modèle doivent être conservés dans la mémoire de l'appareil, tout en laissant un Espace suffisant pour les activations, les caches KV et autres tampons requis pour la passe avant (« forward pass ») à travers le modèle. Les modèles les plus puissants dépassent la capacité des GPU les plus puissants, de sorte que l'inférence doit être répartie sur plusieurs appareils, qui peuvent stocker collectivement les poids du modèle. En fonction des schémas de partitionnement des poids, différents schémas de communication et de calcul doivent être utilisés pour décharger les calculs et synchroniser les appareils.

Agrandir

Le parallélisme des experts (EP), illustré dans la première figure, parallélise uniquement le calcul des experts. Différents experts sont assignés à différents appareils, qui détiennent leurs poids. Après le routage, les tokens sont envoyés à l'appareil correspondant, les résultats étant ensuite rassemblés et accumulés. La complexité du routage dépend du degré de parallélisme dans les autres parties du modèle : la réplication d'autres couches pourrait éliminer complètement le besoin de routage, car chaque « rank » peut sélectionner les tokens à partir d'une table de routage répliquée localement. Cependant, si un seul « rank » exécute le routage, une diffusion (« broadcast ») est nécessaire pour acheminer les tokens, les indices et les poids vers leurs experts respectifs. Enfin, un regroupement global (« all-gather ») ou une diffusion all-to-all synchronise les tokens de sortie avec le « rank » qui continue l'exécution du modèle. Une telle implémentation est relativement simple puisque torch expose déjà les primitives requises, bien que certains avantages puissent être obtenus en fusionnant l'accumulation avec l'opération de regroupement collectant les tokens des experts.

Agrandir

Le parallélisme axé uniquement sur les experts ne s'adapte pas de manière idéale, car les nœuds d'un cluster peuvent être inactifs pendant que le modèle exécute des couches autres que celles des experts, telles que l'attention, la normalisation (« norm ») et l'échantillonnage. Cependant, le calcul de ces couches, principalement l'attention, peut également bénéficier du parallélisme des tenseurs (TP). La plupart des modèles reposent sur l'attention multi-tête, ce qui signifie que les têtes d'attention et leurs projections Q, K et V correspondantes peuvent également être partagées entre les appareils, en répliquant ou en rassemblant les tranches entre les différentes couches. Si l'attention est répartie sur tous les appareils, un « all-gather » peut synchroniser les activations, permettant de répliquer le routage, ce qui nécessite des primitives de synchronisation similaires au cas du parallélisme des experts pour une implémentation efficace. Cependant, il existe des limites au parallélisme à ce niveau, car la réduction du nombre de têtes d'attention en dessous d'un certain seuil entraînera des rendements décroissants.

Agrandir

Pour utiliser au mieux tous les appareils et prendre en charge un très haut degré de parallélisme des experts, allant jusqu'à 128 ou 256 GPU, le parallélisme des données (DP) est requis. Selon ce schéma, les appareils sur lesquels un modèle est divisé sont regroupés pour traiter les demandes simultanément, en calculant l'attention et en maintenant les caches KV partagés au sein de leur groupe local. Plusieurs instances de ces groupes parallèles collaborent à l'évaluation des experts, chacun hébergeant un sous-ensemble différent des poids des experts. En fonction du nombre de têtes d'attention, un groupe peut typiquement s'étendre jusqu'à la taille d'un nœud entier, car la communication intra-nœud est plus rapide. Par exemple, dans la figure ci-dessus, un « rank » DP dessert indépendamment deux requêtes, A et B, gérant l'attention, la normalisation et toute autre comptabilité pour les requêtes. L'autre « rank » DP traite une requête distincte C. Cependant, le premier nœud héberge la moitié des experts et l'autre nœud l'autre moitié ; par conséquent, après le routage, les tokens de A et B peuvent être envoyés au deuxième nœud et vice-versa. Cela conduit à un problème de communication creuse (« sparse communication ») : chaque appareil peut envoyer un nombre différent de tokens vers n'importe quelle autre destination. Les primitives existantes de torch, principalement all_to_all, ne sont pas particulièrement bien adaptées, car elles peuvent nécessiter une forme de remplissage (« padding ») ou une synchronisation GPU-à-UC (CPU) et une diffusion de métadonnées. Pour implémenter la communication efficacement, des noyaux personnalisés sont requis pour traiter les informations de routage et initier la communication à partir des appareils sans intervention de l'UC. Après le routage, un noyau de distribution (« dispatch kernel ») doit envoyer les tokens aux « ranks » vers lesquels ils ont été acheminés, tandis que du côté de la combinaison, les activations appartenant aux requêtes du groupe DP actuel doivent être collectées. De plus, le travail doit être équilibré au sein d'un groupe DP qui peut comporter plusieurs appareils, afin de minimiser la diffusion et de répliquer correctement les tokens sur tous les « ranks » pour les couches suivantes.

NVSHMEM

NVSHMEM est une implémentation OpenSHMEM spécifique à NVIDIA, fournissant des fonctionnalités de communication inter-appareils portables qui font abstraction de la complexité du matériel sous-jacent. L'API peut exprimer des lectures et des écritures de l'appareil à l'appareil, qui sont mappées aux primitives des couches de transport individuelles. NVLink et RDMA sont tous deux pris en charge, accordant un degré de portabilité et permettant aux noyaux NVSHMEM de fonctionner au sein d'un nœud ou à travers un cluster de plusieurs nœuds. Dans nos noyaux, nous utilisons des écritures en mémoire distante et des opérations de signalisation afin de transférer des données et de synchroniser les appareils participants.

Les opérations NVSHMEM sont construites autour du concept de mémoire symétrique : elles opèrent sur des tampons (« buffers ») qui ont été alloués sur tous les appareils participant à la communication inter-appareils. Un appareil local peut utiliser une adresse dérivée de son propre tampon local ainsi que l'index du « rank » cible pour spécifier la destination d'une opération. La figure ci-dessous illustre ce concept : les deux GPU allouent des tampons symétriques de la même taille, conservant des pointeurs src vers ceux-ci. Le premier « rank » souhaite envoyer 3 entiers au second, en les plaçant au début du tampon : nvshmem_int_put_nbi dérive l'adresse de départ du tampon local, en spécifiant l'appareil cible. Le second « rank » dérive un décalage (« offset ») de son propre tampon, envoyant un élément au premier appareil avec un décalage d'un. Bien que les adresses de destination doivent toujours être des tampons symétriques alloués à l'aide de nvshmem_alloc, les tampons source peuvent être des régions arbitraires de la mémoire de l'appareil, à condition qu'ils soient pré-enregistrés auprès de NVSHMEM.

Bien que NVSHMEM offre un large éventail de primitives, nos noyaux reposent sur seulement 3 fonctions, construisant toute la synchronisation et le cloisonnement (« fencing ») sur celles-ci.

Agrandir

nvshmemx_putmem_signal_nbi_warp : transfère un bloc de données d'un appareil à un autre, tout en définissant un indicateur (« flag ») sur l'appareil distant. L'opération définit (NVSHMEM_SIGNAL_SET) ou incrémente (NVSHMEM_SIGNAL_ADD) un emplacement de 64 bits. L'indicateur est mis à jour après le transfert de l'ensemble du bloc de mémoire. Si l'appareil distant observe un changement dans l'indicateur, il peut accéder en toute sécurité au tampon dans sa propre mémoire. Cette fonction est utile pour fusionner le transfert de données et la synchronisation.

nvshmemx_signal_op opère sur un seul emplacement mémoire, généralement un indicateur de 64 bits, en le définissant ou en l'incrémentant de manière atomique. Il est utile pour transmettre des métadonnées et synchroniser les appareils.

nvshmem_uint64_wait_until est utilisé du côté de la réception d'un signal, pour scruter (« poll ») un indicateur jusqu'à ce que l'appareil distant le mette à jour.

Garantir toute forme d'ordre entre les opérations par le biais de nvshmem_quiet ou nvshmem_fence est coûteux, c'est pourquoi notre implémentation des noyaux évite l'utilisation de toute autre primitive NVSHMEM et minimise les dépendances autres que les paires d'envoi-réception individuelles. Lors de l'envoi de données, nous utilisons toujours la version non bloquante des fonctions, sans attendre que les données soient même envoyées hors du « rank » local. D'autre part, une synchronisation implicite à travers le groupe de noyaux garantit que ces tampons ne sont pas écrasés de manière destructive pendant que les cartes réseau ont besoin d'y accéder.

NVSHMEM utilise GPUDirect RDMA afin que la carte réseau puisse lire et écrire directement dans la mémoire du GPU pour l'échange de données sans impliquer les UC. De plus, sur les cartes réseau ConnectX, NVSHMEM prend en charge la communication initiée par le GPU (également connue sous le nom d'Infiniband GPUDirect Async, ou IBGDA), qui permet au GPU de soumettre des opérations RDMA directement à la carte réseau et de scruter directement la file d'attente de complétion de la carte réseau. Sur les plateformes qui ne prennent pas en charge la communication initiée par le GPU, NVSHMEM génère un fil d'exécution (« thread ») « proxy » sur l'UC hôte pour initier la communication et scruter la complétion pour le compte du GPU. Le programme NVSHMEM est portable, que la communication initiée par le GPU soit prise en charge ou non. Cependant, la communication initiée par le GPU réduit considérablement la latence car elle contourne complètement le détour par l'UC.

Noyaux portables

Nous implémentons la communication MoE par le biais d'une paire de noyaux de distribution et de combinaison. Les noyaux de distribution sont responsables de la lecture des tokens et des informations de routage sur chaque « rank », en les distribuant aux experts appropriés. Les noyaux de combinaison collectent les activations produites par les experts et les renvoient à leurs « ranks » d'origine, tout en calculant la moyenne pondérée à partir des experts sélectionnés sur la base des poids calculés par le routeur. Les noyaux sont en outre divisés en un composant d'envoi et un composant de réception, afin de permettre le chevauchement des transferts de données avec le calcul. Les noyaux d'envoi sont tous non bloquants et non synchronisants : ils distribuent simplement toutes les écritures vers les destinataires distants. De l'autre côté, les noyaux de combinaison se contentent de lire en mémoire, en attendant que toutes les données requises aient été transférées. Après avoir distribué le travail aux cartes réseau, pendant que les données sont transférées de manière asynchrone sur le réseau, les GPU peuvent effectuer un autre travail utile localement, comme l'application d'experts partagés ou le calcul de l'attention.

Chaque paire de noyaux possède ses propres tampons de mémoire symétrique, chaque « rank » allouant un stockage privé pour recevoir des données de tous les autres « ranks ». Le nombre maximal de tokens pouvant être distribués à partir d'un groupe DP est fixe pour tous les « ranks », ce qui définit également la limite supérieure que chaque « rank » peut recevoir pour chaque expert local de chaque « rank » DP. Cela permet aux « ranks » émetteurs de dériver une adresse unique sur l'appareil de destination où écrire, sans nécessiter de synchronisation entre eux. Une fois les données reçues, les noyaux réorganisent les tokens pour les disposer en mémoire dans le format optimal requis par les noyaux récepteurs. Bien que les tampons aient une dimension considérable, ils sont réutilisés à travers toutes les couches séquentielles d'un modèle.

Agrandir

La seule forme de synchronisation globale, comme illustré dans la figure ci-dessus, est implémentée dans le noyau de réception et de combinaison (« combine-receive »). Une fois les données distribuées à partir de l'envoi, un « rank » revient et poursuit son exécution lorsqu'il reçoit tous les tokens dont il a besoin des autres « ranks » et entre dans les noyaux d'envoi et de combinaison dès qu'ils sont traités. La barrière dans la réception et la combinaison garantit qu'aucun « rank » ne peut prendre d'avance et démarrer l'envoi et la distribution (« dispatch-send ») pendant qu'un autre « rank » attend toujours de recevoir des données, ce qui pourrait provoquer un chevauchement destructif. La synchronisation est effectuée dans le noyau de combinaison par souci de simplicité ; comme optimisation, la même barrière pourrait être implémentée entre les noyaux de réception et de distribution et les noyaux d'envoi et de combinaison, réduisant ainsi davantage la latence en donnant aux cartes réseau plus de temps pour stabiliser les transferts pendant que le GPU exécute d'autres noyaux.

Les deux noyaux sont divisés entre tous les multiprocesseurs streaming (SM) disponibles des appareils : tandis que les noyaux d'envoi et de distribution et de réception et de combinaison doivent se paralléliser sur un nombre maximal de tokens par « rank » (max_m), les noyaux de réception et de distribution et d'envoi et de combinaison doivent traiter ce même nombre de tokens provenant de chaque « rank » et de chaque expert local, pour un maximum de max_m * (num_experts // EP) * (EP // TP).

Distribution (« Dispatch »)

Le noyau de distribution est responsable de la distribution des tokens et des facteurs de mise à l'échelle par blocs aux pairs qui hébergent les experts vers lesquels les tokens ont été acheminés. Le côté émetteur s'appuie sur la spécialisation des structures warp (« warp specialization ») pour paralléliser deux tâches : agréger les informations de routage pour déterminer combien de tokens sont envoyés à partir de chaque « rank » et regrouper les tokens dans des messages à distribuer à l'appareil distant. Le côté récepteur attend d'abord que tous les décomptes de tokens soient reçus, puis il lit et déballe les charges utiles (« payloads ») de la mémoire partagée dans les tenseurs qui seront transmis aux noyaux GEMM groupés en aval. Il stocke également des informations dans des tampons partagés avec les noyaux de combinaison, pour indiquer où chaque token doit être renvoyé. Ces informations sont nécessaires car le noyau de réception réorganise les tokens dans des tampons contigus selon un ordre non déterministe.

Agrandir

Dans la partie émettrice, un seul warp de 32 fils d'exécution (« threads ») sur chaque bloc est responsable de la lecture des informations de routage et du comptage du nombre de tokens que chaque expert de destination doit recevoir de ce « rank », en envoyant le décompte plus un à l'aide de nvshmem_signal_op. Le décompte est incrémenté de un, car la transition de zéro à une valeur non nulle du côté distant signifie la réception des décomptes d'un « rank » source. En parallèle, les warps restants coopèrent pour copier les tokens dans la mémoire symétrique à travers tous les blocs en parallèle, en emballant les activations, les facteurs de mise à l'échelle et leur index sur le « rank » local dans un bloc contigu. L'index est requis par l'émetteur de combinaison pour déterminer l'adresse où le token sera écrit. Ensuite, après s'être assuré que toutes les données ont été copiées par le biais d'une barrière à travers les groupes de warps, les warps opèrent à nouveau indépendamment, chacun envoyant le même tampon à un expert différent en parallèle. Les tokens sont envoyés à l'aide de nvshmemx_putmem_signal_nbi_warp, qui augmente également de manière atomique le décompte des tokens envoyés du « rank » local sur l'appareil distant. Au sein d'un groupe DP, étant donné que chaque « rank » possède une réplique du token à envoyer, la distribution est équilibrée uniformément, chaque appareil envoyant un sous-ensemble des tokens.

Du côté de la réception, tous les noyaux se synchronisent d'abord avec les « ranks » émetteurs en attendant que les décomptes totaux de tokens soient transmis. Ensuite, ils attendent tous que les décomptes de tokens envoyés incrémentés de manière atomique se stabilisent aux décomptes totaux, indiquant que toutes les charges utiles des « ranks » source ont également été transmises, grâce à la sémantique de l'appel putmem. Les noyaux scrutent les décomptes à l'aide de nvshmem_uint64_wait_until, en parallélisant l'opération sur tous les blocs et fils d'exécution (« threads »). Par la suite, une barrière inter-blocs garantit qu'aucun bloc ne lit les tampons à moins que toutes les données n'aient été correctement reçues. Répartis sur les blocs et synchronisés via un compteur incrémenté de manière atomique, les tokens sont copiés de la mémoire symétrique vers des tampons contigus réguliers qui deviennent les entrées des noyaux en aval implémentant les experts. Le « rank » source, l'index de l'expert et l'index du token sont stockés séparément, indiquant précisément l'emplacement où le noyau de combinaison doit envoyer les activations. Même si les tokens d'un groupe DP sont envoyés à partir d'appareils différents, ils sont tous regroupés pour être transmis à l'expert correspondant.

Combinaison (« Combine »)

Les noyaux de combinaison sont à nouveau divisés en moitiés d'envoi et de réception : les émetteurs copient les activations non quantifiées de 16 bits vers les appareils distants, tandis que les noyaux de réception attendent que les données soient transmises et calculent la moyenne pondérée de toutes les contributions des experts localement. De plus, ils agissent également comme une barrière pour synchroniser la séquence distribution-combinaison : chaque « rank » définit un indicateur sur chaque pair à l'entrée des noyaux d'envoi, les noyaux de réception n'étant pas autorisés à revenir à moins d'observer que l'indicateur est défini. La latence de la synchronisation est minimale, car elle se chevauche avec la communication et le calcul réels.

Agrandir

Du côté de l'émetteur, les noyaux parcourent la liste des tokens assignés à tous les experts locaux en parallèle, les écrivant dans un tampon sur l'appareil de destination. L'appareil cible, l'index de l'expert et l'index du token sont lus à partir des tampons par token peuplés par les noyaux de diffusion lors de la réception des tokens. Chaque émetteur possède sa propre région de mémoire privée par expert où écrire, comme l'indique la figure ci-dessus, évitant ainsi le besoin de synchronisation. De manière similaire à la distribution, la combinaison incrémente de manière atomique les compteurs par token sur l'appareil de destination pour indiquer la réception des données : lorsque le compteur correspond au nombre d'experts auxquels un token a été distribué, le contenu du token peut être consulté.

Dans le noyau de réception, la liste des tokens est parcourue en parallèle sur plusieurs blocs, en attendant que leur contenu arrive par la scrutation de l'indicateur défini par l'opération de signalisation. Dès l'arrivée, les charges utiles sont lues à partir des tampons privés, la table de routage indiquant quel tampon lire et quel poids assigner à chaque expert. Les résultats sont ensuite écrits dans des tenseurs alloués à l'interne, le noyau terminant son exécution une fois que tous les appareils ont passé la barrière. À travers un groupe DP, tous les « ranks » reçoivent une copie de chaque activation d'expert pour calculer leurs propres répliques.

Benchmarks

Nous évaluons nos noyaux sur un cluster de 128 GPU répartis sur 16 nœuds connectés via InfiniBand et des cartes réseau CX-7. Nous essayons à la fois GPUDirect Async (IBGDA) et Reliable Connection (RC) avec un proxy UC. Nous les comparons à l'implémentation hautement optimisée de DeepSeek, ainsi qu'aux primitives denses fournies par PyTorch (via NCCL) ou NVSHMEM.

Agrandir

Bien que sur l'opération dense all-to-all de NVSHMEM, les performances de IBRC et IBGDA soient similaires (6 378 µs vs 6 180 µs), IBGDA est nettement plus rapide avec les noyaux creux (« sparse kernels »). En additionnant la distribution et la combinaison, IBGDA utilise 902 µs tandis que IBRC prend 3 223 µs, soit une amélioration de la latence de 3,6 fois. Alors que all-to-all est limité par la bande passante, les noyaux creux diffusent des ordres de grandeur de données en moins, étant limités par la latence. En déclenchant les transferts réseau directement à partir d'un GPU, sans nécessiter de proxy UC pour coordonner le GPU et la carte réseau, la latence de bout en bout est considérablement réduite.

Bien que nos noyaux portables soient environ 2 fois plus lents que les noyaux DeepSeek hautement optimisés, ils améliorent la latence de 10 fois par rapport aux noyaux denses (902 µs vs 9 944 µs).

De plus, la division en composants émetteur-récepteur permet également de masquer une partie de la latence, contrairement aux primitives de bibliothèque.

Agrandir

Sur un seul nœud (EP8), NVSHMEM utilise NVLINK pour le transport, offrant une latence plus faible et un débit plus élevé que le réseautage inter-nœuds. Nos noyaux portables sont environ 2,5 fois plus rapides que DeepEP sur un seul nœud (186 µs vs 481 µs).

Améliorations supplémentaires

Les noyaux décrits ici surpassent les primitives intégrées des frameworks d'apprentissage automatique (ML) et offrent des performances correctes sans surspécialisation pour des transports inter-appareils particuliers, tels que NVLink ou InfiniBand. Outre les opportunités déjà mentionnées, des gains de performances supplémentaires peuvent être obtenus en remplaçant les primitives de communication par des versions plus spécialisées. Par exemple, à travers NVLink, l'utilisation de la mémoire symétrique pourrait être remplacée par des tampons partagés entre les appareils, éliminant une partie des copies et permettant des schémas de synchronisation plus fins, à l'échelle des tokens individuels au lieu des lots de tokens. À travers InfiniBand, une implémentation pourrait accéder directement aux paires de files d'attente sous-jacentes, éliminant le besoin de transférer à la fois les décomptes de tokens attendus et envoyés et de vider les tokens des files d'attente dans les tampons de sortie dès leur arrivée. Cependant, de telles implémentations se font au détriment de la portabilité, tandis que nos implémentations permettent d'évaluer des clusters dotés de configurations spécifiques avant de s'engager dans un effort de développement pour optimiser pour un matériel particulier.

Conclusion

Nous avons présenté une bibliothèque portable haute performance pour la communication MoE qui offre des performances 10 fois plus rapides que la communication all-to-all standard tout en maintenant la compatibilité entre diverses configurations matérielles. Sur les déploiements à nœud unique avec NVLink, notre solution démontre une latence 2,5 fois inférieure à celle des implémentations précédentes.

Notre approche équilibre les performances et la portabilité grâce à des innovations clés comprenant la prise en charge de la communication initiée par le GPU, une architecture de noyau divisée permettant le chevauchement calcul-communication, et une distribution efficace des tokens utilisant des primitives NVSHMEM minimales. Bien qu'environ 2 fois plus lente que les implémentations hautement spécialisées sur les configurations multi-nœuds, notre bibliothèque offre une flexibilité supérieure dans divers environnements réseau (NVLink, CX-7 et EFA).

À mesure que les modèles MoE continuent d'évoluer, des stratégies de communication efficaces comme la nôtre deviendront de plus en plus importantes pour un déploiement pratique. Notre implémentation entièrement open source est disponible sur https://github.com/ppl-ai/pplx-kernels.

Références

DeepEP

Rapport technique de DeepSeek-V3

NVIDIA NVSHMEM

Citation

@online{perplexity2025EfficientAndPortableMixtureOfExpertsCommunication,
  author = {{Perplexity} and Nandor Licker and Lequn Chen and Phi Hoang},
  title = {{Efficient and portable mixture-of-experts communication}},
  date = {2025-04-02},
  url = {https://www.perplexity.ai/hub/blog/efficient-and-portable-mixture-of-experts-communication}
}