BaGuaLu-2: Adaptive Activation Memory Management for MoE Training at Scale

Problem Background: MoE Training Memory Spikes

Mixture-of-Experts (MoE) models scale to trillion parameters but introduce activation memory spikes. When tokens concentrate on a few experts in a micro-batch, activated memory surges, causing OOM on a single GPU and crashing the entire distributed job — wasting hundreds of iterations. Load balancing mitigates long-term congestion but cannot predict per-batch routing. Measurements show 97.1% of iterations have expert load imbalance below 1.2×, yet 2.9% reach 2.6×, repeatedly interrupting training.

Pipeline parallelism worsens the issue: during warm-up, multiple micro-batches' activations reside simultaneously. If one batch routes heavily to local experts, the extra MoE activations stack on top of existing ones, turning local skew into device-level OOM.

Static remediation — recomputation and activation offloading — forces a dilemma. Aggressive settings yield 30.0% MFU normally but guarantee OOM at spikes, incurring 284-second cold restarts. Conservative full-layer recomputation drops MFU to 24.6%, leaving compute idle. Neither is acceptable.

Design: Flexible Activation Provisioning

BaGuaLu-2 makes memory mitigation adaptive at runtime: light mitigation during stable phases, temporary reinforcement at spikes, automatic recovery afterward — no job termination, no config changes, no restarts.

① Activation Memory Modeling: Extensible Cost Model

The system unifies saving activations, recomputation, and offloading in a single cost model that estimates both memory savings and time overhead, accounting for overlap with existing compute/communication. The same offload at different pipeline stages exposes different critical-path time. Across multiple models and configurations, average prediction error is only 1.56%. The model is extensible for new operators and strategies.

② Adaptive Provisioning: Operator-Level Fine-Grained Control

Decision point: after MoE gating completes but before expert computation starts. Real load is known, risk not yet materialized, leaving time to strengthen recomputation or offloading. Granularity shifts from whole layers to individual operators — each operator independently chooses save, recompute, or offload, picking the minimal combination that fits the memory budget. The system traces activation provenance to eliminate redundant combinations (e.g., same activation saved at different precisions, or one offloaded while another copy remains resident).

③ Dual-Batch Fusion: Dynamic Programming Scheduling with Negligible Overhead

Recomputation and offloading are not free; placing them on the critical path trades throughput for safety. BaGuaLu-2 exploits pipeline bubbles: during stable phases, a device concurrently holds current forward, current backward, previous forward's offload, and next backward's reload. These belong to different micro-batches and can be interleaved while respecting dependencies. A lightweight dynamic programming pass reorders compute and communication streams, packing extra work into hardware gaps. Consecutive operator fusion and budget-state caching let most iterations reuse the scheduling result. Decision overhead stays below 10 ms per iteration (iteration time > 18 s).

Evaluation: Multi-Scale Verification

Experiments on a 10K-card intelligent computing center (Huawei DCU BW1000) with models from 188B to 1.04 trillion parameters, scaling from 128 to 12,288 cards.

Overall performance: vs. Megatron-Kuai and AdaPipe baselines, BaGuaLu-2 averages 1.14× and 1.13× higher throughput (including OOM restart losses). At 12,288 cards, aggregate throughput reaches 2.012 EFLOPS, average MFU 32.9%, peak 35%.

Ablation: Starting from full recomputation baseline: fine-grained activation strategy → 1.16×; + dynamic management → 1.22×; + overlap scheduling → 1.38×. Each module contributes measurable gain.

Open Source

Core code open-sourced at GitHub: https://github.com/thu-pacman/BaGuaLu, including activation memory analysis model, scheduler-integrated offload/reload. Provides integration examples and neutral benchmarks for easy adoption into existing training pipelines. Zero-config usage supports research and rapid validation. The repository does not yet contain the full runtime mechanism described in the paper, but core modules are independently usable.

Paper: arXiv:2610.23064 (accepted at SC26). https://arxiv.org/abs/2610.23064