PaperlayerView exact revision on arXiv

Source-linked AI summary

BaGuaLu-2: Taming Dynamic Memory Fluctuation for Efficient MoE Model Training over 10,000 GPUs

Mingshu Zhai, Minkun Guo, Hongrui Lian, Shuangyu Li, Qi Zhang, Chen Hu, Xiawei Du, Yuyang Jin, Zan Zong, Jidong Zhai

arXiv:2610.23064v1cs.DC

TL;DR

Fine-grained MoE training faces rare load spikes that static activation-memory policies cannot handle without either wasting memory or risking OOMs. BaGuaLu-2 adapts activation provisioning per batch and sustains 2.012 EFLOPs at 12,288 GPUs, with average throughput gains over Megatron-Kuai and AdaPipe.

  • Problem

    Rare, high-magnitude MoE load spikes make static activation-memory policies a poor fit: average-case provisioning risks OOMs, while worst-case provisioning wastes memory.

    Source: §I · p. 2, A. MoE Training at Scale · p. 3

  • Method

    BaGuaLu-2 models activation costs and uses per-batch adaptive provisioning with interleaved execution to control memory and mitigate overhead.

    Source: §I · p. 2, A. Activation Cost Model · p. 5, §V · p. 6

  • Results

    At 12,288 GPUs, BaGuaLu-2 sustains 2.012 EFLOPs; average throughput gains are 1.14× over Megatron-Kuai and 1.13× over AdaPipe.

    Source: §I · p. 2, A. Setup · p. 10

  • Takeaways & Limitations

    Flexible provisioning addresses the static over-provisioning versus under-provisioning dilemma while scaling to a 1040-billion-parameter model across 12,288 GPUs.

    Source: §VIII · p. 11

Abstract

from arXiv · show

Fine-grained Mixture-of-Experts (MoE) is the de facto architecture for scaling large language models. However, dynamic expert load imbalance causes severe memory fluctuations that degrade training efficiency. Existing systems rely on static memory provisioning: they either over-provision, leaving GPU memory idle during normal execution, or under- provision, inviting out-of-memory crashes and costly manual rollbacks when load spikes occur. To address these challenges, we propose BaGuaLu-2, a large-scale MoE training system that eliminates this dilemma through flexible provisioning. By transforming memory allocation from a static, one-time tuning choice into a continuous and fine-grained control. BaGuaLu-2 grants exact activation memory on demand for each training batch. Furthermore, it offsets any dynamic management overhead through system optimizations. BaGuaLu-2 demonstrates exceptional scalability, training a model with 1040 billion parameters across 12,288 GPUs while sustaining over 2 EFlops at 35% MFU. Compared with state-of-the-art systems, BaGuaLu-2 delivers 1.14x and 1.13x average throughput improvements over Megatron-Kuai and AdaPipe, respectively.

I. INTRODUCTION

BaGuaLu-2 replaces static activation-memory provisioning with continuous per-batch control to handle MoE load-driven memory fluctuations. It scales to 12,288 GPUs with higher average throughput than Megatron-Kuai and AdaPipe.

I. INTRODUCTION

  • Dynamic sparse routing creates load imbalance; about 2.9% of iterations show spikes reaching 2.6×, which can push training beyond hardware capacity.Existing approaches such as sharding, recomputation, and offloading do not resolve load-driven memory fluctuation.Source: §I · p. 1
  • Static systems either waste GPU memory by over-provisioning or risk OOM crashes and manual intervention by under-provisioning.The latter can require a restart and a switch to a conservative configuration after a spike.Source: §I · p. 1
  • Activation memory fluctuates with expert imbalance, while weight and optimizer-state memory remain stable and can be distributed across GPUs.Larger clusters require proportionally more input data to maintain compute density, making activation memory difficult to distribute effectively.Source: §I · p. 2
  • Continuous, fine-grained activation provisioning gives each training batch the memory it needs, replacing the static over-provision-or-under-provision choice.Moderate imbalance receives steady mitigation; severe imbalance triggers stronger mitigation to prevent OOMs and avoid cold restarts.Source: §I · p. 2
  • 1.14× and 1.13× average throughput gains over Megatron-Kuai and AdaPipe, respectively, accompany training a 1040-billion-parameter model across 12,288 GPUs at over 2 EFlops and 35% MFU.The reported gains are for Megatron-Kuai and AdaPipe in that order.Source: §I · p. 2

II. BACKGROUND AND MOTIVATION

