LAER-MoE: Dynamic Expert Re-layout

• LAER-MoE is a MoE training system that dynamically optimizes expert placement to alleviate load imbalance and reduce tail latency induced by skewed token distributions.
• Its key innovation, Fully Sharded Expert Parallelism (FSEP), shards expert parameters across all devices, enabling efficient on-demand reconstruction via optimized All-to-All communication.
• The system employs an online, topology-aware planning heuristic to adjust expert replication and routing per iteration, achieving up to 1.69× acceleration over existing distributed training methods.

LAER-MoE is an MoE training system for efficient distributed training under the observation that expert-parallel execution is often limited not by average work but by tail latency induced by dynamic routing. In conventional expert parallelism, experts are statically placed on devices, so skewed token-to-expert assignments create stragglers, inflate synchronization delay, and can drive the All-to-All fraction of iteration time from under 10% to over 40%. LAER-MoE addresses this system-level load-imbalance problem without changing the training algorithm itself. Its two central mechanisms are Fully Sharded Expert Parallelism (FSEP), which shards every expert across all devices and reconstructs only selected complete experts at expert granularity, and an online load-balancing planner that decides expert re-layout and token routing at every iteration. Reported experiments on an A100 cluster show up to $1.69\times$ acceleration over current state-of-the-art training systems (Liu et al., 12 Feb 2026).

1. Problem formulation and system objective

LAER-MoE is motivated by the mismatch between sparse MoE computation and static expert placement. In an MoE layer, each token is routed by a gate to only a few experts, and this routing is data-dependent and changes from iteration to iteration. Under conventional expert parallelism, a device permanently hosts specific experts. When routing concentrates on a few experts, the devices that host them become bottlenecks while other devices idle at synchronization points. The resulting imbalance affects both expert computation and the All-to-All dispatch/combine phases, and the system-level bottleneck manifests as MoE tail latency rather than insufficient average throughput (Liu et al., 12 Feb 2026).

The system therefore targets load balancing at the execution level rather than the optimization level. This is a defining distinction. LAER-MoE is presented as a system-level alternative to stronger algorithmic balancing interventions such as large auxiliary balancing losses. The intent is to preserve training semantics and convergence while reducing the maximum per-device token count and thus the critical-path time of each MoE layer. This suggests a design philosophy in which expert placement becomes a dynamic systems variable rather than a fixed memory mapping.

The paper formalizes the planning problem using expert-placement variables $A_{i,j}$ and token-routing variables $S_{i,j,k}$. Here $R_{i,j}$ denotes the number of tokens originally on device $i$ that are routed by the gate to expert $j$; $A_{i,j}\in\{0,1\}$ indicates whether expert $j$ is placed on device $i$; and $S_{i,j,k}$ is the number of those tokens ultimately sent to device $A_{i,j}$0 for expert-$A_{i,j}$1 computation. The planner minimizes total MoE-layer time

$A_{i,j}$2

The communication term is modeled as

$A_{i,j}$3

and the computation term as

$A_{i,j}$4

with

$A_{i,j}$5

The resulting constrained optimization problem is

$A_{i,j}$6

subject to

$A_{i,j}$7

$A_{i,j}$8

In plain terms, the planner seeks a legal expert layout and token assignment that are topology-aware and avoid making any device the computation bottleneck (Liu et al., 12 Feb 2026).

2. Fully Sharded Expert Parallelism

The architectural core of LAER-MoE is Fully Sharded Expert Parallelism. FSEP combines the memory semantics of FSDP with the execution semantics of expert parallelism. Instead of storing complete experts on designated devices, each expert is evenly partitioned across all $A_{i,j}$9 devices. If there are $S_{i,j,k}$0 experts and $S_{i,j,k}$1 devices, each device stores one shard of every expert, so locally it holds $S_{i,j,k}$2 chunks, each equal to $S_{i,j,k}$3 of an expert. During training, however, a device reconstructs only $S_{i,j,k}$4 complete experts, where $S_{i,j,k}$5 is the expert capacity per device, and these reconstructed experts can be arbitrary. This is what makes expert re-layout possible at iteration granularity (Liu et al., 12 Feb 2026).

FSEP is organized through three operations. During initialization, shard flattens each expert’s parameters, concatenates them, and evenly partitions them across devices. During forward and backward, unshard reconstructs only the selected complete experts needed on a device, at expert granularity, using All-to-All communication rather than gathering every expert. After backward, reshard partitions gradients back into shards and synchronizes them, again via All-to-All. A subtle implementation issue arises because flattening all expert parameters conflicts with restoring only a subset of complete experts while preserving PyTorch autograd metadata. The implementation resolves this by separating flattened parameter storage, total_experts, from expert-level metadata, real_experts, so that unshard can rebuild the correct full tensors for computation while underlying storage remains fully sharded (Liu et al., 12 Feb 2026).

