Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Condense, Don't Just Prune: Enhancing Efficiency and Performance in MoE Layer Pruning

About

Mixture-of-Experts (MoE) has garnered significant attention for its ability to scale up neural networks while utilizing the same or even fewer active parameters. However, MoE does not alleviate the massive memory requirements of networks, which limits their practicality in real-world applications, especially in the era of large language models (LLMs). While recent work explores the possibility of removing entire layers of MoE to reduce memory, the performance degradation is still notable. In this paper, we propose ConDense-MoE (CD-MoE), which, instead of dropping the entire MoE layer, condenses the large, sparse MoE layer into a smaller, denser layer with only a few experts activated for all tokens, while maintaining hardware friendliness. Our approach is specifically designed for fine-grained MoE with shared experts, where Feed-Forward Networks are split into many small experts, with certain experts isolated to serve as shared experts that are always activated, such as DeepSeekMoE and QwenMoE. We demonstrate the effectiveness of our method. Specifically, for the DeepSeekMoE-16B model, our approach maintains 90% of the average accuracy while reducing memory usage by 27.5% and increasing inference speed by 1.26 times. Moreover, we show that by applying lightweight expert fine-tuning -- only to the condensed layers -- and using 5 hours on a single 80G A100 GPU, we can successfully recover 98% of the original performance. Our code is available at: https://github.com/duterscmy/CD-MoE/tree/main.

Mingyu Cao, Gen Li, Jie Ji, Jiaqi Zhang, Ajay Jaiswal, Li Shen, Xiaolong Ma, Shiwei Liu, Lu Yin• 2024

Related benchmarks

TaskDatasetResultRank
Physical Commonsense ReasoningPIQA
Accuracy76.1
696
Natural Language InferenceRTE
Accuracy67.7
590
Commonsense ReasoningWinoGrande
Accuracy73.9
453
Multi-task Language UnderstandingMMLU
MMLU Accuracy39.5
442
Sentence CompletionHellaSwag
Accuracy66.7
364
Boolean Question AnsweringBoolQ
Accuracy79.4
350
Multiple-choice Question AnsweringARC Challenge (test)
Accuracy42.4
57
OpenBook Question AnsweringOBQA
Accuracy40.3
32
Zero-shot Evaluation AggregateARC-C, BoolQ, HellaSwag, MMLU, OBQA, PIQA, RTE, Winogrande Aggregate
Average Accuracy60.4
13
Zero-shot EvaluationDownstream Task Suite (ARC-C, BoolQ, HellaSwag, MMLU, OBQA, PIQA, RTE, WinoGrande) zero-shot Qwen1.5-MoE-A2.7B
ARC-C Accuracy41.1
6
Showing 10 of 11 rows

Other info

Follow for update