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

Radha Gulhane; Quentin Anthony; Beren Millidge

Source-linked AI summary

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

Radha Gulhane, Quentin Anthony, Beren Millidge

arXiv:2610.09372v1Oct 7, 2026cs.CLcs.DC

TL;DR

MoE expert parallelism incurs substantial all-to-all communication to dispatch tokens and return expert outputs. This paper exploits within-layer and cross-layer routing correlations through expert placement, deduplication, and token shuffling. The methods reduce all-to-all time by 1.16×–2.63× and end-to-end step time by up to 1.41× in the reported evaluation.

Problem
At EP32, all-to-all communication accounts for 44.8% of step time under top-2 and 60.4% under top-6 routing.
Source: §I · p. 1
Method
The methods co-locate frequently co-selected experts, deduplicate token destinations, and shuffle tokens toward predicted next-layer expert GPUs during reduce-scatter.
Source: §I · p. 2
Results
All-to-all time falls by 1.16×–2.63× and end-to-end step time by up to 1.41× in the reported evaluation.
Source: §VII · p. 10
Takeaways & Limitations
Both methods improve token–expert locality without changing routing decisions or expert parameters.
Source: §I · p. 2, §VII · p. 10
The evaluation covers one 12-layer model trained for 2.1B tokens; other parallelism layouts and EP inference are left for future work.
Source: §VII · p. 10

Abstract

from arXiv · show

Mixture-of-Experts (MoE) layers replace the feed-forward block of a Transformer with E expert networks, and each token is routed to k of these experts. Under expert parallelism (EP) the experts are distributed across GPUs, and every MoE layer runs all-to-all collectives in the forward and backward passes to dispatch tokens to their experts and then combine the results. On a cluster with 8 AMD Instinct MI300X GPUs per node, these collectives can take 45% of the training step at EP32 with top-2 routing and 60% with top-6 routing. We find that early in pretraining routers have already learned to assign tokens to experts in correlated patterns, both within a layer and across layers. At top-2, 0.8% of the expert pairs in a layer are selected together by 42% of tokens, and the experts a token selects at one layer predict the experts it selects at the next layer. We use these correlations to keep more token--expert assignments on the token's own GPU, which reduces communication across GPUs and across nodes. Correlated expert placement puts experts that are often selected together on the same GPU. Combined with a dispatcher that sends each token to each GPU once, it removes up to 58% of dispatched rows. Token shuffling applies when sequence parallelism shards tokens across the EP group. It moves each token to the GPU predicted to hold its next-layer experts during the reduce-scatter that follows attention. On one node this raises the share of token--expert assignments served on the token's GPU from 12.5% to 59%. In Megatron-LM, across EP degrees from 8 to 64 with top-2 and top-6 routing, the two methods reduce all-to-all time by 1.16-2.63X and end-to-end step time by up to 1.41X. Neither method changes the models' underlying routing decisions or expert parameters.

I. INTRODUCTION

MoE expert parallelism makes dispatch and combine collectives a substantial training cost, while conventional placements and prior approaches do not exploit pretraining routing correlations to reduce total communication. This paper measures those correlations and uses them for expert placement and token shuffling, reducing all-to-all time across evaluated routing settings.

I. INTRODUCTION

At EP32, all-to-all communication accounts for 44.8% of step time under top-2 and 60.4% under top-6 routing.Collectives dispatch tokens to selected experts and return their outputs in both forward and backward passes.Source: §I · p. 1

Characterization of Expert Routing

After about 2.1B training tokens, 0.8% of possible expert pairs carry 42% of top-2 tokens, and routing at one layer predicts routing at the next.Both correlations appear with all three tested load-balancing methods, with differing strengths.Source: §I · p. 2

I. INTRODUCTION

The methods raise token–expert locality without changing routing decisions, expert parameters, or training loss.Token shuffling is limited to TP–EP layouts where the tensor-parallel and expert-parallel groups are the same GPUs.Source: §I · p. 2

