Résumé Technique : DisagMoE

1. Énoncé du Problème

Les architectures de type Mélange d'Experts (MoE) à grande échelle permettent l'entraînement de modèles de langage de grande taille (LLM) comportant des billions de paramètres en activant de manière éparsée un sous-ensemble d'experts. Cependant, l'entraînement de ces modèles se heurte à des goulots d'étranglement sévères dus à la communication tout-à-tout requise pour le parallélisme des experts (EP). À mesure que les modèles s'étendent sur plusieurs nœuds GPU, la bande passante réseau inter-nœuds limitée (par exemple, InfiniBand) exacerbe la latence de communication, qui peut consommer jusqu'à 50 % du temps d'une étape d'entraînement.

Les stratégies d'optimisation existantes se concentrent principalement sur le chevauchement de la communication avec le calcul au niveau de l'opérateur (par exemple, le fractionnement des calculs FFN) ou sur l'utilisation de pipelines de micro-lots. Cependant, ces approches souffrent de deux limitations fondamentales :

1. Ralentissements Résiduels : Les déséquilibres inhérents aux ratios calcul-communication entre les couches d'attention et les couches FFN empêchent un chevauchement parfait. Le calcul d'attention évolue de manière quadratique avec la longueur de séquence, tandis que le FFN et la communication évoluent de manière linéaire. Cela crée des « queues de communication » qui ne peuvent être masquées, en particulier lorsque les longueurs de séquence augmentent.
2. Inefficacité des Ressources : Les approches précédentes de désagrégation (séparant l'attention et le FFN sur des matériels différents) nécessitent souvent que chaque groupe détienne l'intégralité des paramètres du modèle, entraînant des erreurs de mémoire insuffisante (OOM) à grande échelle, ou échouent à adapter l'allocation des ressources aux caractéristiques spécifiques de calcul-communication de chaque composant.

2. Méthodologie : DisagMoE

DisagMoE est un système d'entraînement conçu pour atténuer les goulots d'étranglement tout-à-tout en optimisant conjointement le placement du modèle et la planification grâce à une architecture désagrégée.

2.1 Placement des Composants Désagrégés

Contrairement au parallélisme d'experts traditionnel où les couches d'attention et FFN coexistent sur le même groupe de GPU, DisagMoE partitionne le modèle par type de composant :

• Groupes d'Attention (A-Workers) : Ces groupes contiennent les couches d'auto-attention de plusieurs couches de transformateurs. Ils répliquent les composants denses et distribuent les données d'entraînement via le Parallélisme de Données (DP).
• Groupes FFN (F-Workers) : Ces groupes contiennent les couches FFN (experts) de plusieurs couches de transformateurs. Les experts sont fragmentés (shardés) sur les GPU au sein du groupe via le Parallélisme d'Experts (EP).
• Attribution des Couches : Les couches sont attribuées aux groupes de manière entrelacée (par exemple, le Groupe 0 gère les couches 0, $p$, $2p$ ; le Groupe 1 gère les couches 1, $p+1$, etc.), maximisant l'utilisation des appareils sous les contraintes de mémoire.

2.2 AF-Pipe : Pipeline Multi-Étape avec Chevauchement

DisagMoE introduit AF-Pipe, un plan de pipeline qui traite la communication comme une étape de premier ordre plutôt que comme un effet secondaire.

• Communication Many-to-Many (M2N) : Au lieu des échanges tout-à-tout standard au sein d'un seul groupe de nœuds, AF-Pipe établit des échanges de jetons unidirectionnels et many-to-many entre les groupes d'Attention et FFN.
• Alignement des Étapes : Le pipeline aligne les limites entre les groupes, fusionnant les opérations traditionnelles de « dispatch » (distribution) et de « combine » (combinaison) en une seule étape de communication unifiée. Cela réduit la surcharge de communication d'environ $1/k$ (où $k$ est le nombre d'experts top-$k$) et élimine les transferts de données redondants.
• Exécution Asynchrone : Le système utilise trois flux coordonnés (avant, arrière et communication) au sein de chaque worker. Cela permet au groupe FFN d'envoyer des gradients (N2M) tout en calculant simultanément le passage avant sur les données disponibles, assurant une exécution non bloquante et une utilisation quasi continue du GPU.

2.3 Allocation Adaptative des Workers

