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

arXiv:2610.37720v1 [cs.AI] 27 Oct 2026

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

Zhida Jiang∗

Zhaolong Xing

JD.com Beijing, China

JD.com Beijing, China

JD.com Beijing, China

Yu Li

Zhen Chen∗

Ke Zhang

JD.com Beijing, China

JD.com Beijing, China

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.

CCS Concepts • Networks → Network performance analysis; • Computing methodologies → Machine learning algorithms.

Keywords mixture-of-experts, expert replication, load balancing, topologyaware scheduling, distributed training

1

Introduction

Web-scale search and recommendation services rely on learned models [1, 9]. Training these models is an infrastructure concern: finite GPU memory and interconnect bandwidth constrain the workloads that a cluster can support. Sparse Mixture-of-Experts (MoE) architectures increase parameter capacity while activating only a small subset of experts per token [2, 5, 13, 21, 25]. For MoE-based Web services, this shifts part of the scaling challenge from dense computation to managing routed work across devices. We study this Web-infrastructure problem at the training-system level: how to coordinate expert placement, communication, and backward workload under a fixed per-GPU resource budget, rather than improve a particular search or recommendation model. ∗ Zhida Jiang and Zhen Chen are co-corresponding authors.

Expert Parallelism (EP) distributes experts across GPUs and exchanges routed token–expert pairs through all-to-all communication [6, 13, 21]. Each pair contributes to forward and backward communication, expert computation, and saved-activation storage. Input-dependent routing can concentrate work on GPUs hosting popular experts, delaying synchronous execution [35]. Replication distributes this work but consumes memory and introduces statetransfer and gradient-synchronization costs. Even balanced placement can therefore leave more backward work than the available memory and communication capacity can accommodate. The resulting resource-management problem concerns both where routed pairs execute and how much backward work is retained [28, 35]. Existing systems adapt expert placement and execution to uneven loads. FasterMoE shadows popular experts, FlexMoE adapts placement, and SmartMoE searches hybrid parallelization plans [8, 19, 31]. SYMI decouples model parameters from optimizer states, while LAER-MoE reconstructs expert instances from sharded parameters [17, 26]. UltraEP responds to realized postgating loads, and TopoEP incorporates communication topology into planning [28, 35]. These techniques improve utilization, but replication still consumes expert storage and incurs state-transfer and gradient-synchronization costs. For a fixed admitted token– expert pair set, placement and replication change where work executes without reducing its logical computation. Replicas also compete with token buffers for memory. Under a constrained budget, redistributing that workload can therefore leave total demand above the available service capacity. Selective-Backprop reduces backward work by prioritizing highloss training examples [10]. We instead select routed token–expert pairs by their estimated contributions to the expert update. Once an MoE layer receives its incoming gradients, these can be combined with saved activations and routing weights to characterize pairwise contributions [7]. At a protected expert projection, a pair’s weight-gradient contribution is the outer product of its weighted incoming gradient and saved projection input. We derive the Load Recalibration Metric (LRM) from the squared norm of this contribution, allowing pairs to be ranked without materializing individual weight-gradient matrices. LRM is a projection-level proxy for selecting backward contributions under a limited budget. Because these gradients are unavailable during forward execution, we preserve all routed forward computations and defer contribution-aware selection until backpropagation. MegaBlocks addresses irregular expert execution with block-sparse computation [6]. Our selection must likewise translate into smaller physical workloads. Removing work

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

from an unconstrained path does not accelerate execution, and reducing total work without rerouting can leave imbalances intact. The selected backward update remains approximate, so its effect on learning must be evaluated empirically rather than inferred from pairwise scores alone. We present LRMoE, a memory-aware MoE training system that combines contribution-aware selection with resource-constrained quota planning and load routing. Its two-phase design separates how much backward work the system can afford from which pairs fill that budget. After forward gating reveals the routed load, the controller sets retention quotas for each expert and plans replica placement. These decisions account for full forward execution and expected backward work within memory and communication limits. The full forward pass then preserves the context needed for backpropagation. Once gradients arrive, LRM selects the highestscoring pairs within each expert’s quota. Capacity-aware routing assigns retained pairs to eligible instances, and physical compaction groups their gradients and saved activations by destination and expert. The resulting dense payloads reduce backward all-to-all traffic. Variable-size grouped GEMMs use the actual retained row counts without padding workloads across experts. An inverse mapping restores gradients to the layout expected by surrounding layers. This coordination reduces backward communication, expert computation, and execution-buffer demand, while requiring selection and packing overhead to be included in evaluation. Our contributions are: • We introduce a contribution-aware calibration metric and reduce backward communication and computation through grouped-GEMM workloads. • We develop memory-aware routing that re-balances retained workloads across eligible expert instances. • On Qwen3-30B-A3B and GLM4.5-106B-A12B, LRMoE reduces cumulative MoE-layer training time to fixed-loss targets by 21.2-44.8% versus state-to-the-art baselines, a separate GPU replay shows a 57.8% reduction in the subtotal of expert-placement, replica, and token-buffer memory relative to fixed ownership.

2 Background and Motivation 2.1 MoE and Load Balancing In Transformer models [27], a Mixture-of-Experts (MoE) layer replaces a dense feed-forward block with experts and a learned router. For each token, the router activates only a small subset of experts, increasing model capacity without executing every parameter [5, 13, 25]. Expert parallelism (EP) partitions these experts across GPUs: it dispatches each routed token to the GPU hosting its selected expert and combines the returned expert outputs [6, 21]. We call one token–expert selection a routed pair. Routing makes the workload data dependent. Pair counts change across inputs, layers, and iterations, producing unequal expert computation and all-to-all traffic across GPUs. Because a collective and an MoE layer finish at the pace of their slowest participant, the busiest GPU becomes a straggler while others wait [8, 19, 26, 28, 31]. Training-time routing regularizers encourage balance in expectation, but cannot ensure balance for every realized layer invocation. Runtime balancing is therefore still required for the instantaneous load. BASE layers use

balanced token assignment, whereas Expert Choice fixes expert bucket sizes [14, 34]. Gating Dropout reduces communication by bypassing global routing for some tokens [15]. These approaches modify forward routing, while our selection operates on backward contributions.

2.2

Limitations of Existing Methods

Production MoE training must accommodate model parameters, gradients, optimizer states, and runtime buffers within a fixed GPU budget. When many experts share each GPU, persistent training states leave limited memory for additional expert instances and transient execution buffers. Offloading moves training states to host memory [12, 24]. Input-dependent routing further concentrates demand on selected experts, so sufficient aggregate resources do not necessarily imply that every device can serve its assigned workload. Existing systems provide spatial flexibility through expert placement, replication, and adaptive parallelization [8, 19, 31]. EfficientMoE adapts expert scheduling to predicted loads [30], while NetMoE rearranges samples to exploit routing locality [16]. SYMI and LAER-MoE reduce state-management costs [17, 26], while UltraEP adapts to realized post-gating loads [28]. These mechanisms can improve utilization and communication locality, but they preserve the logical expert computation required by a fixed admitted token–expert pair set. Their ability to accommodate hot experts therefore remains constrained by the capacity available for executing that work. Consider a common backward service window. Let 𝑊𝑛ℓ𝑒 denote expert 𝑒’s backward demand for layer invocation bwd denote GPU 𝑔’s guarded service capacity within (𝑛, ℓ), and let 𝐶𝑛ℓ𝑔 that window. Under an additive model in which demand and capacity are expressed in the same calibrated work units, a necessary Í Í bwd feasibility condition is 𝑒 𝑊𝑛ℓ𝑒 ≤ 𝑔 𝐶𝑛ℓ𝑔 . Violating this bound makes the fixed workload infeasible within the window, regardless of redistribution. Satisfying it is not sufficient, however: legal expert destinations are restricted, replica and staging capacity is bounded, and per-device memory and network-path constraints must also be satisfied. A placement optimizer may therefore fail to find a feasible assignment even when aggregate service capacity is sufficient. Aggregate overload and spatial infeasibility are distinct limitations, neither can be resolved merely by making GPU loads more uniform.

2.3

Key Insight

Our key insight is to treat backward demand as a contributionaware scheduling variable rather than a fixed consequence of forward routing. Calibration replaces demand 𝑊𝑛ℓ𝑒 with retained dee𝑛ℓ𝑒 ≤ 𝑊𝑛ℓ𝑒 , prioritizing estimated update contributions mand 𝑊 under constrained capacity. Load routing distributes retained work among available expert instances. These decisions control which backward work consumes capacity and where it executes. Incoming gradients, saved activations, and routing weights estimate each pair’s contribution to the expert update [7]. The selected update is approximate, so its convergence impact must be evaluated. Selection should target overloaded expert paths, since reducing unconstrained work need not shorten execution. Retained pairs must

