PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.37720v1cs.AIcs.DCcs.LG

TL;DR

MoE training faces uneven routed workloads and limited GPU memory; placement and replication can redistribute that work but compete with token buffers and do not reduce fixed-set computation. LRMoE plans quotas and expert placement, selects and reroutes backward contributions, and reduces cumulative MoE-layer training time to fixed-loss targets by 21.2–44.8% versus state-of-the-art baselines.

  • Problem

    Placement and replication redistribute routed work but consume memory and leave the logical computation of a fixed admitted token–expert pair set unchanged.

    Source: §1 · p. 1, §2.2 · p. 2

  • Method

    LRMoE plans expert quotas and placement, selects token–expert pairs by estimated update contribution during backpropagation, and routes retained pairs to eligible expert instances.

    Source: §1 · p. 2, §3.3 · p. 5

  • Results

    21.2–44.8% lower cumulative MoE-layer training time to fixed-loss targets on Qwen3-30B-A3B and GLM4.5-106B-A12B versus state-of-the-art baselines.

    Source: §1 · p. 2

  • Takeaways & Limitations

    LRMoE jointly controls the size and spatial distribution of backward computation, with compact grouped GEMMs reducing communication payloads and computation costs without zero-padding overhead.

    Source: §6 · p. 9

    Sufficient aggregate service capacity does not guarantee feasible placement because destinations, replica and staging capacity, per-device memory, and network paths constrain assignments.

    Source: §2.2 · p. 2

Abstract

from arXiv · show

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.

1 Introduction

LRMoE addresses MoE training under fixed GPU resource budgets by coordinating expert placement, communication, and backward workload. It selects and routes retained backward work, reducing cumulative MoE-layer training time to fixed-loss targets by 21.2–44.8% versus state-of-the-art baselines.

1 Introduction

  • Expert replication redistributes routed work but consumes memory and adds state-transfer and gradient-synchronization costs, leaving backward demand potentially beyond capacity.For a fixed admitted pair set, placement and replication change where work executes, not its logical computation.Source: §1 · p. 1
  • LRM ranks token–expert pairs by a proxy derived from the squared norm of each pair’s weight-gradient contribution, without materializing individual weight-gradient matrices.Because the needed gradients are unavailable during forward execution, LRMoE preserves all routed forward computations and selects pairs during backpropagation.Source: §1 · p. 1
  • The selected backward update is approximate, so its effect on learning must be evaluated empirically rather than inferred from pairwise scores alone.Reducing work on an unconstrained path does not accelerate execution, and reducing work without rerouting can leave imbalance intact.Source: §1 · p. 2
  • LRMoE plans expert retention quotas and placement under resource constraints, selects high-scoring pairs during backpropagation, and routes retained pairs to eligible instances.Physical compaction forms smaller communication payloads and variable-size grouped GEMMs; an inverse mapping restores gradients to the expected layout.Source: §1 · p. 2
  • 21.2–44.8% lower cumulative MoE-layer training time to fixed-loss targets on Qwen3-30B-A3B and GLM4.5-106B-A12B versus state-of-the-art baselines.A separate GPU replay reports a 57.8% reduction in the subtotal of expert-placement, replica, and token-buffer memory relative to fixed ownership.Source: §1 · p. 2

2 Background and Motivation

MoE routing creates input-dependent expert loads, while placement and replication redistribute work without reducing the computation required by a fixed admitted pair set. LRMoE treats backward demand as adjustable and distributes retained work within capacity constraints.

2.1 MoE and Load Balancing

  • A routed pair is one token–expert selection; input-dependent routing produces unequal expert computation and all-to-all traffic, making the busiest GPU a straggler.Training-time routing regularizers encourage balance in expectation but cannot ensure balance for each realized layer invocation.Source: §2.1 · p. 2

2.2 Limitations of Existing Methods

  • Redistributing a fixed backward workload is constrained by aggregate service capacity and by legal destinations, replica and staging capacity, per-device memory, and network paths.Sufficient aggregate capacity is not sufficient for a feasible placement; overload and spatial infeasibility are distinct limitations.Source: §2.2 · p. 2

