Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration

Simeng Sun; Roger Waleffe

Source-linked AI summary

Training Communication-Efficient Mixture-of-Experts Language Models with Layer Re-Configuration

Simeng Sun, Roger Waleffe

arXiv:2608.28511v1Aug 28, 2026cs.AI

TL;DR

Large MoE models can spend substantial training time on expert-parallel all-to-all communication. The paper proposes CE-MoE, which reduces MoE depth and reallocates capacity into token-mixing-heavy heterogeneous layers; across scaling experiments, it lowers training cost while retaining quality, with 33.3% fewer GPU-hours at 31.5B parameters.

Problem
Expert-parallel all-to-all communication for routing, combining, and gradients can consume a substantial fraction of end-to-end MoE training time.
Method
CE-MoE reduces the number of MoE layers and reallocates capacity into additional token-mixing layers and fewer, wider expert and dense-FFN blocks under comparable parameter budgets.
Results
CE-MoE consistently reduces GPU-hour training cost while keeping validation loss close to matched full-MoE baselines and remaining competitive on downstream benchmarks.
Takeaways & Limitations
At larger scale, CE-MoE can train for more tokens at lower cost while remaining competitive on downstream benchmarks and showing preliminary inference-throughput gains.
Takeaways & Limitations
Measured savings may change with other precisions, larger world sizes, or overlapping communication, and inference results are preliminary from a synthetic static benchmark.

Abstract

from arXiv · show

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.

1 Introduction

Expert-parallel all-to-all communication can consume substantial training time in large MoE models. CE-MoE reduces MoE depth and reallocates capacity into a heterogeneous layer layout to lower communication volume while preserving model quality.

Motivation: Expert-parallel MoE training uses all-to-all collectives for token dispatch, expert-output combination, and routed gradients.These collectives move activations and gradients across ranks hosting selected experts.
Motivation: Communication volume can be reduced through lower precision, lower-dimensional dispatch, smaller routing width, or fewer MoE layers.The paper focuses on reducing the total moved-data volume rather than only improving communication systems.
Approach: CE-MoE substantially reduces MoE depth while adding token-mixing and dense-FFN blocks and concentrating capacity in fewer, wider expert layers.The resulting heterogeneous model is roughly parameter matched to the regular baseline.
Evaluation: The report evaluates CE-MoE on training cost, downstream quality, inference throughput, and residual-stream behavior against parameter-matched hybrid-MoE baselines.Its sections also examine training instability and provide a construction recipe.
Scope: The report concludes with a summary of findings and a discussion of limitations.The stated evaluation scope covers scaling-ladder training and model-quality comparisons.

2 Overview

CE-MoE departs from regular token-mixing/channel-mixing alternation by aggressively reducing MoE depth and skewing the layer composition toward token mixing. Its construction uses a phased-greedy procedure, while consecutive token mixers motivate stability checks favoring Mamba-2 over self-attention.

Layer layout: Standard sparse architectures reduce routed MoE depth but retain regular token-mixing and channel-mixing alternation.The paper contrasts this pattern with more heterogeneous CE-MoE layouts.
Stability motivation: In ∼370M-parameter toy models, attention variants with seven consecutive token-mixing layers are less stable, whereas Mamba-2 variants train smoothly at both learning rates.The variants are matched in per-layer parameter count.
Layer layout: CE-MoE decouples token-mixing and channel-mixing counts, departing from the standard Transformer block design.This more aggressive reduction of MoE depth produces a heterogeneous layer pattern.
Layer layout: CE-MoE assigns four token-mixing layers for every MoE layer or more than two token-mixing layers for every channel-mixing layer.This contrasts with the standard Transformer’s 1:1 token-mixing-to-channel-mixing ratio.
Construction: A phased-greedy procedure spaces attention and channel-mixing blocks along a Mamba-2 spine while permitting consecutive token-mixing blocks.The procedure constructs the heterogeneous layer pattern while avoiding poor local orderings.

3 Rank-collapse risk

Consecutive token-mixing layers can destabilize training by collapsing token representations toward low-rank structure. The experiments find stronger rank collapse and norm growth for self-attention than for Mamba-2, supporting Mamba-2 as the main token-mixing path.

Risk: The token-mixing-heavy CE-MoE layout introduces adjacent token-mixing blocks and raises concerns about training stability and expressivity.Prior work links pure self-attention stacks to convergence toward rank 1, while these experiments test stacks with skip connections.
Measurements: The analysis measures stable rank and effective rank on post-residual token representations using singular-value statistics.The representations contain Ntok sampled tokens and hidden dimension D=1024.
Results: Self-attention networks have much lower stable and effective rank than Mamba-2, with strongest collapse in early layers.At the higher learning rate, early self-attention layers reach stable rank close to 1 and single-digit effective rank.
Results: High-learning-rate self-attention develops per-token ℓ2 norms orders of magnitude larger than Mamba-2, with token differences increasingly carried by magnitude.The paper identifies this magnitude-based representation as a major source of instability.
Implication: The results suggest consecutive token-mixing blocks are viable when the mixer maintains healthy rank under repeated stacking.Mamba-2 is therefore selected as the better candidate for the main token-mixing path in the heterogeneous layout.

