LRMoE: Demand-Side Load Re-Calibration for Efficient MoE Training on Memory-Constrained Nodes

Jiaming Hu; Zhida Jiang; Zhaolong Xing; Yu Li; Zhen Chen; Ke Zhang

cs.AIOct 27, 2026

LRMoE: Demand-Side Load Re-Calibration for Efficient MoE Training on Memory-Constrained Nodes

Authors: Jiaming Hu, Zhida Jiang, Zhaolong Xing, Yu Li, Zhen Chen, Ke Zhang

Organizations: JD.com, Beijing, China

Abstract

Sparse Mixture-of-Experts (MoE) models offer scalable capacity for Web applications, but uneven routing creates computation and communication bottlenecks. Existing placement and replication methods redistribute this workload, yet additional expert copies compete with token-related buffers for limited GPU memory. We present LRMoE, a memory-aware MoE training system that jointly recalibrates backward demand and balances the retained workload. During forward execution, it plans retention quotas and expert placement under resource constraints while preserving all routed computations. During backpropagation, a load recalibration metric estimates each token-expert pair's contribution to the expert update and selects pairs within these quotas. A capacity-aware routing algorithm distributes retained pairs among eligible instances of their original experts. Physical compaction then forms smaller communication payloads and variable-size grouped GEMMs, reducing backward traffic, computation, and execution-buffer demand. Experiments on Qwen3-30B-A3B and GLM4.5-106B-A12B show that LRMoE reduces cumulative MoE-layer training time to fixed-loss targets by 21.2-44.8% relative to SYMI.

Figures & tables

Figure 1. Uncalibrated load. All layers exceeds the maximum capacity for load balancing. Per-layer routed workload compared with the maximum load-balancing capacity.

Figure 2. Selective backward load calibration for uncalibrated load (a) and calibrated load (b). In backward, LRMoE drops low-value token-expert pairs and reassign, ensuring only retained pairs with load fitting within per-GPU capacity enter All-to-All, eliminating congestion. Comparison of uncalibrated and calibrated backward workloads. Low-value token–expert contributions are removed, and retained pairs are reassigned to fit per-GPU capacity.

Figure 3. Two-phase Calibration workflow. Forward plans expert placement and retention quotas under resource constraints. Backward uses LRM scores to select retained token-expert pairs and refines the quotas to satisfy load and resource budgets. Workflow with full GPU forward execution, quota planning, contribution-based backward selection, compact computation, and CPU offload.

Figure 5. PPL comparison for 10000 steps training under different recalibration rate. Both models trained with Alpaca SFT dataset, EP=8, sequence length = 8192, global batch = 16. Ck denotes omitting k% of token–expert backward contributions before backward Combine All-to-All. Perplexity curves over 10,000 training steps for two models at different backward recalibration rates.

Model	Loss	Megatron-LM	SYMI	Ours
Qwen3-30B-A3B	0.05	50.1 / 6,359	48.7 / 6,359	28.0 / 6,630
	0.02	66.4 / 8,440	64.6 / 8,440	36.7 / 8,693
GLM4.5- 106B-A12B	2.66	22.8 / 2,898	22.2 / 2,898	17.5 / 4,151
	2.63	42.4 / 5,388	41.3 / 5,388	22.8 / 5,402

Table 1. Measured cumulative MoE-layer training time to each target loss (minutes / step), including complete forward and backward execution. Bold marks the minimum time in each row.

Figure 6. Effective backward throughput under increasing batch-size pressure (the base batch-size is 512 per-GPU). The vertical line marks the point at which even a perfectly balanced, uncompressed backward workload exceeds the modeled memory budget. Effective backward throughput versus batch-size pressure, with a marked modeled memory-capacity boundary.

Figure 7. GPU-memory breakdown. Bars show the breakdown of each MoE layer. Stacked bars compare memory used by base expert placement, extra replicas, and token buffers across methods.

Figure 8. Measured latency per MoE-layer. Bars show MoE compute, dispatch All-to-All, and combine All-to-All latency breakdown. Stacked bars compare MoE computation, dispatch All-to-All, and combine All-to-All latency across methods.

Metric	Qwen3-30B-A3B	GLM4.5- 106B-A12B
$\alpha_{u}^{2}\lVert g_{t}\rVert_{2}^{2}\lVert\phi_{u}\rVert_{2}^{2}$ (Ours)	0.012474	0.003464
L-2 Norm for Grad $\lVert g\rVert_{2}^{2}$	0.097654	0.011987
Absolute Value for Grad $\|g\|$	0.107485	0.016662
Random Drop	0.775797	0.079282

Table 2. Mean

\Delta

PPL over 10000 steps.

Mixture-of-Experts Models

Mixture-of-experts models use input-dependent routing to activate specialized subnetworks, increasing model capacity without running every parameter for each input.

See all 462 papers

Memory-Efficient Expert Routing for Distributed MoE Training
Oct 5, 2026Arnab Kanti Tarafder, Jaume Guasch-Martí, Gokcen Kestor +1Expert Routing Mixture-of-Experts Models
William & Mary · Barcelona Supercomputing Center
EMO: Frustratingly Easy Progressive Training of Extendable MoE
May 13, 2026Linghao Jin, Chufan Shi, Huijuan Wang +4Sparse Mixture-of-Experts Efficient Language Model Training
USC-ISI · MBZUAI-IFM
ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning
May 9, 2026Chao Jin, Xinming Wei, Yinmin Zhong +6RL for Language Models Expert Load Balancing
School of Computer Science, Peking University · Independent Researcher

Oct 5, 2026J/K move · Enter open · S save

Memory-Efficient Expert Routing for Distributed MoE Training

Arnab Kanti Tarafder, Jaume Guasch-Martí, Gokcen Kestor, Jie Ren

William & Mary · Barcelona Supercomputing Center

Expert Routing Mixture-of-Experts Models Memory-Efficient Training Communication-Efficient Distributed Training

Open paper AskPDF arXiv

As Mixture-of-Experts (MoE) models scale toward hundreds of experts and higher top-

k

routing, memory efficiency in distributed training becomes a critical bottleneck. Peak memory is dominated by the MoE block, not attention: every intermediate buffer in the MoE dispatch pipeline is individually scaled by top-k routing. The standard all-to-all dispatcher sends all routed tokens in a single collective step, requiring the full top-

k

-expanded buffer to be constructed at once. In this work, we propose RelayMoE, a ring-based MoE execution model that computes locally as expert weights or tokens circulate, avoiding full top-

k

-expanded dispatch buffers. RelayMoE selects between expert and token routing according to communication volume and overlaps transfers with computation. The ring structure naturally supports memory-efficient MoE recomputation during backward: each hop reconstructs expert intermediates, uses them to compute gradients, and releases them before the next hop. The saved memory supports longer sequences and larger batches, or retains more attention activations to reduce attention recomputation and improve training throughput. We evaluate RelayMoE on 30B

-

57B production MoE models and varied expert configurations. In single-layer MoE experiments, RelayMoE achieves a

2\times

average speedup over Megatron-LM. In full-model training under the same GPU memory budget, it improves throughput by up to

2.02\times

and extends the largest tested trainable sequence length by up to

2.85\times

Abstract

Mixture-of-Experts Models

Explore similar work

Memory-Efficient Expert Routing for Distributed MoE Training

EMO: Frustratingly Easy Progressive Training of Extendable MoE

ReLibra: Routing-Replay-Guided Load Balancing for MoE Training in Reinforcement Learning