Scalable Training of Mixture-of-Experts Models with Megatron Core

Ce rapport présente les optimisations intégrées de Megatron Core pour le entraînement évolutif des modèles à mélange d'experts (MoE), couvrant la mémoire, la communication et le calcul afin d'atteindre des performances élevées sur des clusters massifs tout en offrant une solution open-source prête pour la production.

L'essentiel

🏗️ Le Grand Projet : Construire une Ville de Génies (Les Modèles MoE)

Imaginez que vous essayez de construire la plus grande bibliothèque de connaissances du monde. Traditionnellement, pour rendre cette bibliothèque plus intelligente, vous deviez ajouter des livres à chaque étagère. C'est ce qu'on appelle les modèles "denses" : tout le monde travaille tout le temps. Mais cela devient trop cher et trop lourd.

NVIDIA a une idée géniale : Le Modèle "Mélange d'Experts" (MoE).
Au lieu d'avoir un seul bibliothécaire géant qui lit tout, vous engagez des milliers de petits experts spécialisés (un expert en cuisine, un en astronomie, un en histoire). Quand un visiteur pose une question, un guichetier intelligent (le "routeur") regarde la question et envoie le visiteur uniquement vers les 2 ou 3 experts qui savent vraiment répondre.

Le problème ? C'est comme si vous aviez une ville de 1 million d'habitants (les paramètres), mais que seulement 100 personnes travaillaient à la fois. Cela crée trois gros problèmes de logistique que ce rapport explique comment résoudre.

---

🚧 Les Trois Murs qui bloquent la construction

Pour faire fonctionner cette ville géante sur des milliers de super-ordinateurs (des GPU), il faut franchir trois obstacles majeurs :

1. Le Mur de la Mémoire (La Boîte à Outils trop pleine)

• Le problème : Même si seuls 100 experts travaillent, vous devez avoir les outils de tous les 1 million d'experts dans votre camion de livraison (la mémoire de la carte graphique). Sinon, vous ne pouvez pas démarrer. C'est comme essayer de transporter une usine entière dans une petite voiture.
• La solution de Megatron-Core :
  • Le "Tetris" intelligent : Ils réorganisent les outils pour qu'ils prennent moins de place sans casser les choses (compression des données).
  • La "Recette" (Recomputation) : Au lieu de stocker tous les outils, ils les jettent et les recréent à la volée quand on en a besoin. C'est comme cuisiner : au lieu de garder tous les ingrédients déjà coupés sur la table (ce qui prend de la place), vous les coupez juste avant de les mettre dans la casserole.
  • Le "Déménagement" (Offloading) : Quand le camion est plein, on met les outils moins utilisés dans un camion de suivi (la mémoire du processeur) et on les ramène seulement quand on en a besoin.

2. Le Mur de la Communication (Les embouteillages sur l'autoroute)

• Le problème : Comme les experts sont répartis sur des milliers de camions (GPU), le guichetier doit constamment envoyer des visiteurs d'un camion à l'autre. Si le guichetier envoie trop de monde en même temps, les routes se bouchent. C'est le "All-to-All" : tout le monde parle à tout le monde.
• La solution de Megatron-Core :
  • Les "Express-Postes" (DeepEP/HybridEP) : Au lieu d'envoyer des lettres une par une, ils construisent des autoroutes ultra-rapides et des camions de livraison optimisés pour transporter des paquets entiers sans s'arrêter.
  • Le "Multitâche" (Recouvrement) : Pendant que les camions roulent pour livrer les visiteurs, les experts sur place commencent déjà à travailler sur le chargement précédent. On ne laisse jamais les experts attendre en regardant par la fenêtre.

3. Le Mur de l'Efficacité (Les ouvriers qui s'ennuient)

• Le problème : Parce que les experts sont si petits et spécialisés, les super-ordinateurs passent leur temps à attendre que le chef (le processeur central) leur dise quoi faire. C'est comme avoir une équipe de 1000 maçons, mais le chef ne leur donne qu'une brique à la fois. Les maçons attendent, et le chantier avance lentement.
• La solution de Megatron-Core :
  • Les "Blocs de travail" (Grouped GEMM) : Au lieu de donner une tâche à un seul maçon, on donne un gros bloc de travail à une équipe entière en même temps.
  • Le "Plan Pré-enregistré" (CUDA Graphs) : Au lieu que le chef crie "Fais ça !", "Fais ça !" à chaque seconde, il enregistre un film de tout le travail à faire et le lance d'un seul coup. Les ouvriers travaillent sans interruption.
  • L'Équilibre (ECHO) : Si un expert est surchargé, le système copie son travail sur un autre expert libre pour que personne n'attende.

