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

> Expert routing in MoE models is highly predictable, enabling correlated expert placement and token shuffling that cut all-to-all communication time by up to 2.63x without altering training.

- **Source:** [arXiv](https://arxiv.org/abs/2610.09372)
- **Published:** 2026-10-10
- **Permalink:** https://picx.dev/p/p3TnE0
- **Whiteboard:** https://picx.dev/p/p3TnE0/image

## Summary

# 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:

| Configuration | EP8 | EP32 |
|---|---|---|
| **Top-2** | 13.5% | 44.8% |
| **Top-6** | 24.0% | 60.4% |

### 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 %):

| EP | Placement | Same GPU (top-2) | GPUs (top-2) | Dedup (top-2) | Local experts moved | Local experts stay |
|---|---|---|---|---|---|---|
| EP8 | Contiguous | 12 | 1.88 | 6 | 33 | 13 |
| EP8 | Correlated | 53 | 1.47 | 26 | 58 | 13 |
| EP16 | Contiguous | 6 | 1.95 | 3 | 28 | 6 |
| EP16 | Correlated | 45 | 1.55 | 23 | 51 | 6 |
| EP32 | Contiguous | 2 | 1.98 | 1 | 25 | 3 |
| EP32 | Correlated | 36 | 1.64 | 18 | 42 | 3 |

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):

| Configuration | Top-2 EP8 | Top-2 EP16 | Top-2 EP32 | Top-2 EP64 |
|---|---|---|---|---|
| **Placement (P)** | 1.36× | 1.29× | 1.22× | 1.12× |

| Configuration | Top-6 EP8 | Top-6 EP16 | Top-6 EP32 | Top-6 EP64 |
|---|---|---|---|---|
| **Placement (P)** | 2.36× | 1.95× | 1.54× | 1.24× |

With token shuffling at TP = EP layouts:

| Configuration | Top-2 EP8 | Top-2 EP16 | Top-6 EP8 | Top-6 EP16 |
|---|---|---|---|---|
| **Placement alone** | 1.4× | 1.3× | 2.4× | 1.9× |
| **Token shuffling (S)** | 2.4× | 2.1× | 3.1× | 2.4× |

### All-to-All Time Reduction

| Configuration | Top-2 | Top-6 |
|---|---|---|
| **EP8 (intra-node)** | 1.16× | 1.58× |
| **EP16 (2 nodes)** | 1.34× / 1.38× (with shuffling: 1.74×) | 1.95× / 1.48× (with shuffling: 2.63×) |
| **EP32 (4 nodes)** | 1.38× | 1.48× |
| **EP64 (8 nodes)** | 1.34× | 1.35× |

### End-to-End Step Time Speedup

| Configuration | Top-2 | Top-6 |
|---|---|---|
| **EP8** | 1.05× | 1.14× |
| **EP16** | 1.13× (1.06× with shuffling) | 1.41× (1.34× with shuffling at TP16 EP16) |
| **EP32** | 1.14× | 1.25× |
| **EP64** | 1.11× | 1.23× |

### 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.

---

_Markdown view of https://picx.dev/p/p3TnE0, served by PicX — AI-generated visual whiteboard summaries of research papers._