Entraînement MoE en Long-Context

Docs stables : @docs/training/moe-optimization.md Card : @skills/nemo-mbridge-perf-moe-long-context/card.yaml

Changements en Long-Context

Une fois que la longueur de séquence dépasse largement le régime 4K, la mémoire d'attention et la résidence des activations deviennent les contraintes dominantes. Pour les modèles MoE, cela signifie généralement que vous avez besoin d'une combinaison de :

context parallelism
selective recompute
précision réduite
CPU offload pour l'état de l'optimiseur
un dispatcher et une disposition PP qui ne gaspillent pas le plus petit budget DP restant

Motifs d'Échelle Arrondis

DSV3 sur H100

Les exécutions long-context DSV3 montrent un motif stable :

selective recompute fonctionne mieux que full recompute une fois que vous dépassez les contexts les plus courts
le débit reste dans une bande assez étroite de la longueur médiane aux très longs contexts si CP est augmenté de manière appropriée
l'arbitrage passe de « memory fit » à « GPU-count feasibility » à mesure que CP augmente

En d'autres termes, le long-context ne réduit pas immédiatement l'utilisation si la disposition est bien choisie, mais il consomme très rapidement le budget DP.

Qwen3-Next sur GB200

Qwen3-Next se comporte davantage comme un modèle medium-scale sensible à la mémoire :

8K et 32K restent pratiques avec un CP modéré
64K est possible, mais la baisse de débit est notable et la mémoire devient beaucoup plus serrée
la disposition pipeline et les améliorations grouped-GEMM comptent presque autant que CP

Qwen3 235B sur GB200

Qwen3 235B montre que le long-context peut rester efficace sur les systèmes NVL72 quand TP, CP et HybridEP sont coordonnés. Les meilleures configurations class-128K ne sont pas juste des recettes « fit-only » ; elles peuvent rester très efficaces si le routing, le parallelism et le recompute sont équilibrés.

Règles Pratiques de Dimensionnement de CP

Partez d'une cible shard 4K : une première hypothèse correcte est CP ~= seq_len / 4096, puis arrondissez à une disposition power-of-two pratique.
Gardez DP en vie si possible : le scaling long-context devient fragile une fois que CP, EP, TP et PP ensemble réduisent DP au minimum.
Préférez selective recompute : recomputez des modules tels que up_proj, norm, moe, moe_act, ou mlp avant de recourir à full recompute.
Évitez SDPA-heavy recompute en très long-context : recomputer les internals d'attention peut ajouter beaucoup de travail pour un bénéfice mémoire inférieur à celui de recomputer les modules MoE et MLP plus petits.
Utilisez TP comme autre levier sur les systèmes NVL72 : les exécutions GB200 et GB300 peuvent parfois échanger un peu de CP contre TP tout en restant efficaces.
Supposez que GBS devra diminuer : à mesure que CP augmente et DP diminue, vous devrez peut-être réduire la taille du batch global ou accepter un GA plus élevé.

Familles de Config Représentatives

DSV3 à 128K sur H100

TP=1  CP=32  EP=32  PP=8  VPP=4
Precision: FP8-class
Dispatcher: DeepEP
Recompute: up_proj, norm, moe, mlp
Extra memory help: optimizer CPU offload

DSV3 à 256K sur H100

TP=1  CP=64  EP=32  PP=8  EDP=2  VPP=4
Precision: FP8-class
Dispatcher: DeepEP
Recompute: up_proj, norm, moe, mlp
Extra memory help: optimizer CPU offload

Qwen3 235B à 128K sur GB200

TP=4  CP=4  EP=32  PP=4  VPP=12
Precision: BF16 or MXFP8
Dispatcher: HybridEP
Recompute: moe_act, norm
CUDA Graph: attn + moe_router + moe_preprocess

Guidance Recompute et CUDA Graph

Pour l'entraînement long-context MoE :

commencez avec selective recompute
n'ajoutez CUDA graphs que quand les shapes et le routing path sont stables
gardez la longueur de séquence et le MBS fixes quand vous utilisez CUDA graphs
si l'exécution dépend de batches très dynamiques, préférez eager execution

Références utiles :

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

Pièges

CP ne remplace pas EP ou PP : il ajoute une autre dimension ; il ne fait pas disparaître les autres.
Un bon baseline 4K peut toujours être un mauvais baseline long-context : le mode routing, le choix recompute et la stratégie offload ont souvent besoin de changer.
GPU-count feasibility devient la vraie contrainte : le très long-context peut sembler correct dans une seule recette, puis devenir impossible une fois que EP et PP sont ajoutés honnêtement sur le modèle complet.
CUDA graphs ont besoin de shapes statiques : les batches de longueur variable et les stratégies de padding opportunistes peuvent silencieusement casser le chemin.
Le support conteneur et kernel compte plus à 128K+ : les chemins long-context tendent à dépendre de kernels plus récents et de corrections de bugs que la bring-up short-context ne le fait.