Pour remédier au déséquilibre entre l'attention (limitée par le calcul) et le FFN (limité par la communication), DisagMoE emploie un allocateur adaptatif guidé par un Modèle Roofline Calcul-Communication.

• Modélisation : Le système modélise le point de basculement ($\hat{I}$) où un composant passe d'une limitation par le réseau à une limitation par le calcul. Les couches d'atteignent plus rapidement le plafond de calcul que les couches FFN à mesure que la longueur de séquence augmente.
• Optimisation : Le système formule le partitionnement des GPU et des cartes d'interface réseau (NIC) comme un problème de Programmation Linéaire en Nombres Entiers Mixtes (MILP). Il alloue dynamiquement plus de bande passante réseau aux groupes FFN (pour soulever leur plafond de communication) et plus de ressources de calcul aux groupes d'Attention, équilibrant les étapes du pipeline pour minimiser les bulles.
• Stratégie de Recherche : Un algorithme utilise le modèle roofline pour générer une configuration de départ, qui est ensuite affinée par une recherche locale utilisant les temps mesurés par étape pour tenir compte des surcharges non modélisées.

3. Contributions Clés

1. Placement de Composants Désagrégés : Une architecture novatrice qui sépare les couches d'attention et FFN en groupes de workers disjoints, permettant une mise à l'échelle et une allocation de ressources indépendantes.
2. Planification AF-Pipe : Un pipeline multi-étapes qui remplace le tout-à-tout conventionnel par des échanges M2N inter-groupes, chevauchant systématiquement la communication avec les calculs d'attention et FFN.
3. Allocation Adaptative des Ressources : Un mécanisme guidé par le roofline qui dimensionne conjointement les ressources GPU et NIC pour les groupes d'attention et FFN, équilibrant efficacement le ratio calcul-communication à travers différentes longueurs de séquence et configurations de modèles.

4. Résultats Expérimentaux

Les auteurs ont évalué DisagMoE sur des clusters de 16 nœuds (128× GPU H800) en utilisant trois modèles MoE : DeepSeek-MoE, GPT-OSS-120B et Qwen3, avec des longueurs de séquence allant de 4K à 32K.

• Amélioration du Débit : DisagMoE a constamment surpassé des bases de référence solides, notamment Megatron-LM (entrelacé 1F1B), Tutel, Comet et DualPipe.
  • Jusqu'à 1,81× d'accélération par rapport à Megatron-1F1B.
  • Jusqu'à 1,34× d'accélération par rapport aux systèmes de chevauchement MoE les plus avancés.
• Réduction de la Communication : Par rapport aux bases de référence, DisagMoE a réduit le temps de communication non chevauché de jusqu'à 88 % (vs Tutel), 75 % (vs Comet) et 45 % (vs DualPipe).
• Passage à l'Échelle : Le système a démontré son efficacité sur différentes valeurs de top-$k$ et tailles de Parallélisme d'Experts (EP).
• Sensibilité aux Ressources : Les expériences ont confirmé que les ratios optimaux de GPU entre les groupes d'attention et FFN varient avec la longueur de séquence (par exemple, des séquences plus longues favorisent davantage de ressources d'attention), validant la nécessité d'une allocation adaptative.

5. Signification et Revendications

L'article revendique que DisagMoE adresse une limitation fondamentale de l'entraînement MoE actuel : l'incapacité de chevaucher parfaitement la communication en raison du déséquilibre structurel entre les couches d'attention et FFN. En désagrégeant ces composants et en allouant adaptativement les ressources basées sur un modèle roofline, le système étend la fenêtre de chevauchement effective au-delà de ce que peut offrir une provision de bande passante symétrique.

Les auteurs soulignent que leur approche est particulièrement significative pour les charges de travail de pré-entraînement à grande échelle avec des longueurs de séquence fixes. Ils notent que, bien que leur conception actuelle suppose des formes stables (longueur de séquence et micro-lot fixes), l'idée centrale — que le matériel et les logiciels futurs devraient cibler des ratios calcul-communication favorables entre les composants plutôt que de simplement augmenter la bande passante brute — offre une nouvelle direction pour l'entraînement efficace des modèles MoE. Le système est implémenté sur Megatron-LM et se révèle être une solution viable pour l'entraînement de modèles comportant des centaines de milliards de paramètres, sans les contraintes de mémoire qui ont entravé les tentatives précédentes d'entraînement désagrégé.