Not All Experts are Equal: Efficient Expert Pruning and Skipping for Mixture-of-Experts Large Language Models (2024)

Abstract: 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 active 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. Code will be made available at https://github.com/Lucky-Lance/Expert_Sparsity.

## introduction large language models (llms) have shown remarkable abilities across various domains @xcite , as evidenced by the widespread use of chatgpt and gemini @xcite . recent notable advancement in this area is the introduction of the opensourced mixture-of-experts (moe) llm, mixtral 8x7b @xcite , which sparsely activates only a portion of its parameters during the training and inference process. this model surpasses the performance of dense transformer-based llms, such as llama-2 70b @xcite , with fewer active parameters (13b) during inference. moe llms achieve a reduction in on-the-fly (active) parameters by choosing only top-k experts for the inference of each token, thereby enhancing inference speed @xcite . however, the static parameters, particularly those required for constructing the moe architecture, still demand considerable memory and storage for deployment. for example, loading the mixtral 8x7b model in bf16 format requires at least two a100-80g gpus. notably, in this moe model, the eight experts constitute around 96% (45b out of 47b) of the total number of parameters. on the other hand, not all experts are equal in the moe model. recent studies, such as @xcite , have demonstrated this discrepancy in expert training outcomes. the differing levels of training among each expert highlight the importance and practicality of identifying and pruning less significant experts, thereby improving the deployment efficiency of moe models. unlike existing post-training weight pruning schemes for llms, which primarily target unstructured sparsity and n:m semi-structured sparsity @xcite , our approach focuses on expert-level sparsity for model sparsification. the aforementioned finegrained weight pruning techniques are effective in reducing the total number of parameters. however, they face challenges in plug-and-play deployment due to the necessity for specific hardware designs (e.g., fpga) @xcite , which demands an extensive co-design of hardware and systems. in this paper, we systematically explore expertlevel sparsity in moe llms and, for the first time to our best knowledge, introduce hardware-friendly post-training methods for either permanently removing unimportant experts (expert pruning) or dynamically skipping experts during inference (dynamic expert skipping). our proposed method significantly reduces memory usage for deploying moe llms and enhances their inference speed. initially, we investigate how to prune less important experts while maintaining satisfactory performance, utilizing an efficient post-training approach. we aim to minimize the token reconstruction loss in a layer-by-layer manner. given the limited number of experts in a single moe layer of the llm, we meticulously enumerate and choose combinations of experts that yield the lowest token reconstruction loss, subsequently, concatenating them to obtain the final pruned moe model. this strategy significantly lowers the memory demands for deploying moe llms. we examine expert-level pruning for both task-agnostic and task-specific (first in literature) models, tailoring our strategies to optimize performance across a wide range of applications. building on this foundation, we further dive into strategies for accelerating the inference speed of moe llms without compromising their robustness. specifically, based on the model's fixed expert count, we introduce an online method for dynamically skipping certain experts. this approach, which is complementary to our expert pruning strategy, allows for on-the-fly adjustment of the number of active experts during inference, thus enhancing the inference speed. by integrating the dynamic (expert) skipping approach with expert pruning, we achieve a more streamlined and efficient deployment for moe llms. experiments on mixtral 8x7b (instruct) models @xcite demonstrate that our methods significantly reduce memory usage and increase inference speed. take the example of posttraining pruning two experts. as shown in fig. 1 , we halve the number of gpus needed, allowing deployment on a single 80g gpu and achieving a 1.2× inference speedup. the pruning also results in mild performance loss, specifically, around 2.9 points for task-agnostic and 6.2 points (reducible to 1.6 with task-specific fine-tuning) for task-specific models. further combination of dynamic skipping with expert pruning can lead to the same inference speedup with dropping 4 experts while achieving much higher model performances. to the best of our knowledge, this study is the first to discuss expert-level sparsity and propose efficient schemes for expert pruning and skipping for moe llms. ## method to enhance the deployment efficiency of moe llms, we concentrate on expert-level model sparsity and innovatively propose post-training techniques designed to reduce memory usage and increase inference speed. in this section, we offer a comprehensive explanation of our proposed methods for expert pruning (sec. 3.2) and dynamic (expert) skipping (sec. 3.3), considering both memory consumption and inference speed. ## experiment in this section, a series of experiments are carried out to evaluate our proposed methods. we introduce experiments of expert pruning for general tasks in sec. 4.1, domain-specific tasks in sec. 4.2, and dynamic expert skipping results in sec. 4.3. ## conclusion and discussion in this paper, based on the structural characteristics of moe llms and the shortcomings of current weight pruning schemes, we focus on expert-level model sparsification and, for the first time, provide post-training expert pruning together with dynamic (expert) skipping methods to enhance the deployment efficiency of moe llms. our methods can significantly reduce memory usage and enhance inference speed while maintaining high model performance. looking ahead, we aim to further refine our pruning/skipping techniques and incorporate them with weight pruning or parameter quantization strategies, achieving more effective deploying approaches for moe llms.