Technique for Simultaneous Control of Four Memory Peaks in Long-Context MoE Training Released

To address the issue where training fails when the memory peak of a single component exceeds GPU capacity during MoE model training with long contexts or large batch sizes, a technique controlling all four major memory peaks has been proposed. The four peaks—expert dispatch, vocabulary projection, gradient checkpoint boundaries, and optimizer states—which were not managed by existing parallelization methods, increase differently depending on the model, context length, and number of devices.

The research team limited all four peaks through scheduling that fixes the GPU working set at startup. Specifically, PipelinedLLEP sets an upper limit on the number of source-specific tokens contributing to dispatch chunks, while Ring-DTP reduces peaks by circulating activation or weight shards in a ring structure during the vocabulary projection stage and folding logit blocks with online log-sum-exp. Selective Checkpoint Offload (SCO) offloads long-lived single tensors at checkpoint boundaries to CPU memory, and OffloadStreamAdamW converts serial CPU Adam updates into a bucket pipeline during optimizer offloading. Since these techniques only change the order and granularity of computation and data movement, the loss function and gradients maintain accuracy.

Performance test results show that the MoE dispatch peak decreased by up to 59.3%, the vocabulary projection peak decreased by 86.6%, and the offloaded optimizer step became 2.05 times faster. When applied to MoE models with 120B to 667B parameters, training is possible at a 1M context length, providing 8 to 32 times the reach and up to 10.4 times the throughput compared to a tuned FSDP2 baseline.