Tsinghua Unveils BaGuaLu-2: Runtime Adaptation Solves Memory Spikes in Large-Scale Training

Runtime-adaptive resolution of memory spikes without configuration changes; code is open-source.

Published by Synced (Machine Heart)

Large model training has scaled to clusters with tens of thousands of GPUs, yet GPU memory remains an invisible barrier.

The MoE architecture has pushed large models to the trillion-parameter level but also planted a time bomb: memory spikes. When tokens in a specific micro-batch concentrate on a few experts, activation memory usage surges instantly. If one card runs out of memory (OOM), the entire distributed task crashes—wasting hundreds of steps.

Recomputation and activation offloading are common techniques, but the problem lies in static configurations. Aggressive settings yield 30% MFU during normal phases but inevitably crash during spikes, requiring a cold restart that takes 284 seconds. Conservative full-layer recomputation is stable but drops MFU to 24.6%, leaving compute resources idle for long periods. Either choice results in loss.

Can the system adapt dynamically to memory pressure—running at full speed when stable, temporarily intensifying mitigation during spikes, and reverting once the spike passes?

Professor Zhai Jidong’s team from Tsinghua University offers the answer: BaGuaLu-2, whose paper has been accepted by SC26. It features runtime adaptation with zero configuration intervention. Within the window after gating computation and before expert computation, the system makes real-time decisions based on actual load. Different operators independently choose between saving, recomputing, or offloading activations, keeping peak memory within device limits. No need to stop tasks or modify configurations.

Real-world tests in a 10,000-card intelligent computing center: Using 128 to 12,288 Hygon DCUs, supporting models from 188B to 1.04T parameters, average throughput is 1.14× higher than the baseline. At 12,288 cards, total throughput reaches 2.012 EFLOPS, with an average MFU of 32.9% and a peak of 35%.

• Paper Link: https://arxiv.org/abs/2610.23064

I. Problem Background: The Pain Point of Memory Spikes in MoE Training

If you have trained MoE models, you have likely encountered this scenario: Training proceeds smoothly for hundreds of iterations, then suddenly, in a specific micro-batch, tokens concentrate on just a few experts. Activation memory usage spikes instantly, causing one card to OOM, which crashes the entire distributed task.

The Problem Lies in "Selection" Itself

The core of MoE is the router selecting a few experts to execute each token. Modern load-balancing methods can suppress long-term expert congestion but cannot predict the specific content of the current batch. Empirical data shows that in 97.1% of iterations, the expert load imbalance ratio is below 1.2×, but in the remaining 2.9% of anomalous iterations, it reaches up to 2.6×. It is precisely these rare iterations that repeatedly interrupt the training process.

Pipelining Exacerbates the Issue

During the pipeline warm-up phase, only forward passes occur without backward passes releasing activations, meaning multiple micro-batches' activations reside on the device simultaneously. If one batch routes a large number of tokens to a few experts on the current device, the additional MoE activations stack on top of already resident activations. Local skewness directly translates into device-level OOM.

The Dilemma of Static Configuration

Recomputation and offloading are standard tools, but traditional approaches hardcode configurations before startup. Aggressive configurations achieve 30.0% MFU during normal phases but inevitably cause OOM during spikes, with a single cold restart taking 284 seconds. Conservative configurations using full-layer recomputation drop MFU to 24.6%, leaving compute resources idle for extended periods. Regardless of the choice, resources are wasted.

Figure 1: MoE load spikes drive up activation memory, forcing static strategies to oscillate between OOM-induced restarts and conservative operation.

II. Design Philosophy: Runtime Adaptation with Zero Configuration Intervention

The core idea behind BaGuaLu-2 is straightforward: As memory pressure fluctuates, strategies should adapt accordingly. During stable phases, minimal intervention preserves high throughput; during spikes, mitigation is temporarily intensified; once the spike passes, the system automatically reverts. This mechanism does not terminate tasks, modify configurations, or trigger restarts. It is called Flexible Activation Provisioning.

Figure 2: Lighter memory mitigation is used during stable phases, while intensity increases temporarily during spikes, ensuring the memory curve always stays within device limits.

① Activation Memory Modeling: A Scalable Cost Model

