A training system adapts memory to rare expert-load spikes

Large-scale mixture-of-experts training routes tokens among many expert networks. Rare routing spikes can raise activation-memory demand, while a fixed memory strategy must either slow ordinary work or leave a run exposed to out-of-memory crashes. How can training respond to those brief changes without paying the same cost on every batch?

BaGuaLu-2 chooses memory-saving steps for each small training batch, then overlaps that work across batches as memory pressure changes.

BaGuaLu-2 estimates activation-memory needs and chooses recomputation, offloading, quantization, or no mitigation separately for each micro-batch. It overlaps that work with operations from neighboring batches to reduce time spent on memory management.

Across evaluated model and cluster configurations from 128 to 12,288 GPUs, the authors report average end-to-end throughput 1.14 times Megatron-Kuai's and 1.13 times AdaPipe's. Throughput here averages Model FLOPs Utilization (MFU), a measure of how much of the GPUs' available computing capacity is used, for each iteration across the entire run, including machine warm-up time after crashes. The results support higher overall throughput in these tests.

The average comparison covers the evaluated configurations, which began with the same offline-selected parallelism; the paper does not specify shared starting checkpoints.

Restart accounting includes machine warm-up but omits human diagnosis and intervention, so it understates that part of the restart cost.