4 Communication-Efficient Hybrid-MoE

CE-MoE constructs heterogeneous hybrid-MoE architectures by reducing routed MoE depth, reallocating depth to token-mixing and dense-FFN blocks, and matching baseline parameter budgets. A phased-greedy ordering and shape search reduce communication while preserving validation performance in the reported ablation.

Determine the layer composition: CE-MoE reduces MoE layers and reallocates depth to additional token-mixing and dense-FFN blocks while preserving total layer count.The design relaxes the standard 1:1 token-mixing-to-channel-mixing ratio.
Determine the layer composition: 65.2% fewer MoE all-to-all calls are incurred by CE-MoE than the 23-layer baseline in the 52-layer running example.CE-MoE uses 8q calls versus 23q for the baseline; the comparison assumes the stated composition and expert-parallel setup.
Generate the layer pattern: The phased-greedy procedure spreads attention and channel-mixing blocks across a Mamba-2 spine using virtual blocks, residue handling, and even partitioning.Its four phases create attention-bearing, non-attention, and residue blocks before distributing them into the final pattern.
Generate the layer pattern: At 2B parameters and 140B tokens, phased-greedy CE-MoE matches the baseline and outperforms all 21 constrained-random layer patterns in validation loss.At the final step, CE-MoE is roughly 1.8 standard deviations below the constrained-random mean.
Determine the model shape: The shape search matches total and activated parameter budgets while allowing wider experts and more experts to compensate for fewer routed MoE layers.The search varies expert count, routed top-K, expert and dense-FFN widths, and related shape parameters.
Determine the model shape: 59.4% lower LEK is achieved by the selected 31.5B CE-MoE shape while changing total and activated parameters by −0.001% and −0.257%.LEK is used as a communication-cost proxy when dispatch dimension and communication precision are fixed.

5 Empirical results

CE-MoE reduces end-to-end training cost across matched scaling experiments while preserving validation loss and downstream quality, and it also improves preliminary inference throughput. Residual-stream analysis suggests this quality is reached through a different, more periodic representational trajectory.

Model quality and scaling-up experiments: 33.3% fewer GPU-hours at 31.5B-A3.5B while improving average downstream score from 58.37 to 58.65.At 14B-A2B, same-token CE-MoE also uses 34.1% fewer GPU-hours and raises the average score from 56.93 to 58.41.
Model quality and scaling-up experiments: 26.1% fewer GPU-hours than the baseline when the larger-GBS CE-MoE run trains on 1.25× as many tokens and reaches an average score of 59.36.The larger-GBS experiment reinvests saved cost into additional training tokens.
Inference speed: 28–36% higher throughput across output lengths from 128 to 2048 for the 31.5B-A3.5B CE-MoE model.The benchmark uses input lengths 2048 and 8192; the authors attribute the gain to fewer MoE layers, reducing router overhead and synchronization points.
Residual stream analysis: CE-MoE follows a staircase-like cosine-to-final trajectory, with fewer expert layers making larger representational updates and intervening token-mixing layers providing smaller refinements.Its effective-rank variation is larger, but both models end with similar effective rank and rank recovers near expert layers.

6 Conclusion

The paper presents CE-MoE as a heterogeneous architecture that reduces expert-parallel communication by reducing MoE depth and reallocating capacity. Across scaling experiments, it lowers GPU-hours while keeping validation loss close to matched baselines, with downstream and preliminary inference gains at larger scale.

Conclusion: CE-MoE reduces expert-parallel communication by using fewer MoE layers and reallocating capacity to token-mixing, dense-FFN, and wider expert blocks.The models use comparable total and activated parameter budgets.
Conclusion: CE-MoE consistently lowers GPU-hours while keeping validation loss close to matched full-MoE baselines across the scaling ladder.At larger scale, it can train for more tokens at lower cost and remain competitive on downstream benchmarks.
Limitations: Measured savings may change with other compute or communication precisions, larger world sizes, or overlapping communication strategies.The authors also report higher sequence-level load-imbalance signals and more early-training loss spikes when global batch size does not scale with expert counts.
Limitations: Inference results are preliminary and come from a synthetic static benchmark rather than a full production serving study.This limits how directly the reported throughput gains transfer to production deployment.