PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Simeng Sun, Roger Waleffe

arXiv:2608.28511v1cs.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.

    Source: §1 · p. 1

  • 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.

    Source: §1 · p. 2, §2 · p. 3

  • 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.

    Source: §5.4 · p. 8, §6 · p. 10

  • 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.

    Source: §5.4 · p. 8, §6 · p. 10

    Measured savings may change with other precisions, larger world sizes, or overlapping communication, and inference results are preliminary from a synthetic static benchmark.

    Source: §6 · p. 10

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.Source: §1 · p. 1
  • 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.Source: §1 · pp. 1–2

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.Source: §1 · p. 2

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.Source: §1 · p. 2

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.Source: §1 · p. 2

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.Source: §2 · pp. 2–3

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.Source: Figure 2 · p. 3

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.Source: §2 · p. 3
  • 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.Source: §2 · p. 3

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.Source: §2 · p. 3

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.Source: §3 · p. 3

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.Source: §3 · p. 3

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.Source: §3 · p. 3, Figure 3 · p. 4
  • 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.Source: §3 · p. 4

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.Source: §3 · p. 4

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.Source: §4.1 · p. 4
  • 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.Source: §4.1 · pp. 4–5

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.Source: §4.2 · p. 5, §4.1 · p. 5
  • 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.Source: §4.2 · p. 6, Figure 4 · p. 6

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.Source: §4.3 · p. 6
  • 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.Source: §4.3 · pp. 6–7

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.Source: §5.4 · p. 8
  • 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.Source: §5.4 · p. 8

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.Source: §5.5 · pp. 8–9, Figure 6 · p. 9

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.Source: §5.6 · pp. 9–10

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.Source: §6 · p. 10
  • 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.Source: §6 · p. 10

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.Source: §6 · p. 10
  • 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.Source: §6 · p. 10

Related papers

Loading 2608.28511v1…