Flux de travail d'optimisation MoE

Docs stables : @docs/training/moe-optimization.md Card : @skills/nemo-mbridge-perf-moe-optimization-workflow/card.yaml Source : Scalable Training of MoE Models with Megatron Core

Référence rapide

Pensez en termes des Trois Murs du papier :

• memory wall
• communication wall
• compute and host-overhead wall

Le tuning MoE est itératif. Corriger un mur expose généralement le suivant, donc le meilleur flux de travail est : adapter d'abord, mettre à l'échelle deuxièmement, profiler troisièmement, puis réajuster.

Phase 1 : Rendre l'exécution réalisable en mémoire

Commencez par une configuration qui s'adapte de manière fiable avant de chercher le débit.

Ordre recommandé :

1. Utilisez la plus petite quantité de parallélisme de modèle qui tient encore.
2. Activez le selective recompute avant de revenir au full recompute.
3. Ajoutez l'offloading uniquement quand le recompute et le parallélisme sont toujours insuffisants.
4. Utilisez --fake-init-process-group pour vérifier les grandes topologies parallèles sur un seul GPU avant de consommer du temps cluster.

Guidance sur le recompute

Préférez le selective recompute pour les exécutions MoE :

• bons premiers choix : layernorm, core_attn, moe_act, mlp, ou modules spécifiques au modèle (shared_experts, mla_up_proj)
• utilisez le full recompute uniquement quand l'exécution ne tient toujours pas
• revisitez le recompute après l'activation des CUDA graphs, car certaines portées de graphes et les chemins full recompute ne se mélangent pas bien

En règle générale, le recompute fin-grained récupère souvent la majorité de la mémoire nécessaire tout en gardant le débit beaucoup plus proche de la baseline sans recompute que le recompute au niveau des couches entières.

Phase 2 : Choisir le parallélisme pour la mise à l'échelle

Ordre de priorité :

1. Maximisez DP une fois que le modèle tient.
2. Gardez le chemin de communication chaud à l'intérieur de l'interconnexion rapide quand c'est possible.
3. Utilisez PP, plus VPP si nécessaire, pour la mise à l'échelle multi-nœud.
4. Préférez EP à du TP supplémentaire pour les couches expert.
5. Ajoutez CP pour le contexte long une fois que la longueur de séquence rend la mémoire d'attention dominante.

Parallel Folding

Parallel Folding découple le parallélisme de l'attention et du MoE pour que vous n'ayez pas à choisir une seule topologie de compromis :

Attention: TP × CP × DP × PP
MoE:       ETP × EP × EDP × PP

Leviers clés :

• --expert-model-parallel-size
• --expert-tensor-parallel-size

Utilisez-le quand l'attention préfère du TP ou CP, mais les couches expert bénéficient d'un degré EP plus grand que les couches denses ne peuvent le supporter.

Phase 3 : Profiler le goulot d'étranglement dominant

Guidance sur Dispatcher et Overlap

Utilisez le choix de dispatcher comme correctif de goulot d'étranglement, pas comme premier levier de tuning.

• moe_token_dispatcher_type="alltoall" : chemin de mise en route le plus sûr, bien pour les tailles EP plus petites
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="deepep" : forte valeur par défaut pour les déploiements de style H100 et B200
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="hybridep" : point de départ le plus fort sur les systèmes GB200 ou GB300 NVL72

Si le chemin all-to-all est visible dans les profils, combinez le tuning du dispatcher avec :

• --overlap-moe-expert-parallel-comm
• --overlap-grad-reduce
• --tp-comm-overlap

Décision rapide sur recette FP8

Gardez le router en FP32. Les plus grands gains viennent généralement des GEMMs expert et autre grosse arithmétique matricielle, pas de l'essai de quantifier chaque petit composant MoE.

CUDA Graphs pour MoE

Pour le MoE dropless, commencez par des graphes à portée TE partielle :

• attn
• moe_router
• moe_preprocess

Ce chemin donne généralement un gain step-time significatif tout en gardant le travail expert dynamique en dehors du graphe. Attendez-vous à un speedup modéré quand l'overhead de lancement est visible, mais réservez plusieurs Go supplémentaires de mémoire et vérifiez que les formes restent statiques.

Utilisez les graphes d'itération complète uniquement pour des workloads adaptés aux graphes comme drop-and-pad ou des expériences static-shape étroitement contrôlées.

Références connexes :

• @skills/nemo-mbridge-perf-cuda-graphs/SKILL.md
• @docs/training/cuda-graphs.md
• @docs/training/activation-recomputation.md

Pièges

1. Ne pas optimiser dans le mauvais ordre : adapter le modèle et sélectionner un parallélisme sensé importent plus que les micro-optimisations.

2. Les changements de plateforme modifient le mur limitant : les exécutions de classe H100 semblent souvent plus liées à la communication, tandis que les exécutions GB200 ou GB300 exposent souvent le CPU ou l'overhead de lancement plus tôt.

3. Le MFU FP8 peut paraître étrangement bas : comparez le débit absolu ainsi que le MFU en changeant les modes de précision.

4. Les CUDA graphs et le recompute interagissent : les graphes à portée TE sont généralement appairés avec le selective recompute, pas le full recompute systématique.

5. Parallel Folding n'est pas optionnel à grande échelle : une fois que l'attention et les couches expert veulent des topologies clairement différentes, un plan TP ou EP partagé unique devient une taxe sur les deux.