Entraînement MoE Long-Context

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

Ce qui change à long contexte

Une fois que la longueur de séquence dépasse nettement 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 budget DP plus réduit restant

Motifs de scaling arrondi

DSV3 sur H100

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

le selective recompute fonctionne mieux que le full recompute une fois qu'on dépasse les contextes les plus courts
le throughput reste dans une bande assez étroite du contexte de longueur moyenne aux contextes très longs 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 contexte ne s'effondre pas immédiatement l'utilisation si la disposition est bien choisie, mais consomme très rapidement le budget DP.

Qwen3-Next sur GB200

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

8K et 32K restent pratiques avec un CP modéré
64K est possible, mais la baisse de throughput 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 contexte peut rester efficace sur les systèmes NVL72 lorsque TP, CP et HybridEP sont coordonnés. Les meilleures configurations classe 128K ne sont pas simplement des recettes « fit-only » ; elles peuvent rester hautement efficaces si le routage, le parallelism et le recompute sont équilibrés.

Règles pratiques de dimensionnement CP

Partir d'une cible de shard 4K : une bonne première estimation est CP ~= seq_len / 4096, puis arrondir à une disposition pratique en puissance de deux.
Garder 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érer le selective recompute : recompute les modules tels que up_proj, norm, moe, moe_act ou mlp avant de recourir au full recompute.
Éviter le recompute lourd en SDPA à très long contexte : recompute les internals d'attention peut ajouter beaucoup de travail pour moins de bénéfice mémoire que de recompute les modules MoE et MLP plus petits.
Utiliser TP comme autre levier sur les systèmes NVL72 : les exécutions GB200 et GB300 peuvent parfois échanger du CP contre du TP tout en restant efficaces.
Supposer 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 configuration 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

Recommandations pour recompute et CUDA Graph

Pour l'entraînement MoE long-context :

commencer par le selective recompute
ajouter les CUDA graphs seulement après que les shapes et le chemin de routage soient stables
garder la longueur de séquence et MBS fixes quand vous utilisez les CUDA graphs
si l'exécution dépend de batches hautement dynamiques, préférer l'eager execution

Références utiles :

@docs/training/activation-recomputation.md
@skills/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.
Une bonne baseline 4K peut toujours être une mauvaise baseline long-context : le routing mode, le choix de recompute et la stratégie d'offload ont souvent besoin de changer.
GPU-count feasibility devient la véritable contrainte : le très long contexte peut sembler correct dans une seule recette, puis devenir impossible une fois que EP et PP sont ajoutés honnêtement sur l'ensemble du modèle.
Les CUDA graphs ont besoin de shapes statiques : les batches de longueur variable et les stratégies de padding opportuniste peuvent casser silencieusement le chemin.
Le support container et kernel compte plus à 128K+ : les chemins long-context tendent à s'appuyer sur des kernels plus récents et des correctifs de bugs que ne le fait la mise en place short-context.