Fine-grained MoE increases activation-memory pressure, and rare, high-magnitude load spikes expose the limits of static mitigation. Measurements and a routing analysis connect these spikes to changing input distributions and show the cost of fixed memory policies.

A. MoE Training at Scale

  • Fine-grained MoE increases activated-expert count and activation-memory demands, while expert imbalance has shifted from persistent load to rare spikes.The passage contrasts early coarse-grained MoE designs with newer fine-grained architectures.Source: A. MoE Training at Scale · p. 3
  • 97.1% of iterations have imbalance below 1.2×, while 2.9% exceed that rate with a maximum of 2.6×, making rare spikes—not persistent imbalance—the key memory risk.Existing MoE systems emphasize compute load balancing and are poorly equipped to prevent memory failures from occasional spikes.Source: A. MoE Training at Scale · p. 3
  • Static mitigation assumes activation memory is stable across iterations, so average-sized policies risk OOMs on spikes and worst-case policies waste memory in ordinary iterations.Existing workflows profile or model a configuration, choose a strategy for a fixed budget, and apply it uniformly throughout training.Source: A. MoE Training at Scale · p. 3

B. Dynamic Memory Pressure: A Quantitative Case

  • 30.0% normal-iteration MFU under average-case provisioning gives way to OOMs in 2.9% of spike iterations; each restart costs 284 seconds, while full-layer recompute sustains 24.6% MFU.The 284-second recovery overhead is 16× a normal iteration and excludes human response latency.Source: B. Dynamic Memory Pressure: A… · p. 3
  • Routing randomness and shifting token distributions can produce up to 2.6× imbalance; few local experts and warmup accumulation can push device memory peaks beyond capacity.The passage argues that the spike mechanism is intrinsic to fine-grained token-choice routing, rather than an artifact of a particular balancing mechanism.Source: B. Dynamic Memory Pressure: A… · p. 3
  • Across SWE domains, both models show systematic changes in max/mean expert-load balance and hottest-expert identity at the same layer, implicating input token mix.Figure 3 reports the balance rate above and hottest expert below for DeepSeek-V3.1-Terminus and GLM-4.7.Source: B. Dynamic Memory Pressure: A… · p. 4, Figure 3 · p. 4
  • Static activation mitigation mismatches dynamic pressure, producing both conservative-policy performance deficits and aggressive-policy stability failures.The section motivates runtime adjustment of mitigation in response to actual memory pressure.Source: B. Dynamic Memory Pressure: A… · p. 4

III. OVERVIEW

BaGuaLu-2 plans activation-memory use at runtime for each upcoming micro-batch. Its pipeline-aware provisioning and interleaved execution aim to keep memory within device limits while reducing mitigation overhead.

III. OVERVIEW

  • At each pipeline scheduling point, BaGuaLu-2 uses current pipeline state and MoE imbalance to plan memory use and operator order for the next micro-batch.The plan combines activation-memory modeling, adaptive provisioning, and fused dual-batch execution.Source: §III · p. 4
  • Adaptive provisioning sets a pipeline-aware budget and selects per-operator recomputation, offloading, or quantization to keep peak memory within GPU capacity while minimizing overhead.The selection uses an iterative algorithm for each batch.Source: §III · p. 4
  • A lightweight solver interleaves operators across computation and communication streams to hide mitigation overhead in hardware bubbles.It examines mitigation strategies for the current and adjacent micro-batches.Source: §III · pp. 4–5

IV. ACTIVATION MEMORY MODELING

BaGuaLu-2 models activation memory and execution time together, from individual operators through micro-batches to full pipeline iterations. This model explains how pipeline phases shape memory occupancy and where peaks arise.

IV. ACTIVATION MEMORY MODELING

  • The activation cost model is developed from operator-level costs to a full iteration, providing the basis for quantifying activation mitigation strategies.Source: §IV · p. 5

A. Activation Cost Model

  • An operator’s activation cost is (∆t, ∆m): execution time and activation-memory change, represented as a trajectory over time and GPU memory occupancy.A micro-batch concatenates operator costs into an approximate triangle with total height P ∆m and width P ∆t.Source: A. Activation Cost Model · p. 5
  • Pipeline warmup stacks forward-batch activations, steady state interleaves forward and backward batches around a plateau, and cooldown releases activations.The iteration-level model shows that severe MoE imbalance during steady-state forwards can push memory beyond device capacity; cooldown follows the peak.Source: A. Activation Cost Model · p. 5, Figure 5 · p. 5