This differs from conventional expert parallelism in two ways. First, expert placement is no longer static: every device persistently stores shards of all experts and materializes a chosen subset on demand. Second, unlike a conventional FSDP+EP hybrid, FSEP does not treat experts as whole units assigned to EP groups and then gathered within FSDP groups. Instead, expert restoration itself becomes a flexible All-to-All problem over the full device set. That flexibility is the enabler for per-iteration re-layout and effective expert replication.

In LAER-MoE, “re-layout” means changing which complete experts are restored on which devices during unshard. Because only shards are stored persistently, placing an expert on multiple devices in the same iteration becomes equivalent to restoring it multiple times, so those devices act as replicas for that iteration. If routing statistics indicate that some experts are hot, the next iteration can restore those experts on more devices and distribute their tokens across replicas. If an expert cools down, fewer devices need restore it. A plausible implication is that FSEP turns replication from a heavyweight state-migration event into a lightweight restoration choice.

The paper also analyzes communication and memory. With $S_{i,j,k}$6, expert size $S_{i,j,k}$7, and $S_{i,j,k}$8 complete experts restored per device, the communication volume per device is given as

$S_{i,j,k}$9

For an FSDP+EP baseline with EP size $R_{i,j}$0, FSDP size $R_{i,j}$1, $R_{i,j}$2, and $R_{i,j}$3, the unshard communication is

$R_{i,j}$4

The ratio is reported as

$R_{i,j}$5

which approaches $R_{i,j}$6 as cluster size grows; the paper gives about $R_{i,j}$7 for $R_{i,j}$8, $R_{i,j}$9, $i$0. Memory-wise, optimizer state remains fully sharded, and the extra overhead comes primarily from overlap-oriented communication optimization rather than from FSEP itself (Liu et al., 12 Feb 2026).

3. Dynamic planning, token dispatch, and expert re-layout

The full optimization problem is a nonlinear integer program with $i$1 variables, so LAER-MoE does not solve it directly online. Instead, it decomposes planning into a fast synchronous GPU token dispatcher and a slower asynchronous CPU expert layout tuner. This split is central to the system’s practicality: the token-routing path must incur negligible overhead, whereas the placement search can tolerate slightly more computation so long as it runs asynchronously and feeds the next iteration (Liu et al., 12 Feb 2026).

The token dispatcher uses a topology-aware “lite routing” heuristic. It assumes knowledge of the current global expert layout but does not require expensive global coordination over token assignments. For each expert, if replicas exist within the same node, tokens are evenly distributed among intra-node replicas; otherwise they are evenly distributed among all global replicas. The appendix states

$i$2

for intra-node replicas, and otherwise

$i$3

This rule is greedy and topology-aware, intended to reduce inter-node transfers while avoiding cross-device coordination overhead. Its measured overhead is reported as $i$4 ms for Mixtral-8x7B e8k2 and $i$5 ms for e16k4, both less than $i$6 of total training time (Liu et al., 12 Feb 2026).

The expert layout tuner operates on routing statistics from the current and prior iterations. It first estimates replica counts per expert through a priority-queue-based proportional allocation heuristic. Starting from one replica per expert,

$i$7

it repeatedly selects the expert with maximum average load $i$8, increments its replica count, and reinserts it until the replica budget satisfies $i$9. It then places those replicas using a greedy topology-aware relocation procedure: replicas are ordered by per-replica load, and each one is placed by preferring nodes with fewer existing replicas of that expert and then choosing the least-loaded device with remaining expert capacity.

The full layout algorithm does not commit to a single replica-allocation heuristic. It evaluates a small candidate set that includes proportional allocation, even allocation, and random perturbations of them. For each candidate it performs expert relocation, invokes lite routing, estimates total time with the cost model, and picks the best layout. In the experiments, the number of candidate schemes $j$0 is fixed to $j$1, specifically proportional and even allocation, for efficiency. The planner therefore updates expert placement essentially every iteration rather than every hundreds of iterations, which is the operational distinction the paper emphasizes relative to earlier relocation or replication systems (Liu et al., 12 Feb 2026).

