Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training
Four scheduling techniques bound all memory peaks in long-context MoE training, enabling 120B-667B models at 1M-token context with up to 10.4x throughput.
The paper addresses memory peaks in long-context Mixture-of-Experts training by bounding four unbounded components: expert dispatch with the routing matrix, vocabulary projection, gradient checkpoint boundaries, and optimizer state. It introduces PipelinedLLEP (capped token contributions per dispatch chunk), Ring-DTP (ring circulation of activations or weight shards with online log-sum-exp), Selective Checkpoint Offload (SCO), and OffloadStreamAdamW, all preserving exact loss and gradients. Composed on MoE models from 120B to 667B parameters, the methods enable training at 1M context length, 8-32x the context reach of a tuned FSDP2 baseline, and up to 10.4x its throughput.
- Bounds four memory peaks: expert dispatch, vocabulary projection, checkpoint boundaries, optimizer state.
- Introduces PipelinedLLEP, Ring-DTP, Selective Checkpoint Offload, and OffloadStreamAdamW schedules.
- Cuts dispatch peak by up to 59.3% and vocabulary projection peak by 86.6% without throughput loss.
- Enables 120B-667B MoE training at 1M context, up to 10.4x FSDP2 baseline throughput.
Full article268 words · extracted from huggingface.co · click to collapse
Training a Mixture-of-Experts (MoE) model at long context or large batch size fails as soon as any one component's peak allocation exceeds device memory, so the target is every peak at once, not the average footprint. Four are left unbounded by the parallelism plans in common use, and each grows differently: expert dispatch with the routing matrix, the vocabulary projection with tokens times vocabulary, gradient checkpoint boundaries with depth times sequence length, and optimizer state with parameter count. Which one runs out first changes with the model, the context length, and the device count, so lowering the largest only exposes the next. We bound all four with schedules whose GPU working set is fixed at launch: PipelinedLLEP extends least-loaded expert parallelism with a cap on the tokens each source contributes to a dispatch chunk, Ring-DTP circulates activations or weight shards around a ring at the vocabulary projection and folds each block of logits into an online log-sum-exp, Selective checkpoint offload (SCO) keeps the one long-lived tensor of each checkpoint boundary in CPU memory, and OffloadStreamAdamW turns the serial CPU Adam update of optimizer offload into a bucket pipeline. All four change only the order and granularity of computation and data movement, so the loss and gradients stay exact. In matched component tests, they cut the MoE dispatch peak by up to 59.3% without losing throughput, the vocabulary projection peak by 86.6%, and the offloaded optimizer step by 2.05times faster. Composed on MoE models from 120B to 667B parameters, they train at 1M context length, 8--32times the reach of a tuned FSDP2 baseline, and up to 10.4times its throughput.
Text extracted automatically; images, tables and formatting may be missing. Original: https://huggingface.co/papers/2609.14306