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

2024-05-23

Abstract: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.

Machine Learning,Artificial Intelligence,Computation and Language

What problem does this paper attempt to address?

The problem this paper attempts to address is the performance bottleneck encountered in training large-scale language models (LLMs) using existing Mixture of Experts (MoE) models. Specifically, these issues include load imbalance, high-latency All-to-All communication, and relatively redundant computation due to large expert capacity. These problems not only affect the efficiency of model training but may also lead to overtraining of some experts while others are undertrained, thereby impacting the balance of learning efficiency and load distribution.

To solve the aforementioned issues, the paper proposes a new routing strategy—LocMoE (Low Overhead MoE). This strategy combines load balancing and locality principles by converting some inter-node communication to intra-node communication to reduce communication latency. It also introduces locality loss as a regularization term, encouraging local experts to compete with specialized experts, reducing communication time consumption, and avoiding undertraining of some experts. Additionally, the paper explores the minimum threshold of MoE expert capacity and demonstrates the relationship between this threshold and the characteristics of the input corpus, finding that reducing expert capacity does not affect model accuracy as long as it does not exceed this critical limit.

Through these improvements, LocMoE was experimented on Ascend clusters of different scales. The results show that compared to traditional routers (such as hash routers and switch routers), LocMoE can reduce the training time per epoch by 12.68% to 22.24% without affecting model accuracy. This indicates that LocMoE has significant advantages in improving the efficiency of large-scale language model training.