Expert Coupling in MoE Pretraining: Reducing All-to-All Overhead with Correlated Placement and Token Shuffling

Authors: Radha Gulhane, Quentin Anthony, Beren Millidge (Zyphra)

---

Summary (Overview)

• Key Problem: All-to-all (A2A) communication in Mixture-of-Experts (MoE) pretraining under expert parallelism (EP) can consume up to 60.4% of training step time at EP32 with top-6 routing on 8-GPU AMD MI300X nodes.
• Core Finding: Expert routing in pretrained MoE models exhibits strong correlations—both within a layer (certain expert pairs are co-selected far more often than chance) and across layers (experts chosen at layer ℓ predict experts at layer ℓ+1).
• Two Proposed Methods: (1) Correlated expert placement + deduplicating dispatcher co-locates frequently co-selected experts on the same GPU; (2) Token shuffling moves tokens to predicted next-layer expert GPUs during the existing reduce-scatter collective.
• Results: A2A time reduced by 1.16–2.63× and end-to-end step time by up to 1.41×, with no changes to routing decisions, expert parameters, or training loss.
• Key Advantage: Correlations form early (~1 billion tokens) and require only ~4,096 tokens to estimate reliably, making the approach practical for pretraining from scratch.

---

Introduction and Theoretical Foundation

Background

Mixture-of-Experts (MoE) layers replace the single feed-forward network (FFN) of a dense Transformer with $E$ parallel expert networks and a lightweight router. Each token is routed to its top-$k$ experts, growing model capacity with $E$ while FLOPs per token grow only with $k$.

The Communication Bottleneck

Under expert parallelism (EP), experts are sharded across GPUs, requiring two all-to-all (A2A) collectives per MoE layer (dispatch and combine) in both forward and backward passes. The communication cost scales with:

• EP degree: A larger EP group spreads experts over more GPUs, reducing the fraction of a token's experts on its own GPU as $1/\text{EP}$.
• Top-k: Higher $k$ increases the number of activation rows per token (top-6 moves 3× the volume of top-2).

Measured A2A share of step time:

Key Theoretical Insight: Expert Correlation

Within-layer correlation: Some expert pairs are selected together far more often than independent routing predicts. After ~2.1B training tokens with top-2 routing:

• 42% of tokens select one of only 0.8% of possible expert pairs (64 of 8,128 pairs)
• 14 pairs are chosen with ≥50% confidence and carry 16% of all tokens
• Independent routing with the same loads would give only 1.6%

Cross-layer correlation: The experts a token selects at layer ℓ predict its experts at layer ℓ+1:

• 112 of 128 layer-8 experts send ≥30% of their tokens to one specific layer-9 expert
• Median share is 48%

These correlations enable two complementary strategies: placing correlated experts together (reducing dispatch rows) and moving tokens to predicted expert locations (increasing locality).

---

Methodology

Model Configuration

• 12-layer MoE models with $d_{\text{model}} = 4096$
• 128 experts of FFN width 1024 per layer
• Top-2 and top-6 routing variants
• Sequence-level auxiliary loss (coefficient $10^{-2}$) for load balancing
• Trained on 2.1B tokens (2000 steps) of FineWeb-Edu, sequence length 8192, global batch size 128

Hardware

• Nodes with 8 AMD Instinct MI300X GPUs each
• Intra-node: fully connected xGMI (direct link between each GPU pair)
• Inter-node: eight 100 Gb/s RoCE NICs per node (one per GPU)

Method 1: Correlated Expert Placement + Deduplicating Dispatcher

Deduplicating Dispatcher: Conventional top-k dispatch sends one copy of a token's hidden state per (token, expert) assignment, even when multiple selected experts share a GPU. The deduplicating dispatcher sends one row per unique (token, destination rank) pair, carrying metadata (expert IDs and routing probabilities) for all selected experts on that rank.

