Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

TL;DR

LocMoE addresses key bottlenecks in MoE-based LLM training, notably load imbalance and high inter-node communication, by introducing a locality-aware routing strategy and a GrAP-based gating mechanism. It combines a locality regularizer with an auxiliary load-balancing loss, and derives a theoretical lower bound on expert capacity linked to token-gating geometry, enabling reduced capacity without sacrificing accuracy. Group-wise All-to-All and communication overlap further cut latency, and experiments on PanGu-$\Sigma$ across Ascend clusters demonstrate up to $22.24\%$ reduction in time per epoch while preserving or improving task performance. The work offers practical improvements for scalable sparse transformers and lays groundwork for more efficient, domain-local MoE routing in large-scale NLP systems.

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.

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

TL;DR

LocMoE addresses key bottlenecks in MoE-based LLM training, notably load imbalance and high inter-node communication, by introducing a locality-aware routing strategy and a GrAP-based gating mechanism. It combines a locality regularizer with an auxiliary load-balancing loss, and derives a theoretical lower bound on expert capacity linked to token-gating geometry, enabling reduced capacity without sacrificing accuracy. Group-wise All-to-All and communication overlap further cut latency, and experiments on PanGu- across Ascend clusters demonstrate up to reduction in time per epoch while preserving or improving task performance. The work offers practical improvements for scalable sparse transformers and lays groundwork for more efficient, domain-local MoE routing in large-scale NLP systems.

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.
Paper Structure (19 sections, 4 theorems, 16 equations, 15 figures, 1 table)

This paper contains 19 sections, 4 theorems, 16 equations, 15 figures, 1 table.

Table of Contents

  1. Introduction
  2. Related Work
  3. MoE.
  4. Ascend Architecture.
  5. PanGu Series Model.
  6. Methodology
  7. PanGu-$\Sigma$
  8. MoE With Local Routing Strategy
  9. MoE in Encoder Layers of Transformers
  10. Localized Bias Weighting Loss
  11. Critical Value of Expert Capacity
  12. Group-Wise All-to-All and Communication Overlap
  13. Experiment Results and Analysis
  14. Analysis for Expert Capacity
  15. Ablation Analysis
  16. ...and 4 more sections

Key Result

Lemma 1

The top-1 router is essentially the mechanism to select the expert $i^*$ with the minimum angle $\theta_{i^*, j}$ to the gating weight $\omega_i$.

Figures (15)

  • Figure 1: The networking scheme applied in the Ascend cluster.
  • Figure 2: The algorithm bandwidth of each communication operator in HCCL under 64N, 128N, and 256N, respectively.
  • Figure 3: The architecture of sparse Transformer layers in PanGu-$\Sigma$.
  • Figure 4: Difference between feature extraction via the dense layer and the GrAP layer.
  • Figure 5: The action principle of locality loss.
  • ...and 10 more figures

Theorems & Definitions (7)

  • Lemma 1: Minimum Angle of Expert
  • Lemma 2: Equivalent Probability for Assignment
  • Lemma 3: Assignment Probability for Unit Vector
  • Theorem 1: Lower Bound of Expert Capacity
  • proof
  • proof
  • proof