BaGuaLu-2: Runtime-Adaptive Memory Management for MoE Training on 10K+ GPU Clusters

Problem: Activation Memory Spikes in Large-Scale MoE Training

Mixture-of-Experts (MoE) models scale to trillion parameters but introduce unpredictable activation memory spikes. When tokens in a micro-batch concentrate on a few experts, activated memory surges, causing OOM on a single GPU and crashing the entire distributed job — wasting hundreds of iterations. Load-balancing methods reduce long-term imbalance (97.1% of iterations below 1.2×), but 2.9% of iterations still hit 2.6× imbalance, repeatedly interrupting training.

Pipeline parallelism worsens the issue: during warm-up, multiple micro-batches' activations reside simultaneously. A skewed batch adds MoE activations on top of existing ones, turning local skew into device-level OOM.

Static Configurations Force a Trade-off

Traditional remedies — recomputation and activation offloading — are configured statically before launch:

Aggressive config : normal-phase MFU 30.0%, but spikes guarantee OOM; each cold restart costs 284 seconds.

Conservative full recomputation : MFU drops to 24.6%, leaving compute idle most of the time.

Neither adapts to actual runtime pressure.

BaGuaLu-2: Flexible Activation Provisioning

The core idea: memory pressure changes, so strategy should change with it. Light intervention during stable phases for high throughput; temporary reinforcement during spikes; automatic recovery afterward — no job termination, no config edits, no restarts.

① Scalable Activation Memory Cost Model

The system unifies three actions — save activation, recompute, offload — into one 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 policies.

② Per-Operator Adaptive Provisioning

Decision point: after MoE gating finishes, before expert computation starts. Real load is known; risk has not materialized; there is still time to strengthen recomputation or offload. Granularity is per-operator, not per-layer. 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 stays resident).

③ Dual-Batch Fusion via Dynamic Programming

Recomputation and offloading are not free; naively placing them on the critical path trades throughput for safety. BaGuaLu-2 exploits pipeline bubbles: during steady state, 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 scheduler reorders compute and communication streams, packing extra work into hardware gaps. 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).

Experimental Validation

Full-Scale Results

Tested in a 10K-GPU intelligent computing center on Hygon DCU (BW1000), models from 188B to 1.04 trillion parameters, clusters from 128 to 12,288 cards. Compared to Megatron-Kuai and AdaPipe (including their OOM restart penalties), BaGuaLu-2 delivers 1.14× and 1.13× average throughput gains. At 12,288 cards: 2.012 EFLOPS aggregate throughput, average MFU 32.9%, peak 35%.

Ablation Study

Starting from full recomputation baseline:

Fine-grained activation policy → 1.16×

+ Dynamic memory management → 1.22×

+ Overlap scheduling → 1.38×

Each module contributes measurable improvement.

Open Source & Usability

Repository: https://github.com/thu-pacman/BaGuaLu Released components: activation memory analysis model, scheduler-integrated offload/reload core.

Low-barrier integration: adaptation examples and neutral benchmarks provided.

Zero-config usage: works out-of-the-box for research and rapid validation.

Note: the public repo does not yet contain the full runtime mechanism described in the paper, but core modules are independently usable.

Team & Community

Led by Prof. Zhai Jidong (Tsinghua University), jointly developed with Chitu Engine team, University of Science and Technology Beijing, and independent researchers. The group focuses on open-source co-construction of large-model training and inference systems. Related projects: Chitu inference engine ( https://github.com/thu-pacman/chitu) and BaGuaLu training system ( https://github.com/thu-pacman/BaGuaLu). Their evaluation technology is also applied to energy control in intelligent computing centers, with pronounced benefits for finance clouds using domestic compute.

Full paper: arXiv:2610.23064 (accepted at SC26).