---

🧠 L'Innovation Magique : Le "Pliage Parallèle" (Parallel Folding)

C'est la pièce maîtresse du rapport.
Imaginez que vous avez deux types de travaux dans votre ville :

1. La circulation (Attention) : Besoin de routes larges et rapides.
2. Les ateliers (Experts) : Besoin de petits ateliers spécialisés.

Avant, on était obligé de construire la ville avec le même plan pour les routes et les ateliers. Résultat : les routes étaient trop étroites pour la circulation, ou les ateliers étaient trop grands et vides.

Megatron-Core invente le "Pliage Parallèle" :
C'est comme si on pouvait plier la carte de la ville différemment selon le quartier !

• Pour la circulation, on utilise un plan avec de grandes avenues.
• Pour les ateliers, on plie la carte pour avoir des ruelles spécialisées.
  Cela permet d'utiliser chaque route et chaque atelier exactement comme il faut, sans gaspillage. C'est ce qui permet de faire tenir des modèles de 1 000 milliards de paramètres sur des milliers de cartes graphiques.

---

🚀 Les Résultats : La Ville de Demain

Grâce à toutes ces astuces, NVIDIA a pu entraîner des modèles comme DeepSeek-V3 et Qwen3 à une vitesse incroyable.

• Sur les nouvelles puces GB200/GB300 (les camions de livraison de dernière génération), ils atteignent des vitesses de calcul qui étaient inimaginables il y a un an.
• Ils peuvent même gérer des contextes très longs (lire un livre entier d'un coup) sans que le camion ne se renverse.

En résumé

Ce rapport ne parle pas juste de maths compliquées. Il explique comment NVIDIA a réorganisé toute la logistique pour construire des intelligences artificielles géantes.

• Ils ont appris à gérer l'espace (mémoire) comme un expert du Tetris.
• Ils ont lissé le trafic (communication) pour éviter les embouteillages.
• Ils ont gardé les ouvriers occupés (calcul) pour qu'ils ne s'ennuient jamais.

Le résultat ? Des modèles plus intelligents, entraînés plus vite, et à un coût plus bas. C'est la clé pour passer de l'IA de laboratoire à l'IA qui change le monde.

Résumé technique

1. Problématique : Les Défis du Passage à l'Échelle des MoE

Le rapport identifie que l'entraînement des modèles Mixture-of-Experts (MoE) à très grande échelle (milliards à billions de paramètres) introduit des défis systèmes fondamentaux absents dans les modèles denses traditionnels. Ces défis découlent de la sparsité inhérente aux MoE, où seuls un sous-ensemble d'experts est activé par token, créant un décalage entre le nombre total de paramètres et le calcul effectif par token.

Ce décalage engendre trois "Murs" (Three Walls) interconnectés qui limitent l'efficacité de l'entraînement :

1. Le Mur de la Mémoire (Memory Wall) :
   • Tous les paramètres des experts (même ceux non activés) doivent résider en mémoire GPU, ainsi que leurs états d'optimiseur et leurs gradients.
   • Pour des modèles comme DeepSeek-V3 (685B paramètres totaux, 37B actifs par token), la pression mémoire dépasse largement la capacité des GPU (ex: >199 Go par GPU sur H100 sans optimisation), rendant l'entraînement impossible sans techniques avancées.
2. Le Mur de la Communication (Communication Wall) :
   • Le parallélisme d'experts (Expert Parallelism - EP) nécessite des communications All-to-All pour router les tokens vers les GPU appropriés.
   • Ce volume de communication est proportionnel au nombre de tokens et à la dimension cachée, et augmente avec le nombre d'experts. Sur des topologies inter-nœuds, la bande passante limitée peut consommer jusqu'à 60 % du temps d'entraînement.
3. Le Mur de l'Efficacité de Calcul (Compute Efficiency Wall) :
   • Les experts fins génèrent de nombreuses petites opérations GEMM (matrice-matrice) qui sous-utilisent les unités de calcul des GPU (Tensor Cores).
   • La surcharge de l'hôte (CPU) due au lancement de nombreux petits noyaux (kernels) et à la synchronisation pour les shapes dynamiques (routage sans perte de tokens) crée des goulots d'étranglement, laissant les GPU en attente.

De plus, il existe un décalage dense-sparse : les couches d'attention (denses) et les couches MoE (sparse) ont des configurations de parallélisme optimales contradictoires (ex: TP élevé pour l'attention vs EP élevé pour les MoE), ce que les frameworks traditionnels ne gèrent pas bien.

