LocMoE: A Low-overhead MoE for Large Language Model Training
Jing Li Zhijie Sun Xuan He Li Zeng Yi Lin Entong Li Binfan Zheng Rongqian Zhao Xin Chen
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.