2.3 Key Insight

  • LRMoE treats backward demand as a contribution-aware scheduling variable, selecting retained pairs within quotas and routing them to eligible instances of their original experts.Physical compaction reduces communication payloads and expert-GEMM token dimensions; masking full-sized tensors alone need not provide these savings.Source: §2.3 · pp. 2–3

3 LRMoE Design

LRMoE plans resource-feasible backward quotas and expert placement during forward execution, then selects and reroutes retained pairs during backpropagation. It physically compacts retained work into smaller communication payloads and grouped GEMMs.

3.1 Design Overview

  • The design preserves every routed pair in the complete forward pass, then uses recorded activations, routing weights, pair identities, and parameter versions to support backward execution.Retained gradients and matching saved states are compacted into destination-specific payloads and expert-major segments; returned input-gradient contributions are scattered to original positions, while omitted contributions remain discarded.Source: §3.1 · p. 3

3.2 Load Recalibration Metric

  • LRM scores each routed pair by the squared-norm perturbation its omission causes to a protected expert weight gradient, combining routing strength, backward signal, and saved expert state.For fixed per-expert quotas and equal row costs, selecting the highest-scoring pairs minimizes the additive sum of isolated omission perturbations, not necessarily aggregate gradient error or subsequent training effects.Source: §3.2 · pp. 3–4

3.3 Layerwise Load Calibration

  • Forward planning tests retained-count candidates from full retention downward against GPU-memory, replica-placement, communication, and complete-forward-workload constraints, then fixes feasible per-expert quotas.Backward selects the highest-scoring pairs within each expert's quota and reroutes each retained pair exactly once to an eligible instance of its original expert.Source: §3.3 · pp. 4–5

3.4 Repacking GEMM with Load Calibration

  • Retained rows are packed into variable-size grouped GEMMs without padding, reducing expert-backward computation and the physical replicas needed to meet capacity.The logical expert set remains unchanged; LRMoE reduces redundant physical replicas and their associated parameter memory, temporary buffers, and communication.Source: Figure 4 · p. 5, §3.4 · p. 6

4 LRMoE Evaluation

LRMoE is evaluated for convergence, loss-matched MoE-layer training time, throughput, memory use, latency, and backward-selection quality. Across the reported settings, it reduces training time to fixed-loss targets while preserving convergence closely at recalibration rates up to C20.

4.1 Experimental Setup

  • Evaluation uses eight Hopper GPUs with EP8, BF16, an 80 GiB per-GPU memory budget, and complete forward and backward passes.The workload uses Qwen3-30B-A3B and GLM4.5-106B-A12B, 128 experts, top-8 routing, and sequence length 8,192.Source: §4.1 · p. 6

4.2 Convergence Analysis

  • At C20, Qwen3-30B-A3B’s final-1,000-step mean PPL is 1.015141 versus 1.015190 for C0; C30 reaches 1.016270.The passage reports that recalibration up to C20 closely preserves convergence; C30 is approximately 0.106% above C0.Source: Figure 5 · p. 6, §4.2 · p. 6

4.3 Convergence Time Analysis to Target Loss

  • 21.2–44.8% lower cumulative MoE-layer training time than SYMI across four fixed-loss targets, despite requiring more steps to reach each target.The reported times include complete forward and backward MoE-layer execution, not whole-model end-to-end training time.Source: Table 1 · p. 6, §4.3 · p. 7

4.4 Throughput Analysis

  • At 20× batch-size pressure, LRMoE maintains 100% effective throughput using 7–9 additional replicas under the modeled memory budget.Effective throughput is the fraction of the offered batch whose backward workload fits the modeled per-GPU memory budget.Source: Figure 6 · p. 7, §4.4 · p. 7

4.5 GPU Memory Breakdown

  • LRMoE’s displayed placement, replica, and token-buffer memory subtotal is 5.44 GiB, 57.8% below Megatron-LM.The comparison concerns displayed GPU allocations, not total device memory or host-memory optimizer state.Source: Figure 7 · p. 7, §4.5 · p. 7