Correlated Expert Placement and a Deduplicating Dispatcher

Correlated placement co-locates frequently co-selected experts, while a deduplicating dispatcher sends each token to each destination GPU at most once.Together they reduce intra-node and inter-node communication volume.Source: §I · p. 2

Token Shuffling

Token shuffling predicts a token’s next-layer experts from its previous-layer selections and moves it to their predicted GPU during the post-attention reduce-scatter.It applies when sequence parallelism uses the same GPU group as expert parallelism.Source: §I · p. 2

Evaluation in Megatron-LM

1.16×–2.63× lower all-to-all time with both methods, across top-2 and top-6 routing, with gains varying by configuration.Placement and deduplication alone yield 1.16×–1.38× under top-2 and 1.35×–1.95× under top-6; adding token shuffling yields 1.40×–1.74× and 1.46×–2.63×, respectively.Source: §I · p. 2

II. MOTIVATION

All-to-all communication grows costly with expert-parallel degree and routing top-k, while pretraining routes show exploitable correlations both within and across layers. These findings motivate expert placement and token prediction as ways to improve locality and reduce dispatched rows.

A. All-to-All Cost in MoE Pretraining

At EP32 with top-6 routing, all-to-all takes 60.4% of step time; widening EP from 8 to 32 raises its share from 13.5% to 44.8% under top-2.At EP8, top-6 raises the share from 13.5% to 24.0% compared with top-2.Source: A. All-to-All Cost in MoE Pretraining · p. 2, Figure 1 · p. 3

B. Expert Correlation Within and Across Layers

Among the 8,128 possible pairs, the passage reports a 42% concentration, versus 1.6% under independent routing with the same loads; frequent pairs can reduce duplicated token-state traffic.At EP8, correlated placement raises the share of top-2 tokens whose experts share a GPU from 12% to 53%. Rows removed by deduplication rise from 6% to 26% under top-2 and from 25% to 58% under top-6.Source: B. Expert Correlation Within and… · p. 3
Across layers, 112 of 128 layer-8 experts send at least 30% of their tokens to one specific layer-9 expert, and the median share is 48%.Source: B. Expert Correlation Within and… · p. 3

GPU GPUs

Under sequence parallelism, cross-layer routing predicts where many next-layer experts reside, enabling tokens to be moved there during attention.

GPU GPUs

At EP8 under top-2, 58% of a moved token’s next-layer experts are local on its predicted GPU.The share is 51% at EP16 and 42% at EP32 under top-2; the corresponding top-6 shares are 60%, 49%, and 33%.Source: GPU GPUs · p. 3
A token that stays on its home GPU has 1/EP of its next-layer experts local.Under sequence parallelism, the token can move to its predicted GPU during attention.Source: GPU GPUs · p. 3
At EP8, correlated placement gives a moved token 58% of next-layer experts local, compared with 33% under contiguous placement.Source: GPU GPUs · p. 3

III. BACKGROUND

Expert parallelism distributes experts across GPUs, while hybrid tensor–expert parallelism uses collectives to move between attention and MoE token layouts. Related affinity-based systems share some placement ideas, but differ in training setting and mechanisms.

A. Expert Parallelism

Each token–expert assignment crossing a GPU boundary adds a row to dispatch and combine all-to-all collectives in both forward and backward passes.EP partitions experts across a GPU group, typically alongside data parallelism.Source: A. Expert Parallelism · p. 4

B. Hybrid Parallelism: Tensor–Expert Parallelism

In hybrid TP–EP, reduce-scatter partitions tokens after attention, dispatch and combine serve the MoE layer, and all-gather restores the full sequence for the next Transformer block.Token shuffling replaces reduce-scatter and all-gather with shuffled versions, requiring TP and EP to use the same GPUs.Source: B. Hybrid Parallelism: Tensor–Expert… · p. 4

IV. RELATED WORK

