Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training

Flattening Every Memory Peak in Long-Context Mixture-of-Experts Training. By Shrey Pandit and three others, from Salesforce AI Research. Published on the thirteenth of September. No code repository is linked to it.

Training a Mixture-of-Experts model, MoE, at long context or large batch size fails the moment any single component's peak memory allocation exceeds what the device holds. The average footprint does not matter. What matters is every peak, all at once. The authors identify four peaks that the parallelism plans in common use leave unbounded, and each grows in its own way. It is a game of whack-a-mole, and the only way to win is to hold every peak down at the same time.

Expert dispatch grows with the routing matrix. The vocabulary projection grows with the number of tokens times the size of the vocabulary. The boundaries where gradients are checkpointed grow with depth times sequence length. And optimizer state grows with the parameter count. Which of these runs out first changes with the model, the context length, and the number of devices. So lowering the largest peak only exposes the next one behind it.

The answer is to bound all four together, each with a schedule whose GPU working set is fixed at launch. However long the context or large the model, the memory each component needs on the device is decided in advance and does not grow. Crucially, all four schedules change only the order and granularity of computation and data movement. None changes the arithmetic, so the loss and the gradients stay exact.

PipelinedLLEP extends least-loaded expert parallelism with a cap on how many tokens each source may contribute to one dispatch chunk. Ring-DTP handles the vocabulary projection by circulating activations or weight shards around a ring, and folds each block of logits into an online log-sum-exp, so the full set of logits never has to exist at once. No single source can then flood a chunk with its tokens.

Selective checkpoint offload, SCO, keeps the one long-lived tensor at each checkpoint boundary in CPU memory rather than on the GPU. And OffloadStreamAdamW takes the serial CPU Adam update used in optimizer offload and turns it into a pipeline over buckets, so the update streams instead of waiting in line. Between them, the four schedules match the four peaks one to one: dispatch, vocabulary, checkpoints, and optimizer.

In matched tests of each component, the schedules cut the MoE dispatch peak by up to fifty-nine point three percent without losing throughput, and cut the vocabulary projection peak by eighty-six point six percent. The offloaded optimizer step runs two point zero five times faster.

Composed together, on MoE models from one hundred and twenty billion parameters up to six hundred and sixty-seven billion, they train at a context length of one million tokens. That is eight to thirty-two times the reach of a tuned FSDP2 baseline, with up to ten point four times its throughput.

The framing is what does the work here. Memory failure in long-context MoE training is treated not as one bottleneck to be squeezed but as four independent ceilings, any of which can be the one that breaks. Fixing the working set of each in advance, while leaving the mathematics untouched, is what lets the whole run scale without trading away exactness.