The system integrates activation saving, recomputation, and offloading into a unified cost model, estimating both memory savings and time overheads while considering whether extra operations can overlap with existing computation and communication. For example, offloading the same amount of data at different pipeline stages may expose vastly different amounts of time on the critical path. Across multiple models and configurations, the average prediction error is only 1.56%. The model design is scalable, facilitating the integration of new operators and strategies.

② Adaptive Provisioning: Operator-Level Fine-Grained Control

The decision point is placed in a critical window: after MoE gating computation completes but before expert computation begins. At this stage, the actual load is known, risks have not yet materialized, and there is still time to temporarily intensify recomputation or offloading. Strategies are no longer applied at the layer level; instead, different operators independently choose between saving, recomputing, or offloading, selecting only the minimal combination that satisfies the budget. The system also traces activation sources to eliminate redundant combinations, such as those saved in different precisions or where one copy is offloaded while another remains resident.

Figure 3: Based on the current pipeline state and MoE load, the system sequentially performs activation memory modeling, adaptive provisioning, and dual-batch fused execution.

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

Recomputation and offloading are not free; placing them directly on the critical path trades throughput for safety. BaGuaLu-2 exploits pipeline bubbles. In stable phases, the device simultaneously hosts the current forward pass, current backward pass, previous forward offload, and next backward recovery. These operations belong to different micro-batches and can be interleaved while maintaining dependency order. The system uses lightweight dynamic programming to reorder computation and communication streams, squeezing extra work into hardware bubbles. Consecutive operators are merged, and budget states are cached, allowing most iterations to reuse scheduling results directly. Ultimately, the decision overhead is less than 10 ms, compared to an iteration time exceeding 18 seconds.

Figure 4: First, dynamic strategies are used to lower peak memory usage; then, adjacent batches are interleaved and scheduled to hide additional computation and data movement within hardware bubbles.

III. Empirical Results: Validated Across Multiple Scales, Code Open-Sourced

Leading Performance at Full Scale

In a large-scale AI computing center with tens of thousands of cards, models ranging from 188B to 1.04 trillion parameters were tested on clusters scaling from 128 to 12,288 Hygon DCUs (BW1000). Compared to Megatron-Kuai and AdaPipe, the system achieved average speedups of 1.14× and 1.13× respectively (including losses from OOM restarts). At the 12,288-card scale, total throughput reached 2.012 EFLOPS, with an average Model FLOPs Utilization (MFU) of 32.9% and a peak of 35%.

Figure 5: From 128 to 12,288 accelerator cards, BaGuaLu 2 maintains leading throughput across all tested scales.

Ablation Studies

Starting from a full recomputation baseline: fine-grained activation strategy yields 1.16× → + dynamic management yields 1.22× → + overlapped scheduling yields 1.38×. Each module contributes clear performance gains.

Figure 6: The sequential addition of fine-grained activation strategies, dynamic memory management, and overlapped scheduling results in an overall gain of 1.38×.

Open Source and Quick Start

• Code Repository: GitHub — thu-pacman/BaGuaLu
• Open-Sourced Components: Core code including the activation memory analysis model, scheduler-integrated activation offloading and reloading mechanisms.
• Low-Barrier Integration: Provides adaptation examples and neutral evaluation benchmarks to facilitate integration into existing training pipelines.
• Zero-Configuration Usage: Ready-to-use out of the box, supporting research experiments and rapid validation.

We welcome Issues and Pull Requests for technical collaboration. Note that the public repository does not include all runtime mechanisms described in the paper, but core modules are independently usable.

IV. Team and Community

BaGuaLu 2 is led by Professor Zhai Jidong's team at Tsinghua University, co-developed with Qingcheng Jizhi, University of Science and Technology Beijing, and independent researchers. The team is dedicated to open-source collaboration for large model training and inference systems, welcoming developers and researchers to explore cutting-edge technologies together.

• Chitu Inference Engine: https://github.com/thu-pacman/chitu
• BaGuaLu Training System: https://github.com/thu-pacman/BaGuaLu

The team's evaluation techniques are also applied to energy consumption control in AI computing centers, yielding particularly significant benefits for financial clouds heavily utilizing domestic compute resources.

Click "Read Original" or visit https://arxiv.org/abs/2610.23064 to read the full paper!

Qingcheng Jizhi is currently hiring AI Infrastructure Engineers with competitive salaries. Applications are welcome!

© THE END

Please contact this official account for authorization before reprinting.

For submissions or media inquiries: [email protected]

Read Original

Open in WeChat