Occult and GRACE-MoE use within-layer expert affinity for trained-model inference or fine-tuning, whereas this work studies pretraining from scratch without changing routing or adding replicas.This work also introduces token shuffling, which neither system uses.Source: §IV · p. 4

V. IMPLEMENTATION

Correlated expert placement co-locates frequently co-selected experts, while a deduplicating dispatcher sends each token to each destination rank only once. Together, they increase locality and reduce duplicate communication rows without changing router assignments or expert computation.

A. Within-Layer Correlation: Correlated Expert Placement

Correlated expert placement co-locates frequently co-selected experts, and deduplication sends each token to each destination rank at most once.The methods increase token–expert locality and give the dispatcher more rows to merge.Source: A. Within-Layer Correlation:… · p. 4
The deduplicating dispatcher sends one row per unique (token, destination rank) pair, carrying expert IDs and routing probabilities for that rank.The receiving rank distributes copies to its selected experts; the combine similarly returns one partial output per token-rank pair.Source: A. Within-Layer Correlation:… · p. 4
Placement maximizes the total weight of co-selected expert pairs sharing a GPU while keeping the same number of experts on every GPU.A Kernighan–Lin-style local search starts from contiguous placement and swaps experts when the objective improves.Source: A. Within-Layer Correlation:… · p. 5
Across multiple nodes, experts are partitioned first across nodes and then across GPUs to keep co-selected experts close where possible.The resulting per-layer permutation leaves router decisions and expert parameters unchanged.Source: A. Within-Layer Correlation:… · p. 5

B. Cross-Layer Correlation: Token Shuffling

With fixed sequence-shard ownership, co-located experts alone do not keep tokens local. Token shuffling uses prior-layer routes to predict an owner and fuses token movement into existing collectives, subject to TP = EP.

B. Cross-Layer Correlation: Token Shuffling

With fixed sequence-shard ownership, a token is on the GPU holding its co-located experts with probability only 1/EP, so its row usually still crosses the network.The dispatch and combine move the row when token and expert owners differ.Source: B. Cross-Layer Correlation: Token… · p. 5
Token shuffling uses prior-layer routes to predict the next-layer owner and sends each token there during the reduce-scatter after attention.Cross-layer routing predicts the next experts above chance, with the prediction available before the next-layer router runs.Source: B. Cross-Layer Correlation: Token… · p. 5
The fused reduce-scatter sends tokens to planned owners, while a fused all-gather restores canonical order before attention; this requires TP = EP and adds no collective.The token’s residual-stream row moves with it, and the MoE layer sees tokens whose experts are mostly on the same GPU.Source: B. Cross-Layer Correlation: Token… · p. 6
The planner predicts each token’s next owner from its preceding expert routes, using offline route counts and confidence-ranked assignment to meet equal per-rank token quotas.At runtime, previous layers’ routes are the planner’s only input, and it adds no collective.Source: B. Cross-Layer Correlation: Token… · p. 6
The planner runs concurrently with expert computation and the next attention block, with the main stream waiting about one millisecond per step for its plan.A side stream hides prediction and quota assignment behind existing computation.Source: B. Cross-Layer Correlation: Token… · p. 6

VI. EXPERIMENTAL RESULTS

The experiments evaluate routing-correlation tables and methods on trained top-2 and top-6 models. Small fitting samples capture much of the placement and prediction benefit, and the correlations form during training and later stabilize.

VI. EXPERIMENTAL RESULTS

Experiments use 12-layer models with 128 experts per layer, top-2 or top-6 routing, and measurements from trained step-2000 checkpoints on held-out data.The models were trained for 2.1 billion tokens; timing runs were repeated and reported as medians over the last 10 of 14 steps.Source: §VI · pp. 6–7

A. Building Expert Correlation Within and Across Layers

