Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration
2026-08-28 • Artificial Intelligence
Artificial Intelligence
AI summaryⓘ
The authors studied a way to make training large Mixture-of-Experts (MoE) language models faster by changing how the model layers are organized. Instead of spreading the expert parts across many layers, they focused expert capacity into fewer layers and added other types of layers to keep the model's depth. Their approach, called Communication-Efficient MoE (CE-MoE), uses less training time but performs just as well as traditional models across different sizes. At the largest model scale tested, CE-MoE trained faster, performed better on tests, and ran inference more efficiently.
Mixture-of-Experts (MoE)language modelsexpert parallelismtoken dispatchattention layersfeed-forward networks (FFN)model scalinginference throughputtraining efficiency
Authors
Simeng Sun, Roger Waleffe
Abstract
When training Mixture-of-Experts (MoE) language models with expert parallelism, all-to-all token dispatch and combine collectives can consume a substantial fraction of end-to-end training time. In this work, we study communication-efficient MoE models (CE-MoE), in which we adopt a heterogeneous layer pattern that decouples token-mixing and channel-mixing depth. Compared to conventional models which interleave MoE layers after each token-mixing layer (e.g., attention, Mamba-2), CE-MoE models concentrate expert capacity in a select few routed MoE layers, while maintaining depth by adding additional token-mixing and dense-FFN layers. Across a scaling ladder from 2B to 31.5B total parameters, under matched total and activated parameters, CE-MoE models consistently reduce training cost while matching validation loss and downstream benchmarks with full-MoE baselines. At the 31.5B scale, CE-MoE uses 33.3\% fewer GPU-hours while improving average downstream score and inference throughput.