LRMoE

Peak GPU Load

MILP: 48/48 exceed the limit

MILP-balanced

Load (pairs/GPU)

800k

600k

400k GPU limit = 262k

200k

0

0

4

8

12

16

20

24

28

MoE layer

32

36

40

44

Figure 1: Uncalibrated load. All layers exceeds the maximum capacity for load balancing.

reach eligible instances of their original experts. Physical compaction reduces communication payloads and expert-GEMM token dimensions while preserving token–expert correspondence. Masking full-sized tensors alone need not yield these savings. Net benefits must account for scoring, packing, and routing overhead and the execution efficiency of the reduced workloads.

3 LRMoE Design 3.1 Design Overview Figure 2 presents the design of LRMoE, which combines layer-wise load calibration, replica-aware load routing, and compact groupedGEMM execution. The workflow separates forward-time resource planning from backward-time contribution selection. The model router first determines each token’s logical experts and routing weights, while the placement plan identifies their physical execution instances. After resource admission, the complete forward pass dispatches every routed token–expert pair, executes the expert computations, and combines their outputs into the dense representation passed to the next layer. The runtime preserves the activations, required expert intermediates, routing weights, and pair identities, together with the forward executor and parameter version, establishing the correspondence needed for subsequent backward execution. Before expert forward begins, the controller observes the routed demand 𝑛𝑒 of each expert and determines the backward retention quotas 𝐾𝑒 under the admitted placement and guarded resource capacities. Quota planning accommodates as much backward demand as the declared capacity model permits, while accounting separately for the complete forward workload, saved-state residency, replica Í storage, and execution buffers. With 𝑁 = 𝑒 𝑛𝑒 , the retention rate Í is 𝜌 = 𝑒 𝐾𝑒 /𝑁 , expert 𝑒 omits 𝑛𝑒 − 𝐾𝑒 backward contributions. Quotas remain fixed for the invocation. At backward Combine, LRM uses incoming gradients and saved state to select the highest-valued 𝐾𝑒 pairs per expert, without another quota search. Load routing determines how these retained contributions occupy the admitted expert instances. During forward planning, a capacity-constrained integer allocation assigns 𝑋𝑒𝑔 retained pairs Í of expert 𝑒 to GPU 𝑔, with 𝑔 𝑋𝑒𝑔 = 𝐾𝑒 and nonzero assignments permitted only at eligible instances of that expert. The planner minimizes the peak assigned GPU load for the fixed placement

rather than dividing each expert’s demand equally among replicas. During backward, rerouting realizes these instance-level quotas using the selected pair identities, preserving their logical experts and assigning every retained pair exactly once. Eligibility requires the correct expert parameters and access to the matching forward state at the recorded version. Because the selected identities are unknown during planning, the admitted state-access and communication envelope must remain valid for any selected set satisfying the quotas. Backward realization therefore neither introduces new replicas nor changes the previously admitted workload. The compact runtime translates the retention mask and routing assignments into physical communication and computation. It gathers retained gradient rows and their matching activations and intermediates in a consistent pair order, grouping them by destination and expert instance. Counts and offsets define compact All-to-All payloads and variable-size grouped GEMMs. FC2, the activation function, FC1, and backward Dispatch use the same retained identities, omitted rows are not processed in full-sized masked tensors. Returned input-gradient contributions are scatter-added at their original token and routing-slot positions in zero-initialized dense buffers. This restores the layout, not discarded contributions, router gradients follow a separate path. Replica parameter gradients are summed at logical owners before the required accumulation and distributed reductions, one optimizer update per logical expert, and weight synchronization.

3.2

Load Recalibration Metric

