LocMoE: A Low-Overhead MoE for Large Language Model Training

📄 arXiv: 2401.13920v3 📥 PDF

作者: Jing Li, Zhijie Sun, Xuan He, Li Zeng, Yi Lin, Entong Li, Binfan Zheng, Rongqian Zhao, Xin Chen

分类: cs.LG, cs.AI, cs.CL

发布日期: 2024-01-25 (更新: 2024-05-23)

备注: 1. Update the font size of all figures. 2. Update the name of the proposed layer Grouped Average Pooling (GrAP). 3. Change the order of the Section Contribution Statement

💡 一句话要点

提出LocMoE以解决大语言模型训练中的负载不均和延迟问题

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: Mixtures-of-Experts 大语言模型 负载均衡 通信优化 训练效率 深度学习 路由策略

📋 核心要点

  1. 现有的MoE模型在负载不均衡和高延迟的All-to-All通信方面存在显著性能瓶颈，影响了训练效率。
  2. 本文提出了一种新型路由策略，通过优化节点间和节点内的通信方式，改善负载均衡和降低延迟。
  3. 实验结果显示，LocMoE在训练时间上较传统路由器有显著提升，且保持了模型的准确性，具有良好的实用性。

📝 摘要（中文）

Mixtures-of-Experts (MoE) 模型是一种广泛应用于大语言模型（LLM）的分布式学习方法，因其高效的稀疏化和扩展能力而受到青睐。然而，MoE的性能受到负载不均衡和All-to-All通信高延迟的限制，同时由于专家容量过大导致计算冗余。为了解决这些问题，本文提出了一种新颖的路由策略，通过将部分节点间通信转换为节点内通信来实现负载均衡和局部性。我们还阐明了专家容量的最小阈值，该阈值通过专家的门控权重与分配的标记之间的最大角度偏差计算得出。我们在MindSpore框架下的PanGu-Sigma模型上实现了这些修改，并在Ascend集群上进行了实验。实验结果表明，所提出的LocMoE相比于传统路由器（如哈希路由器和交换路由器），每个epoch的训练时间减少了12.68%至22.24%，且不影响模型的准确性。

🔬 方法详解

问题定义：本文旨在解决大语言模型训练中MoE模型的负载不均衡和高延迟问题。现有的路由策略导致某些专家被频繁选择，造成负载不均，且All-to-All通信的频繁进行显著延长了训练时间。

核心思路：论文提出了一种结合负载均衡和局部性的路由策略，通过将部分节点间通信转化为节点内通信，来减少延迟并提高训练效率。这种设计旨在优化资源的使用，降低冗余计算。

技术框架：整体架构基于PanGu-Sigma模型，采用MindSpore框架，实施多级路由策略。主要模块包括路由选择、专家调用和通信优化，确保在训练过程中有效分配计算资源。

关键创新：最重要的技术创新在于提出了新的路由策略和专家容量的最小阈值计算方法。这一创新与传统方法的本质区别在于通过优化通信方式和专家选择，显著提升了训练效率。

关键设计：在设计中，专家容量的最小阈值通过专家的门控权重与分配标记之间的最大角度偏差进行计算。此外，采用了多级路由策略以优化节点间和节点内的通信，确保了负载均衡。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，LocMoE在训练时间上相比于传统的哈希路由器和交换路由器减少了12.68%至22.24%。这一显著的性能提升在不影响模型准确性的前提下，展示了新路由策略的有效性和实用性。

🎯 应用场景

该研究具有广泛的应用潜力，尤其在大规模自然语言处理任务中，如机器翻译、文本生成和对话系统等领域。通过提高训练效率，LocMoE能够加速模型的迭代和部署，推动智能应用的发展。未来，该方法还可能扩展到其他类型的深度学习模型中，进一步提升其性能和应用价值。

📄 摘要（原文）

The Mixtures-of-Experts (MoE) model is a widespread distributed and integrated learning method for large language models (LLM), which is favored due to its ability to sparsify and expand models efficiently. However, the performance of MoE is limited by load imbalance and high latency of All-to-All communication, along with relatively redundant computation owing to large expert capacity. Load imbalance may result from existing routing policies that consistently tend to select certain experts. The frequent inter-node communication in the All-to-All procedure also significantly prolongs the training time. To alleviate the above performance problems, we propose a novel routing strategy that combines load balance and locality by converting partial inter-node communication to that of intra-node. Notably, we elucidate that there is a minimum threshold for expert capacity, calculated through the maximal angular deviation between the gating weights of the experts and the assigned tokens. We port these modifications on the PanGu-Sigma model based on the MindSpore framework with multi-level routing and conduct experiments on Ascend clusters. The experiment results demonstrate that the proposed LocMoE reduces training time per epoch by 12.68% to 22.24% compared to classical routers, such as hash router and switch router, without impacting the model accuracy.