MoE Long-Context Training

MoE Long-Context Training

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

What Changes At Long Context

Once sequence length moves well past the 4K-class regime, attention memory and activation residency become the dominant constraints. For MoE models, that usually means you need some combination of:

• context parallelism
• selective recompute
• lower precision
• CPU offload for optimizer state
• a dispatcher and PP layout that do not waste the smaller remaining DP budget

Rounded Scaling Patterns

DSV3 on H100

The DSV3 long-context runs show a stable pattern:

• selective recompute works better than full recompute once you move past the shortest contexts
• throughput stays in a fairly narrow band from mid-length through very long contexts if CP is increased appropriately
• the trade shifts from "memory fit" to "GPU-count feasibility" as CP grows

In other words, long context does not immediately collapse utilization if the layout is chosen well, but it does consume the DP budget very quickly.

Qwen3-Next on GB200

Qwen3-Next behaves more like a memory-sensitive medium-scale model:

• 8K and 32K remain practical with moderate CP
• 64K is possible, but the throughput drop is noticeable and memory becomes much tighter
• pipeline layout and grouped-GEMM improvements matter almost as much as CP