With tables fitted on 1,024 tokens, correlated placement removes 23.6% of dispatched rows at top-2 EP8 and 57.9% at top-6.The tables were evaluated on held-out iteration 1999.Source: A. Building Expert Correlation Within… · p. 7
With 4,096 tokens, fitted-versus-held-out co-selection correlations are 0.93 for top-2 and 0.98 for top-6, with placement and prediction within one percentage point of 524k-token tables.The authors report that the tables capture a trained-router property rather than depending on the particular fitting tokens.Source: A. Building Expert Correlation Within… · p. 7
At step 199, placement removes 14.4% of dispatched rows at top-2 EP8 and 42.7% at top-6, versus 26.5% and 57.7% for end-of-training tables.Tables fitted at step 1000 come within 2–3 percentage points of final values; later tables give nearly the same results.Source: A. Building Expert Correlation Within… · p. 7

B. Correlated Expert Placement

Correlated placement reduces all-to-all communication and step time by grouping co-selected experts, while token shuffling raises locality by moving tokens toward their predicted next-layer experts. The gains vary with routing, EP layout, and the share of step time spent in all-to-all communication.

B. Correlated Expert Placement

Correlated placement cuts communication volume by 1.36× at top-2 and 2.36× at top-6 on EP8, with smaller reductions as EP widens.A2A time falls by 1.16× at top-2 and 1.58× at top-6 on EP8; across EP16–EP64, reductions range from 1.34× to 1.38× at top-2 and 1.35× to 1.95× at top-6.Source: B. Correlated Expert Placement · p. 7, Figure 7 · p. 8
1.41× at top-6 EP16 is the largest correlated-placement step speedup reported; at top-2, speedups range from 1.05× to 1.14× across EP8–EP64.At top-6, speedups are 1.14×, 1.41×, 1.25×, and 1.23× at EP8, EP16, EP32, and EP64, respectively.Source: B. Correlated Expert Placement · pp. 7–8

C. Token Shuffling

Token shuffling raises token–expert locality from 12.5% to 59% at top-2 and from 12.5% to 57% at top-6 on EP8.Wire bytes fall by 2.4× at EP8 and 2.1× at EP16 for top-2, and by 3.1× and 2.4× for top-6.Source: C. Token Shuffling · p. 8
At TP16 EP16, token shuffling raises the A2A reduction to 1.74× at top-2 and 2.63× at top-6, from 1.55× and 1.95× with placement alone.On TP8 EP8, A2A reductions with shuffling are 1.40× at top-2 and 1.46× at top-6; the latter is slightly below placement alone’s 1.48×.Source: C. Token Shuffling · p. 8
On TP=EP layouts, token shuffling yields step speedups of 1.01–1.06× at top-2 and 1.08× at TP8 EP8 to 1.34× at TP16 EP16 at top-6.A2A occupies 9–21% of step time at top-2 and 18–38% at top-6; at TP16 EP16, shuffling raises the top-6 speedup from 1.28× to 1.34×.Source: C. Token Shuffling · p. 9, Figure 10 · p. 9

VII. CONCLUSION

Expert routing correlations emerge early and remain stable, enabling placement and token shuffling to exploit within-layer and cross-layer patterns. On the evaluated MI300X cluster, the methods reduce all-to-all time by 1.16–2.63× and end-to-end step time by up to 1.41×.

VII. CONCLUSION

Correlated placement groups co-selected experts on a GPU and deduplicates token dispatch; token shuffling moves tokens toward the predicted next-layer experts under per-rank capacity constraints.Shuffling is folded into sequence parallelism’s reduce-scatter, and neither method changes routing decisions or expert parameters.Source: §VII · p. 9
On a cluster with eight MI300X GPUs per node, the methods reduce all-to-all time by 1.16–2.63× and end-to-end step time by up to 1.41×.The largest gains occur at top-6 when the EP group spans several nodes.Source: §VII · p. 10
The evaluation covers one 12-layer model trained for 2.1B tokens; other parallelism layouts and EP inference remain future work.The authors expect expert correlation to hold across model families and scales, but present that as a belief rather than an evaluated result.Source: §VII · p. 10

