阿里云通义大模型新技术：MoE模型训练专家平衡的关键细节

为 MoE 模型优化提供新视角。

原标题：阿里云通义大模型新技术：MoE模型训练专家平衡的关键细节
文章来源：机器之心
内容字数：7053字

机器之心AIxiv专栏：阿里云通义千问团队解决MoE模型训练关键问题

机器之心AIxiv专栏长期关注并报道全球AI领域的学术前沿，近日报道了阿里云通义千问Qwen团队的一篇论文，该论文揭示并解决了MoE（混合专家模型）训练中的一个关键问题，并提出了一种显著提升MoE模型性能和专家特异性的新方法。

1. MoE模型训练中的关键问题：局部负载均衡的局限性

MoE模型通过路由机制动态激活部分模型参数，从而高效地扩展模型规模。然而，基于TopK机制的稀疏激活会导致专家激活不均衡：部分专家被频繁选择，而其他专家被闲置，造成资源浪费。为了解决这个问题，通常引入负载均衡损失（LBL）来鼓励专家激活均衡。然而，现有MoE训练框架大多实现的是局部（micro-batch）负载均衡，即在每个小批量数据内进行均衡分配。这会导致一个问题：如果小批量数据来自同一领域，则模型会将这些数据均匀分配给所有专家，阻碍专家在特定领域的分化，从而限制模型性能。

2. 从局部均衡到全局均衡：轻量级通信的解决方案

为了克服局部均衡的局限性，阿里云通义千问团队提出了一种全新的方法：将局部均衡放松为全局均衡。通过轻量级的通信机制，该方法能够同步所有micro-batch的专家选择频率，并在全局范围内计算负载均衡损失。这种方法的关键在于，只需要在各个节点之间同步一个专家数大小的向量，通信开销极小。即使对于需要梯度积累的情况，通过缓存机制也能有效降低通信成本。

3. 实验结果：显著提升模型性能和专家特异性

研究人员在不同规模的模型（3.4B、15B、43B参数）上进行了实验，结果表明，将均衡范围从局部扩展到全局后，模型的性能（Benchmark指标和PPL）均得到了显著提升。实验还表明，增加全局均衡范围能持续提升模型效果，并在一定范围后达到饱和。此外，通过对比实验，研究人员验证了全局均衡的提升主要源于在更通用、多样化的token集合上计算损失，而不是单纯的token数量增加。

4. 效率提升：少量局部均衡损失的补充

完全依赖全局均衡可能会降低局部均衡，影响计算效率。因此，研究人员进一步实验了在主要使用全局均衡的同时，添加少量局部均衡损失。结果表明，这种方法既能提高模型训练速度，又能保持模型性能。

5. 与现有工作的比较

与现有的一些工作相比，该研究更系统地验证了全局负载均衡的有效性，详细分析了均衡范围对性能的影响，并通过消融实验证明了全局均衡通过纳入更多样化的领域信息来提升性能。

6. 结论

阿里云通义千问团队的研究解决了MoE模型训练中的一个关键问题，通过轻量级通信实现全局负载均衡，显著提升了MoE模型的性能和专家特异性。这项工作为MoE模型的优化提供了新的视角，并有助于构建更大规模、更有效的MoE模型。

联系作者

文章来源：机器之心
作者微信：
作者简介：专业的人工智能媒体和产业服务平台