Source-linked AI summary
Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models
Xudong Lu, Qi Liu, Yuhui Xu, Aojun Zhou, Siyuan Huang, Bo Zhang, Junchi Yan, Hongsheng Li
TL;DR
MoE LLMs offer strong performance with fewer active parameters but remain difficult to deploy because their full expert collections require substantial memory. The paper proposes plug-and-play post-training expert pruning and dynamic skipping, achieving lower memory use and faster inference while maintaining satisfactory performance across tasks.
Problem
MoE LLMs reduce active parameters but remain difficult to deploy because their static expert parameters impose substantial memory and storage demands.
Method
The paper introduces hardware-friendly post-training task-agnostic and task-specific expert pruning together with dynamic expert skipping for MoE LLMs.
Results
The methods reduce memory usage and increase inference speed while maintaining satisfactory performance, including nearly 90% of Mixtral 8x7B Instruct performance with half the parameters and 1.33× token-generation speedup.
Takeaways & Limitations
Expert-level sparsification provides a more efficient MoE deployment approach that complements existing model-sparsification strategies.
Takeaways & Limitations
The pruning algorithm relies on expert-combination enumeration, which becomes cumbersome as each MoE layer scales beyond the currently evaluated 4- or 8-expert settings.
Abstract
from arXiv · showhide
A pivotal advancement in the progress of large language models (LLMs) is the emergence of the Mixture-of-Experts (MoE) LLMs. Compared to traditional LLMs, MoE LLMs can achieve higher performance with fewer parameters, but it is still hard to deploy them due to their immense parameter sizes. Different from previous weight pruning methods that rely on specifically designed hardware, this paper mainly aims to enhance the deployment efficiency of MoE LLMs by introducing plug-and-play expert-level sparsification techniques. Specifically, we propose, for the first time to our best knowledge, post-training approaches for task-agnostic and task-specific expert pruning and skipping of MoE LLMs, tailored to improve deployment efficiency while maintaining model performance across a wide range of tasks. Extensive experiments show that our proposed methods can simultaneously reduce model sizes and increase the inference speed, while maintaining satisfactory performance. Data and code will be available at https://github.com/Lucky-Lance/Expert_Sparsity.
1 Introduction
MoE LLMs improve parameter efficiency but remain difficult to deploy because experts dominate static memory requirements. This paper introduces hardware-friendly post-training expert pruning and dynamic skipping to reduce deployment cost while preserving performance.
- 1 Introduction: Mixtral 8x7B requires at least two A100-80G GPUs in bf16 because its eight experts contain around 96% of its parameters.The experts account for 45B of 47B total parameters.
- 1 Introduction: Unequal expert training outcomes motivate identifying and removing less important experts to improve MoE deployment efficiency.
- 1 Introduction: The paper proposes hardware-friendly post-training expert pruning and dynamic expert skipping as expert-level sparsification methods for MoE LLMs.Pruning permanently removes unimportant experts, while skipping dynamically omits experts during inference.
- 1 Introduction: Expert pruning searches layer-wise expert combinations with low token reconstruction loss, covering both task-agnostic and task-specific models.The approach concatenates selected expert combinations into the final pruned model.
- 1 Introduction: Dynamic skipping adjusts the number of active experts online and complements pruning to streamline MoE deployment.
- 1 Introduction: Pruning two experts halves GPU requirements and achieves a 1.2× inference speedup, with performance losses of around 2.9 points task-agnostically and 6.2 points task-specifically.Task-specific fine-tuning reduces the latter loss to 1.6 points.
2 Related Works
Prior work established MoE architectures and expert pruning in specialized settings, while LLM pruning mainly targeted weights and often required dedicated hardware. This paper addresses the resulting gap with post-training expert-level sparsification for decoder-only MoE LLMs.
- 2 Related Works: MoE models use multiple separate networks, each processing a subset of the data, and have been extended from recurrent networks to Transformer-based models.
- 2 Related Works: MoE LLMs sparsely activate experts for each token, but this paper focuses specifically on post-training expert pruning and skipping.
- 2 Related Works: Prior NLP expert-pruning studies targeted machine translation using activation patterns or gate statistics to remove unnecessary language specialists.
- 2 Related Works: Existing LLM post-training methods primarily sparsify linear-layer weights and require dedicated hardware, leaving efficient decoder-only MoE expert pruning unaddressed.
- 2 Related Works: The proposed approach targets expert-level sparsity to reduce memory usage and increase inference speed without relying on specialized deployment hardware.
3 Method
The paper introduces post-training expert pruning and dynamic skipping to reduce MoE deployment costs while preserving task performance. Pruning permanently removes low-contribution experts, whereas dynamic skipping selectively omits experts during inference based on routing weights.
- 3.1 Preliminary: Each token is routed to the top-k experts, whose SwiGLU outputs are combined using normalized routing weights.Mixtral 8x7B uses k = 2 selected experts.
- 3.2 Post-training Expert Pruning: Task-specific pruning differs from general-task pruning because C4 and MATH select substantially different expert combinations, matching in only four layers.The comparison uses r = 6 for Mixtral 8x7B.
- 3 Method: The deployment method targets expert-level sparsity because existing fine-grained weight pruning requires specialized hardware for efficient MoE deployment.The proposed expert pruning is post-training and does not require parameter updates.
- 3 Method: The method combines permanent expert pruning with dynamic skipping to reduce MoE memory use and inference cost.Pruning changes the model structure, while skipping adjusts active experts per token during inference.
- 3.2 Post-training Expert Pruning: Expert pruning uses calibration data, enumerates preserved subsets per layer, and retains the subset with minimum reconstruction loss.The remaining n − r experts are discarded, and pruned layers are concatenated into the final model.
- 3.3 Dynamic Skipping During Inference: Dynamic skipping omits the lower-weighted selected expert when we1 < βwe0, using a layer-specific threshold calibrated from routing-weight statistics.The scheme can operate alongside expert pruning and reports a 1.2× to 1.3× speedup with r = 6.
4 Experiment
Experiments evaluate expert pruning across general and domain-specific tasks, compare it with weight-pruning and other pruning baselines, and combine pruning with dynamic expert skipping. The methods reduce deployment memory and increase inference speed while retaining performance, especially when calibration data matches the target task and fine-tuning is allowed.
- 4.1 Expert Pruning for General Tasks: At roughly 50% parameter reduction, expert pruning outperforms Wanda’s structured 2:4 sparsity in memory usage and benchmark performance.The comparison uses r = 4 for expert pruning and the commonly used 2:4 pattern for Wanda.
- 4.1 Expert Pruning for General Tasks: 2.9-point and 7.1-point average performance drops result from discarding two and four experts, respectively, with the proposed method achieving the least degradation among pruning baselines.Random pruning and activation-frequency pruning perform worse, while activation frequency alone does not reliably identify unimportant experts.
- 4.1 Expert Pruning for General Tasks: 1.20× and 1.27× token-generation speedups are obtained after pruning two and four experts, while inference memory falls from two A100-80G GPUs to one.The speed improvement is attributed to reduced GPU intercommunication after lowering the number of GPUs needed to load the model.
- 4.2 Expert Pruning for Domain-Specific Tasks: Domain-specific calibration data substantially improves 5-shot GSM8K performance after expert pruning compared with pre-training calibration data.The expert-pruning strategy also outperforms Wanda with 2:4 structured sparsity on this task.
- 4.2 Expert Pruning for Domain-Specific Tasks: Fine-tuning greatly reduces pruning-induced performance gaps, and the pruned seven-expert Mixtral 8x7B Instruct model exceeds the eight-expert model on GSM8K.Models pruned using MATH calibration data outperform those using C4 after tuning.
- 4.3 Dynamic Expert Skipping Results: 1.33× token generation speedup and nearly 90% of Mixtral 8x7B Instruct performance are achieved with half the parameters by combining pruning and dynamic skipping.Dynamic skipping further improves inference speed with negligible performance drops; at equal speed, pruning fewer experts plus skipping yields higher LM-eval accuracy than pruning more experts alone.
5 Conclusion and Discussion
The paper focuses on expert-level sparsification and introduces post-training expert pruning and dynamic skipping to improve MoE deployment efficiency while maintaining model performance.
- Post-training expert pruning and dynamic expert skipping significantly reduce memory usage and enhance inference speed while maintaining high model performance.
Limitations
The method has limitations in scalability and evaluation scope: expert-combination enumeration becomes cumbersome with many experts, and experiments cover only Mixtral models.
- Expert pruning relies on enumerating expert combinations, which is feasible for four or eight experts but cumbersome for layers containing many experts.
- Experiments evaluate only Mixtral 8x7B and Mixtral 8x7B Instruct, leaving generalizability and scalability to other MoE models for future work.
Ethics Statement
The paper acknowledges ethical considerations surrounding broader LLM deployment and commits to responsible use and code transparency.
- The research aims to make advanced LLM deployment more efficient and broadly accessible through reduced model sizes and faster inference.
- The authors encourage responsible use of their findings while considering the societal impacts of deploying large language models.
A.1 Expert Selection Tendency in MoE Models
The appendix analyzes expert-selection frequencies in Mixtral 8x7B across general and task-specific data, showing that selection patterns vary by layer and downstream task.
- Expert-selection frequencies are visualized for layers 0, 15, and 31 using C4 and MATH samples to compare general and task-specific inference behavior.
- Each grid cell records the frequency with which a pair of top-2 experts is selected during inference.
- The model exhibits selection tendencies, particularly when inference is tailored to specific downstream tasks.
A.2 Theoretical Insight and Broader Application of Dynamic Skipping
The analysis models dynamic skipping as retaining only the highest-weight experts while bounding reconstruction loss. Under a fixed inter-expert distance assumption, the bound determines how many experts to reserve.
- Output Reconstruction: The layer output changes from combining all routed expert features to combining only the retained top-i experts.The formulation begins with routing weights and expert outputs for each token, then compares original and skipped outputs.
- Reconstruction Objective: The reconstruction loss is defined as the squared distance between the original and approximated output features.The analysis uses this loss to quantify the effect of skipping experts.
- Top-2 Simplification: Assuming pairwise expert-output distances equal a fixed D enables a simplified skipping criterion for the top-2 setting.The fixed-distance assumption is motivated by a concentrated ||f1 − f2||2 distribution observed on Mixtral 8x7B with C4 calibration data.
- Generalized Dynamic Skipping: Dynamic skipping reserves the top-i experts, where i is chosen to keep reconstruction loss within the upper bound H.The bound trades off accuracy and inference speed.
A.3 Experiments on the Sizes of Calibration Datasets
Experiments examine calibration-set size and dynamic skipping for domain-specific mathematical tasks. The results indicate robustness to calibration-set size, while task-specific skipping can incur larger performance drops but improve deployment trade-offs when combined with pruning.
- Calibration-Set Size: Using 64 or 128 calibration sequences produced the highest overall results, while very small sets, especially one sequence, could degrade performance.The calibration sequences contain 2048 tokens and are sampled from C4 for pruning Mixtral 8x7B to r = 6.
- Domain-Specific Tasks: For mathematical reasoning tasks, dynamic expert skipping caused larger performance drops, whereas combining it with pruning improved the speed–accuracy trade-off for Mixtral 8x7B Instruct.Pruning two experts plus dynamic skipping matched the speedup of pruning four experts while achieving higher evaluation accuracy.
- Domain-Specific Tasks: Dynamic skipping was evaluated using layer-wise β calibration from MATH training samples and 5-shot GSM8K accuracy alongside token-generation speed.The reported β values differ substantially across layers and between C4 and MATH calibration data.
- Broader Application: Expert pruning and dynamic skipping are plug-and-play with other lightweighting methods and compatible with weight quantization.The paper positions these techniques as orthogonal to weight and token pruning.