B. Activation Mitigation Cost

Activation mitigation trades memory savings for extra execution time, while MoE routing makes activation costs fluctuate across micro-batches. BaGuaLu-2 uses runtime provisioning to keep memory within capacity and reduce execution time.

B. Activation Mitigation Cost

  • Recomputation, offloading, and quantization reduce forward activation memory while adding execution-time costs at different points in the forward and backward passes.Recomputation regenerates discarded activations during backward; offloading transfers them to CPU and reloads them; quantization stores them at lower precision and casts them back.Source: B. Activation Mitigation Cost · p. 5
  • A roofline-style model compares each mitigation strategy’s memory reduction rate against target activation size to guide time-minimizing choices.Compute-bound operators usually favor offloading, while memory-bound operators such as layer normalization usually favor recomputation.Source: Figure 6 · p. 6

V. MODEL-GUIDED RUNTIME MEMORY PROVISIONING

  • For dense models, batches have identical activation costs, making a static mitigation strategy sufficient to prevent out-of-memory errors.Source: §V · p. 6
  • Dynamic token-to-expert routing makes MoE activation costs fluctuate, so fixed mitigation can be wasteful in ordinary batches or unsafe during spikes.The runtime objective is to choose per-operator mitigation while keeping peak memory within GPU capacity and minimizing total execution time.Source: §V · p. 6
  • BaGuaLu-2 applies dynamic mitigation to bring peak memory within capacity, then uses fused execution to hide mitigation operations behind useful computation.Source: §V · p. 6

A. Dynamic Activation Memory Mitigation

Dynamic activation mitigation assigns each micro-batch a memory budget and selects strategies to meet it with minimal added time. Deferred planning addresses unpredictable MoE spikes, while conflict checks prevent combinations that fail to save memory.

A. Dynamic Activation Memory Mitigation

  • Unpredictable MoE load is known only after the gate executes, so provisioning defers some decisions to subsequent micro-batches and can reduce the next batch’s quota or recompute the current MoE module.During pipeline steady state, the system can switch to full recomputation after the gate but before expert execution to avoid an unexpected peak.Source: A. Dynamic Activation Memory Mitigation · p. 6, Figure 7 · p. 7
  • Pipeline-aware budgeting distributes remaining capacity across warmup batches and assigns steady-phase batches memory freed by an earlier backward pass.Budgets are updated to actual consumption so one batch’s over- or under-estimation does not cascade to later batches.Source: A. Dynamic Activation Memory Mitigation · p. 7
  • Given a batch budget, the greedy algorithm repeatedly applies the available mitigation with the largest current memory reduction until the budget is satisfied.It reevaluates reductions after each choice because one strategy can change the benefit of later strategies.Source: A. Dynamic Activation Memory Mitigation · pp. 7–8
  • Mitigation strategies can conflict when consumers of the same activation retain inconsistent precision or device placement, leaving duplicate copies or preventing GPU memory from being freed.Examples include separate FP32 and FP8 copies, or an offloaded tensor whose GPU copy remains saved by another consumer.Source: A. Dynamic Activation Memory Mitigation · p. 8
  • After each candidate strategy, the system checks whether consumers of a shared source tensor agree on precision and placement, rejecting conflicts rather than undoing a more beneficial prior choice.Source: A. Dynamic Activation Memory Mitigation · p. 8

B. Fused Dual-Batch Execution

Fused dual-batch execution reorders independent forward, backward, offload, and recovery operations to overlap work and reduce execution time. A dynamic program finds the schedule, with compaction and caching reducing its runtime cost.

B. Fused Dual-Batch Execution

  • During pipeline steady state, BaGuaLu-2 plans the order of forward, backward, offloading, and recovery operators to improve overlap and minimize total execution time.Dependencies form a chain within each operator type, while the four types are independent across different batches.Source: B. Fused Dual-Batch Execution · p. 8
  • The dynamic program defines each state as the minimum time for prefixes of the four operator types and records optimal transfers to reconstruct the execution order.Its stated time complexity is O(|F|^2|B|^2|O|^2|R|^2).Source: B. Fused Dual-Batch Execution · p. 8
  • Stage compaction reduces DP states, while caching transfer points by discretized memory budget exploits temporal locality and makes amortized overhead nearly zero in most iterations.An MoE transformer forward batch’s more than 20 operators can be compacted into 4 stages.Source: B. Fused Dual-Batch Execution · p. 9
  • The reordering preserves each operator class’s program order, and pipeline barriers protect cross-batch dependencies; the method requires pipeline-parallel degree greater than two.The evaluated configurations use degree at least 4.Source: B. Fused Dual-Batch Execution · p. 9

