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

Researchers Simeng Sun and Roger Waleffe introduce Communication-Efficient Mixture-of-Experts (CE-MoE) models to address the high communication overhead associated with expert parallelism in large language models. By reconfiguring how layers are structured, the authors aim to reduce the time spent on "all-to-all" token dispatch and combine collectives, which often account for a significant portion of training time in conventional Mixture-of-Experts (MoE) architectures.

The Communication Bottleneck

In standard MoE models, expert layers are typically interleaved after every token-mixing layer, such as attention or Mamba-2 blocks. This design requires frequent communication between processors to route tokens to the correct experts. The authors identify these communication collectives as a primary bottleneck that slows down end-to-end training.

Reconfiguring Model Layers

CE-MoE addresses this by decoupling token-mixing depth from channel-mixing depth. Instead of placing MoE layers throughout the entire model, the researchers concentrate expert capacity into a select few routed MoE layers. To maintain the model's overall depth and capacity, they supplement these routed layers with additional token-mixing and dense feed-forward network (FFN) layers. This heterogeneous pattern reduces the frequency of communication-heavy operations while preserving the model's ability to process information.

Training Efficiency and Performance

The authors evaluated CE-MoE across a scaling range from 2B to 31.5B total parameters. When compared to full-MoE baselines with matched total and activated parameters, CE-MoE models consistently reduced training costs while maintaining similar validation loss and performance on downstream benchmarks.
At the 31.5B parameter scale, the researchers report that CE-MoE achieved:

• A 33.3% reduction in GPU-hours required for training.

• Improved inference throughput.

• Higher average scores on downstream benchmarks.

Franklin Analysis

The evidence suggests that the primary value of the CE-MoE approach is its ability to maintain model performance while lowering the hardware requirements for training. By shifting the architecture away from a uniform distribution of MoE layers, the authors successfully mitigate the communication overhead that typically limits the efficiency of large-scale MoE training. The 33.3% reduction in GPU-hours at the 31.5B scale provides a concrete metric for the practical benefit of this layer re-configuration strategy.