---

2. Méthodologie : Megatron-Core MoE

NVIDIA propose Megatron-Core MoE, une pile logicielle open-source intégrée conçue pour résoudre simultanément ces trois murs et le décalage de parallélisme. L'approche repose sur une co-conception système couvrant la mémoire, la communication et le calcul.

A. Parallélisme Multi-Dimensionnel et "Parallel Folding"

Pour résoudre le décalage dense-sparse, Megatron-Core introduit le Parallel Folding :

• Découplage des configurations : Il permet aux couches d'attention et aux couches MoE d'utiliser des schémas de parallélisme indépendants (TP, CP, DP, PP pour l'attention ; EP, ETP, EDP pour les MoE).
• Flexibilité : Cela brise la contrainte traditionnelle où le degré d'EP doit être inférieur ou égal au degré de DP (EP ≤ DP). Le "folding" permet de plier le parallélisme d'experts sur des sous-groupes arbitraires, optimisant l'utilisation de la bande passante NVLink intra-nœud.
• Gestion des groupes : Une gestion fine des groupes de processus permet d'appliquer le FSDP (Fully Sharded Data Parallelism) spécifiquement aux experts, réduisant la mémoire des états d'optimiseur.

B. Optimisations Mémoire (Briser le Mur de la Mémoire)

• Permutation Économe en Mémoire (Memory-Efficient Permutation) : Réarrange les calculs pour éliminer les tenseurs intermédiaires redondants sans coût de calcul supplémentaire.
• Re-calcul Granulaire (Fine-grained Recomputation) : Ne re-calculent que les opérations les plus coûteuses en mémoire mais peu coûteuses en calcul (ex: fonctions d'activation, LayerNorm) plutôt que des couches entières.
• Déchargement (Offloading) : Déplace les activations et les états d'optimiseur vers la mémoire CPU de manière asynchrone, en chevauchant les transferts PCIe avec le calcul GPU.
• Précision Réduite : Utilisation de formats FP8/FP4 pour les activations, réduisant la mémoire de 50 % à 75 %.

C. Optimisations de Communication (Briser le Mur de la Communication)

• Dispatchers Optimisés (DeepEP et HybridEP) : Des noyaux de communication spécialisés qui fusionnent les étapes de permutation et de réduction, exploitant les primitives matérielles (TMA, IBGDA) pour maximiser la bande passante, notamment sur les topologies NVLink (NVL72).
• Chevauchement Calcul-Communication : Mise en œuvre d'un schéma 1F1B (1 Forward, 1 Backward) fusionné avec des opérations All-to-All, permettant de cacher la latence de communication derrière le calcul des experts.
• Réduction de la synchronisation : Techniques pour éviter que le CPU ne bloque le GPU en attendant les shapes dynamiques des tokens.

D. Optimisations de Calcul (Briser le Mur de l'Efficacité)

• GEMM Groupés (Grouped GEMM) : Regroupe les calculs de multiples experts en un seul noyau pour améliorer l'utilisation des Tensor Cores.
• Fusion de Noyaux : Fusionne le routage, la permutation et les pertes auxiliaires en moins de noyaux pour réduire la surcharge du CPU.
• CUDA Graphs : Capture les séquences de noyaux pour éliminer la surcharge de lancement par itération. Pour les MoE sans perte de tokens (dropless), l'approche utilise des noyaux sans synchronisation (Sync-Free) initiés par le GPU, ECHO (clonage élastique des experts populaires pour équilibrer la charge) et Paged Stashing (gestion de mémoire paginée) pour gérer les shapes dynamiques sans bloquer le CPU.

E. Entraînement en Précision Réduite (FP8/FP4)

Une stratégie de précision sélective est adoptée :

• Le routeur, les embeddings et les états d'optimiseur restent en haute précision (FP32/BF16) pour la stabilité.
• Les calculs massifs (GEMM des experts) et les communications sont effectués en FP8 (Blockwise ou MXFP8) ou FP4 (NVFP4 sur Blackwell), offrant des gains de performance et de mémoire sans dégrader la convergence.

F. Support des Contextes Longs et RL

• Context Parallelism (CP) : Essentiel pour les séquences longues (16K-64K+), où l'attention devient le goulot d'étranglement dominant ($O(S^2)$).
• Packed Sequences & Dynamic-CP : Gestion efficace des séquences de longueurs variables (typique en RL) en évitant le padding et en adaptant dynamiquement le degré de CP par micro-lot.

---

3. Contributions Clés

1. Architecture Parallèle Unifiée : Introduction du Parallel Folding pour découpler les stratégies de parallélisme des couches denses et sparse, permettant des configurations optimisées pour chaque type de couche.
2. Stack d'Optimisation Systémique : Une approche holistique qui traite la mémoire, la communication et le calcul comme un système couplé, où les solutions à un mur (ex: FP8) sont conçues pour ne pas aggraver les autres.
3. Techniques de Bas Niveau Avancées : Développement de noyaux spécifiques (DeepEP, HybridEP, Grouped GEMM, Sync-Free kernels) et de mécanismes de gestion de mémoire (Paged Stashing, ECHO) pour les MoE dynamiques.
4. Support Production et RL : Intégration de fonctionnalités critiques pour le déploiement industriel : checkpointing distribué, upcycling (conversion dense vers MoE), et support natif pour le Reinforcement Learning (Router Replay, Dynamic-CP).

---

4. Résultats de Performance

Les benchmarks ont été réalisés sur des clusters NVIDIA GB300, GB200 et H100 avec des modèles de pointe :

• DeepSeek-V3 (685B paramètres, 256 experts) :
  • GB300 (256 GPU) : 1 233 TFLOPS/GPU (MXFP8).
  • GB200 (256 GPU) : 1 048 TFLOPS/GPU (MXFP8).
  • H100 (1024 GPU) : 368 TFLOPS/GPU (FP8-Blockwise).
• Qwen3-235B (235B paramètres) :
  • GB300 (256 GPU) : 974 TFLOPS/GPU.
  • GB200 (256 GPU) : 919 TFLOPS/GPU.
  • H100 (256 GPU) : 320 TFLOPS/GPU.

Observations :

• Les plateformes Blackwell (GB200/GB300) offrent environ 3x plus de débit (tokens/seconde) que H100 pour des configurations similaires, grâce à la bande passante mémoire supérieure et au support natif MXFP8/FP4.
• L'optimisation complète permet de réduire la mémoire par GPU de ~199 Go à moins de 80 Go pour DeepSeek-V3, rendant l'entraînement possible sur du matériel actuel.
• Le surcoût de communication All-to-All est réduit à moins de 5-10 % du temps d'itération grâce aux optimisations de chevauchement et aux dispatchers avancés.

---

5. Signification et Impact

Ce rapport marque une étape majeure dans l'entraînement des modèles de langage à l'échelle du trillion de paramètres :

• Démocratisation des MoE à grande échelle : En fournissant une solution open-source (Megatron-Core) prête pour la production, NVIDIA permet à la communauté académique et industrielle d'entraîner des modèles MoE complexes sans développer ses propres infrastructures systèmes.
• Preuve de concept Hardware-Software Co-design : Les résultats démontrent que l'exploitation maximale du matériel de nouvelle génération (Blackwell) nécessite des optimisations logicielles profondes (CUDA Graphs, Sync-Free, FP4) qui vont au-delà des simples ajustements de hyperparamètres.
• Adaptabilité aux nouveaux paradigmes : La capacité à gérer efficacement les contextes longs et les charges de travail de Reinforcement Learning (RL) positionne Megatron-Core comme un moteur essentiel pour la prochaine génération de modèles de raisonnement (type o1, R1).
• Efficacité Économique : En réduisant le coût de calcul par token et en permettant l'entraînement sur des clusters plus petits ou plus efficaces, cette technologie accélère le cycle de développement des IA génératives.

En résumé, Megatron-Core MoE transforme l'entraînement des modèles MoE d'un défi de faisabilité en une tâche industrielle maîtrisée, en résolvant systématiquement les goulots d'étranglement mémoire, communicationnels et computationnels.