A. Setup

BaGuaLu-2 outperforms both baselines across 128–12,288 GPUs, while evaluation uses shared parallelism configurations to compare activation-memory management. The setup uses a 1,040-billion-parameter model at the largest scale.

A. Setup

  • For a fair comparison, all systems share the same offline-optimal parallelism configuration at training start, so differences in memory pressure, stability, and performance stem from activation-memory management under that configuration.Evaluations use interleaved 1F1B scheduling on a cluster with 12,288 GPUs.Source: A. Setup · p. 9
  • Megatron-Kuai is evaluated with MoE-specific memory-model terms, while only AdaPipe’s adaptive recomputation is enabled.For baseline runs, OOM crashes prompt manual adjustment of the reduction strategy and a training restart.Source: A. Setup · p. 9
  • 1.14× and 1.13× average throughput gains over Megatron-Kuai and AdaPipe, respectively, with 2.012 EFLOPs sustained at 12,288 GPUs.The average MFU across scales is 32.9%.Source: A. Setup · p. 10, Figure 8 · p. 10

C. Robustness under Varying Memory Pressure

BaGuaLu-2 maintains its advantage under both runtime imbalance spikes and changes in workload configuration. Its MFU remains highest across varying MoE sparsity and degrades more gracefully as sequence length grows.

C. Robustness under Varying Memory Pressure

  • 1.28× overall average speedup over baselines during imbalance spikes, when Megatron-Kuai’s MFU drops 21.4% and AdaPipe’s drops 18.2% on average.Spikes occur in approximately 2.9% of iterations; BaGuaLu-2 transitions between offloading and selective strategies without switching overhead.Source: C. Robustness under Varying Memory… · p. 10
  • BaGuaLu-2 achieves the highest MFU in all five varying-sparsity configurations, averaging 1.11× speedup over the best baseline.Its advantage grows with activation pressure as more parameters are activated.Source: C. Robustness under Varying Memory… · pp. 10–11
  • All systems’ MFU decreases as sequence length rises from 4K to 8K and 16K, but BaGuaLu-2 degrades more gracefully.Longer sequences increase per-token activation footprint and total micro-batch memory demand; fine-grained mitigation avoids the baselines’ all-or-nothing overhead.Source: C. Robustness under Varying Memory… · p. 11

D. Ablation Study

BaGuaLu-2’s optimizations contribute cumulative speedups, while its activation-memory predictions remain accurate across configurations. Runtime overhead is small relative to iteration time.

D. Ablation Study

  • 1.16×, 1.05×, and 1.13× gains come respectively from static memory optimization, dynamic activation mitigation, and runtime overlap scheduling.Together, the three optimizations yield the full BaGuaLu-2 performance.Source: D. Ablation Study · p. 11
  • 1.56% mean prediction error across configurations for BaGuaLu-2’s activation-memory model.Runtime decisions depend on accurately modeling the actual memory footprint.Source: D. Ablation Study · p. 11
  • Less than 10 ms of total runtime overhead per iteration, compared with more than 18 seconds of training time.Temporal locality lets most iterations hit the cache and skip the algorithms, reducing amortized overhead to near zero.Source: D. Ablation Study · p. 11

VII. RELATED WORK

Prior systems address MoE computation, communication, or training memory, but the cited approaches do not account for activation-memory variation from fine-grained MoE routing. BaGuaLu-2 introduces continuous per-batch activation-memory control and reports large-scale training results.

VII. RELATED WORK

  • Prior MoE systems optimize expert parallelism or token-dispatch communication, while memory-efficient systems target static graphs or parameters and optimizer states, leaving dynamic MoE activation footprints unaddressed.Related work includes SmartMoE, FlexMoE, Lina, FasterMoE, Janus, and activation-memory systems,, –.Source: §VII · p. 11

VIII. CONCLUSION

  • BaGuaLu-2 provides continuous, per-batch activation-memory control and scales to 12,288 GPUs on a 1,040-billion-parameter model, sustaining over 2 EFlops at 35% MFU.Average throughput gains are 1.14× over Megatron-Kuai and 1.13× over AdaPipe.Source: §VIII · p. 11
Loading 2610.23064v1…