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.