Entraînement MoE à long contexte

Documentation stable : @docs/training/moe-optimization.md Fiche : @skills/nemo-mbridge-perf-moe-long-context/card.yaml

Quels changements dans le cadre des contextes longs ?

Dès que la longueur de la séquence dépasse largement le régime de l’ordre de 4K, la mémoire d’attention et la résidence d’activation deviennent les contraintes dominantes. Pour les modèles MoE, cela signifie généralement qu’il faut combiner :

• le parallélisme contextuel
• recalcul sélectif
• précision réduite
• déchargement du CPU pour l'état de l'optimiseur
• un répartiteur et une disposition PP qui ne gaspillent pas le budget DP restant, plus restreint

Modèles de mise à l'échelle arrondis

DSV3 sur H100

Les exécutions à long contexte du DSV3 montrent un schéma stable :

• le recalcul sélectif fonctionne mieux que le recalcul complet dès lors que l’on dépasse les contextes les plus courts
• le débit reste dans une fourchette assez étroite, des contextes de longueur moyenne aux très longs, si le CP est augmenté de manière appropriée
• le compromis passe de l’« adaptation à la mémoire » à la « faisabilité en nombre de GPU » à mesure que le CP augmente

En d’autres termes, un contexte long n’entraîne pas immédiatement une chute du taux d’utilisation si la disposition est bien choisie, mais il épuise très rapidement le budget DP.

Qwen3-Next sur GB200

Qwen3-Next se comporte davantage comme un modèle de taille moyenne sensible à la mémoire :

• 8 Ko et 32 Ko restent viables avec un CP modéré
• 64K est possible, mais la baisse de débit est perceptible et la mémoire devient beaucoup plus limitée
• les améliorations apportées à la disposition du pipeline et aux GEMM groupées ont presque autant d’importance que le CP

Qwen3 235B sur GB200

Qwen3 235B montre qu'un contexte long peut encore être efficace sur les systèmes NVL72 lorsque le TP, le CP et l'HybridEP sont coordonnés. Les meilleures configurations de classe 128K ne sont pas seulement des recettes « adaptées » ; elles peuvent rester très efficaces si le routage, le parallélisme et le recalcul sont équilibrés.

Règles empiriques pour le dimensionnement du CP

1. Commencez par une cible de fragment de 4 Ko: une bonne première estimation est CP ~= seq_len / 4096, puis arrondissez à une configuration pratique de type « puissance de deux ».

2. Maintenez le DP actif si possible: la scalabilité des contextes longs devient fragile dès que le CP, l’EP, le TP et le PP, combinés, réduisent le DP à son minimum.

3. Privilégiez le recalcul sélectif: recalculez des modules tels que up_proj, norm, moe, moe_act ou mlp avant d’en arriver à un recalcul complet.

4. Évitez les recalculs gourmands en SDPA avec des contextes très longs: le recalcul des composants internes de l’attention peut ajouter beaucoup de travail pour un gain de mémoire moindre que le recalcul de modules plus petits côté MoE et MLP.

5. Utilisez TP comme levier supplémentaire sur les systèmes NVL72: les exécutions sur GB200 et GB300 peuvent parfois échanger une partie de CP contre TP tout en restant efficaces.

6. Partez du principe que le GBS devra être réduit: à mesure que le CP augmente et que le DP diminue, vous devrez peut-être réduire la taille globale des lots ou accepter un GA plus élevé.

Familles de configurations représentatives

DSV3 à 128K sur H100

TP=1  CP=32  EP=32  PP=8  VPP=4
Précision : classe FP8
Dispatcher : DeepEP
Recalcul : up_proj, norm, moe, mlp
Aide mémoire supplémentaire : déchargement du processeur par l’optimiseur

DSV3 à 256K sur H100

TP=1  CP=64  EP=32  PP=8  EDP=2  VPP=4
Précision : classe FP8
Dispatcher : DeepEP
Recalcul : up_proj, norm, moe, mlp
Aide mémoire supplémentaire : déchargement du CPU par l’optimiseur

Qwen3 235B à 128K sur GB200

TP=4  CP=4  EP=32  PP=4  VPP=12
Précision : BF16 ou MXFP8
Dispatcher : HybridEP
Recalcul : moe_act, norm
Graphique CUDA : attn + moe_router + moe_preprocess

Conseils sur le recalcul et le graphe CUDA

Pour l’entraînement MoE à contexte long :

• commencez par un recalcul sélectif
• n'ajoutez des graphes CUDA qu'une fois que les formes et le chemin de routage sont stables
• maintenez la longueur de séquence et le MBS fixes lorsque vous utilisez des graphes CUDA
• si l'exécution dépend de lots très dynamiques, privilégiez l'exécution « eager »

Références utiles :

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

Pièges

1. Le CP ne remplace ni l’EP ni le PP: il ajoute une dimension supplémentaire ; il ne fait pas disparaître les autres.

2. Une bonne base de référence 4K peut tout de même s’avérer médiocre pour les contextes longs: le mode de routage, le choix de recalcul et la stratégie de déchargement doivent souvent être modifiés.

3. La faisabilité en termes de nombre de GPU devient la véritable contrainte: un contexte très long peut sembler convenir dans une recette unique, puis devenir impossible une fois que l’EP et le PP sont ajoutés de manière cohérente à l’ensemble du modèle.

4. Les graphes CUDA nécessitent des formes statiques: les lots de longueur variable et les stratégies de remplissage opportunistes peuvent rompre le chemin de manière silencieuse.

5. La prise en charge des conteneurs et des noyaux revêt davantage d’importance à partir de 128K+: les chemins à long contexte ont tendance à s’appuyer davantage sur des noyaux plus récents et des corrections de bogues que ne le fait la mise en route à court contexte.