Flux de travail d'optimisation MoE

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

Référence rapide

Réfléchissez en termes des trois murs du papier :

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

L'ajustement MoE est itératif. Corriger un mur expose généralement le suivant, donc le meilleur flux est : ajuster d'abord, augmenter l'échelle ensuite, profiler en troisième, puis réajuster.

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

Commencez par une configuration qui tient de manière fiable avant de chasser le débit.

Ordre recommandé :

1. Utilisez la plus petite quantité de parallélisme de modèle qui tient encore.
2. Activez la recalculation sélective avant de revenir à la recalculation complète.
3. Ajoutez l'offloading uniquement quand la recalculation et le parallélisme restent 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.

Guide de recalculation

Préférez la recalculation sélective 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 la recalculation complète uniquement si l'exécution ne tient toujours pas
• révisitez la recalculation après l'activation des graphes CUDA, car certains scopes et chemins de recalculation complète ne se mélangent pas bien

En règle générale, la recalculation à grain fin récupère souvent la plupart de la mémoire nécessaire tout en maintenant le débit bien plus proche de la ligne de base sans recalculation que la recalculation par couche complète.

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

Ordre de priorité :

1. Maximisez DP une fois que le modèle tient.
2. Gardez le chemin de communication actif à l'intérieur de l'interconnexion rapide si possible.
3. Utilisez PP, plus VPP si nécessaire, pour la montée en charge multi-nœuds.
4. Préférez EP plutôt que TP supplémentaire pour les couches d'experts.
5. Ajoutez CP pour un long contexte une fois que la longueur de séquence rend la mémoire d'attention dominante.

Parallel Folding

Parallel Folding découple le parallélisme d'attention et de MoE pour que vous n'ayez pas à choisir une seule topologie 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 du CP, mais les couches d'experts bénéficient d'un degré EP plus large que les couches denses ne peuvent tolérer.

Phase 3 : Profiler le goulot d'étranglement dominant

Dispatcher et guide de chevauchement

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

• moe_token_dispatcher_type="alltoall" : chemin de mise en place le plus sûr, bon pour les petites tailles EP
• moe_token_dispatcher_type="flex" + moe_flex_dispatcher_backend="deepep" : défaut robuste 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 de recette FP8

Gardez le routeur en FP32. Les plus grands gains proviennent généralement des GEMM d'experts et d'autres mathématiques matricielles lourdes, pas de la tentative de quantifier chaque petit composant MoE.

Graphes CUDA pour MoE

Pour un MoE sans perte, commencez par des graphes partiellement scopés TE :

• attn
• moe_router
• moe_preprocess

Ce chemin donne généralement un gain de temps d'étape significatif tout en gardant le travail d'experts dynamiques hors du graphe. Attendez-vous à un accélération modérée quand l'overhead de lancement est visible, mais budgétisez plusieurs GB supplémentaires de mémoire et vérifiez que les formes restent statiques.

Utilisez les graphes de pleine itération uniquement pour les charges de travail compatibles avec les graphes comme drop-and-pad ou les expériences à formes statiques fortement contrôlées.

Références connexes :

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

Pièges

1. N'optimisez pas dans le mauvais ordre : ajuster le modèle et choisir un parallélisme sensé importent plus que les micro-optimisations.

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

3. MFU FP8 peut sembler faible de manière trompeuse : comparez le débit absolu ainsi que la MFU quand vous changez les modes de précision.

4. Les graphes CUDA et la recalculation interagissent : les graphes scopés TE sont généralement appairés avec une recalculation sélective, pas une recalculation complète généralisée.

5. Parallel Folding n'est pas optionnel à grande échelle : une fois que les couches d'attention et d'experts veulent clairement des topologies différentes, un plan TP ou EP partagé unique devient un coût pour les deux.