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

May 10, 2026cs.LG

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

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.

Zhichen Zeng, Chi-Chih Chang, Jiayi Wang +10

ByteDance Seed · University of Washington · Cornell University

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

Ask PDF

Apr 25, 2026cs.LG

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

Most recent state-of-the-art (SOTA) large language models (LLMs) use Mixture-of-Experts (MoE) architectures to scale model capacity without proportional per-token compute, enabling higher-quality outputs at manageable serving costs. However, MoE inference at scale is fundamentally bottlenecked by expert load imbalance and inefficient token routing, especially in multi-node deployments where tokens are not guaranteed to be routed to local experts, resulting in significant inter-node all-to-all communication overhead. To systematically characterize these challenges, we profile SOTA open-source MoE models, including Llama 4 Maverick, DeepSeek V3-671B, and Qwen3-230B-A22B, on various datasets and collected over 100k real expert activation traces. Upon studying the expert activation patterns, we uncover various persistent properties across all the frontier MoE models: variable expert load imbalance, domain-specific expert activation where expert popularity shifts across task families (code, math, chat, general), and a strong correlation between prefill and decode expert activations. Motivated by these findings, we propose workload-aware micro-batch grouping and an expert placement strategy to maximize token locality to the destination expert, thereby reducing inter-node communication. Across models and datasets, these optimizations help reduce all2all communication data up to 20, resulting in lower MoE decode latency and better accelerator utilization.

Abhimanyu Bambhaniya, Geonhwa Jeong, Jason Park +6

1Georgia Institute of Technology, USA · 2Meta Platforms, Inc.

Mixture-Of-Expert Inference Mixture-Of-Experts LLM Inference Optimization Model Activations Class Imbalance

Ask PDF

Aug 11, 2026cs.LG

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

Mixture-of-experts (MoE) models have recently moved beyond routing a fixed number of complete experts. Shared-expert designs preserve reusable knowledge, fine-grained methods vary computation within experts, and dynamic routers adapt the number of active experts. Yet these decisions are usually made independently, overlooking a basic dependency: extracting reusable computation changes both what remains and how much expert capacity the remainder needs. We study this dependency by decomposing sparsely upcycled feed-forward experts into key-value channels. Co-activated experts align at a subset of value positions; removing these positions changes expert preference; and greater shared coverage is associated with lower residual expert demand. These observations lead to one principle: share first, then route what remains. We instantiate it in UniF-MoE, a unified framework for token-adaptive MoE computation. Each expert is partitioned into aligned blocks. A shared-demand score sets the shared block count and pathway weight, key prototypes select the shared content, and the complementary demand determines the residual expert count through cumulative routing mass. A Gram regularizer separates and normalizes router embeddings, promoting diverse routing directions, sparse expert overlap, and a simple routing geometry. Experiments on DomainBed and GLUE show that this unified design improves predictive performance over representative static and dynamic MoEs while reducing activated computation, inference latency, and memory. Code is available at https://github.com/existence0420/UniF-MoE.

Gongli Zhang, Zhulin Liu, C. L. Philip Chen

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

Mixture-Of-Experts Unified Framework

Ask PDF