We derive a pairwise metric that quantifies the local effect of omitting an already-routed token–expert pair from the protected backward update. Recalibration keeps the logical expert assignment and routing weight fixed, and determines which contributions remain active during backward communication and computation. To compare these contributions, we first identify the parameter-gradient contribution associated with each pair, then measure the perturbation caused by removing it. This connects the selection criterion directly to the protected expert-weight gradient rather than using routing strength or incoming gradient magnitude alone. For a pair 𝑢 = (𝑡, 𝑒), the chain rule expresses its contribution to the protected projection gradient as the outer product 𝐺𝑢 = (𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ , where 𝑔𝑡 is the incoming gradient before multiplication by the routing weight and 𝜙𝑢 is the saved projection input. Omitting this pair while leaving the other contributions unchanged produces the perturbation Δ𝑢 = −𝐺𝑢 . We measure its magnitude using the squared Frobenius norm. Because 𝐺𝑢 is rank one, this norm factorizes into the squared norms of the weighted gradient and saved input, yielding a score that combines routing strength, backward-signal energy, and expert-state magnitude without materializing the pairwise gradient matrix. The following proposition formalizes this projection-level calibration metric. Proposition 3.1 (Token-expert calibration metric). Consider a completed forward route and a protected expert projection with weight matrix 𝑊𝑒 . For a routed token-expert pair 𝑢 = (𝑡, 𝑒), let 𝛼𝑢 denote its fixed routing weight, let 𝑔𝑡 denote the gradient arriving at the routed expert output, and let 𝜙𝑢 denote the saved input to the protected projection. The contribution of pair 𝑢 to the weight gradient

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

(a) Forward congestion (unbalanced) Backward gradient dY

Pair importance

? ? ?

? ? ? ?

Not available yet Token owners

(b) Backward: selective drop + balance

Retain every pair

Backward gradient dY

Pair importance

Available

Ranked by contribution

Cannot rank by contribution No forward compression

GPU 0

GPU 1

GPU 2

Token owners

GPU 0

GPU 1

Value-aware token drop Select → compact → send

From GPU (token owners)

All-to-All

From GPU (token owners)

To GPU (expert

Hot-expert overload

executors)

GPU 2 / E2

GPU 1 / E1

Balanced within physical capacity GPU 2 / E2 GPU 0 / E0 GPU 1 / E1

3 queued

Drop 2

Exceed

memory capacity

Per-GPU

Drop 2

Only retained pairs

All-to-All

GPU 0 / E0

Keep high-value work

GPU 2 Choose drop to meet capacity

Send all pairs

To GPU (expert executors)

Drop low-value pairs

+ E1 replica

Rebalance 1

Per-GPU

capacity

capacity

3 ranks

6 ranks (3 active + 3 queued)

Retained workload

2 ranks

Dropped workload

3 ranks Reassigned workload

Physical capacity

6 − 2 − 1 = 3 ranks

2 + 1 = 3 ranks

Drop 2 + rebalance 1: (3, 6, 2) → (3, 3, 3)

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. is 𝜕L𝑢 𝐺𝑢 = = (𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ . (1) 𝜕𝑊𝑒 If pair 𝑢 is omitted while all other routed pairs remain unchanged, the protected weight gradient changes by Δ𝑢 = −𝐺𝑢 = −(𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ .

(2)

We therefore define the pairwise calibration score as the squared Frobenius norm of this perturbation: 𝑅𝑢 := ∥Δ𝑢 ∥ 2𝐹 2

= (𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ 𝐹 = 𝛼𝑢2 ∥𝑔𝑡 ∥ 22 ∥𝜙𝑢 ∥ 22 .

(3)

• Routing strength 𝛼𝑢2 . This term measures how strongly token 𝑡 was assigned to expert 𝑒 during the completed forward pass. Because routing is already fixed, 𝛼𝑢 is treated as a constant during recalibration. • Backward signal ∥𝑔𝑡 ∥ 22 . This term measures the magnitude of the gradient propagated from the subsequent computation. It distinguishes routed pairs that receive strong optimization signals from pairs that contribute little to the current update. • Expert-state magnitude ∥𝜙𝑢 ∥ 22 . This term measures the magnitude of the saved projection input. Two pairs with similar routing weights and incoming gradients may still have different effects on the expert update because their saved expert states differ. Equation 3 follows from the rank-one identity ∥𝑎𝑏 ⊤ ∥ 2𝐹 = ∥𝑎∥ 22 ∥𝑏 ∥ 22 . Thus, 𝑅𝑢 equals the squared change in the protected projection’s weight gradient when pair 𝑢 is omitted and all other contributions remain unchanged. A larger score identifies a pair whose omission induces a larger isolated gradient perturbation. The projection-level factorization and detailed derivation are

provided in Appendix 7.1. LRMoE separates how much backward demand can be admitted from which contributions should be retained. During forward planning, the resource controller chooses Í the largest feasible retained count 𝐾ℓ = 𝑒 𝐾ℓ,𝑒 from its tested candidates under the quota-allocation policy, HBM budget, replicaplacement constraints, and communication window, then fixes the per-expert quotas {𝐾ℓ,𝑒 }. Once backward gradients become available, the metric ranks pairs within each expert, and the 𝐾ℓ,𝑒 highest-scoring pairs fill its quota. For fixed per-expert quotas and equal row costs, this selection minimizes the additive sum of isolated omission perturbations, as established in Appendix 7.2. The selected gradient rows and their matching saved activations are physically compacted before backward communication and grouped GEMMs, turning contribution selection into a smaller workload for load routing and execution. Merely masking full-sized tensors would not reduce their communication payloads or matrix dimensions. The selection guarantee concerns the additive local calibration objective, rather than the exact aggregate gradient error or subsequent training behavior, these distinctions and the applicable assumptions are detailed in Appendices 7.3 and 7.4.

3.3

Layerwise Load Calibration

Figure 3 details the quota admission and execution procedure in Algorithm 1. For layer ℓ, the controller counts the full routed load 𝑁 ℓ,𝑒 of each expert and allocates a candidate retained count 𝐾ℓ Í into quotas satisfying 0 ≤ 𝐾ℓ,𝑒 ≤ 𝑁 ℓ,𝑒 and 𝑒 𝐾ℓ,𝑒 = 𝐾ℓ . It jointly evaluates these quotas and a proposed replica placement against the load-balancing target, GPU-memory budget, and communication window. Candidates are tested in decreasing order, starting with full retention. The first feasible candidate fixes the quotas and placement Πℓ𝐹 , feasibility includes the complete forward workload

LRMoE

Algorithm 1 Resource-Constrained Layer-wise Load Calibration Ð Require: Routed pairs Pℓ = 𝑒 Pℓ,𝑒 , forward placement Πℓ𝐹 , candidate rates R ℓ , available HBM {𝑀𝑟avail }, bandwidth 𝐵 eff , window ovlp 𝑇ℓ Ensure: Retained pairs Sℓ∗ and backward plan Π𝐵ℓ Forward-time admission 1: 𝑁 ℓ,𝑒 ← |Pℓ,𝑒 | for each expert 𝑒 2: for 𝜌 ∈ R ℓ in decreasing order do b 𝑇bcomm ) 3: ({𝐾ℓ,𝑒 }, Π, M, ← PlanQuotasAndReplicas({𝑁 ℓ,𝑒 }, 𝜌, Πℓ𝐹 , 𝐵 eff ) b𝑟 ≤ 𝑀𝑟avail for all 𝑟 , and 𝑇bcomm ≤ 𝑇 ovlp if Π is feasible, 𝑀 4: ℓ then ∗ ∗ 𝐴 5: (𝜌 ℓ , {𝐾ℓ,𝑒 }, Πℓ ) ← (𝜌, {𝐾ℓ,𝑒 }, Π) 6: break 7: end if 8: end for 9: if Π𝐴 ℓ is undefined then 10: return Infeasible 11: end if 12: Cℓ ← FullForwardAndSaveContext(Pℓ , Π𝐴 ℓ ) Backward-time calibration 13: 𝑣 𝑝 ← 𝛼 𝑝2 ∥𝑔𝑡 ∥ 22 ∥𝜙 𝑝 ∥ 22 for each 𝑝 = (𝑡, 𝑒) ∈ Pℓ ∗ ← TopK(P , {𝑣 }, 𝐾 ∗ ) for each expert 𝑒 14: Sℓ,𝑒 ℓ,𝑒 𝑝 ℓ,𝑒 ∗ }) ← PruneAndRoute( Ð S ∗ , Π𝐴 , C ) 15: (Π𝐵 , {S ℓ 𝑒 ℓ,𝑒 ℓ,𝑒,𝑟 ℓ ℓ ∗ }, C ) 16: {(𝐺 ℓ,𝑒,𝑟 , Φℓ,𝑒,𝑟 )} ← CompactAndAllToAll({Sℓ,𝑒,𝑟 ℓ 17: (Δ𝑊ℓ,𝑒,𝑟 , 𝐷 ℓ,𝑒,𝑟 ) ← PackedExpertBackward(𝐺 ℓ,𝑒,𝑟 , Φℓ,𝑒,𝑟 ) for each active (𝑒, 𝑟 ) Í 18: QCPU ← QCPU ∪ {(𝑒, 𝑟 Δ𝑊ℓ,𝑒,𝑟 )}𝑒 ∗ } 19: ScatterWithZeros(){𝐷 ℓ,𝑒,𝑟 }, {Sℓ,𝑒,𝑟 Ð ∗ ∗ 𝐵 20: return ( 𝑒 Sℓ,𝑒 , 𝜌 ℓ , Π ℓ )

GPU forward

Count routed load

Limits: HBM and Bandwidth

Plan quotas & replicas

Admit & materialize

Full forward & save context

Route & compact

GEMM & reduce

Lower ρ

GPU backward Compute LRM scores

CPU offload

Select Top-K

Queue & accumulate gradients

Update host-memory optimizer states

Figure 3: Two-phase Calibration workflow. Forward plans expert placement and retention quotas under resource constraints. Backward uses LRM scores to select retained tokenexpert pairs and refines the quotas to satisfy load and resource budgets.

and its required saved state, not only the reduced backward demand. The runtime then materializes the admitted replicas, executes the full forward pass, and records pair identities, parameter versions, and backward context. If no candidate is feasible, admission fails.

1 Mask

2 Pack

dY

dY

e0 e2 e0 e1 e2 e0 e1 e2 e0

X, φ dropped

3 Counts

X, φ

4 Grouped GEMM dY e0: M = 3

e0

R = [3, 1, 2]

e1

offset = [0, 3, 4, 6]

e1: M = 1

dropped dropped

X, φ

Grouped GEMM

Variable M No padding rows

e2

Me = R e

e2: M = 2

Figure 4: Expert replication reduction process. Dropped token-expert backward contributions are removed before expert-backward computation, while retained gradients and saved expert states are packed into contiguous expert segments. Their counts and offsets define variable-𝑀 grouped GEMMs without padding. The reduced per-expert row counts lower both expert-backward computation and the number of physical expert replicas required to satisfy the capacity bound.

These admitted quotas remain fixed when backward begins. Using the incoming gradients, saved activations, and routing weights, the controller computes the LRM score 𝑣 𝑝 and selects the 𝐾ℓ,𝑒 highest-scoring pairs within each logical expert. Per-expert selection preserves the planned load distribution, whereas an unrestricted layer-wide ranking could concentrate retained work on overloaded experts. Thus, forward planning determines the admitted counts, and backward LRM selection determines which pair identities fill them. The selected pairs then become the input to capacity-aware rerouting. Each pair is assigned exactly once to an eligible instance of its original expert within Πℓ𝐹 , using the retained workload and the hosting GPUs’ available capacity. If 𝑚 ℓ,𝑒,𝑖 denotes the count assigned to eligible instance 𝑖, routing preserves Í 𝑖 𝑚 ℓ,𝑒,𝑖 = 𝐾ℓ,𝑒 . A reassignment changes only the physical executor: the destination must provide the recorded parameter version and access to the matching saved context, with any required state movement included in the resource accounting. Rerouting therefore redistributes the selected workload without changing logical expert assignments or creating additional replicas. The resulting per-instance counts directly determine compact communication and computation. Retained gradient rows and their corresponding saved activations are packed in a consistent pair order, with counts and prefix offsets defining destination-specific payloads and expert-major tensor segments. Variable-size grouped GEMMs consume these segments: 𝑚 ℓ,𝑒,𝑖 determines the row dimension of input-gradient products and the reduction dimension of weightgradient products, while parameter dimensions remain unchanged. The inverse mapping restores input-gradient contributions to their original token positions, and parameter-gradient contributions are summed across instances of each logical expert before the optimizer update. In this way, the admitted quotas, selected identities, and routing assignments jointly determine the workload physically executed.

3.4

Repacking GEMM with Load Calibration

Figure 4 shows how selective dropping is converted into an actual reduction in expert computation. After identifying the retained token-expert pairs, we remove the dropped rows and repack the

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

C0 · 0%

4.2

Convergence Analysis

Figure 5 compares training perplexity under the uncompressed baseline (C0) and three LRM-guided recalibration settings: C10, C20, and C30, corresponding to nominal reductions of 10%, 20%, and 30% in the backward token-expert workload. Both models exhibit similar overall convergence trends. The late-training region reveals the differences between recalibration settings more clearly. For Qwen3-30B-A3B, the mean PPL over the final 1,000 steps is 1.015190, 1.015156, and 1.015141 for C0, C10, and C20, respectively. Meanwhile, C30 reaches 1.016270, approximately 0.106% above the C0 baseline. These observations confirm that recalibration up to C20 closely preserves the original training convergence, whereas

1.015

10

0

2k 4k 6k 8k 10k Training Step

C30 · 30%

(b) GLM4.5-106B-A12B 60 8k–9k 16 40 13 20

1

0

2k 4k 6k 8k 10k Training Step

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. 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. Model

Loss

Megatron-LM

SYMI

Ours

Qwen3-30B-A3B

0.05 0.02

50.1 / 6,359 66.4 / 8,440

48.7 / 6,359 64.6 / 8,440

28.0 / 6,630 36.7 / 8,693

2.66 2.63

22.8 / 2,898 42.4 / 5,388

22.2 / 2,898 41.3 / 5,388

17.5 / 4,151 22.8 / 5,402

GLM4.5106B-A12B

4 LRMoE Evaluation 4.1 Experimental Setup We evaluate MoE-layer execution on eight NVIDIA Hopper GPUs with EP8, PyTorch 2.11.0, CUDA 13.0, and BF16 precision. We use Qwen3-30B-A3B and GLM4.5-106B-A12B with a locally preprocessed Alpaca corpus, with text records concatenated and cyclically reused to construct long inputs. The sequence length is 8,192 tokens and the per-GPU microbatch size is four, yielding 32,768 source tokens per GPU. We conduct identical controlled tensors across compared methods, using 128 experts, top-8 routing, hidden dimension 𝐻 = 2048, and expert intermediate dimension 𝐹 = 768. We compare the Megatron-LM fixed-ownership strategy, native UltraEP, an independent implementation of SYMI, and LRMoE under the same 80 GiB per-GPU memory budget. We execute complete forward and backward passes, with 50 warm-up iterations followed by 1000 measured iterations.

C20 · 20%

PPL

15 5 1

C10 · 10%

(a) Qwen3-30B-A3B 8k–9k 1.029

20 PPL

corresponding gradients and saved expert states into contiguous, expert-specific buffers. The retained row counts and offsets are then passed to a variable-size grouped GEMM. Each expert therefore processes only its retained rows, without computing over zero-filled rows or padding the input back to its original size. This distinction is essential. Zero-masking the dropped rows while preserving the original tensor shape leaves the GEMM size and its memory requirement largely unchanged. Repacking instead reduces the row dimension of all four expert-backward GEMMs: the FC1 and FC2 data-gradient and weight-gradient operations. The resulting gradients are equivalent to those produced by the zeromasked implementation for the same retained pairs, apart from ordinary floating-point differences. After the expert computation, the retained results are scattered back to their original positions, and zeros are restored only before the next dense-layer boundary. Smaller packed GEMMs also reduce the physical capacity required for overloaded experts. Expert replication is provisioned to accommodate the peak number of rows assigned to each expert. Once selective dropping reduces this workload, fewer physical instances are needed to keep the per-replica load within the GPU capacity. In Figure 4, the original expert loads are reduced from [4, 2, 3] to [3, 1, 2]. An expert whose retained workload now fits on one GPU no longer requires an additional replica. The set of logical experts remains unchanged, our method reduces only the redundant physical replicas, along with their parameter memory, temporary buffers, and communication.

C30 introduces a negligible late-training gap on Qwen. This minor gap raises a complementary systems question regarding whether the substantial savings in MoE-layer execution time outweigh the cost of any additional optimization steps required to reach the exact same loss target. We address this question directly through the rigorous loss-matched comparison below.

4.3

Convergence Time Analysis to Target Loss

To relate the convergence behavior above to training efficiency, we compare measured cumulative MoE-layer training time at common target losses rather than at a fixed training step. This metric accumulates the time spent in complete MoE-layer execution, including forward and backward computation and communication, up to the first step at which each target loss is reached. It therefore measures the full MoE-layer training path rather than communication alone, but does not include execution outside the MoE layers or represent whole-model end-to-end training time. Taking the uncompressed trajectory as a reference, we use two target cross-entropy losses for each model: 0.05 and 0.02 for Qwen3-30B-A3B, and 2.66 and 2.63 for GLM4.5-106B-A12B. These thresholds sample two optimization levels after the initial rapid loss decrease. Both the uncompressed and C30 trajectories reach these targets within the observed training horizon. The thresholds provide common comparison points selected from Figure 5 for convergence. We primarily benchmark long-trajectory convergence against SYMI, as it represents the closest state-of-the-art framework decoupling expert training states. UltraEP primarily targets rack-scale

LRMoE

Effective throughput (%)

Megatron-LM

UltraEP

SYMI

Ideal full

Ours

100

Base placement

50 0

Token buffers 12.88

1

4

8.03

UltraEP

Full ideal limit 8 12 16 20 24 Batch-size multiplier

12.88

SYMI

5.44

Ours

0

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. communication primitives and hits restrictive memory boundaries under our sustained long-sequence training setup. Table 1 reports the cumulative MoE-layer training time and the step at which each target is first reached. LRMoE achieves the lowest cumulative time at all four targets. For Qwen3-30B-A3B, it requires 28.0 and 36.7 minutes, compared with 48.7 and 64.6 minutes for SYMI. For GLM4.5-106B-A12B, the corresponding times are 17.5 and 22.8 minutes, compared with 22.2 and 41.3 minutes. These results represent 21.2 to 44.8% lower cumulative MoE-layer training time relative to SYMI, despite LRMoE requiring more training steps to reach each target. The comparison thus accounts for both MoE-layer execution cost and the additional optimization steps associated with recalibration. While SYMI uses current-step expert loads to select placements, LRMoE actively reduces the backward token to expert workload and balances the retained pairs across eligible expert instances. The measured reductions show that these execution savings substantially outweigh the additional steps at the evaluated loss targets.

4.4

Extra replicas

Megatron-LM

Throughput Analysis

Figure 6 evaluates how memory pressure affects effective throughput as batch size increases. We define effective throughput as the fraction of the offered batch whose corresponding backward workload fits within the modeled per-GPU memory budget. The experiment assumes an 80 GiB GPU with a 64 GiB resident footprint, leaving 16 GiB for token buffers and additional expert replicas. Megatron-LM quickly loses throughput because its fixed placement concentrates tokens on overloaded GPUs. UltraEP and SYMI improve this placement by utilizing additional replicas, but their overall benefit remains strictly constrained by residual load imbalance and the memory required to host those replicas. The Ideal full curve illustrates an uncompressed backward workload with perfectly balanced token traffic, optimistically charging zero memory for additional replicas. It reaches the memory limit at a batch-size multiplier of 16×. Beyond this threshold, placement strategies alone cannot accommodate the full backward workload within the given memory model. In contrast, our method first systematically reduces the backward token working set and then dynamically balances the retained pairs, requiring only 7 to 9 additional replicas across eight GPUs. Consequently, it preserves a 100% effective throughput even at 20×, compared with 58.4% for UltraEP, 44.9% for SYMI, 40.0% for

2

4 6 8 10 Dynamic Memory (GiB / GPU)

12

14

Figure 7: GPU-memory breakdown. Bars show the breakdown of each MoE layer.

Megatron-LM, and 80.0% for the optimistic full-workload bound. This result directly characterizes the backward-capacity benefit of contribution-aware selection under the declared memory model.

4.5

GPU Memory Breakdown

Figure 7 separates base expert placement, additional replica storage, and token buffers. Their measured subtotal is 12.88 GiB for Megatron-LM, 8.03 GiB for UltraEP, 12.88 GiB for SYMI, and 5.44 GiB for LRMoE. Forward load balancing distributes the complete routed workload across admitted expert instances, while backward recalibration and rerouting reduce and redistribute the retained workload. Together with compact execution, these mechanisms reduce the reported token-buffer component from 12.46 GiB under fixed ownership and 7.55 GiB under UltraEP to 4.97 GiB under LRMoE. The three-component subtotal is therefore 57.8% lower than Megatron-LM and 32.2% lower than UltraEP. These reductions concern the displayed GPU allocations, not total device memory or host-memory optimizer state. These results demonstrate the concrete runtime advantages of our design. By combining contributionaware selection with physical compaction, LRMoE eliminates unselected rows prior to communication, directly trimming All-to-All transfer volumes. Furthermore, our variable-size grouped GEMMs execute compact per-instance workloads without incurring padding overhead. This substantial reduction along the critical GPU execution path forms the direct mechanical foundation for the overall training speedups reported in our end-to-end evaluations.

4.6

Latency Breakdown

Using the same workload and resource constraints, Figure 8 examines the computation and token-communication costs associated with these allocations. Dispatch transfers incoming gradients to expert executors, whereas combine returns input-gradient contributions to token owners. For each iteration, all three components are measured from the critical rank with the longest complete backward execution to preserve a consistent critical-path evaluation. Their mean subtotal is 29.64,ms for Megatron-LM, 16.41,ms for UltraEP, 30.11,ms for SYMI, and 12.69,ms for LRMoE. Specifically, LRMoE achieves an expert computation latency of 7.26,ms, a dispatch latency of 2.58,ms, and a combine latency of 2.85,ms.

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

Latency (ms)

MoE Compute 30 20 10 0

Dispatch All-to-All

Combine All-to-All

30.11

29.64 14.54

16.41

17.33

12.69

5.19

3.47 3.61

5.26

2.85 2.58

9.91

9.33

7.52

7.26

Megatron-LM

UltraEP

SYMI

Ours

Figure 8: Measured latency per MoE-layer. Bars show MoE compute, dispatch All-to-All, and combine All-to-All latency breakdown.

5

Table 2: Mean ΔPPL over 10000 steps.

Metric 𝛼𝑢2 ∥𝑔𝑡 ∥ 22 ∥𝜙𝑢 ∥ 22 (Ours) L-2 Norm for Grad ∥𝑔∥ 22 Absolute Value for Grad |𝑔| Random Drop

Qwen3-30B-A3B

GLM4.5106B-A12B

0.012474 0.097654 0.107485 0.775797

0.003464 0.011987 0.016662 0.079282

These results demonstrate the concrete runtime advantages of our design. By combining contribution-aware selection with physical compaction, LRMoE eliminates unselected rows prior to communication, directly trimming All-to-All transfer volumes. Furthermore, our variable-size grouped GEMMs execute compact perinstance workloads without incurring padding overhead. Consequently, LRMoE reduces the backward execution subtotal by 57.2% relative to Megatron-LM, 57.9% relative to SYMI, and 22.7% relative to UltraEP. This substantial reduction along the critical GPU execution path forms the direct mechanical foundation for the overall training speedups reported in Table 1.

4.7

In distributed Expert Parallelism, two routed pairs can receive nearly identical upstream gradient norms but possess entirely different initial routing weights or saved activation magnitudes. A gradient-only rule assigns them identical importance, even when their actual mathematical contributions to an expert weight projection differ drastically. Our LRM accurately and efficiently accounts for these diverse factors. Crucially, it achieves this comprehensive scoring without requiring the computationally prohibitive materialization of full per-pair weight-gradient matrices. The tight PPL comparison provides strong empirical validation that this efficient rank-one local approximation preserves global model convergence and training quality over extensive training horizons.

Ablation Study

Table 2 compares our LRM against alternative backward selection criteria over a sustained 10,000-step training trajectory, contrasting it with a squared gradient norm approach, an absolute gradient magnitude baseline, and a pure random selection mechanism. Our recalibration metric achieves the lowest mean ΔPPL across both evaluated architectures, registering an average deviation of just 0.012474 for Qwen3-30B-A3B and an exceptionally tight 0.003464 for GLM4.5-106B-A12B. Both gradient-only metrics outperform random selection but produce larger PPL deviations than LRM. The ordering is consistent across both models: the squared gradient norm is the stronger gradient-only baseline, followed by the absolute gradient magnitude, while random selection yields the largest deviations. This empirical result strongly supports our core design hypothesis: properly selecting retained token to expert pairs requires considering the established routing strength and the saved expert-state magnitude alongside the raw backward gradient signal.

Related Work

MoE routing and load balancing. BASE layers balance token assignment, while Expert Choice fixes expert buckets and varies the number of selected experts per token [14, 34]. StableMoE distills and freezes a router to stabilize assignments [3]. Gating Dropout bypasses global routing for some tokens to reduce communication [15]. These methods change forward routing. FasterMoE shadows popular experts and can alter expert selection to avoid network congestion [8]. FlexMoE adapts placement, while SmartMoE searches hybrid parallelization plans [19, 31]. SYMI separates parameters from optimizer states, and LAER-MoE reconstructs expert instances from sharded parameters [17, 26]. UltraEP and topology-aware TopoEP respond to realized post-gating loads [28, 35]. LRMoE instead selects backward pairs under quotas and balances them across instances of their original experts. Selective backward computation. Selective-Backprop prioritizes high-loss training examples [10]. In contrast, LRMoE ranks routed pairs within each expert’s quota. Its projection-level score uses the factorization underlying efficient per-example gradientnorm computation [7], applied to saved projection inputs and weighted incoming gradients. Memory and communication reduction. ZeRO partitions training states to remove redundancy [22]. ZeRO-Offload and ZeRO-Infinity extend memory management across GPU, CPU, and NVMe [23, 24]. ES-MoE offloads expert parameters and pipelines execution [12]. FlashAttention reduces attention HBM accesses through IO-aware tiling [4]. These methods do not select expertupdate contributions. RCMoE quantizes all-to-all exchanges and sparsifies all-reduce gradients [29]. Our row selection instead reduces token-dependent communication and expert computation without changing parameter dimensions. Parallel execution and overlap. PipeDream pipelines minibatches across partitioned layers, while Alpa automatically plans parallel strategies [18, 33]. Within MoE layers, MegaBlocks uses block-sparse computation for irregular loads without capacitybased token dropping [6]. Lancet, COMET, and FSMoE co-schedule and overlap communication with computation across different granularities [11, 20, 32]. LRMoE changes which backward rows enter execution and packs retained rows into contiguous expert segments. Unlike masking full-sized tensors, this compaction reduces physical GEMM row counts.

LRMoE

6

Conclusion

In this work, we present LRMoE to tackle load imbalances and memory bottlenecks in Expert Parallelism. Unlike balancers that merely redistribute a fixed workload, LRMoE tightly coordinates contribution-aware load recalibration with capacity-aware routing under strict memory and bandwidth constraints. It selectively retains high-value token-expert pairs to fit within planned quotas and routes them across admitted expert instances. Compact grouped GEMMs then translate this reduced workload into smaller communication payloads and lower computation costs without zero-padding overhead. By keeping the forward pass intact for accurate contribution scoring, LRMoE advances MoE balancing by jointly controlling the size and spatial distribution of backward computation.

References [1] Paul Covington, Jay Adams, and Emre Sargin. 2016. Deep Neural Networks for YouTube Recommendations. In Proceedings of the 10th ACM Conference on Recommender Systems, Boston, MA, USA, September 15-19, 2016, Shilad Sen, Werner Geyer, Jill Freyne, and Pablo Castells (Eds.). ACM, 191–198. doi:10.1145/2959100. 2959190 [2] Damai Dai, Chengqi Deng, Chenggang Zhao, R. X. Xu, Huazuo Gao, Deli Chen, Jiashi Li, Wangding Zeng, Xingkai Yu, Y. Wu, Zhenda Xie, Y. K. Li, Panpan Huang, Fuli Luo, Chong Ruan, Zhifang Sui, and Wenfeng Liang. 2024. DeepSeekMoE: Towards Ultimate Expert Specialization in Mixture-of-Experts Language Models. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2024, Bangkok, Thailand, August 11-16, 2024, Lun-Wei Ku, Andre Martins, and Vivek Srikumar (Eds.). Association for Computational Linguistics, 1280–1297. doi:10.18653/V1/2024.ACL-LONG.70 [3] Damai Dai, Li Dong, Shuming Ma, Bo Zheng, Zhifang Sui, Baobao Chang, and Furu Wei. 2022. StableMoE: Stable Routing Strategy for Mixture of Experts. In Proceedings of the 60th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), ACL 2022, Dublin, Ireland, May 22-27, 2022, Smaranda Muresan, Preslav Nakov, and Aline Villavicencio (Eds.). Association for Computational Linguistics, 7085–7095. doi:10.18653/V1/2022.ACL-LONG.489 [4] Tri Dao, Daniel Y. Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. 2022. FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (Eds.). http://papers.nips.cc/paper_files/paper/2022/ hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html [5] William Fedus, Barret Zoph, and Noam Shazeer. 2022. Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity. J. Mach. Learn. Res. 23 (2022), 120:1–120:39. https://jmlr.org/papers/v23/21-0998.html [6] Trevor Gale, Deepak Narayanan, Cliff Young, and Matei Zaharia. 2023. MegaBlocks: Efficient Sparse Training with Mixture-of-Experts. In Proceedings of the Sixth Conference on Machine Learning and Systems, MLSys 2023, Miami, FL, USA, June 4-8, 2023, Dawn Song, Michael Carbin, and Tianqi Chen (Eds.). mlsys.org. https://proceedings.mlsys.org/paper_files/paper/2023/hash/ 5a54f79333768effe7e8927bcccffe40-Abstract-mlsys2023.html [7] Ian J. Goodfellow. 2015. Efficient Per-Example Gradient Computations. CoRR abs/1510.01799 (2015). arXiv:1510.01799 doi:10.48550/ARXIV.1510.01799 [8] Jiaao He, Jidong Zhai, Tiago Antunes, Haojie Wang, Fuwen Luo, Shangfeng Shi, and Qin Li. 2022. FasterMoE: modeling and optimizing training of large-scale dynamic pre-trained models. In PPoPP ’22: 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, Seoul, Republic of Korea, April 2 6, 2022, Jaejin Lee, Kunal Agrawal, and Michael F. Spear (Eds.). ACM, 120–134. doi:10.1145/3503221.3508418 [9] Po-Sen Huang, Xiaodong He, Jianfeng Gao, Li Deng, Alex Acero, and Larry P. Heck. 2013. Learning deep structured semantic models for web search using clickthrough data. In 22nd ACM International Conference on Information and Knowledge Management, CIKM’13, San Francisco, CA, USA, October 27 - November 1, 2013, Qi He, Arun Iyengar, Wolfgang Nejdl, Jian Pei, and Rajeev Rastogi (Eds.). ACM, 2333–2338. doi:10.1145/2505515.2505665 [10] Angela H. Jiang, Daniel L.-K. Wong, Giulio Zhou, David G. Andersen, Jeffrey Dean, Gregory R. Ganger, Gauri Joshi, Michael Kaminsky, Michael Kozuch, Zachary C. Lipton, and Padmanabhan Pillai. 2019. Accelerating Deep Learning by Focusing on the Biggest Losers. CoRR abs/1910.00762 (2019). arXiv:1910.00762 doi:10.48550/ARXIV.1910.00762

[11] Chenyu Jiang, Ye Tian, Zhen Jia, Shuai Zheng, Chuan Wu, and Yida Wang. 2024. Lancet: Accelerating Mixture-of-Experts Training via Whole Graph ComputationCommunication Overlapping. In Proceedings of the Seventh Annual Conference on Machine Learning and Systems, MLSys 2024, Santa Clara, CA, USA, May 13-16, 2024, Phillip B. Gibbons, Gennady Pekhimenko, and Christopher De Sa (Eds.). mlsys.org. https://proceedings.mlsys.org/paper_files/paper/2024/hash/ 339caf45a6fa281cae8adc6465343464-Abstract-Conference.html [12] Yechan Kim, Hwijoon Lim, and Dongsu Han. 2024. Scaling Beyond the GPU Memory Limit for Large Mixture-of-Experts Model Training. In Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024 (Proceedings of Machine Learning Research, Vol. 235), Ruslan Salakhutdinov, Zico Kolter, Katherine A. Heller, Adrian Weller, Nuria Oliver, Jonathan Scarlett, and Felix Berkenkamp (Eds.). PMLR / OpenReview.net, 24342–24353. https://proceedings.mlr.press/v235/kim24w.html [13] Dmitry Lepikhin, HyoukJoong Lee, Yuanzhong Xu, Dehao Chen, Orhan Firat, Yanping Huang, Maxim Krikun, Noam Shazeer, and Zhifeng Chen. 2021. GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net. https://openreview.net/forum? id=qrwe7XHTmYb [14] Mike Lewis, Shruti Bhosale, Tim Dettmers, Naman Goyal, and Luke Zettlemoyer. 2021. BASE Layers: Simplifying Training of Large, Sparse Models. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event (Proceedings of Machine Learning Research, Vol. 139), Marina Meila and Tong Zhang (Eds.). PMLR, 6265–6274. http://proceedings.mlr.press/ v139/lewis21a.html [15] Rui Liu, Young Jin Kim, Alexandre Muzio, and Hany Hassan. 2022. Gating Dropout: Communication-efficient Regularization for Sparsely Activated Transformers. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA (Proceedings of Machine Learning Research, Vol. 162), Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (Eds.). PMLR, 13782–13792. https://proceedings.mlr.press/ v162/liu22g.html [16] Xinyi Liu, Yujie Wang, Fangcheng Fu, Xupeng Miao, Shenhan Zhu, Xiaonan Nie, and Bin Cui. 2025. NetMoE: Accelerating MoE Training through Dynamic Sample Placement. In The Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24-28, 2025. OpenReview.net. https://openreview.net/ forum?id=1qP3lsatCR [17] Xinyi Liu, Yujie Wang, Fangcheng Fu, Xuefeng Xiao, Huixia Li, Jiashi Li, and Bin Cui. 2026. LAER-MoE: Load-Adaptive Expert Re-layout for Efficient Mixtureof-Experts Training. In Proceedings of the 31st ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 2, ASPLOS 2026, Pittsburgh, PA, USA, March 22-26, 2026, Benjamin C. Lee, Harry Xu, Mark Silberstein, and Bingyao Li (Eds.). ACM, 1055–1072. doi:10.1145/ 3779212.3790180 [18] Deepak Narayanan, Aaron Harlap, Amar Phanishayee, Vivek Seshadri, Nikhil R. Devanur, Gregory R. Ganger, Phillip B. Gibbons, and Matei Zaharia. 2019. PipeDream: generalized pipeline parallelism for DNN training. In Proceedings of the 27th ACM Symposium on Operating Systems Principles, SOSP 2019, Huntsville, ON, Canada, October 27-30, 2019, Tim Brecht and Carey Williamson (Eds.). ACM, 1–15. doi:10.1145/3341301.3359646 [19] Xiaonan Nie, Xupeng Miao, Zilong Wang, Zichao Yang, Jilong Xue, Lingxiao Ma, Gang Cao, and Bin Cui. 2023. FlexMoE: Scaling Large-scale Sparse Pre-trained Model Training via Dynamic Device Placement. Proc. ACM Manag. Data 1, 1 (2023), 110:1–110:19. doi:10.1145/3588964 [20] Xinglin Pan, Wenxiang Lin, Lin Zhang, Shaohuai Shi, Zhenheng Tang, Rui Wang, Bo Li, and Xiaowen Chu. 2025. FSMoE: A Flexible and Scalable Training System for Sparse Mixture-of-Experts Models. In Proceedings of the 30th ACM International Conference on Architectural Support for Programming Languages and Operating Systems, Volume 1, ASPLOS 2025, Rotterdam, The Netherlands, 30 March 2025 - 3 April 2025, Lieven Eeckhout, Georgios Smaragdakis, Kaitai Liang, Adrian Sampson, Martha A. Kim, and Christopher J. Rossbach (Eds.). ACM, 524–539. doi:10.1145/3669940.3707272 [21] Samyam Rajbhandari, Conglong Li, Zhewei Yao, Minjia Zhang, Reza Yazdani Aminabadi, Ammar Ahmad Awan, Jeff Rasley, and Yuxiong He. 2022. DeepSpeedMoE: Advancing Mixture-of-Experts Inference and Training to Power NextGeneration AI Scale. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA (Proceedings of Machine Learning Research, Vol. 162), Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvári, Gang Niu, and Sivan Sabato (Eds.). PMLR, 18332–18346. https: //proceedings.mlr.press/v162/rajbhandari22a.html [22] Samyam Rajbhandari, Jeff Rasley, Olatunji Ruwase, and Yuxiong He. 2020. ZeRO: memory optimizations toward training trillion parameter models. In Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9-19, 2020, Christine Cuicchi, Irene Qualters, and William T. Kramer (Eds.). IEEE/ACM, 20. doi:10.1109/SC41405.2020.00024

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

[23] Samyam Rajbhandari, Olatunji Ruwase, Jeff Rasley, Shaden Smith, and Yuxiong He. 2021. ZeRO-infinity: breaking the GPU memory wall for extreme scale deep learning. In International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2021, St. Louis, Missouri, USA, November 14-19, 2021, Bronis R. de Supinski, Mary W. Hall, and Todd Gamblin (Eds.). ACM, 59. doi:10. 1145/3458817.3476205 [24] Jie Ren, Samyam Rajbhandari, Reza Yazdani Aminabadi, Olatunji Ruwase, Shuangyan Yang, Minjia Zhang, Dong Li, and Yuxiong He. 2021. ZeRO-Offload: Democratizing Billion-Scale Model Training. In Proceedings of the 2021 USENIX Annual Technical Conference, USENIX ATC 2021, July 14-16, 2021, Irina Calciu and Geoff Kuenning (Eds.). USENIX Association, 551–564. https://www.usenix. org/conference/atc21/presentation/ren-jie [25] Noam Shazeer, Azalia Mirhoseini, Krzysztof Maziarz, Andy Davis, Quoc V. Le, Geoffrey E. Hinton, and Jeff Dean. 2017. Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layer. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Conference Track Proceedings. OpenReview.net. https://openreview.net/forum?id= B1ckMDqlg [26] Athinagoras Skiadopoulos, Mark Zhao, Swapnil Gandhi, Thomas Norrie, Shrijeet Mukherjee, and Christos Kozyrakis. 2026. SYMI: Efficient Mixture-of-Experts Training via Model and Optimizer State Decoupling. In 23rd USENIX Symposium on Networked Systems Design and Implementation, NSDI 2026, Renton, WA, May 46, 2026, Srikanth Kandula and Hakim Weatherspoon (Eds.). USENIX Association, 75–92. https://www.usenix.org/conference/nsdi26/presentation/skiadopoulos [27] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attention is All you Need. In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 49, 2017, Long Beach, CA, USA, Isabelle Guyon, Ulrike von Luxburg, Samy Bengio, Hanna M. Wallach, Rob Fergus, S. V. N. Vishwanathan, and Roman Garnett (Eds.). 5998–6008. https://proceedings.neurips.cc/paper/2017/hash/ 3f5ee243547dee91fbd053c1c4a845aa-Abstract.html [28] Xinming Wei, Chao Jin, Tuo Dai, Yinmin Zhong, Shan Yu, Chengxu Yang, Bingyang Wu, Zili Zhang, Jing Mai, Qianchao Zhu, Zhouyang Li, Yuliang Liu, and Guojie Luo. 2026. UltraEP: Unleash MoE Training and Inference on RackScale Nodes with Near-Optimal Load Balancing. CoRR abs/2606.04101 (2026). arXiv:2606.04101 doi:10.48550/ARXIV.2606.04101 [29] Donglei Wu, Xiao Cai, Jinglei Tan, Jinda Jia, Guangming Tan, Dingwen Tao, Wen Xia, and Zhihong Tian. 2026. RCMoE: A Communication-Efficient Random Compression Framework for Resource-Constrained Mixture-of-Experts Training. In Fortieth AAAI Conference on Artificial Intelligence, Thirty-Eighth Conference on Innovative Applications of Artificial Intelligence, Sixteenth Symposium on Educational Advances in Artificial Intelligence, AAAI 2026, Singapore, January 20-27, 2026, Sven Koenig, Chad Jenkins, and Matthew E. Taylor (Eds.). AAAI Press, 26876–26884. doi:10.1609/AAAI.V40I32.39899 [30] Yan Zeng, Chengchuang Huang, Yipeng Mei, Lifu Zhang, Teng Su, Wei Ye, Wenqi Shi, and Shengnan Wang. 2025. EfficientMoE: Optimizing Mixture-of-Experts Model Training With Adaptive Load Balance. IEEE Trans. Parallel Distributed Syst. 36, 4 (2025), 677–688. doi:10.1109/TPDS.2025.3539297 [31] Mingshu Zhai, Jiaao He, Zixuan Ma, Zan Zong, Runqing Zhang, and Jidong Zhai. 2023. SmartMoE: Efficiently Training Sparsely-Activated Models through Combining Offline and Online Parallelization. In Proceedings of the 2023 USENIX Annual Technical Conference, USENIX ATC 2023, Boston, MA, USA, July 10-12, 2023, Julia Lawall and Dan Williams (Eds.). USENIX Association, 961–975. https: //www.usenix.org/conference/atc23/presentation/zhai [32] Shulai Zhang, Ningxin Zheng, Haibin Lin, Ziheng Jiang, Wenlei Bao, Chengquan Jiang, Qi Hou, Weihao Cui, Size Zheng, Li-Wen Chang, Quan Chen, and Xin Liu. 2025. COMET: Fine-grained Computation-communication Overlapping for Mixture-of-Experts. In Proceedings of the Eighth Conference on Machine Learning and Systems, MLSys 2025, Santa Clara, CA, USA, May 12-15, 2025, Matei Zaharia, Gauri Joshi, and Yingyan (Celine) Lin (Eds.). OpenReview.net/mlsys.org. https: //openreview.net/forum?id=fGgQS5VW09 [33] Lianmin Zheng, Zhuohan Li, Hao Zhang, Yonghao Zhuang, Zhifeng Chen, Yanping Huang, Yida Wang, Yuanzhong Xu, Danyang Zhuo, Eric P. Xing, Joseph E. Gonzalez, and Ion Stoica. 2022. Alpa: Automating Inter- and Intra-Operator Parallelism for Distributed Deep Learning. In 16th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA, USA, July 11-13, 2022, Marcos K. Aguilera and Hakim Weatherspoon (Eds.). USENIX Association, 559–578. https://www.usenix.org/conference/osdi22/presentation/zhenglianmin [34] Yanqi Zhou, Tao Lei, Hanxiao Liu, Nan Du, Yanping Huang, Vincent Y. Zhao, Andrew M. Dai, Zhifeng Chen, Quoc V. Le, and James Laudon. 2022. Mixture-of-Experts with Expert Choice Routing. In Advances in Neural Information Processing Systems 35: Annual Conference on Neural Information Processing Systems 2022, NeurIPS 2022, New Orleans, LA, USA, November 28 - December 9, 2022, Sanmi Koyejo, S. Mohamed, A. Agarwal, Danielle Belgrave, K. Cho, and A. Oh (Eds.). http://papers.nips.cc/paper_files/paper/2022/hash/ 2f00ecd787b432c1d36f3de9800728eb-Abstract-Conference.html

[35] Jiacheng Zhu, Xie Zhao, Gongming Zhao, Hongli Xu, Yao Fei, and Jin Fang. 2026. TopoEP: Topology-Aware Load Balancing for Expert-Parallel MoE Training. CoRR abs/2609.35481 (2026). arXiv:2609.35481 doi:10.48550/ARXIV.2609.35481

Appendices 7

Derivation and Optimality of Token–Expert Recalibration

This appendix proves Proposition 3.1 and the optimality of the resulting drop rule under the fidelity objective used by our recalibration controller. The first result is an exact, projection-level identity for one routed pair. The second result concerns selection: for a fixed number of equal-size rows to remove, dropping the rows with the smallest scores minimizes the sum of their isolated projectiongradient perturbations. We state the assumptions explicitly because this selection result is not a claim about the nonlinear loss or the entire optimizer trajectory.

7.1

Pairwise projection-gradient contribution

Consider one completed forward route and the protected output projection (FC2) of expert 𝑒. Let 𝑊𝑒 ∈ R𝑑out ×𝑑in ,

𝜙𝑢 ∈ R𝑑in ,

𝑔𝑡 ∈ R𝑑out ,

(4)

where 𝑢 = (𝑡, 𝑒) is a routed token–expert pair. The saved vector 𝜙𝑢 is the actual input to the protected projection. We take 𝑔𝑡 to be the upstream cotangent before multiplication by the fixed routing weight 𝛼𝑢 . Equivalently, an implementation may expose the already weighted cotangent 𝛿𝑢 := 𝛼𝑢 𝑔𝑡 , (5) in that representation, the routing weight must not be applied again. The contribution of pair 𝑢 to the routed expert output is 𝑦𝑢 = 𝛼𝑢𝑊𝑒 𝜙𝑢 .

(6)

The forward route, 𝛼𝑢 , and 𝜙𝑢 are fixed during recalibration. For an arbitrary perturbation 𝑑𝑊𝑒 , the differential of the loss through this pair is 𝑑L𝑢 = 𝑔𝑡⊤𝑑𝑦𝑢 = 𝑔𝑡⊤ 𝛼𝑢 (𝑑𝑊𝑒 )𝜙𝑢  
⊤ = tr (𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ 𝑑𝑊𝑒 .

(7)

By the defining identity 𝑑L = tr((∇𝑊 L) ⊤𝑑𝑊 ), the isolated contribution of pair 𝑢 to the protected weight gradient is therefore   𝐺𝑢 := ∇𝑊𝑒 L 𝑢 = (𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ = 𝛿𝑢 𝜙𝑢⊤ . (8) Let U𝑒 denote all pairs contributing to the same protected projection. Before omission, its routed-pair gradient is ∑︁ 𝐺𝑒 = 𝐺𝑣 . (9) 𝑣 ∈ U𝑒

If only pair 𝑢 is omitted and reconstructed as zero while every other pair is unchanged, then 𝐺𝑒(−𝑢 ) = 𝐺𝑒 − 𝐺𝑢 ,

Δ𝑢 := 𝐺𝑒(−𝑢 ) − 𝐺𝑒 = −𝐺𝑢 .

(10)

The sign is immaterial after taking a norm, but Eq. (10) makes clear that the metric measures a backward contribution removed from a fixed completed route, it does not rerun or alter the forward computation.

LRMoE

For vectors 𝑎 and 𝑏, the squared Frobenius norm of their outer product factorizes as
∥𝑎𝑏 ⊤ ∥ 2𝐹 = tr (𝑎𝑏 ⊤ ) ⊤ (𝑎𝑏 ⊤ )
= tr 𝑏 (𝑎 ⊤𝑎)𝑏 ⊤ = (𝑎 ⊤𝑎)(𝑏 ⊤𝑏) = ∥𝑎∥ 22 ∥𝑏 ∥ 22 . (11) Applying Eq. (11) with 𝑎 = 𝛼𝑢 𝑔𝑡 and 𝑏 = 𝜙𝑢 proves 2

∥Δ𝑢 ∥ 2𝐹 = (𝛼𝑢 𝑔𝑡 )𝜙𝑢⊤ 𝐹 = 𝛼𝑢2 ∥𝑔𝑡 ∥ 22 ∥𝜙𝑢 ∥ 22 = 𝑅𝑢 ,

(12)

which is exactly Eq. (3). Hence the three factors in 𝑅𝑢 are not an ad hoc weighted sum: together they form the exact squared norm of the isolated gradient contribution to the protected projection.

7.2

Optimality of the drop rule

Proof. Let 𝐷 be any dropped set with |𝐷 | = 𝑑. Suppose that 𝐷 contains a pair 𝑖 while retaining a pair 𝑗 such that 𝑅𝑖 > 𝑅 𝑗 . Replacing 𝑖 by 𝑗 gives 𝐷 ′ = (𝐷 \ {𝑖}) ∪ { 𝑗 },

(14)

and changes the objective by b(𝐷 ′ ) − E b(𝐷) = 𝑅 𝑗 − 𝑅𝑖 < 0. E

(15)

Thus a set with such an inversion cannot be optimal. Repeatedly applying this exchange terminates at a set containing the 𝑑 smallest scores, because each exchange strictly decreases a finite-set objective. Conversely, a set containing the 𝑑 smallest scores has no improving exchange: every retained score is at least as large as every dropped score. It therefore minimizes the objective. A strict boundary gap fixes the set, ties yield multiple sets with the same sum. □ The equivalent retained-set form follows directly from ∑︁ ∑︁ ∑︁ 𝑅𝑢 = 𝑅𝑢 − 𝑅𝑢 . 𝑢∈U

(16)

𝑢 ∈𝑆

The first term is constant, so minimizing dropped value is identical to maximizing retained value. This is the precise sense in which our drop is optimal: once the load controller determines that 𝑑 equal-size rows must be removed, no other 𝑑-row mask preserves more of the declared pairwise calibration value.

𝑢 ∈𝐷

For the fixed per-expert quotas used by our controller, this objective separates across experts and is minimized by selecting the smallest scores within each expert.

Relation to the aggregate expert gradient

The single-pair identity in Eq. (12) is exact. For several dropped pairs from the same expert, their contributions to a shared 𝑊𝑒 add. Here let 𝐷 ⊆ U𝑒 and 𝑑 = |𝐷 |. The actual change in the protected expert gradient is ∑︁ Δ𝐷 = −

(18)

𝐺𝑢 ,

𝑢 ∈𝐷

and therefore

𝑢 ∈𝐷

Proposition 7.1 (Optimal score-based dropping). Under a fixed completed route, equal row costs, and a fixed drop count 𝑑, any set containing the 𝑑 smallest scores 𝑅𝑢 minimizes Eq. (13). Equivalently, any retained set containing the 𝐾 = 𝑁 − 𝑑 largest scores maximizes preserved calibration value. The solution is unique when there is a strict gap between the 𝑑th and (𝑑 + 1)st scores in ascending order, otherwise every difference between optimal solutions is confined to tied boundary scores.

𝑢 ∈𝐷

𝐷∈F

7.3

Let U be the causally available token–expert pairs at the recalibration boundary, with 𝑁 = |U|. Let 𝐷 ⊆ U be the dropped set and 𝑆 = U \ 𝐷 the retained set. All rows at this boundary have the same shape and encoding, and therefore the same communication cost. For a prescribed drop count 𝑑, or equivalently a retained count 𝐾 = 𝑁 − 𝑑, define the additive local calibration distortion ∑︁ ∑︁ b(𝐷) := E ∥Δ𝑢 ∥ 2𝐹 = 𝑅𝑢 . (13) 𝑢 ∈𝐷

When the controller imposes a fixed drop count 𝑑𝑒 separately on each expert or overloaded path, the same proof applies independently inside every partition: dropping the 𝑑𝑒 smallest scores in each partition minimizes the sum of Eq. (13) over all masks obeying those cardinalities. More general placement and capacity constraints define a restricted feasible family F . The corresponding selection problem is ∑︁ 𝐷 ∗ ∈ arg min 𝑅𝑢 . (17)

∥Δ𝐷 ∥ 2𝐹 =

∑︁

𝑅𝑢 + 2

𝑢 ∈𝐷

∑︁

⟨𝐺𝑢 , 𝐺 𝑣 ⟩𝐹 .

(19)

𝑢<𝑣 𝑢,𝑣 ∈𝐷

The cross terms can be positive or negative. Consequently, Proposition 7.1 exactly optimizes the additive sum of isolated perturbations, it does not, without further assumptions, claim to minimize the realized norm of the summed shared-parameter gradient error. Two useful special cases recover an aggregate interpretation. First, if the dropped contributions are pairwise Frobeniusb(𝐷). orthogonal, all cross terms in Eq. (19) vanish and ∥Δ𝐷 ∥ 2𝐹 = E Second, if the cross terms of the selected dropped set have zero expectation, the same equality holds in expectation. Without either condition, the triangle and Cauchy–Schwarz inequalities still give the deterministic bound !2 ∑︁ √︁ ∑︁ 2 ∥Δ𝐷 ∥ 𝐹 ≤ 𝑅𝑢 ≤ 𝑑 𝑅𝑢 . (20) 𝑢 ∈𝐷

𝑢 ∈𝐷

For fixed 𝑑, the smallest-score rule also minimizes the final upper bound in Eq. (20).

7.4

Scope of the result

The proof relies on four declared conditions. First, the forward route and saved projection inputs are fixed, recalibration acts only on the completed route’s backward contributions. Second, 𝜙𝑢 is the actual input to the protected linear projection. For a nonlinear multilayer expert, the same formula is exact for the selected projection but is not the squared norm of the expert’s complete parametergradient contribution. Third, the closed-form drop rule assumes equal row costs and a fixed cardinality. Unequal byte costs lead to a knapsack problem, while replica, topology, and HBM constraints lead to Eq. (17). Fourth, the objective protects a local projectiongradient quantity. It is neither an identity for nonlinear loss change nor a guarantee about future optimizer states or model quality, those effects are evaluated empirically.

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

Within these conditions, the result is exact: 𝑅𝑢 is the squared isolated perturbation caused by dropping pair 𝑢, and our drop rule

is optimal for the declared additive calibration objective at the required recalibration rate.

Related documents

• LRMoE: Demand-Side Load Re-Calibration for Efficient MoE Training on Memory-Constrained Nodes#1488512
• Accelerating MoE Training through Expert Placement and Replica Coordination#1325694
• CRISP: Co-designing Pruning and Scheduling for Efficient MoE Inference on Edge Servers#1473611
• CRISP: Co-designing Pruning and Scheduling for Efficient MoE Inference on Edge Servers#1469045
• MoEless: Efficient MoE LLM Serving with Serverless Experts#1142637