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.

翻译：混合专家模型（MoE）是一种广泛用于大语言模型（LLM）的分布式集成学习方法，因其能够高效地稀疏化和扩展模型而受到青睐。然而，MoE的性能受到负载不均衡和全到全通信高延迟的限制，同时由于较大的专家容量导致计算冗余。负载不均衡可能源于现有路由策略持续倾向于选择特定专家。全到全过程中频繁的节点间通信也显著延长了训练时间。为缓解上述性能问题，我们提出一种新型路由策略，通过将部分节点间通信转换为节点内通信，实现了负载均衡与局部性的结合。值得注意的是，我们阐明了专家容量存在最小阈值，该阈值通过计算专家门控权重与分配令牌之间的最大角度偏差得出。我们基于MindSpore框架在多级路由棋盘上将这些改进移植到盘古-西格玛模型中，并在Ascend集群上进行实验。实验结果表明，与经典路由算法（如哈希路由和开关路由）相比，所提出的LocMoE在不影响模型准确率的情况下，将每轮训练时间减少了12.68%至22.24%。