Training Mixture-of-Experts (MoE) models at long context or large batch sizes often fails when any single component's peak memory allocation exceeds device capacity. A new research paper argues that the real target is not average memory footprint but every peak at once, since four components remain unbounded under common parallelism plans: expert dispatch, vocabulary projection, gradient checkpoint boundaries, and optimizer state.
Each of these peaks grows differently—expert dispatch scales with the routing matrix, vocabulary projection with tokens times vocabulary size, checkpoint boundaries with depth times sequence length, and optimizer state with parameter count. Because the first component to run out of memory changes with the model, context length, and device count, simply lowering the largest peak only exposes the next one.
The paper proposes four techniques that bound all four peaks with GPU working sets fixed at launch. PipelinedLLEP extends least-loaded expert parallelism with a cap on 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. OffloadStreamAdamW converts the serial CPU Adam update of optimizer offload into a bucket pipeline.
All four methods change only the order and granularity of computation and data movement, so loss and gradients remain exact. In matched component tests, they reduced the MoE dispatch peak by up to 59.3% without losing throughput, cut the vocabulary projection peak by 86.6%, and made the offloaded optimizer step 2.05 times faster.
When composed on MoE models ranging from 120B to 667B parameters, the techniques enabled training at 1M context length—8 to 32 times the reach of a tuned FSDP2 baseline—and delivered up to 10.4 times its throughput.