通过层重配置训练通信高效的混合专家语言模型

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