Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Layer-adaptive Expert Pruning for Pre-Training of Mixture-of-Experts Large Language Models

YuanLab. ai, Shawn Wu, Jiangang Luo, Tong Yu, Darcy Chen, Sean Wang, Xudong Zhao, Louie Li, Claire Wang, Hunter He, Carol Wang, Allen Wang

TL;DR

This work addresses pre-training inefficiencies in Mixture-of-Experts large language models caused by skewed expert utilization. It introduces Layer-Adaptive Expert Pruning (LAEP), which prunes underutilized experts and rearranges the remaining ones across devices based on stable token-distribution statistics, avoiding auxiliary load-balancing losses. The approach yields substantial parameter reductions and training-efficiency gains, demonstrated on models up to 1010B base and 1515B MoE configurations, while preserving or improving accuracy across diverse benchmarks. LAEP thus offers a practical path to more memory- and compute-efficient MoE pre-training at extreme scales, with implications for deployment and real-world use.

Abstract

Although Mixture-of-Experts (MoE) Large Language Models (LLMs) deliver superior accuracy with a reduced number of active parameters, their pre-training represents a significant computationally bottleneck due to underutilized experts and limited training efficiency. This work introduces a Layer-Adaptive Expert Pruning (LAEP) algorithm designed for the pre-training stage of MoE LLMs. In contrast to previous expert pruning approaches that operate primarily in the post-training phase, the proposed algorithm enhances training efficiency by selectively pruning underutilized experts and reorganizing experts across computing devices according to token distribution statistics. Comprehensive experiments demonstrate that LAEP effectively reduces model size and substantially improves pre-training efficiency. In particular, when pre-training the 1010B Base model from scratch, LAEP achieves a 48.3\% improvement in training efficiency alongside a 33.3% parameter reduction, while still delivering excellent performance across multiple domains.

Layer-adaptive Expert Pruning for Pre-Training of Mixture-of-Experts Large Language Models

TL;DR

This work addresses pre-training inefficiencies in Mixture-of-Experts large language models caused by skewed expert utilization. It introduces Layer-Adaptive Expert Pruning (LAEP), which prunes underutilized experts and rearranges the remaining ones across devices based on stable token-distribution statistics, avoiding auxiliary load-balancing losses. The approach yields substantial parameter reductions and training-efficiency gains, demonstrated on models up to 1010B base and 1515B MoE configurations, while preserving or improving accuracy across diverse benchmarks. LAEP thus offers a practical path to more memory- and compute-efficient MoE pre-training at extreme scales, with implications for deployment and real-world use.

Abstract

Although Mixture-of-Experts (MoE) Large Language Models (LLMs) deliver superior accuracy with a reduced number of active parameters, their pre-training represents a significant computationally bottleneck due to underutilized experts and limited training efficiency. This work introduces a Layer-Adaptive Expert Pruning (LAEP) algorithm designed for the pre-training stage of MoE LLMs. In contrast to previous expert pruning approaches that operate primarily in the post-training phase, the proposed algorithm enhances training efficiency by selectively pruning underutilized experts and reorganizing experts across computing devices according to token distribution statistics. Comprehensive experiments demonstrate that LAEP effectively reduces model size and substantially improves pre-training efficiency. In particular, when pre-training the 1010B Base model from scratch, LAEP achieves a 48.3\% improvement in training efficiency alongside a 33.3% parameter reduction, while still delivering excellent performance across multiple domains.
Paper Structure (21 sections, 2 equations, 4 figures, 7 tables, 2 algorithms)

This paper contains 21 sections, 2 equations, 4 figures, 7 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Expert Pruning for MoE models
  4. Expert load balance method for MoE models
  5. Method
  6. Expert Token distribution in Pre-training
  7. Expert Pruning
  8. Expert Rearrangement
  9. Ablation study
  10. Sensitivity Analysis of Pruning Hyper-parameters
  11. Effect of pruning on different Attention Architectures
  12. Impact of LAEP on Downstream Task Performance
  13. Scalability of LAEP
  14. Pre-training of 1010B Base model with LAEP
  15. Conclusion
  16. ...and 6 more sections

Figures (4)

  • Figure 1: The left column (a-c) displays the trend of token distribution among experts across different layers changing over the course of training. The right column (d-f) shows token distribution of individual experts in the stable stage of pre-training (80B tokens). The red dashed line denotes the average number of tokens per expert per layer. Max/Min represents the ratio between the token count of the most heavily utilized expert and that of the least utilized expert.
  • Figure 2: (a) Number of experts under different token loads; (b) Accumulated tokens of experts that are sorted by token loads from smallest to largest. The area to the left of the red dashed line represents experts that need to be pruned according to hyperparameter of $\alpha$ and $\beta$, while the green dashed line represents the average token load per expert.
  • Figure 3: Illustration of the Expert Rearrangement algorithm for load balancing among computing devices in MoE LLM pre-training.
  • Figure 4: The trend of token distribution among experts across different layers over the course of training. The top two columns display the results with auxiliary loss from Deepseek V3 (auxiliary loss coefficient c = 0.0001 and 0.01), while the bottom two columns show the results with Mixtral auxiliary loss (auxiliary loss coefficient c = 0.0001 and 0.01).