Planner complexity is reported as $j$2 for replica allocation and $j$3 for placement, yielding $j$4 overall. Planner-performance experiments report that the GPU lite routing overhead is negligible and that the CPU expert-layout solver remains fast even when analytically scaled to $j$5 GPUs. This suggests that online per-iteration re-layout is viable primarily because placement is treated as a lightweight heuristic search over already-sharded parameters rather than as migration of large optimizer and parameter states.

4. Communication scheduling and iteration execution

A major reason LAER-MoE can afford per-iteration re-layout is its fine-grained communication scheduling. FSEP introduces three major communications: expert-parameter restoration during forward unshard, parameter restoration during backward, and gradient resharing or synchronization after backward. LAER-MoE overlaps these with computation using multiple CUDA streams. Similar to FSDP, it uses parameter prefetching, but it relaxes the standard next-layer prefetch rule. Rather than overlapping expert prefetch with attention computation, it overlaps prefetch of the next MoE layer’s experts with the current layer’s expert computation, which is heavier and therefore hides communication more effectively. It also serializes prefetch after token-dispatch All-to-All to reduce channel contention. On the backward side, gradient synchronization is delayed so that it overlaps with backward computation of the next expert layer, rather than being launched immediately by the autograd engine (Liu et al., 12 Feb 2026).

The paper gives an overlap condition for an MoE layer with SwiGLU MLP, bfloat16, hidden size $j$6, intermediate size $j$7, top-$j$8 routing, and $j$9 tokens per device. Under balanced load, compute per device is

$A_{i,j}\in\{0,1\}$0

and prefetch communication volume per device is

$A_{i,j}\in\{0,1\}$1

The overlap condition is

$A_{i,j}\in\{0,1\}$2

In the reported setup, this is theoretically satisfied when $A_{i,j}\in\{0,1\}$3, while empirically $A_{i,j}\in\{0,1\}$4 already suffices (Liu et al., 12 Feb 2026).

A training iteration proceeds as follows. Each device first executes the attention part of the Transformer layer under the chosen non-MoE parallel strategy. In the MoE layer, the gate computes token-to-expert assignments, producing routing counts $A_{i,j}\in\{0,1\}$5. The GPU-side token dispatcher converts the current global layout into a concrete routing plan $A_{i,j}\in\{0,1\}$6, preferring intra-node replicas. Tokens are sent with All-to-All dispatch to the devices that will execute their experts. In parallel, the CPU-side planner uses current routing statistics and historical information to compute the expert re-layout for the next iteration of the layer. After token dispatch, the executor fetches the planner’s chosen layout for the next MoE layer and prefetches those experts using FSEP unshard, overlapping this with current expert MLP computation. Each device reconstructs only its assigned complete experts, computes expert FFNs for the received tokens, and returns outputs through All-to-All combine. Backward follows the same routing structure; required experts are restored for backward, expert backward compute is executed, and gradient resharing or synchronization is delayed to overlap with later expert-layer backward computation. Because expert states are fully sharded, optimizer updates apply to local shards in FSDP style (Liu et al., 12 Feb 2026).

Several implementation details support this schedule. LAER-MoE supports heterogeneous parallelism, allowing attention and MoE submodules to use different strategies. It supports fine-grained recomputation at attention and MLP granularity and, for MoE, recomputes only expert computation to avoid extra All-to-All during checkpointing. It implements a custom All-to-All kernel for expert-parameter prefetch and gradient synchronization to avoid extra send/receive buffers and memory rearrangement overhead from PyTorch’s interface; the stated purpose is memory efficiency rather than higher bandwidth. It also makes H2D and D2H interactions with the CPU planner asynchronous on dedicated CUDA streams and replaces blocking token rearrangement operations with a Triton kernel (Liu et al., 12 Feb 2026).

5. Experimental evaluation and observed behavior

LAER-MoE is implemented on top of Galvatron and PyTorch with about 11k additional lines of Python, CUDA, and C++ code. The evaluation uses a 4-node cluster with 8 NVIDIA A100 80GB GPUs per node. Intra-node connectivity is NVLink at 300 GB/s peak unidirectional bandwidth, and inter-node connectivity is 800 Gbps InfiniBand. The software stack includes Python 3.9.2, CUDA 12.1, PyTorch 2.1.0, NCCL 2.18.1, FlashAttention v2.5.8, Apex, and Transformer-Engine for some baselines. The models are Mixtral-8x7B, Mixtral-8x22B, and Qwen-8x7B, each in $A_{i,j}\in\{0,1\}$7 and $A_{i,j}\in\{0,1\}$8 settings. Training is dropless on WikiText-103 and C4. Baselines are Megatron with heterogeneous expert-parallel support, an optimized FSDP+EP baseline that already includes the paper’s communication optimizations, and FlexMoE as the strongest prior dynamic-placement baseline (Liu et al., 12 Feb 2026).

