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

Radha Gulhane; Quentin Anthony; Beren Millidge

cs.CLOct 7, 2026

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

Authors: Radha Gulhane, Quentin Anthony, Beren Millidge

Organizations: Zyphra

Abstract

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.

Figures & tables

Fig. 1 : Share of training step time by component at EP8 and EP32 with top-2 and top-6 routing, using the Megatron-LM all-to-all dispatcher and contiguous placement. Dashed boxes mark the all-to-all share.

Panel of Figure 2 — Fig. 2 : Expert correlation in the top-2 model, computed from the saved routing of 524k tokens at training steps 1983–1995 (every fourth step). (a) $P(e_{2}\in\text{top-2}\mid e_{1}\in\text{top-2})$ within layer 8; (b) $P(e_{2}\text{ at }9\mid e_{1}\text{ at }8)$ across layers. Blue marks strongly correlated pairs (conditional probability above 0.3).

		Within layer			Across layers: local experts
EP	Placement	Same GPU	GPUs	Dedupli- cation	Token moved	Token stays
top-2 routing
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

TABLE I : Offline placement and prediction statistics on held-out iterations 1983 and 1987 (fit on 1991–1999); all values in % (rounded) except GPUs . Same GPU : tokens whose top-

k

experts share a GPU. GPUs : distinct destinations per token. Deduplication : rows removed as duplicates. Local experts : share of a token’s next-layer experts on its own GPU when it is moved to the GPU predicted from its current experts (

T/\mathrm{EP}

tokens per GPU) and when it stays on its home GPU (

1/\mathrm{EP}

Fig. 3 : Reducing A2A with within-layer expert correlation. Each token $T_{i}$ and its selected experts have the same color.

Fig. 4 : TP–EP baseline with correlated placement. Each token stays on the rank of its canonical sequence slice, so its rows still cross the network even when its experts share one GPU.

Fig. 5 : Token shuffling end to end. After the attention of layer $\ell$ 1 , the shuffle planner predicts each token’s owner rank from its layer- $(\ell{-}1)$ routes, and the shuffled reduce-scatter 2 delivers each row to that rank. The dispatch 3 , the expert computation 4 and the combine 5 then find most of the token’s experts local. The shuffle-aware all-gather 6 restores canonical order before the next attention block 7 .

Fig. 6 : The planner runs on a side stream between the layer- $i$ routes and the layer- $(i{+}1)$ reduce-scatter.

Fig. 7 : All-to-All Cost Under Expert Parallelism. A2A time (top) and communication volume per rank (bottom) with EP8–EP64, for top-2 and top-6 routing. B is the baseline, D adds deduplication, and P adds correlated placement to deduplication.

Fig. 8 : Step-time breakdown at EP8–EP64 for top-2 and top-6 routing; dashed boxes indicate A2A share.

Fig. 9 : All-to-All Cost on the $\mathrm{TP}=\mathrm{EP}$ Layouts. Same encoding as Figure 7 , with token shuffling as the fourth configuration (S). The shuffle places tokens on the GPUs that hold their experts, so more of the dispatch stays on the device and less traffic crosses the network.

Fig. 10 : Step-time breakdown on the $\mathrm{TP}=\mathrm{EP}$ layouts (TP8 EP8, TP16 EP16) with token shuffling added as S. Same encoding as Figure 8 .

Appendix figures & tables4 assets

Supplementary material from the paper’s appendix.

Appendix

	Symbol	Value
Transformer
Layers	$L$	12, every layer an MoE layer
Model width	$d_{\text{model}}$	4096
Attention heads	$n_{\text{head}}$	16
Head dimension	$d_{\text{head}}$	32
Attention width	$n_{\text{head}}\,d_{\text{head}}$	512

Appendix

TABLE II: Model and pre-training configuration used in all experiments.

Fig. 11 : Router ablation on three top-6 models that differ only in load balancing: the main recipe (sequence auxiliary loss $10^{-2}$ ), a $10\times$ stronger auxiliary loss ( $10^{-1}$ ), and the DeepSeek-V3 auxiliary-loss-free router. Left: step speedup of deduplication + correlated placement over that model’s baseline, per layout. Right: the structure each router leaves to exploit, as the deduplication rate under correlated placement.

Fig. 12 : Step time (left) and A2A volume per rank (right) for the ten measured steps of each configuration, after the four warm-up steps. Data point labels give the ratio of the baseline median to the median of that configuration.

Symbol	Meaning
$t$ , $T$	a token; number of tokens in the global sequence
$E$ , $k$	experts per MoE layer; experts selected per token (top- $k$ )
$P$ , $N$	EP ranks (in Alg. 2 also the TP size); nodes
$\ell$ , $m$	MoE layer index; MoE position whose token owners are being planned
$\mathcal{S}_{\ell}(t)$	the $k$ experts the router selects for token $t$ at layer $\ell$
$W[e,e^{\prime}]$	number of tokens that select both $e$ and $e^{\prime}$ (co-selection count)

Appendix

TABLE III: Notation.

Explore similar work

DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism
May 10, 2026Zhichen Zeng, Chi-Chih Chang, Jiayi Wang +10Mixture-Of-Expert Inference Mixture-Of-Experts
ByteDance Seed · University of Washington · Cornell University
Scaling Multi-Node Mixture-of-Experts Inference Using Expert Activation Patterns
Apr 25, 2026Abhimanyu Bambhaniya, Geonhwa Jeong, Jason Park +6Mixture-Of-Expert Inference Mixture-Of-Experts
1Georgia Institute of Technology, USA · 2Meta Platforms, Inc.
Share First, Route What Remains: A Unified Framework for Token-Adaptive MoE Computation
Aug 11, 2026Gongli Zhang, Zhulin Liu, C. L. Philip ChenMixture-Of-Experts Unified Framework
Guangdong Provincial Key Laboratory of Computational AI Models and Cognitive Intelligence, School of Computer Science and Engineering, South China University of Technology, Guangzhou 510006, China · Pazhou Lab, Guangzhou 510335, China · Engineering Research Center of the Ministry of Education on Health Intelligent Perception and Paralleled Digital-Human, Guangzhou 510641, China

May 10, 2026, cs.LGJ/K move · Enter open · S save

DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism

Zhichen Zeng, Chi-Chih Chang, Jiayi Wang, Zezhou Wang+9

ByteDance Seed · University of Washington · Cornell University

Mixture-Of-Expert Inference Mixture-Of-Experts Cpu-Gpu Hybrid Designs Bottlenecks Disagreement

Open paper AskPDF arXiv

Mixture-of-experts (MoE) architectures enable trillion-parameter LLMs with sparsely activated experts. Expert parallelism (EP) is a widely adopted MoE training strategy, but it suffers from severe all-to-all communication bottlenecks, which is exaggerated by the limited inter-node network bandwidth as the growing model size requires distributing experts across GPU nodes. Prior work focused on overlapping these all-to-all communications with feed-forward network (FFN) and self-attention computations, which often leaves residual network-bound stalls due to inherent imbalance in attention and FFN layers' computation-communication ratios. We present DisagMoE, a disaggregated MoE training system that jointly optimizes model placement and scheduling for maximal efficiency. DisagMoE separates attention and FFN layers into disjoint GPU groups, introduces a multi-stage pipeline with uni-directional, many-to-many communications, and employs a computation-communication roofline model to balance GPU and network bandwidth allocation among the attention and FFN groups. DisagMoE is implemented on Megatron-LM, and evaluation shows that DisagMoE improves training efficiency across multiple MoE models with up to 1.8x speedup on 16-node 8xH800 clusters.

Abstract

Explore similar work

DisagMoE: Computation-Communication overlapped MoE Training via Disaggregated AF-Pipe Parallelism

Scaling Multi-Node Mixture-of-Experts Inference Using Expert Activation Patterns

Share First, Route What Remains: A Unified Framework for Token-Adaptive MoE Computation