4.6 Latency Breakdown

  • LRMoE reduces backward execution subtotal by 57.2% versus Megatron-LM, 57.9% versus SYMI, and 22.7% versus UltraEP.Physical compaction removes unselected rows before communication, and variable-size grouped GEMMs avoid padding overhead.Source: Figure 8 · p. 8, §4.6 · p. 8

4.7 Ablation Study

  • LRM yields the lowest mean ΔPPL on both models: 0.012474 for Qwen3-30B-A3B and 0.003464 for GLM4.5-106B-A12B.Both gradient-only criteria outperform random selection but yield larger deviations than LRM.Source: Table 2 · p. 8, §4.7 · p. 8

5 Related Work

Prior work addresses MoE routing, placement, memory, and parallel execution through varied strategies. LRMoE differs by selecting backward token–expert pairs and compacting retained rows rather than changing forward routing or merely redistributing a fixed workload.

MoE routing and load balancing

  • Routing approaches alter forward assignments, while LRMoE selects backward pairs within quotas and balances them across instances of their original experts.The passage contrasts LRMoE with BASE, Expert Choice, StableMoE, Gating Dropout, FasterMoE, FlexMoE, SmartMoE, SYMI, LAER-MoE, UltraEP, and TopoEP.Source: §5 · p. 8

Selective backward computation

  • Unlike Selective-Backprop’s high-loss example prioritization, LRMoE ranks routed pairs within each expert’s quota using saved projection inputs and weighted incoming gradients.Its projection-level score applies the factorization used for efficient per-example gradient-norm computation [7].Source: §5 · p. 8

Memory and communication reduction

  • LRMoE reduces token-dependent communication and expert computation through row selection without changing parameter dimensions.The passage distinguishes this from state partitioning, offloading, attention IO tiling, quantized exchanges, and sparse all-reduce gradients.Source: §5 · p. 8

Parallel execution and overlap

  • LRMoE packs retained backward rows into contiguous expert segments, reducing physical GEMM row counts rather than masking full-sized tensors.This contrasts with systems that use block-sparse computation or co-schedule and overlap communication with computation.Source: §5 · p. 8

6 Conclusion

LRMoE jointly controls the size and distribution of backward computation under memory and bandwidth constraints while retaining the full forward pass for contribution scoring.

6 Conclusion

  • LRMoE coordinates contribution-aware recalibration with capacity-aware routing to fit high-value backward pairs within planned quotas and admitted expert instances.The quotas and routing address both workload size and its spatial distribution.Source: §6 · p. 9
  • Compact grouped GEMMs translate retained pairs into smaller communication payloads and lower computation without zero-padding overhead.Source: §6 · p. 9
  • The full forward pass remains intact to support accurate contribution scoring during backward-work selection.Source: §6 · p. 9

7 Derivation and Optimality of Token–Expert Recalibration

The appendix derives an exact isolated-pair gradient contribution and proves when score-based dropping is optimal for its declared additive calibration objective. It also specifies the assumptions and limits of that result.

7.1 Pairwise projection-gradient contribution

  • The metric is the exact squared norm of a routed pair’s isolated contribution to the protected projection’s weight gradient.The derivation holds the completed forward route and saved projection input fixed during recalibration.Source: §7.1 · pp. 10–11

7.2 Optimality of the drop rule

  • Under a fixed completed route, equal row costs, and a fixed drop count, dropping the lowest-score pairs minimizes additive calibration distortion.For fixed per-expert quotas, the objective separates across experts, so the lowest scores are dropped within each expert.Source: §7.2 · p. 11

7.3 Relation to the aggregate expert gradient

  • The rule optimizes the sum of isolated perturbations, but does not generally minimize the realized norm of the shared expert-gradient error.Dropped-pair contributions share parameters, and their cross terms can be positive or negative; orthogonality or zero-mean cross terms recover an aggregate interpretation.Source: §7.3 · p. 11

7.4 Scope of the result

  • The result is exact for the protected projection and declared additive objective under its assumptions, not a guarantee about nonlinear loss, future optimizer states, or model quality.Unequal row costs, or replica, topology, and HBM constraints, require a more general selection problem.Source: §7.4 · pp. 11–12

Related papers

Loading 2610.37720v1…