Across all tested settings at 8K context, after 20 warmup steps and averaging the next 50 iterations, LAER-MoE consistently outperforms the baselines. The headline result is up to $A_{i,j}\in\{0,1\}$9 speedup over Megatron and up to $j$0 over FSDP+EP. Compared with FlexMoE, it achieves up to $j$1 speedup, with $j$2 average improvement. In a Mixtral-8x7B case study, it reduces the All-to-All share of iteration time to below 20% and achieves up to $j$3 speedup in All-to-All communication versus the baseline. Token-count analysis shows that LAER-MoE keeps per-device maximum token counts much closer to the perfect-balance line than either FSDP+EP or FlexMoE (Liu et al., 12 Feb 2026).

The convergence study is used to support the claim that LAER-MoE is a systems optimization rather than a change in learning behavior. On Mixtral-8x7B e8k2 at 4K context, LAER-MoE with auxiliary loss weight $j$4 converges at the same rate as Megatron with the same auxiliary loss, with relative error below $j$5. The paper interprets this as confirming correctness and the absence of numerical degradation from FSEP. Because iterations are faster, the system also delivers the best convergence speed in wall-clock time.

Ablation results show that both the planner and communication scheduling matter. If the planner uses only one replica-allocation scheme, whether proportional-only or even-only, performance becomes less robust because different routing patterns prefer different layouts. Using both candidate schemes yields the best overall results. Disabling communication-scheduling optimizations also degrades performance, indicating that overlap and launch ordering are crucial to hiding FSEP overhead. For scalability, the authors use trace-driven simulation with real Mixtral-8x7B-e8k2 routing traces and report stable simulated MLP speedups from 8 to 128 GPUs: $j$6, $j$7, $j$8, $j$9, and $i$0, respectively (Liu et al., 12 Feb 2026).

6. Relation to prior load-balancing systems, applicability, and limitations

LAER-MoE’s distinguishing claim is not merely that it rebalances experts, but that it changes the execution model so that re-layout is cheap enough to perform every iteration. The paper contrasts this with prior load-balancing systems such as FasterMoE, SmartMoE, Prophet, and FlexMoE. Those systems replicate or relocate experts under communication, synchronization, or memory costs that force them to limit the frequency of layout change. By contrast, LAER-MoE uses FSEP so that placement becomes a decision about which experts to restore from already-sharded parameters, and most of the associated communication can be hidden behind computation (Liu et al., 12 Feb 2026).

In practical terms, the system is most beneficial when routing is highly skewed and changes rapidly across iterations. That is precisely the regime in which static expert placement or infrequent expert relocation creates stragglers. It is also useful when one wishes to keep auxiliary balancing losses small while preserving throughput. In balanced scenarios, the paper states that LAER-MoE performs similarly to FSDP+EP because the communication volume is nearly the same; the advantage emerges when imbalance exists. This suggests that LAER-MoE should be interpreted less as a universal replacement for expert parallelism than as a specialized solution for dynamic skew.

The main trade-off is system complexity. FSEP requires custom sharding logic, multi-stream communication scheduling, and an online planner. The planner is heuristic rather than globally optimal. The architecture is nonetheless described as decoupled enough to support better planners in the future. Compatibility is presented as a strength: LAER-MoE is described as orthogonal to many prior communication-computation overlap optimizations and implementable atop existing PyTorch-based MoE stacks. The released code is in the Hetu-Galvatron tree, with artifact scripts for end-to-end, convergence, case, planner, and ablation experiments (Liu et al., 12 Feb 2026).

A broader context comes from neighboring MoE research. Work on parameter-efficient or layer-aware LoRA-MoE systems has argued that expert budgeting and allocation matter, including non-uniform expert counts across depth (Gao et al., 2024) and sensitivity-driven placement of LoRA-MoE experts under fixed budgets (Xu et al., 6 May 2025). Those systems operate at the model-design or PEFT-allocation level rather than at the distributed-training systems level. LAER-MoE instead treats expert placement as an online execution decision for full MoE training. This suggests that “allocation” in LAER-MoE refers not to architectural capacity assignment across layers, but to runtime restoration, replication, and placement of expert parameters across devices.

Within distributed MoE training, LAER-MoE therefore occupies a specific position: it reframes load balancing as a first-class per-iteration optimization over expert restoration and token routing. By combining FSEP, online topology-aware planning, and communication scheduling, it reduces MoE tail latency and improves throughput without altering the training algorithm itself.