[论文] Training Communication-Efficient Mixture-of-Experts Language Models wi...

论文概要

研究领域: ML 作者: Simeng Sun, Roger Waleffe 发布时间: 2026-08-28 arXiv: 2608.28511

中文摘要

在使用专家并行训练混合专家（MoE）语言模型时，all-to-all token分派和组合集合通信可能消耗端到端训练时间的相当大一部分。在这项工作中，我们研究了通信高效的MoE模型（CE-MoE），其中我们采用异构层模式，将token混合和通道混合深度解耦。与传统模型在每个token混合层后交错MoE层（如注意力、Mamba-2）不同，CE-MoE模型将专家容量集中在少数几个路由MoE层中，同时通过添加额外的token混合和密集FFN层保持深度。在从2B到31.5B总参数的扩展阶梯上，在匹配的总参数和激活参数下，CE-MoE模型一致地降低训练成本，同时与全MoE基线匹配验证损失和下游基准。在31.5B规模下，CE-MoE使用33.3%更少的GPU小时，同时提高平均下游得分和推理吞吐量。

原文摘要

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

#论文 #arXiv #ML #小凯