LRMoE:面向内存受限节点上高效 MoE 训练的需求侧负载重校准

arXiv cs.AI 论文

摘要

LRMoE 是一种内存感知的 MoE 训练系统,可在内存受限的 GPU 节点上联合重新校准反向传播需求并均衡保留的工作负载。与 SYMI 相比,LRMoE 在 Qwen3-30B-A3B 和 GLM4.5-106B-A12B 上将 MoE 层的累计训练时间缩短了 21.2-44.8%。

arXiv:2610.37720v1 Announce Type: new 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.
查看原文
查看缓存全文

缓存时间: 2026/10/31 09:10

# LRMoE: Demand-Side Load Re-Calibration for Efficient MoE Training on Memory-Constrained Nodes
Source: [https://arxiv.org/html/2610.37720](https://arxiv.org/html/2610.37720)
CCS:Networks Network performance analysisCCS:Computing methodologies Machine learning algorithms,Zhida JiangNote:Zhida Jiang and Zhen Chen are co\-corresponding authors\.Affiliation:JD\.com,Beijing,China,Zhaolong XingAffiliation:JD\.com,Beijing,China,Yu LiAffiliation:JD\.com,Beijing,China,Zhen ChenAffiliation:JD\.com,Beijing,ChinaandKe ZhangAffiliation:JD\.com,Beijing,China

© none

###### 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 presentLRMoE, 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 thatLRMoEreduces cumulative MoE\-layer training time to fixed\-loss targets by 21\.2\-44\.8% relative to SYMI\.

###### Keywords:

mixture\-of\-experts, expert replication, load balancing, topology\-aware scheduling, distributed training

## 1\.Introduction

Web\-scale search and recommendation services rely on learned models\([Huang et al\., 2013](https://arxiv.org/html/2610.37720#bib.bib25);[Covington et al\., 2016](https://arxiv.org/html/2610.37720#bib.bib26)\)\. 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\([Shazeer et al\., 2017](https://arxiv.org/html/2610.37720#bib.bib2);[Lepikhin et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib3);[Fedus et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib1);[Rajbhandari et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib6);[Dai et al\., 2024](https://arxiv.org/html/2610.37720#bib.bib28)\)\. 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\.

Expert Parallelism \(EP\) distributes experts across GPUs and exchanges routed token–expert pairs through all\-to\-all communication\([Lepikhin et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib3);[Gale et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib8);[Rajbhandari et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib6)\)\. 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\([Zhu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib24)\)\. Replication distributes this work but consumes memory and introduces state\-transfer 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\([Wei et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib14);[Zhu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib24)\)\.

Existing systems adapt expert placement and execution to uneven loads\. FasterMoE shadows popular experts, FlexMoE adapts placement, and SmartMoE searches hybrid parallelization plans\([He et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib9);[Nie et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib10);[Zhai et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib11)\)\. SYMI decouples model parameters from optimizer states, while LAER\-MoE reconstructs expert instances from sharded parameters\([Skiadopoulos et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib13);[Liu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib12)\)\. UltraEP responds to realized post\-gating loads, and TopoEP incorporates communication topology into planning\([Wei et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib14);[Zhu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib24)\)\. 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 high\-loss training examples\([Jiang et al\., 2019](https://arxiv.org/html/2610.37720#bib.bib19)\)\. 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\([Goodfellow, 2015](https://arxiv.org/html/2610.37720#bib.bib17)\)\. 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\([Gale et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib8)\)\. Our selection must likewise translate into smaller physical workloads\. Removing work 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 presentLRMoE, 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 highest\-scoring 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,LRMoEreduces 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\([Vaswani et al\., 2017](https://arxiv.org/html/2610.37720#bib.bib35)\), 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\([Shazeer et al\., 2017](https://arxiv.org/html/2610.37720#bib.bib2);[Lepikhin et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib3);[Fedus et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib1)\)\. 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\([Rajbhandari et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib6);[Gale et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib8)\)\. 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\([He et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib9);[Nie et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib10);[Zhai et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib11);[Skiadopoulos et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib13);[Wei et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib14)\)\. 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\([Lewis et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib4);[Zhou et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib5)\)\. Gating Dropout reduces communication by bypassing global routing for some tokens\([Liu et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib20)\)\. 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\([Ren et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib23);[Kim et al\., 2024](https://arxiv.org/html/2610.37720#bib.bib7)\)\. 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\([He et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib9);[Nie et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib10);[Zhai et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib11)\)\. EfficientMoE adapts expert scheduling to predicted loads\([Zeng et al\., 2025](https://arxiv.org/html/2610.37720#bib.bib21)\), while NetMoE rearranges samples to exploit routing locality\([Liu et al\., 2025](https://arxiv.org/html/2610.37720#bib.bib22)\)\. SYMI and LAER\-MoE reduce state\-management costs\([Skiadopoulos et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib13);[Liu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib12)\), while UltraEP adapts to realized post\-gating loads\([Wei et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib14)\)\. 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\. LetWn​ℓ​eW\_\{n\\ell e\}denote expertee’s backward demand for layer invocation\(n,ℓ\)\(n,\\ell\), and letCn​ℓ​gbwdC^\{\\mathrm\{bwd\}\}\_\{n\\ell g\}denote GPUgg’s guarded service capacity within that window\. Under an additive model in which demand and capacity are expressed in the same calibrated work units, a necessary feasibility condition is∑eWn​ℓ​e≤∑gCn​ℓ​gbwd\\sum\_\{e\}W\_\{n\\ell e\}\\leq\\sum\_\{g\}C^\{\\mathrm\{bwd\}\}\_\{n\\ell g\}\. 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 contribution\-aware scheduling variable rather than a fixed consequence of forward routing\. Calibration replaces demandWn​ℓ​eW\_\{n\\ell e\}with retained demandW~n​ℓ​e≤Wn​ℓ​e\\widetilde\{W\}\_\{n\\ell e\}\\leq W\_\{n\\ell e\}, prioritizing estimated update contributions 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\([Goodfellow, 2015](https://arxiv.org/html/2610.37720#bib.bib17)\)\. 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 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\.LRMoEDesign

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

Figure 2\.Selective backward load calibration for uncalibrated load \(a\) and calibrated load \(b\)\. In backward,LRMoEdrops 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[2](https://arxiv.org/html/2610.37720#S3.F2)presents the design ofLRMoE, which combines layer\-wise load calibration, replica\-aware load routing, and compact grouped\-GEMM 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 demandnen\_\{e\}of each expert and determines the backward retention quotasKeK\_\{e\}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\. WithN=∑eneN=\\sum\_\{e\}n\_\{e\}, the retention rate isρ=∑eKe/N\\rho=\\sum\_\{e\}K\_\{e\}/N, experteeomitsne−Ken\_\{e\}\-K\_\{e\}backward contributions\. Quotas remain fixed for the invocation\. At backward Combine, LRM uses incoming gradients and saved state to select the highest\-valuedKeK\_\{e\}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 assignsXe​gX\_\{eg\}retained pairs of experteeto GPUgg, with∑gXe​g=Ke\\sum\_\{g\}X\_\{eg\}=K\_\{e\}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 pairu=\(t,e\)u=\(t,e\), the chain rule expresses its contribution to the protected projection gradient as the outer productGu=\(αu​gt\)​ϕu⊤G\_\{u\}=\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}, wheregtg\_\{t\}is the incoming gradient before multiplication by the routing weight andϕu\\phi\_\{u\}is the saved projection input\. Omitting this pair while leaving the other contributions unchanged produces the perturbationΔu=−Gu\\Delta\_\{u\}=\-G\_\{u\}\. We measure its magnitude using the squared Frobenius norm\. BecauseGuG\_\{u\}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 matrixWeW\_\{e\}\. For a routed token\-expert pairu=\(t,e\)u=\(t,e\), letαu\\alpha\_\{u\}denote its fixed routing weight, letgtg\_\{t\}denote the gradient arriving at the routed expert output, and letϕu\\phi\_\{u\}denote the saved input to the protected projection\. The contribution of pairuuto the weight gradient is

\(1\)Gu=∂ℒu∂We=\(αu​gt\)​ϕu⊤\.G\_\{u\}=\\frac\{\\partial\\mathcal\{L\}\_\{u\}\}\{\\partial W\_\{e\}\}=\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}\.If pairuuis omitted while all other routed pairs remain unchanged, the protected weight gradient changes by

\(2\)Δu=−Gu=−\(αu​gt\)​ϕu⊤\.\\Delta\_\{u\}=\-G\_\{u\}=\-\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}\.We therefore define the pairwise calibration score as the squared Frobenius norm of this perturbation:

Ru\\displaystyle R\_\{u\}:=∥Δu∥F2\\displaystyle:=\\lVert\\Delta\_\{u\}\\rVert\_\{F\}^\{2\}=‖\(αu​gt\)​ϕu⊤‖F2\\displaystyle=\\left\\lVert\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}\\right\\rVert\_\{F\}^\{2\}\(3\)=αu2​∥gt∥22​∥ϕu∥22\.\\displaystyle=\\alpha\_\{u\}^\{2\}\\lVert g\_\{t\}\\rVert\_\{2\}^\{2\}\\lVert\\phi\_\{u\}\\rVert\_\{2\}^\{2\}\.

- •Routing strengthαu2\\alpha\_\{u\}^\{2\}\.This term measures how strongly tokenttwas assigned to experteeduring the completed forward pass\. Because routing is already fixed,αu\\alpha\_\{u\}is treated as a constant during recalibration\.
- •Backward signal∥gt∥22\\lVert g\_\{t\}\\rVert\_\{2\}^\{2\}\.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∥ϕu∥22\\lVert\\phi\_\{u\}\\rVert\_\{2\}^\{2\}\.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](https://arxiv.org/html/2610.37720#S3.E3)follows from the rank\-one identity∥a​b⊤∥F2=∥a∥22​∥b∥22\\lVert ab^\{\\top\}\\rVert\_\{F\}^\{2\}=\\lVert a\\rVert\_\{2\}^\{2\}\\lVert b\\rVert\_\{2\}^\{2\}\. Thus,RuR\_\{u\}equals the squared change in the protected projection’s weight gradient when pairuuis 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](https://arxiv.org/html/2610.37720#S7.SS1)\.LRMoEseparates*how much*backward demand can be admitted from*which*contributions should be retained\. During forward planning, the resource controller chooses the largest feasible retained countKℓ=∑eKℓ,eK\_\{\\ell\}=\\sum\_\{e\}K\_\{\\ell,e\}from its tested candidates under the quota\-allocation policy, HBM budget, replica\-placement constraints, and communication window, then fixes the per\-expert quotas\{Kℓ,e\}\\\{K\_\{\\ell,e\}\\\}\. Once backward gradients become available, the metric ranks pairs within each expert, and theKℓ,eK\_\{\\ell,e\}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](https://arxiv.org/html/2610.37720#S7.SS2)\. 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](https://arxiv.org/html/2610.37720#S7.SS3)and[7\.4](https://arxiv.org/html/2610.37720#S7.SS4)\.

### 3\.3\.Layerwise Load Calibration

Algorithm 1Resource\-Constrained Layer\-wise Load Calibration1:Routed pairs

𝒫ℓ=⋃e𝒫ℓ,e\\mathcal\{P\}\_\{\\ell\}=\\bigcup\_\{e\}\\mathcal\{P\}\_\{\\ell,e\}, forward placement

ΠℓF\\Pi\_\{\\ell\}^\{F\}, candidate rates

ℛℓ\\mathcal\{R\}\_\{\\ell\}, available HBM

\{Mravail\}\\\{M\_\{r\}^\{\\mathrm\{avail\}\}\\\}, bandwidth

BeffB\_\{\\mathrm\{eff\}\}, window

TℓovlpT\_\{\\ell\}^\{\\mathrm\{ovlp\}\}
2:Retained pairs

𝒮ℓ∗\\mathcal\{S\}\_\{\\ell\}^\{\*\}and backward plan

ΠℓB\\Pi\_\{\\ell\}^\{B\}
3:Forward\-time admission

4:

Nℓ,e←\|𝒫ℓ,e\|N\_\{\\ell,e\}\\leftarrow\|\\mathcal\{P\}\_\{\\ell,e\}\|for each expert

ee
5:for

ρ∈ℛℓ\\rho\\in\\mathcal\{R\}\_\{\\ell\}in decreasing orderdo

6:

\(\{Kℓ,e\},Π,𝐌^,T^comm\)←PlanQuotasAndReplicas​\(\{Nℓ,e\},ρ,ΠℓF,Beff\)\(\\\{K\_\{\\ell,e\}\\\},\\Pi,\\widehat\{\\mathbf\{M\}\},\\widehat\{T\}\_\{\\mathrm\{comm\}\}\)\\leftarrow\\textsc\{PlanQuotasAndReplicas\}\(\\\{N\_\{\\ell,e\}\\\},\\rho,\\Pi\_\{\\ell\}^\{F\},B\_\{\\mathrm\{eff\}\}\)
7:if

Π\\Piis feasible,

M^r≤Mravail\\widehat\{M\}\_\{r\}\\leq M\_\{r\}^\{\\mathrm\{avail\}\}for all

rr, and

T^comm≤Tℓovlp\\widehat\{T\}\_\{\\mathrm\{comm\}\}\\leq T\_\{\\ell\}^\{\\mathrm\{ovlp\}\}then

8:

\(ρℓ∗,\{Kℓ,e∗\},ΠℓA\)←\(ρ,\{Kℓ,e\},Π\)\(\\rho\_\{\\ell\}^\{\*\},\\\{K\_\{\\ell,e\}^\{\*\}\\\},\\Pi\_\{\\ell\}^\{A\}\)\\leftarrow\(\\rho,\\\{K\_\{\\ell,e\}\\\},\\Pi\)
9:break

10:endif

11:endfor

12:if

ΠℓA\\Pi\_\{\\ell\}^\{A\}is undefinedthen

13:returnInfeasible

14:endif

15:

𝒞ℓ←FullForwardAndSaveContext​\(𝒫ℓ,ΠℓA\)\\mathcal\{C\}\_\{\\ell\}\\leftarrow\\textsc\{FullForwardAndSaveContext\}\(\\mathcal\{P\}\_\{\\ell\},\\Pi\_\{\\ell\}^\{A\}\)
16:Backward\-time calibration

17:

vp←αp2​∥gt∥22​∥ϕp∥22v\_\{p\}\\leftarrow\\alpha\_\{p\}^\{2\}\\lVert g\_\{t\}\\rVert\_\{2\}^\{2\}\\lVert\\phi\_\{p\}\\rVert\_\{2\}^\{2\}for each

p=\(t,e\)∈𝒫ℓp=\(t,e\)\\in\\mathcal\{P\}\_\{\\ell\}
18:

𝒮ℓ,e∗←TopK​\(𝒫ℓ,e,\{vp\},Kℓ,e∗\)\\mathcal\{S\}\_\{\\ell,e\}^\{\*\}\\leftarrow\\textsc\{TopK\}\(\\mathcal\{P\}\_\{\\ell,e\},\\\{v\_\{p\}\\\},K\_\{\\ell,e\}^\{\*\}\)for each expert

ee
19:

\(ΠℓB,\{𝒮ℓ,e,r∗\}\)←PruneAndRoute​\(⋃e𝒮ℓ,e∗,ΠℓA,𝒞ℓ\)\(\\Pi\_\{\\ell\}^\{B\},\\\{\\mathcal\{S\}\_\{\\ell,e,r\}^\{\*\}\\\}\)\\leftarrow\\textsc\{PruneAndRoute\}\(\\bigcup\_\{e\}\\mathcal\{S\}\_\{\\ell,e\}^\{\*\},\\Pi\_\{\\ell\}^\{A\},\\mathcal\{C\}\_\{\\ell\}\)
20:

\{\(Gℓ,e,r,Φℓ,e,r\)\}←CompactAndAllToAll​\(\{𝒮ℓ,e,r∗\},𝒞ℓ\)\\\{\(G\_\{\\ell,e,r\},\\Phi\_\{\\ell,e,r\}\)\\\}\\leftarrow\\textsc\{CompactAndAllToAll\}\(\\\{\\mathcal\{S\}\_\{\\ell,e,r\}^\{\*\}\\\},\\mathcal\{C\}\_\{\\ell\}\)
21:

\(Δ​Wℓ,e,r,Dℓ,e,r\)←PackedExpertBackward​\(Gℓ,e,r,Φℓ,e,r\)\(\\Delta W\_\{\\ell,e,r\},D\_\{\\ell,e,r\}\)\\leftarrow\\textsc\{PackedExpertBackward\}\(G\_\{\\ell,e,r\},\\Phi\_\{\\ell,e,r\}\)for each active

\(e,r\)\(e,r\)
22:

𝒬CPU←𝒬CPU∪\{\(e,∑rΔ​Wℓ,e,r\)\}e\\mathcal\{Q\}\_\{\\mathrm\{CPU\}\}\\leftarrow\\mathcal\{Q\}\_\{\\mathrm\{CPU\}\}\\cup\\\{\(e,\\sum\_\{r\}\\Delta W\_\{\\ell,e,r\}\)\\\}\_\{e\}
23:ScatterWithZeros\(

OPEN\)​\{Dℓ,e,r\},\{𝒮ℓ,e,r∗\}\)\{\\\{D\_\{\\ell,e,r\}\\\},\\\{\\mathcal\{S\}\_\{\\ell,e,r\}^\{\*\}\\\}\}
24:return

\(⋃e𝒮ℓ,e∗,ρℓ∗,ΠℓB\)\(\\bigcup\_\{e\}\\mathcal\{S\}\_\{\\ell,e\}^\{\*\},\\rho\_\{\\ell\}^\{\*\},\\Pi\_\{\\ell\}^\{B\}\)

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[3](https://arxiv.org/html/2610.37720#S3.F3)details the quota admission and execution procedure in Algorithm[1](https://arxiv.org/html/2610.37720#alg1)\. For layerℓ\\ell, the controller counts the full routed loadNℓ,eN\_\{\\ell,e\}of each expert and allocates a candidate retained countKℓK\_\{\\ell\}into quotas satisfying0≤Kℓ,e≤Nℓ,e0\\leq K\_\{\\ell,e\}\\leq N\_\{\\ell,e\}and∑eKℓ,e=Kℓ\\sum\_\{e\}K\_\{\\ell,e\}=K\_\{\\ell\}\. 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ΠℓF\\Pi\_\{\\ell\}^\{F\}, feasibility includes the complete forward workload 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\.

These admitted quotas remain fixed when backward begins\. Using the incoming gradients, saved activations, and routing weights, the controller computes the LRM scorevpv\_\{p\}and selects theKℓ,eK\_\{\\ell,e\}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ΠℓF\\Pi\_\{\\ell\}^\{F\}, using the retained workload and the hosting GPUs’ available capacity\. Ifmℓ,e,im\_\{\\ell,e,i\}denotes the count assigned to eligible instanceii, routing preserves∑imℓ,e,i=Kℓ,e\\sum\_\{i\}m\_\{\\ell,e,i\}=K\_\{\\ell,e\}\. 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:mℓ,e,im\_\{\\ell,e,i\}determines the row dimension of input\-gradient products and the reduction dimension of weight\-gradient 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

![Backward token--expert contributions are selected and packed into contiguous expert segments for variable-size grouped matrix multiplication.](https://arxiv.org/html/2610.37720v1/figure4_block_arrows.png)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\-MMgrouped 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\.Backward token–expert contributions are selected and packed into contiguous expert segments for variable\-size grouped matrix multiplication\.Figure[4](https://arxiv.org/html/2610.37720#S3.F4)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 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 zero\-masked 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](https://arxiv.org/html/2610.37720#S3.F4), the original expert loads are reduced from\[4,2,3\]\[4,2,3\]to\[3,1,2\]\[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\.

## 4\.LRMoEEvaluation

### 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 dimensionH=2048H=2048, and expert intermediate dimensionF=768F=768\. We compare the Megatron\-LM fixed\-ownership strategy, native UltraEP, an independent implementation of SYMI, andLRMoEunder 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\.

### 4\.2\.Convergence Analysis

![Perplexity curves over 10,000 training steps for two models at different backward recalibration rates.](https://arxiv.org/html/2610.37720v1/figure5_font_only.png)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\.Figure[5](https://arxiv.org/html/2610.37720#S4.F5)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 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

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\.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](https://arxiv.org/html/2610.37720#S4.F5)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 communication primitives and hits restrictive memory boundaries under our sustained long\-sequence training setup\. Table[1](https://arxiv.org/html/2610.37720#S4.T1)reports the cumulative MoE\-layer training time and the step at which each target is first reached\.LRMoEachieves 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, despiteLRMoErequiring 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,LRMoEactively 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\.Throughput Analysis

![Effective backward throughput versus batch-size pressure, with a marked modeled memory-capacity boundary.](https://arxiv.org/html/2610.37720v1/figure6_font_only.png)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[6](https://arxiv.org/html/2610.37720#S4.F6)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 of16×16\\times\. 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 at20×20\\times, compared with 58\.4% for UltraEP, 44\.9% for SYMI, 40\.0% for 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

![Stacked bars compare memory used by base expert placement, extra replicas, and token buffers across methods.](https://arxiv.org/html/2610.37720v1/figure7_dynamic_memory_clean.png)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 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 forLRMoE\. 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 underLRMoE\. 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 contribution\-aware selection with physical compaction,LRMoEeliminates 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

![Stacked bars compare MoE computation, dispatch All-to-All, and combine All-to-All latency across methods.](https://arxiv.org/html/2610.37720v1/figure8_latency_three_components_clean.png)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\.Using the same workload and resource constraints, Figure[8](https://arxiv.org/html/2610.37720#S4.F8)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 forLRMoE\. Specifically,LRMoEachieves an expert computation latency of 7\.26,ms, a dispatch latency of 2\.58,ms, and a combine latency of 2\.85,ms\.

These results demonstrate the concrete runtime advantages of our design\. By combining contribution\-aware selection with physical compaction,LRMoEeliminates 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\. Consequently,LRMoEreduces 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](https://arxiv.org/html/2610.37720#S4.T1)\.

### 4\.7\.Ablation Study

Table 2\.MeanΔ\\DeltaPPL over 10000 steps\.Table[2](https://arxiv.org/html/2610.37720#S4.T2)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Δ\\DeltaPPL 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\.

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

## 5\.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\([Lewis et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib4);[Zhou et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib5)\)\. StableMoE distills and freezes a router to stabilize assignments\([Dai et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib27)\)\. Gating Dropout bypasses global routing for some tokens to reduce communication\([Liu et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib20)\)\. These methods change forward routing\. FasterMoE shadows popular experts and can alter expert selection to avoid network congestion\([He et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib9)\)\. FlexMoE adapts placement, while SmartMoE searches hybrid parallelization plans\([Nie et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib10);[Zhai et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib11)\)\. SYMI separates parameters from optimizer states, and LAER\-MoE reconstructs expert instances from sharded parameters\([Skiadopoulos et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib13);[Liu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib12)\)\. UltraEP and topology\-aware TopoEP respond to realized post\-gating loads\([Wei et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib14);[Zhu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib24)\)\.LRMoEinstead selects backward pairs under quotas and balances them across instances of their original experts\.

Selective backward computation\.Selective\-Backprop prioritizes high\-loss training examples\([Jiang et al\., 2019](https://arxiv.org/html/2610.37720#bib.bib19)\)\. In contrast,LRMoEranks routed pairs within each expert’s quota\. Its projection\-level score uses the factorization underlying efficient per\-example gradient\-norm computation\([Goodfellow, 2015](https://arxiv.org/html/2610.37720#bib.bib17)\), applied to saved projection inputs and weighted incoming gradients\.

Memory and communication reduction\.ZeRO partitions training states to remove redundancy\([Rajbhandari et al\., 2020](https://arxiv.org/html/2610.37720#bib.bib31)\)\. ZeRO\-Offload and ZeRO\-Infinity extend memory management across GPU, CPU, and NVMe\([Ren et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib23);[Rajbhandari et al\., 2021](https://arxiv.org/html/2610.37720#bib.bib34)\)\. ES\-MoE offloads expert parameters and pipelines execution\([Kim et al\., 2024](https://arxiv.org/html/2610.37720#bib.bib7)\)\. FlashAttention reduces attention HBM accesses through IO\-aware tiling\([Dao et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib33)\)\. These methods do not select expert\-update contributions\. RCMoE quantizes all\-to\-all exchanges and sparsifies all\-reduce gradients\([Wu et al\., 2026](https://arxiv.org/html/2610.37720#bib.bib18)\)\. 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\([Narayanan et al\., 2019](https://arxiv.org/html/2610.37720#bib.bib30);[Zheng et al\., 2022](https://arxiv.org/html/2610.37720#bib.bib29)\)\. Within MoE layers, MegaBlocks uses block\-sparse computation for irregular loads without capacity\-based token dropping\([Gale et al\., 2023](https://arxiv.org/html/2610.37720#bib.bib8)\)\. Lancet, COMET, and FSMoE co\-schedule and overlap communication with computation across different granularities\([Jiang et al\., 2024](https://arxiv.org/html/2610.37720#bib.bib15);[Zhang et al\., 2025](https://arxiv.org/html/2610.37720#bib.bib16);[Pan et al\., 2025](https://arxiv.org/html/2610.37720#bib.bib32)\)\.LRMoEchanges 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\.

## 6\.Conclusion

In this work, we presentLRMoEto tackle load imbalances and memory bottlenecks in Expert Parallelism\. Unlike balancers that merely redistribute a fixed workload,LRMoEtightly 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,LRMoEadvances MoE balancing by jointly controlling the size and spatial distribution of backward computation\.

## References

- Covingtonet al\.\(2016\)P\. Covington, J\. Adams, and E\. SarginDeep neural networks for youtube recommendations\.InProceedings of the 10th ACM Conference on Recommender Systems, Boston, MA, USA, September 15\-19, 2016,S\. Sen, W\. Geyer, J\. Freyne, and P\. Castells \(Eds\.\),pp\. 191–198\.External Links:[Link](https://doi.org/10.1145/2959100.2959190),[Document](https://dx.doi.org/10.1145/2959100.2959190)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1)\.
- Daiet al\.\(2024\)D\. Dai, C\. Deng, C\. Zhao, R\. X\. Xu, H\. Gao, D\. Chen, J\. Li, W\. Zeng, X\. Yu, Y\. Wu, Z\. Xie, Y\. K\. Li, P\. Huang, F\. Luo, C\. Ruan, Z\. Sui, and W\. LiangDeepSeekMoE: towards ultimate expert specialization in mixture\-of\-experts language models\.InProceedings of the 62nd Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\), ACL 2024, Bangkok, Thailand, August 11\-16, 2024,L\. Ku, A\. Martins, and V\. Srikumar \(Eds\.\),pp\. 1280–1297\.External Links:[Link](https://doi.org/10.18653/v1/2024.acl-long.70),[Document](https://dx.doi.org/10.18653/V1/2024.ACL-LONG.70)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1)\.
- Daiet al\.\(2022\)D\. Dai, L\. Dong, S\. Ma, B\. Zheng, Z\. Sui, B\. Chang, and F\. WeiStableMoE: stable routing strategy for mixture of experts\.InProceedings of the 60th Annual Meeting of the Association for Computational Linguistics \(Volume 1: Long Papers\), ACL 2022, Dublin, Ireland, May 22\-27, 2022,S\. Muresan, P\. Nakov, and A\. Villavicencio \(Eds\.\),pp\. 7085–7095\.External Links:[Link](https://doi.org/10.18653/v1/2022.acl-long.489),[Document](https://dx.doi.org/10.18653/V1/2022.ACL-LONG.489)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Daoet al\.\(2022\)T\. Dao, D\. Y\. Fu, S\. Ermon, A\. Rudra, and C\. RéFlashAttention: fast and memory\-efficient exact attention with io\-awareness\.InAdvances 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,S\. Koyejo, S\. Mohamed, A\. Agarwal, D\. Belgrave, K\. Cho, and A\. Oh \(Eds\.\),External Links:[Link](http://papers.nips.cc/paper_files/paper/2022/hash/67d57c32e20fd0a7a302cb81d36e40d5-Abstract-Conference.html)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p3.1)\.
- Feduset al\.\(2022\)W\. Fedus, B\. Zoph, and N\. ShazeerSwitch transformers: scaling to trillion parameter models with simple and efficient sparsity\.J\. Mach\. Learn\. Res\.23,pp\. 120:1–120:39\.External Links:[Link](https://jmlr.org/papers/v23/21-0998.html)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1)\.
- Galeet al\.\(2023\)T\. Gale, D\. Narayanan, C\. Young, and M\. ZahariaMegaBlocks: efficient sparse training with mixture\-of\-experts\.InProceedings of the Sixth Conference on Machine Learning and Systems, MLSys 2023, Miami, FL, USA, June 4\-8, 2023,D\. Song, M\. Carbin, and T\. Chen \(Eds\.\),External Links:[Link](https://proceedings.mlsys.org/paper_files/paper/2023/hash/5a54f79333768effe7e8927bcccffe40-Abstract-mlsys2023.html)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p2.1),[§1](https://arxiv.org/html/2610.37720#S1.p4.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p4.1)\.
- Goodfellow \(2015\)I\. J\. GoodfellowEfficient per\-example gradient computations\.CoRRabs/1510\.01799\.External Links:[Link](https://doi.org/10.48550/arXiv.1510.01799),[Document](https://dx.doi.org/10.48550/ARXIV.1510.01799),1510\.01799Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p4.1),[§2\.3](https://arxiv.org/html/2610.37720#S2.SS3.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p2.1)\.
- Heet al\.\(2022\)J\. He, J\. Zhai, T\. Antunes, H\. Wang, F\. Luo, S\. Shi, and Q\. LiFasterMoE: modeling and optimizing training of large\-scale dynamic pre\-trained models\.InPPoPP ’22: 27th ACM SIGPLAN Symposium on Principles and Practice of Parallel Programming, Seoul, Republic of Korea, April 2 \- 6, 2022,J\. Lee, K\. Agrawal, and M\. F\. Spear \(Eds\.\),pp\. 120–134\.External Links:[Link](https://doi.org/10.1145/3503221.3508418),[Document](https://dx.doi.org/10.1145/3503221.3508418)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Huanget al\.\(2013\)P\. Huang, X\. He, J\. Gao, L\. Deng, A\. Acero, and L\. P\. HeckLearning deep structured semantic models for web search using clickthrough data\.In22nd ACM International Conference on Information and Knowledge Management, CIKM’13, San Francisco, CA, USA, October 27 \- November 1, 2013,Q\. He, A\. Iyengar, W\. Nejdl, J\. Pei, and R\. Rastogi \(Eds\.\),pp\. 2333–2338\.External Links:[Link](https://doi.org/10.1145/2505515.2505665),[Document](https://dx.doi.org/10.1145/2505515.2505665)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1)\.
- Jianget al\.\(2019\)A\. H\. Jiang, D\. L\.\-K\. Wong, G\. Zhou, D\. G\. Andersen, J\. Dean, G\. R\. Ganger, G\. Joshi, M\. Kaminsky, M\. Kozuch, Z\. C\. Lipton, and P\. PillaiAccelerating deep learning by focusing on the biggest losers\.CoRRabs/1910\.00762\.External Links:[Link](https://doi.org/10.48550/arXiv.1910.00762),[Document](https://dx.doi.org/10.48550/ARXIV.1910.00762),1910\.00762Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p4.1),[§5](https://arxiv.org/html/2610.37720#S5.p2.1)\.
- Jianget al\.\(2024\)C\. Jiang, Y\. Tian, Z\. Jia, S\. Zheng, C\. Wu, and Y\. WangLancet: accelerating mixture\-of\-experts training via whole graph computation\-communication overlapping\.InProceedings of the Seventh Annual Conference on Machine Learning and Systems, MLSys 2024, Santa Clara, CA, USA, May 13\-16, 2024,P\. B\. Gibbons, G\. Pekhimenko, and C\. D\. Sa \(Eds\.\),External Links:[Link](https://proceedings.mlsys.org/paper_files/paper/2024/hash/339caf45a6fa281cae8adc6465343464-Abstract-Conference.html)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p4.1)\.
- Kimet al\.\(2024\)Y\. Kim, H\. Lim, and D\. HanScaling beyond the GPU memory limit for large mixture\-of\-experts model training\.InForty\-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21\-27, 2024,R\. Salakhutdinov, Z\. Kolter, K\. A\. Heller, A\. Weller, N\. Oliver, J\. Scarlett, and F\. Berkenkamp \(Eds\.\),Proceedings of Machine Learning Research, Vol\.235,pp\. 24342–24353\.External Links:[Link](https://proceedings.mlr.press/v235/kim24w.html)Cited by:[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p3.1)\.
- Lepikhinet al\.\(2021\)D\. Lepikhin, H\. Lee, Y\. Xu, D\. Chen, O\. Firat, Y\. Huang, M\. Krikun, N\. Shazeer, and Z\. ChenGShard: scaling giant models with conditional computation and automatic sharding\.In9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3\-7, 2021,External Links:[Link](https://openreview.net/forum?id=qrwe7XHTmYb)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1),[§1](https://arxiv.org/html/2610.37720#S1.p2.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1)\.
- Lewiset al\.\(2021\)M\. Lewis, S\. Bhosale, T\. Dettmers, N\. Goyal, and L\. ZettlemoyerBASE layers: simplifying training of large, sparse models\.InProceedings of the 38th International Conference on Machine Learning, ICML 2021, 18\-24 July 2021, Virtual Event,M\. Meila and T\. Zhang \(Eds\.\),Proceedings of Machine Learning Research, Vol\.139,pp\. 6265–6274\.External Links:[Link](http://proceedings.mlr.press/v139/lewis21a.html)Cited by:[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Liuet al\.\(2022\)R\. Liu, Y\. J\. Kim, A\. Muzio, and H\. HassanGating dropout: communication\-efficient regularization for sparsely activated transformers\.InInternational Conference on Machine Learning, ICML 2022, 17\-23 July 2022, Baltimore, Maryland, USA,K\. Chaudhuri, S\. Jegelka, L\. Song, C\. Szepesvári, G\. Niu, and S\. Sabato \(Eds\.\),Proceedings of Machine Learning Research, Vol\.162,pp\. 13782–13792\.External Links:[Link](https://proceedings.mlr.press/v162/liu22g.html)Cited by:[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Liuet al\.\(2025\)X\. Liu, Y\. Wang, F\. Fu, X\. Miao, S\. Zhu, X\. Nie, and B\. CuiNetMoE: accelerating moe training through dynamic sample placement\.InThe Thirteenth International Conference on Learning Representations, ICLR 2025, Singapore, April 24\-28, 2025,External Links:[Link](https://openreview.net/forum?id=1qP3lsatCR)Cited by:[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1)\.
- Liuet al\.\(2026\)X\. Liu, Y\. Wang, F\. Fu, X\. Xiao, H\. Li, J\. Li, and B\. CuiLAER\-moe: load\-adaptive expert re\-layout for efficient mixture\-of\-experts training\.InProceedings 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,B\. C\. Lee, H\. Xu, M\. Silberstein, and B\. Li \(Eds\.\),pp\. 1055–1072\.External Links:[Link](https://doi.org/10.1145/3779212.3790180),[Document](https://dx.doi.org/10.1145/3779212.3790180)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Narayananet al\.\(2019\)D\. Narayanan, A\. Harlap, A\. Phanishayee, V\. Seshadri, N\. R\. Devanur, G\. R\. Ganger, P\. B\. Gibbons, and M\. ZahariaPipeDream: generalized pipeline parallelism for DNN training\.InProceedings of the 27th ACM Symposium on Operating Systems Principles, SOSP 2019, Huntsville, ON, Canada, October 27\-30, 2019,T\. Brecht and C\. Williamson \(Eds\.\),pp\. 1–15\.External Links:[Link](https://doi.org/10.1145/3341301.3359646),[Document](https://dx.doi.org/10.1145/3341301.3359646)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p4.1)\.
- Nieet al\.\(2023\)X\. Nie, X\. Miao, Z\. Wang, Z\. Yang, J\. Xue, L\. Ma, G\. Cao, and B\. CuiFlexMoE: scaling large\-scale sparse pre\-trained model training via dynamic device placement\.Proc\. ACM Manag\. Data1\(1\),pp\. 110:1–110:19\.External Links:[Link](https://doi.org/10.1145/3588964),[Document](https://dx.doi.org/10.1145/3588964)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Panet al\.\(2025\)X\. Pan, W\. Lin, L\. Zhang, S\. Shi, Z\. Tang, R\. Wang, B\. Li, and X\. ChuFSMoE: A flexible and scalable training system for sparse mixture\-of\-experts models\.InProceedings 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,L\. Eeckhout, G\. Smaragdakis, K\. Liang, A\. Sampson, M\. A\. Kim, and C\. J\. Rossbach \(Eds\.\),pp\. 524–539\.External Links:[Link](https://doi.org/10.1145/3669940.3707272),[Document](https://dx.doi.org/10.1145/3669940.3707272)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p4.1)\.
- Rajbhandariet al\.\(2022\)S\. Rajbhandari, C\. Li, Z\. Yao, M\. Zhang, R\. Y\. Aminabadi, A\. A\. Awan, J\. Rasley, and Y\. HeDeepSpeed\-moe: advancing mixture\-of\-experts inference and training to power next\-generation AI scale\.InInternational Conference on Machine Learning, ICML 2022, 17\-23 July 2022, Baltimore, Maryland, USA,K\. Chaudhuri, S\. Jegelka, L\. Song, C\. Szepesvári, G\. Niu, and S\. Sabato \(Eds\.\),Proceedings of Machine Learning Research, Vol\.162,pp\. 18332–18346\.External Links:[Link](https://proceedings.mlr.press/v162/rajbhandari22a.html)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1),[§1](https://arxiv.org/html/2610.37720#S1.p2.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1)\.
- Rajbhandariet al\.\(2020\)S\. Rajbhandari, J\. Rasley, O\. Ruwase, and Y\. HeZeRO: memory optimizations toward training trillion parameter models\.InProceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, SC 2020, Virtual Event / Atlanta, Georgia, USA, November 9\-19, 2020,C\. Cuicchi, I\. Qualters, and W\. T\. Kramer \(Eds\.\),pp\. 20\.External Links:[Link](https://doi.org/10.1109/SC41405.2020.00024),[Document](https://dx.doi.org/10.1109/SC41405.2020.00024)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p3.1)\.
- Rajbhandariet al\.\(2021\)S\. Rajbhandari, O\. Ruwase, J\. Rasley, S\. Smith, and Y\. HeZeRO\-infinity: breaking the GPU memory wall for extreme scale deep learning\.InInternational Conference for High Performance Computing, Networking, Storage and Analysis, SC 2021, St\. Louis, Missouri, USA, November 14\-19, 2021,B\. R\. de Supinski, M\. W\. Hall, and T\. Gamblin \(Eds\.\),pp\. 59\.External Links:[Link](https://doi.org/10.1145/3458817.3476205),[Document](https://dx.doi.org/10.1145/3458817.3476205)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p3.1)\.
- Renet al\.\(2021\)J\. Ren, S\. Rajbhandari, R\. Y\. Aminabadi, O\. Ruwase, S\. Yang, M\. Zhang, D\. Li, and Y\. HeZeRO\-offload: democratizing billion\-scale model training\.InProceedings of the 2021 USENIX Annual Technical Conference, USENIX ATC 2021, July 14\-16, 2021,I\. Calciu and G\. Kuenning \(Eds\.\),pp\. 551–564\.External Links:[Link](https://www.usenix.org/conference/atc21/presentation/ren-jie)Cited by:[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p3.1)\.
- Shazeeret al\.\(2017\)N\. Shazeer, A\. Mirhoseini, K\. Maziarz, A\. Davis, Q\. V\. Le, G\. E\. Hinton, and J\. DeanOutrageously large neural networks: the sparsely\-gated mixture\-of\-experts layer\.In5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24\-26, 2017, Conference Track Proceedings,External Links:[Link](https://openreview.net/forum?id=B1ckMDqlg)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p1.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1)\.
- Skiadopouloset al\.\(2026\)A\. Skiadopoulos, M\. Zhao, S\. Gandhi, T\. Norrie, S\. Mukherjee, and C\. KozyrakisSYMI: efficient mixture\-of\-experts training via model and optimizer state decoupling\.In23rd USENIX Symposium on Networked Systems Design and Implementation, NSDI 2026, Renton, WA, May 4\-6, 2026,S\. Kandula and H\. Weatherspoon \(Eds\.\),pp\. 75–92\.External Links:[Link](https://www.usenix.org/conference/nsdi26/presentation/skiadopoulos)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Vaswaniet al\.\(2017\)A\. Vaswani, N\. Shazeer, N\. Parmar, J\. Uszkoreit, L\. Jones, A\. N\. Gomez, L\. Kaiser, and I\. PolosukhinAttention is all you need\.InAdvances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4\-9, 2017, Long Beach, CA, USA,I\. Guyon, U\. von Luxburg, S\. Bengio, H\. M\. Wallach, R\. Fergus, S\. V\. N\. Vishwanathan, and R\. Garnett \(Eds\.\),pp\. 5998–6008\.External Links:[Link](https://proceedings.neurips.cc/paper/2017/hash/3f5ee243547dee91fbd053c1c4a845aa-Abstract.html)Cited by:[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1)\.
- Weiet al\.\(2026\)X\. Wei, C\. Jin, T\. Dai, Y\. Zhong, S\. Yu, C\. Yang, B\. Wu, Z\. Zhang, J\. Mai, Q\. Zhu, Z\. Li, Y\. Liu, and G\. LuoUltraEP: unleash moe training and inference on rack\-scale nodes with near\-optimal load balancing\.CoRRabs/2606\.04101\.External Links:[Link](https://doi.org/10.48550/arXiv.2606.04101),[Document](https://dx.doi.org/10.48550/ARXIV.2606.04101),2606\.04101Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p2.1),[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Wuet al\.\(2026\)D\. Wu, X\. Cai, J\. Tan, J\. Jia, G\. Tan, D\. Tao, W\. Xia, and Z\. TianRCMoE: A communication\-efficient random compression framework for resource\-constrained mixture\-of\-experts training\.InFortieth 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,S\. Koenig, C\. Jenkins, and M\. E\. Taylor \(Eds\.\),pp\. 26876–26884\.External Links:[Link](https://doi.org/10.1609/aaai.v40i32.39899),[Document](https://dx.doi.org/10.1609/AAAI.V40I32.39899)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p3.1)\.
- Zenget al\.\(2025\)Y\. Zeng, C\. Huang, Y\. Mei, L\. Zhang, T\. Su, W\. Ye, W\. Shi, and S\. WangEfficientMoE: optimizing mixture\-of\-experts model training with adaptive load balance\.IEEE Trans\. Parallel Distributed Syst\.36\(4\),pp\. 677–688\.External Links:[Link](https://doi.org/10.1109/TPDS.2025.3539297),[Document](https://dx.doi.org/10.1109/TPDS.2025.3539297)Cited by:[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1)\.
- Zhaiet al\.\(2023\)M\. Zhai, J\. He, Z\. Ma, Z\. Zong, R\. Zhang, and J\. ZhaiSmartMoE: efficiently training sparsely\-activated models through combining offline and online parallelization\.InProceedings of the 2023 USENIX Annual Technical Conference, USENIX ATC 2023, Boston, MA, USA, July 10\-12, 2023,J\. Lawall and D\. Williams \(Eds\.\),pp\. 961–975\.External Links:[Link](https://www.usenix.org/conference/atc23/presentation/zhai)Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§2\.2](https://arxiv.org/html/2610.37720#S2.SS2.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Zhanget al\.\(2025\)S\. Zhang, N\. Zheng, H\. Lin, Z\. Jiang, W\. Bao, C\. Jiang, Q\. Hou, W\. Cui, S\. Zheng, L\. Chang, Q\. Chen, and X\. LiuCOMET: fine\-grained computation\-communication overlapping for mixture\-of\-experts\.InProceedings of the Eighth Conference on Machine Learning and Systems, MLSys 2025, Santa Clara, CA, USA, May 12\-15, 2025,M\. Zaharia, G\. Joshi, and Y\. \(\. Lin \(Eds\.\),External Links:[Link](https://openreview.net/forum?id=fGgQS5VW09)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p4.1)\.
- Zhenget al\.\(2022\)L\. Zheng, Z\. Li, H\. Zhang, Y\. Zhuang, Z\. Chen, Y\. Huang, Y\. Wang, Y\. Xu, D\. Zhuo, E\. P\. Xing, J\. E\. Gonzalez, and I\. StoicaAlpa: automating inter\- and intra\-operator parallelism for distributed deep learning\.In16th USENIX Symposium on Operating Systems Design and Implementation, OSDI 2022, Carlsbad, CA, USA, July 11\-13, 2022,M\. K\. Aguilera and H\. Weatherspoon \(Eds\.\),pp\. 559–578\.External Links:[Link](https://www.usenix.org/conference/osdi22/presentation/zheng-lianmin)Cited by:[§5](https://arxiv.org/html/2610.37720#S5.p4.1)\.
- Zhouet al\.\(2022\)Y\. Zhou, T\. Lei, H\. Liu, N\. Du, Y\. Huang, V\. Y\. Zhao, A\. M\. Dai, Z\. Chen, Q\. V\. Le, and J\. LaudonMixture\-of\-experts with expert choice routing\.InAdvances 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,S\. Koyejo, S\. Mohamed, A\. Agarwal, D\. Belgrave, K\. Cho, and A\. Oh \(Eds\.\),External Links:[Link](http://papers.nips.cc/paper_files/paper/2022/hash/2f00ecd787b432c1d36f3de9800728eb-Abstract-Conference.html)Cited by:[§2\.1](https://arxiv.org/html/2610.37720#S2.SS1.p1.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.
- Zhuet al\.\(2026\)J\. Zhu, X\. Zhao, G\. Zhao, H\. Xu, Y\. Fei, and J\. FangTopoEP: topology\-aware load balancing for expert\-parallel moe training\.CoRRabs/2609\.35481\.External Links:[Link](https://doi.org/10.48550/arXiv.2609.35481),[Document](https://dx.doi.org/10.48550/ARXIV.2609.35481),2609\.35481Cited by:[§1](https://arxiv.org/html/2610.37720#S1.p2.1),[§1](https://arxiv.org/html/2610.37720#S1.p3.1),[§5](https://arxiv.org/html/2610.37720#S5.p1.1)\.

## 7\.Derivation and Optimality of Token–Expert Recalibration

This appendix proves Proposition[3\.1](https://arxiv.org/html/2610.37720#S3.Thmtheorem1)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 projection\-gradient 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 expertee\. Let

\(4\)We∈ℝdout×din,ϕu∈ℝdin,gt∈ℝdout,W\_\{e\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{out\}\}\\times d\_\{\\mathrm\{in\}\}\},\\qquad\\phi\_\{u\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{in\}\}\},\\qquad g\_\{t\}\\in\\mathbb\{R\}^\{d\_\{\\mathrm\{out\}\}\},whereu=\(t,e\)u=\(t,e\)is a routed token–expert pair\. The saved vectorϕu\\phi\_\{u\}is the actual input to the protected projection\. We takegtg\_\{t\}to be the upstream cotangent before multiplication by the fixed routing weightαu\\alpha\_\{u\}\. Equivalently, an implementation may expose the already weighted cotangent

\(5\)δu:=αu​gt,\\delta\_\{u\}:=\\alpha\_\{u\}g\_\{t\},in that representation, the routing weight must not be applied again\.

The contribution of pairuuto the routed expert output is

\(6\)yu=αu​We​ϕu\.y\_\{u\}=\\alpha\_\{u\}W\_\{e\}\\phi\_\{u\}\.The forward route,αu\\alpha\_\{u\}, andϕu\\phi\_\{u\}are fixed during recalibration\. For an arbitrary perturbationd​WedW\_\{e\}, the differential of the loss through this pair is

d​ℒu\\displaystyle d\\mathcal\{L\}\_\{u\}=gt⊤​d​yu=gt⊤​αu​\(d​We\)​ϕu\\displaystyle=g\_\{t\}^\{\\top\}dy\_\{u\}=g\_\{t\}^\{\\top\}\\alpha\_\{u\}\(dW\_\{e\}\)\\phi\_\{u\}\(7\)=tr⁡\(\(\(αu​gt\)​ϕu⊤\)⊤​d​We\)\.\\displaystyle=\\operatorname\{tr\}\\\!\\left\(\\bigl\(\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}\\bigr\)^\{\\top\}dW\_\{e\}\\right\)\.By the defining identityd​ℒ=tr⁡\(\(∇Wℒ\)⊤​d​W\)d\\mathcal\{L\}=\\operatorname\{tr\}\(\(\\nabla\_\{W\}\\mathcal\{L\}\)^\{\\top\}dW\), the isolated contribution of pairuuto the protected weight gradient is therefore

\(8\)Gu:=\[∇Weℒ\]u=\(αu​gt\)​ϕu⊤=δu​ϕu⊤\.G\_\{u\}:=\\left\[\\nabla\_\{W\_\{e\}\}\\mathcal\{L\}\\right\]\_\{u\}=\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}=\\delta\_\{u\}\\phi\_\{u\}^\{\\top\}\.
Let𝒰e\\mathcal\{U\}\_\{e\}denote all pairs contributing to the same protected projection\. Before omission, its routed\-pair gradient is

\(9\)Ge=∑v∈𝒰eGv\.G\_\{e\}=\\sum\_\{v\\in\\mathcal\{U\}\_\{e\}\}G\_\{v\}\.If only pairuuis omitted and reconstructed as zero while every other pair is unchanged, then

\(10\)Ge\(−u\)=Ge−Gu,Δu:=Ge\(−u\)−Ge=−Gu\.G\_\{e\}^\{\(\-u\)\}=G\_\{e\}\-G\_\{u\},\\qquad\\Delta\_\{u\}:=G\_\{e\}^\{\(\-u\)\}\-G\_\{e\}=\-G\_\{u\}\.The sign is immaterial after taking a norm, but Eq\. \([10](https://arxiv.org/html/2610.37720#S7.E10)\) makes clear that the metric measures a backward contribution removed from a fixed completed route, it does not rerun or alter the forward computation\.

For vectorsaaandbb, the squared Frobenius norm of their outer product factorizes as

∥a​b⊤∥F2\\displaystyle\\lVert ab^\{\\top\}\\rVert\_\{F\}^\{2\}=tr⁡\(\(a​b⊤\)⊤​\(a​b⊤\)\)\\displaystyle=\\operatorname\{tr\}\\\!\\left\(\(ab^\{\\top\}\)^\{\\top\}\(ab^\{\\top\}\)\\right\)\(11\)=tr⁡\(b⁡\(a⊤​a\)​b⊤\)=\(a⊤​a\)​\(b⊤​b\)=∥a∥22​∥b∥22\.\\displaystyle=\\operatorname\{tr\}\\\!\\left\(b\(a^\{\\top\}a\)b^\{\\top\}\\right\)=\(a^\{\\top\}a\)\(b^\{\\top\}b\)=\\lVert a\\rVert\_\{2\}^\{2\}\\lVert b\\rVert\_\{2\}^\{2\}\.Applying Eq\. \([11](https://arxiv.org/html/2610.37720#S7.E11)\) witha=αu​gta=\\alpha\_\{u\}g\_\{t\}andb=ϕub=\\phi\_\{u\}proves

∥Δu∥F2\\displaystyle\\lVert\\Delta\_\{u\}\\rVert\_\{F\}^\{2\}=‖\(αu​gt\)​ϕu⊤‖F2\\displaystyle=\\left\\lVert\(\\alpha\_\{u\}g\_\{t\}\)\\phi\_\{u\}^\{\\top\}\\right\\rVert\_\{F\}^\{2\}\(12\)=αu2​∥gt∥22​∥ϕu∥22=Ru,\\displaystyle=\\alpha\_\{u\}^\{2\}\\lVert g\_\{t\}\\rVert\_\{2\}^\{2\}\\lVert\\phi\_\{u\}\\rVert\_\{2\}^\{2\}=R\_\{u\},which is exactly Eq\. \([3](https://arxiv.org/html/2610.37720#S3.E3)\)\. Hence the three factors inRuR\_\{u\}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

Let𝒰\\mathcal\{U\}be the causally available token–expert pairs at the recalibration boundary, withN=\|𝒰\|N=\|\\mathcal\{U\}\|\. LetD⊆𝒰D\\subseteq\\mathcal\{U\}be the dropped set andS=𝒰∖DS=\\mathcal\{U\}\\setminus Dthe retained set\. All rows at this boundary have the same shape and encoding, and therefore the same communication cost\. For a prescribed drop countdd, or equivalently a retained countK=N−dK=N\-d, define the additive local calibration distortion

\(13\)ℰ^​\(D\):=∑u∈D∥Δu∥F2=∑u∈DRu\.\\widehat\{\\mathcal\{E\}\}\(D\):=\\sum\_\{u\\in D\}\\lVert\\Delta\_\{u\}\\rVert\_\{F\}^\{2\}=\\sum\_\{u\\in D\}R\_\{u\}\.
###### Proposition 7\.1 \(Optimal score\-based dropping\)\.

Under a fixed completed route, equal row costs, and a fixed drop countdd, any set containing theddsmallest scoresRuR\_\{u\}minimizes Eq\. \([13](https://arxiv.org/html/2610.37720#S7.E13)\)\. Equivalently, any retained set containing theK=N−dK=N\-dlargest scores maximizes preserved calibration value\. The solution is unique when there is a strict gap between theddth and\(d\+1\)\(d\+1\)st scores in ascending order, otherwise every difference between optimal solutions is confined to tied boundary scores\.

###### Proof\.

LetDDbe any dropped set with\|D\|=d\|D\|=d\. Suppose thatDDcontains a pairiiwhile retaining a pairjjsuch thatRi\>RjR\_\{i\}\>R\_\{j\}\. Replacingiibyjjgives

\(14\)D′=\(D∖\{i\}\)∪\{j\},D^\{\\prime\}=\(D\\setminus\\\{i\\\}\)\\cup\\\{j\\\},and changes the objective by

\(15\)ℰ^​\(D′\)−ℰ^​\(D\)=Rj−Ri<0\.\\widehat\{\\mathcal\{E\}\}\(D^\{\\prime\}\)\-\\widehat\{\\mathcal\{E\}\}\(D\)=R\_\{j\}\-R\_\{i\}<0\.Thus a set with such an inversion cannot be optimal\. Repeatedly applying this exchange terminates at a set containing theddsmallest scores, because each exchange strictly decreases a finite\-set objective\. Conversely, a set containing theddsmallest 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

\(16\)∑u∈DRu=∑u∈𝒰Ru−∑u∈SRu\.\\sum\_\{u\\in D\}R\_\{u\}=\\sum\_\{u\\in\\mathcal\{U\}\}R\_\{u\}\-\\sum\_\{u\\in S\}R\_\{u\}\.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 thatddequal\-size rows must be removed, no otherdd\-row mask preserves more of the declared pairwise calibration value\.

When the controller imposes a fixed drop countded\_\{e\}separately on each expert or overloaded path, the same proof applies independently inside every partition: dropping theded\_\{e\}smallest scores in each partition minimizes the sum of Eq\. \([13](https://arxiv.org/html/2610.37720#S7.E13)\) over all masks obeying those cardinalities\. More general placement and capacity constraints define a restricted feasible familyℱ\\mathcal\{F\}\. The corresponding selection problem is

\(17\)D∗∈arg⁡min⁡∑u∈DD∈ℱ⁡Ru\.D^\{\*\}\\in\\arg\\min\_\{D\\in\\mathcal\{F\}\}\\sum\_\{u\\in D\}R\_\{u\}\.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\.

### 7\.3\.Relation to the aggregate expert gradient

The single\-pair identity in Eq\. \([12](https://arxiv.org/html/2610.37720#S7.E12)\) is exact\. For several dropped pairs from the same expert, their contributions to a sharedWeW\_\{e\}add\. Here letD⊆𝒰eD\\subseteq\\mathcal\{U\}\_\{e\}andd=\|D\|d=\|D\|\. The actual change in the protected expert gradient is

\(18\)ΔD=−∑u∈DGu,\\Delta\_\{D\}=\-\\sum\_\{u\\in D\}G\_\{u\},and therefore

\(19\)∥ΔD∥F2\\displaystyle\\lVert\\Delta\_\{D\}\\rVert\_\{F\}^\{2\}=∑u∈DRu\+2​∑u<vu,v∈D⟨Gu,Gv⟩F\.\\displaystyle=\\sum\_\{u\\in D\}R\_\{u\}\+2\\sum\_\{\\begin\{subarray\}\{c\}u<v\\\\ u,v\\in D\\end\{subarray\}\}\\langle G\_\{u\},G\_\{v\}\\rangle\_\{F\}\.The cross terms can be positive or negative\. Consequently, Proposition[7\.1](https://arxiv.org/html/2610.37720#S7.Thmtheorem1)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 Frobenius\-orthogonal, all cross terms in Eq\. \([19](https://arxiv.org/html/2610.37720#S7.E19)\) vanish and∥ΔD∥F2=ℰ^​\(D\)\\lVert\\Delta\_\{D\}\\rVert\_\{F\}^\{2\}=\\widehat\{\\mathcal\{E\}\}\(D\)\. 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

\(20\)∥ΔD∥F2≤\(∑u∈DRu\)2≤d​∑u∈DRu\.\\lVert\\Delta\_\{D\}\\rVert\_\{F\}^\{2\}\\leq\\left\(\\sum\_\{u\\in D\}\\sqrt\{R\_\{u\}\}\\right\)^\{2\}\\leq d\\sum\_\{u\\in D\}R\_\{u\}\.For fixeddd, the smallest\-score rule also minimizes the final upper bound in Eq\. \([20](https://arxiv.org/html/2610.37720#S7.E20)\)\.

### 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,ϕu\\phi\_\{u\}is the actual input to the protected linear projection\. For a nonlinear multi\-layer expert, the same formula is exact for the selected projection but is not the squared norm of the expert’s complete parameter\-gradient 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](https://arxiv.org/html/2610.37720#S7.E17)\)\. Fourth, the objective protects a local projection\-gradient 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\.

Within these conditions, the result is exact:RuR\_\{u\}is the squared isolated perturbation caused by dropping pairuu, and our drop rule is optimal for the declared additive calibration objective at the required recalibration rate\.

相似文章

MoE$^2$-LoRA:当MoE模型遇上MoE风格的低秩适配

arXiv cs.CL

MoE2-LoRA 引入了一种双通道路由条件投影(Routing-Conditioned Projection)和一个全局 LoRA 专家池,以实现用于微调 MoE 模型的 MoE 风格低秩适配,在保留通用能力的同时取得了最先进的准确率。

EasyBalance:分布式MoE推理中的跨层负载均衡

arXiv cs.LG

本文提出了EasyBalance,一种用于分布式混合专家(MoE)推理的跨层负载均衡策略。该策略在不修改专家-设备映射的情况下,调度并联合执行来自不同层的工作负载,以减少GPU空闲时间。实验表明,该策略可将空闲时间减少40%以上。