Correlated Placement: For each MoE layer, build a correlation graph $W$ where $W[e, e']$ counts tokens selecting both experts $e$ and $e'$. A Kernighan–Lin-style local search (Algorithm 1) starts from contiguous placement and swaps expert pairs between GPUs to maximize:

$$\sum_{\text{pairs } (e, e') \text{ sharing a GPU}} W[e, e']$$

Each GPU holds exactly $E/\text{EP}$ experts, preserving memory footprint. For multi-node EP groups, partitioning proceeds in two steps: first across nodes (reducing inter-node traffic), then across GPUs within each node.

Method 2: Token Shuffling

Token shuffling applies to TP = EP layouts (tensor parallelism group = expert parallelism group), where sequence parallelism shards tokens across the same GPU set.

Mechanism:

1. Prediction: The shuffle planner predicts each token's owner rank at layer ℓ+1 from the experts it selected at the two preceding MoE layers (pair-based lookup with fallback to single-expert prediction for sparse observations).
2. Fused Shuffle: The reduce-scatter that already follows attention is modified to deliver each token to its predicted owner instead of its canonical sequence position. The all-gather before the next attention block restores canonical order.
3. Capacity Constraint: Each GPU must own exactly $T/\text{EP}$ tokens, so tokens are assigned in decreasing confidence to their most preferred rank with capacity left.
4. Overhead Hiding: The planner runs on a side stream, concurrent with expert computation and the next attention block, joined with a single event wait (~1 ms per step).

Key constraint: Token shuffling requires TP = EP (e.g., TP8 EP8, TP16 EP16) and does not apply to layouts without tensor parallelism (TP1).

---

Empirical Validation / Results

Offline Placement and Prediction Statistics

Table I (held-out iterations, all values in %):

For top-6 at EP8: deduplication rises from 25% (contiguous) to 58% (correlated); at EP32 from 5.5% to 35.1%.

Correlation Formation and Stability

• Few tokens suffice: With 4,096 tokens, Pearson correlation between fitted and held-out co-selection matrices is 0.93 (top-2) and 0.98 (top-6); results within 1 percentage point of using all 524k tokens.
• Early formation: Tables fitted at step 199 give ~half the final benefit (14.4% vs 26.5% row removal at top-2 EP8). Tables fitted at step 1000 come within 2–3 percentage points of final values.
• Stability: Tables fitted at step 1400+ give nearly identical results to final tables. Refitting every ~1000 steps is expected to suffice for longer runs.

All-to-All Communication Volume Reduction

Relative to baseline (contiguous placement, Megatron-LM dispatcher):

With token shuffling at TP = EP layouts:

All-to-All Time Reduction

End-to-End Step Time Speedup

Token–Expert Locality

Token shuffling raises the share of token–expert assignments served on the token's GPU:

• EP8: from 12.5% (baseline) to 59% (top-2) and 57% (top-6)
• EP16: from 6.3% to 53% (top-2) and 46% (top-6)

---

Theoretical and Practical Implications

Theoretical Significance

1. Routing structure is learnable and stable: The demonstration that expert co-selection patterns form within ~1 billion tokens and remain stable challenges the view of MoE routing as chaotic or input-dependent. This suggests expert specialization is more predictable than previously assumed.

2. Locality as a design principle: The paper establishes that raising token–expert locality (the fraction of assignments served on the token's GPU) is a principled way to reduce A2A overhead, complementing prior approaches that focused on load balancing or traffic distribution.

3. Cross-layer predictability: The strong cross-layer correlation (median 48% of tokens follow a single expert path) suggests MoE models learn sequential processing pipelines, which could inform future architecture design.

Practical Implications

1. No architectural changes required: Both methods work within existing MoE frameworks (implemented in Megatron-LM) without modifying routing decisions, expert parameters, or memory footprint.

2. Cost-effective implementation: Only ~4,096 tokens (one microbatch from one rank) are needed to fit correlation tables, and refitting every ~1000 steps suffices for long training runs.

3. Broad applicability: Correlated placement and deduplication apply to any expert-parallel layout (training or inference). Token shuffling applies to TP = EP layouts, which are common in production systems.

4. Largest gains at scale: The benefits grow with top-k and when EP groups span multiple nodes—exactly the configurations where A2A overhead is most severe.

---

Conclusion

Main Takeaways

• Expert routing in MoE pretraining exhibits stable within-layer and cross-layer correlations that form early (~1B tokens) and require minimal data to estimate.
• Correlated expert placement + deduplicating dispatcher co-locates frequently co-selected experts, reducing communication volume by 1.12–2.36×.
• Token shuffling exploits cross-layer predictability to move tokens to their predicted expert GPUs inside existing collectives, raising token–expert locality from 12.5% to 59% at EP8.
• Combined, the methods reduce A2A time by 1.16–2.63× and end-to-end step time by up to 1.41×, with the largest gains at top-6 routing across multiple nodes.

Future Directions

The authors identify several open avenues:

• Validating expert correlation across other model families and scales
• Applying the methods to other parallelism layouts (e.g., context parallelism)
• Extending the affinity-based approaches to EP inference
• Exploring whether the observed routing structure can inform architectural improvements to MoE layers themselves

Key Quote: "Neither method changes the models' underlying routing decisions or expert parameters." — The methods are purely communication optimizations that preserve training semantics exactly.