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.

通过转换部分节点间通信为节点内通信的路由策略，我们提出了一种新颖的负载均衡和局部性相结合的路由策略，并在 PanGu-Sigma 模型上进行了改进，实验证明该路由策略减少了每轮训练时间约 12.68% 至 22.24%，而不影响模型的准确性。