APPENDIX

The appendix describes the fixed experimental model and implementation, then tests correlated placement across three top-6 routing strategies. Placement remains effective across the routers, while its benefit tracks their co-selection structure and does not directly address load skew.

1) Model Architecture:

The fixed model has 12 layers, and experiments use pre-softmax top-k routing with k=2 or k=6 and a sequence-level auxiliary loss of 10−2 unless stated otherwise.The router ablation keeps architecture fixed and compares that recipe with a 10−1 auxiliary loss and DeepSeek-V3’s auxiliary-loss-free router.Source: 1) Model Architecture · p. 12, Table II · p. 12

2) Training Framework:

The Megatron-LM implementation changes the dispatcher, expert placement, and token shuffling while leaving the router, expert computation, optimizer, and training loop unchanged.Routing traces build offline placement and affinity tables; all variants train the same model and produce the same loss.Source: 2) Training Framework · pp. 12–13

B. Ablation Over Routing Strategies

The auxiliary-loss-free router reaches the lowest loss of the three (4.39), while the stronger auxiliary loss gives the most balanced expert load (maximum-to-mean ratio 1.20).Their losses are 4.39 for auxiliary-loss-free, 4.46 for the main model, and 4.64 for stronger auxiliary loss; load ratios are 1.37 for auxiliary-loss-free and 1.52 for the main model.Source: B. Ablation Over Routing Strategies · p. 13
At TP1 EP16, correlated placement gives step speedups of 1.41×, 1.30×, and 1.26× for the main, stronger-auxiliary-loss, and auxiliary-loss-free routers, respectively.At EP8 it removes 58%, 51%, and 47% of dispatched rows for those routers; the reported EP16 speedups are at least 1.20× on every layout.Source: B. Ablation Over Routing Strategies · p. 13, Figure 11 · p. 14
Correlated placement and deduplication target communication volume through co-selection, not communication skew; the all-to-all still waits for the most heavily loaded rank.The authors describe relocation or replication of heavily loaded experts as complementary ways to address skew.Source: B. Ablation Over Routing Strategies · p. 13

C. Per-Step Measurements

The per-step measurements report step time and per-rank all-to-all volume after warm-up, with labels comparing configuration medians against the baseline median. The appendix algorithm describes how token shuffling predicts owner ranks and assigns tokens within a per-rank cap.

C. Per-Step Measurements

Figure 12 reports step time and A2A volume per rank for ten measured steps after four warm-up steps, with labels showing baseline-median to configuration-median ratios.The figure’s stated measures are per configuration; it does not provide outcome values in the supplied passage.Source: Figure 12 · p. 15, C. Per-Step Measurements · p. 15

D. Algorithms

The shuffle planner predicts each token’s owner-rank distribution from its prior-layer expert choices, mixing first- and second-order tables according to trace support.The planner runs at each MoE position, using prediction tables fit offline on routing traces.Source: D. Algorithms · pp. 16–17
Tokens are assigned greedily from most to least confident under a per-rank cap of T/P, then reduce-scatter writes rows to their owners.The planner runs on a side stream during attention; the main stream waits at the attention output projection before the reduce-scatter.Source: D. Algorithms · p. 17

E. Author Contributions

Radha Gulhane implemented the methods and ran all experiments, while Quentin Anthony and Beren Millidge contributed to ideation, technical discussion, and writing.

Radha Gulhane implemented the methods, carried out all experiments, and led initial ideation and exploration.Source: E. Author Contributions · p. 17
Quentin Anthony contributed to initial ideation and shaping the idea, helped write the paper, and provided substantial guidance throughout.Source: E. Author Contributions · p. 17
Beren Millidge participated in technical discussions, offered suggestions throughout the project, and helped write the paper.Source: E. Author Contributions · p. 17