Source-linked AI summary
Memory-Efficient Expert Routing for Distributed MoE Training
Arnab Kanti Tarafder, Jaume Guasch-Martí, Gokcen Kestor, Jie Ren
TL;DR
Distributed MoE training faces a memory bottleneck because all-to-all dispatch requires full top-k-expanded buffers. RelayMoE computes locally as expert weights or tokens circulate through a ring, avoiding those buffers. It reports faster MoE layer execution, reduced peak memory, and longer trainable sequences.
Problem
All-to-all dispatch constructs full top-k-expanded buffers, making MoE dispatch the dominant peak-memory contributor and limiting trainable sequence length and batch size.
Source: §1 · p. 1Method
RelayMoE circulates expert weights or token activations through a GPU ring and computes locally without materializing full top-k-expanded dispatch buffers.
Source: §1 · p. 2Results
2× average speedup in MoE layer execution at EP=8, 7.6× lower peak GPU memory on average at EP=64, and up to 2.85× longer trainable sequences.
Source: §1 · p. 2Takeaways & Limitations
The saved memory supports longer sequences and larger micro-batches and can reduce attention recomputation.
Source: §1 · p. 2, §6 · p. 13All-to-all is faster on E384 at all three tested sequence lengths; on H2048 it is faster at 40K and 65K, with comparable throughput at 98K.
Source: §4.4 · p. 10
Abstract
from arXiv · showhide
As Mixture-of-Experts (MoE) models scale toward hundreds of experts and higher top-$k$ routing, memory efficiency in distributed training becomes a critical bottleneck. Peak memory is dominated by the MoE block, not attention: every intermediate buffer in the MoE dispatch pipeline is individually scaled by top-k routing. The standard all-to-all dispatcher sends all routed tokens in a single collective step, requiring the full top-$k$-expanded buffer to be constructed at once. In this work, we propose RelayMoE, a ring-based MoE execution model that computes locally as expert weights or tokens circulate, avoiding full top-$k$-expanded dispatch buffers. RelayMoE selects between expert and token routing according to communication volume and overlaps transfers with computation. The ring structure naturally supports memory-efficient MoE recomputation during backward: each hop reconstructs expert intermediates, uses them to compute gradients, and releases them before the next hop. The saved memory supports longer sequences and larger batches, or retains more attention activations to reduce attention recomputation and improve training throughput. We evaluate RelayMoE on 30B$-$57B production MoE models and varied expert configurations. In single-layer MoE experiments, RelayMoE achieves a $2\times$ average speedup over Megatron-LM. In full-model training under the same GPU memory budget, it improves throughput by up to $2.02\times$ and extends the largest tested trainable sequence length by up to $2.85\times$.
1 Introduction
As MoE models grow across experts, top-k routing, sequence length, and FFN width, dispatch memory becomes a training bottleneck. RelayMoE replaces full-buffer all-to-all execution with ring-based computation and reports lower memory use and faster execution.
1 Introduction
- All-to-all requires a full top-k-expanded dispatch buffer, with every intermediate buffer scaled by k; MoE dispatch now dominates peak memory over attention.The growing memory footprint limits maximum trainable sequence length and batch size.Source: §1 · p. 1
- Prior dispatch optimizations target communication efficiency, but token dispatch and expert computation can still create memory peaks that limit training, especially during backward.The cited approaches include fused permutation and communication, stream pipelining, fine-grained scheduling, and fetching expert weights.Source: §1 · pp. 1–2
- RelayMoE circulates expert weights or token activations through a GPU ring, computing locally with memory bounded by local data and a fixed number of in-flight partitions.It processes routed token–expert pairs partition by partition, enabling dropless, padding-free expert computation.Source: §1 · p. 2
- During backward recomputation, each ring hop reconstructs intermediates for local gradient computation and releases them immediately, leaving more memory for attention activations.Retaining more attention activations reduces attention recomputation.Source: §1 · p. 2
- 2× average speedup in MoE layer execution at EP=8, 7.6× lower peak GPU memory on average at EP=64, and up to 2.85× longer trainable sequences.The speedup is compared with all-to-all; sequence length refers to full-model training.Source: §1 · p. 2
2 Background and Motivation
MoE routes tokens to selected experts, which expert parallelism distributes across GPUs using cross-GPU communication. In the profiled training setup, top-k-expanded dispatch buffers create large MoE memory spikes that set the out-of-memory boundary.
2.1 MoE Training with Expert Parallelism
- Each token selects its top-k experts; expert parallelism distributes experts across GPUs, requiring tokens and results to move between GPUs.The standard all-to-all dispatcher expands and permutes tokens before sending them to local experts, then returns and recombines their outputs.Source: §2.1 · pp. 2–3
2.2 Peak Memory Consumption in MoE Training
- In a four-layer, top-8 profiled setup, memory spikes reach 51 GB from an approximately 14 GB baseline, and MoE blocks—not attention—set the OOM boundary.The profile uses micro-batch size 24, sequence length 16,384, and 8 H100 GPUs.Source: §2.2 · p. 3
- Top-k expansion turns the 336 MiB input activation into dispatch tensors of approximately 2.6 GiB each.The reported 336 MiB input activation is expanded 8× under top-8 routing.Source: §2.2 · p. 3
- Coexisting expanded tensors from dispatch and combine stages push the per-layer backward peak above 42 GB.New outputs are allocated while their inputs remain in memory.Source: §2.2 · p. 3
3 Design
RelayMoE replaces single-step all-to-all dispatch with ring execution that computes locally as experts or tokens circulate. It combines adaptive routing, hop-local MoE recomputation, and selective attention retention.
3.1 Ring-Based MoE Dispatch
- At each ring hop, local tokens are processed with the visiting expert partition, partial outputs accumulate, and weights move to the next GPU.Overlapping computation with asynchronous weight transfers can bring per-hop time closer to the larger of compute and transfer time than their sum.Source: §3.1 · pp. 3–4
- Ring dispatch avoids full expanded dispatch buffers, but each hop’s local intermediates still depend on top-k and routing, and uneven token–expert matches can leave ranks idle.The design also avoids token dropping and capacity-factor padding to bound dispatch buffers.Source: §3.1 · p. 5
- 13.9 GiB lower peak memory is reported for RelayMoE under the same training configuration, with the sharp MoE spikes substantially reduced.The coexisting 2.6 GiB dispatch buffers no longer appear in the single-layer profile.Source: Figure 6 · p. 5, §3.1 · p. 5
3.2 Adaptive Routing: Experts vs Tokens
- RelayMoE chooses expert or token routing by comparing per-hop communication volumes under the full training configuration.Expert routing keeps tokens resident; token routing keeps expert weights resident, and the preferred mode can shift with sequence partitioning and expert tensor parallelism.Source: §3.2 · p. 6, Table 2 · p. 6
3.3 Memory-Efficient MoE Recomputation
- During backward recomputation, each hop reconstructs and consumes its expert intermediates, then releases them before the next hop.The peak intermediate memory is determined by the largest hop, and reconstruction uses the data movement already required for backward without another ring traversal.Source: §3.3 · pp. 6–7
3.4 Selective Attention Recomputation
- RelayMoE uses memory saved by MoE recomputation to retain selected attention activations, reducing attention recomputation and improving training throughput.It ranks options by compute saving per unit of memory and selects within the available memory budget.Source: §3.4 · p. 7
4 Evaluation
RelayMoE is evaluated on H100 systems using single-layer and full-model MoE tests across expert configurations and production models. The results cover latency, memory, throughput, capacity, and routing design.
4.2 MoE Layer Performance
- RelayMoE averages 2× the all-to-all speed on single-layer MoE tests; E128 reaches 2.4× at 98K, while E384 and H2048 cross from slower at 40K to faster at 98K.The larger expert counts or wider FFNs raise expert-weight transfer costs, making the benefit depend on available token computation to amortize transfers.Source: §4.2 · p. 8
- At EP=64, RelayMoE reaches 5× speedup over all-to-all for K16; its peak-memory reduction averages 7.6× at EP=64 and reaches 13.4× for K16.Against FasterMoE and Tutel, speedups reach 27× and 16×, respectively, at EP=64; Tutel encounters OOM in some configurations.Source: Figure 10 · p. 8, §4.2 · pp. 9–10
4.3 Model Training Performance
- The largest tested trainable sequence length rises from 393K to 786K on Qwen3-30B and Qwen3.5-35B, while RelayMoE trains Qwen2-57B at 384K where all-to-all runs out of memory.RelayMoE also has higher throughput at every tested sequence length where both dispatchers run on the two Qwen3 models; DeepEP runs out of memory on several configurations at 98K.Source: Figure 11 · p. 9, §4.3 · p. 10
4.4 Impact of Memory Savings
- Up to 2.02× speedup on K16 at 40K accompanies higher throughput for RelayMoE in most configurations when saved memory retains more attention activations.That retention reduces attention recomputation; K4 is near parity, consistent with less memory saved from its lower top-k expansion.Source: §4.4 · p. 10, Figure 13 · p. 11
- At 40,960 sequence length, RelayMoE reduces full-model peak GPU memory by up to 35.7% and fits micro-batches of 18–32, but the largest batch is not always fastest.H2048 is slower at batch size 12 but faster at 30; attention-checkpoint schedules also change, and batch size and checkpointing need joint consideration.Source: §4.4 · p. 11, Figure 14 · p. 11
- With micro-batch size 12 and full recomputation, RelayMoE reaches 90K–150K across configurations and extends the largest tested sequence length by up to 2.85×.Only K4 trains beyond 65K with all-to-all; the tested limits also depend on layer count, retained attention and checkpoint inputs, and model state.Source: §4.4 · p. 11, Figure 15 · p. 12
4.5 Design Analysis
- Token routing outperforms expert routing on both 384-expert configurations, while all-to-all encounters OOM at tested lengths on K12 and K16.On K12, token-routing speedup over expert routing falls from 3.7× at 20K to 2.1× at 40K; token routing avoids circulating expert weights.Source: §4.5 · pp. 11–12, Figure 16 · p. 12
- Communication–computation overlap lowers single-layer latency from 295.9 ms to 203.4 ms per iteration in the measured E128 configuration.The naive ring implementation executes communication and computation sequentially without GroupedGEMM; the overlap configuration is measured in the second bar.Source: §4.5 · p. 12, Figure 17 · p. 12
5 Related Work
Related work addresses MoE communication, memory, expert computation, and attention through different techniques. RelayMoE targets top-k-expanded dispatch buffers and expert intermediates during backward, complementing these approaches.
Communication optimization
- MoE dispatch libraries, communication overlap, and network design optimize expert-parallel communication, while RelayMoE targets top-k-expanded inter-GPU dispatch buffers and can incorporate optimized kernels and overlap.RelayMoE also complements attention parallelism.Source: §5 · p. 12
Expert placement and fetching
- FasterMoE replicates overloaded experts near their tokens; Janus bounds expert-fetch buffers and caches fetched weights in host memory for backward.RelayMoE instead reconstructs and releases expert intermediates within each hop using the weight transfers also used for gradient computation.Source: §5 · pp. 12–13
Memory-efficient training
- MemFine uses token chunking and chunk-level recomputation while retaining dispatch–compute–combine; RelayMoE keeps tokens local and circulates expert partitions.RelayMoE limits reconstructed-intermediate lifetimes and avoids inter-GPU token-dispatch buffers.Source: §5 · p. 13
Attention memory
- FlashAttention and sequence/context parallel systems reduce attention memory, while RelayMoE addresses MoE dispatch buffers and expert activations.The paper states that the approaches can be combined to reduce memory in both sublayers and support longer-context training under the same hardware budget.Source: §5 · p. 13
6 Conclusion
RelayMoE avoids full top-k-expanded dispatch buffers and reduces memory for expert recomputation during backward. The saved memory supports longer sequences, larger micro-batches, and reduced attention recomputation, alongside lower peak memory and improved throughput over all-to-all.
Conclusion
- RelayMoE avoids full top-k-expanded dispatch buffers and reduces the memory required for expert recomputation during backward.The saved memory supports longer sequences and larger micro-batches, or can retain more attention activations to reduce attention recomputation.Source: §6 · p. 13