LRMoE trims the MoE backward pass and speeds up training by 21–45%

A team of academic authors has proposed LRMoE, a system for training MoE models on nodes where GPU memory is short. MoE (Mixture-of-Experts) is an architecture in which a layer consists of dozens or hundreds of small expert subnetworks, and a router activates only a few of them for each token. The authors call one such token–expert selection a pair. The headline result: on Qwen3-30B-A3B and GLM4.5-106B-A12B, the cumulative training time of the MoE layers to reach a preset loss value fell by 21.2–44.8% compared with SYMI. SYMI is a previous system that separates the storage of model parameters from optimizer states in order to place experts more flexibly.

Why expert replicas don't fix load imbalance

Experts are spread across GPUs, and tokens travel to the right cards via all-to-all, an exchange in which every device talks to every other. Popular experts receive more tokens, so their GPUs become the bottleneck while the rest wait. A synchronous step runs at the speed of its most heavily loaded participant.

The usual answers are to move experts around or to replicate the most popular ones. FasterMoE, FlexMoE, SmartMoE, LAER-MoE, UltraEP and TopoEP all work this way. But replicas come at a cost:

• a replica takes up memory that is needed for token buffers and stored activations;
• copies have to be synchronized, and their gradients summed;
• the total amount of work does not shrink; it just moves to other cards.

The authors show that under a hard per-GPU limit, even a perfectly even layout can leave the total load higher than memory and network can handle. Imbalance and overall overload are different problems, and balancing solves only the first.

What LRMoE does: a full forward pass, a trimmed backward pass

The idea is to treat the backward-pass load not as an unavoidable consequence of routing but as a quantity that can be controlled. The process runs in two stages.

• The forward pass stays complete. All token–expert pairs are computed as usual, because no gradients exist yet at this stage and there is nothing to judge the value of pairs by.
• Quotas before computation. A controller sees how many tokens have arrived at each expert and, taking memory and network into account, decides how many pairs of each expert to keep in the backward pass. It also plans replica placement. The quota is fixed for the entire layer call.
• Selection by the LRM metric. When gradients reach the layer, each pair is scored by its contribution to the update of the expert's weights. A pair's contribution is the outer product of the incoming gradient and the saved input of the projection. Its squared norm therefore factors into a product of three terms: the squared routing weight, the squared gradient norm and the squared input norm. There is no need to compute full gradient matrices for this. Within the quota, the pairs with the highest scores are kept.
• Reassignment. An integer-programming algorithm distributes the remaining pairs across the permitted copies of their expert. The goal is to minimize the peak load on a single GPU. No new replicas appear in the backward pass.
• Physical compaction. The remaining rows of gradients and activations are gathered into dense blocks by recipient and expert. The all-to-all exchange becomes more compact, and matrix multiplications run as variable-size grouped GEMM without zero padding. The result is then returned to the original layout.

An important detail: simply masking the extra rows in full-size tensors gains nothing, in either time or memory. The benefit appears only when the actual sizes of transfers and matrices shrink.

What the experiments showed, and the catch

Besides the headline 21.2–44.8% figure, the authors ran a separate replay on GPUs and measured a 57.8% reduction in total memory for expert placement, replicas and token buffers compared with fixed expert ownership. This is a memory measurement, not a training-time one.

The catch is that the backward pass becomes approximate: part of the gradient contributions is simply thrown away. The LRM metric is a score at the level of a single projection, and it offers no convergence guarantees. That is why the authors compare systems by time to a fixed loss rather than by step speed: a speedup counts only if the model reaches the same quality. The overhead of computing scores, packing and routing is included in the measurements.

The takeaway for engineers: if MoE training is bound by GPU memory and network, and there is no room left for more replicas, it makes sense to look at reducing the backward-pass work itself, not just shuffling it around. The caveats matter too. The results come from two models and are compared against a single baseline system, so transfer to other architectures and to long runs needs separate verification. The cost of approximate gradients may show up at stricter quality targets.