先结构化再无结构剪枝,让大模型专家系统更高效
STUN: Structured-Then-Unstructured Pruning for Scalable MoE Pruning
- 利用专家行为相似性构建潜在结构,实现高效剪枝
- 4800亿参数模型剪掉40%专家,性能几乎无损,仅需2小时
- 适合需要高效部署MoE模型的研究与工程人员
混合专家(MoE)通过稀疏激活专家来降低大语言模型的推理成本。尽管如此,专家数量庞大仍导致服务开销高。本文研究如何对MoE进行剪枝。传统上,无结构剪枝因灵活性高而表现优于结构化剪枝。但本文反直觉发现:先进行专家级结构化剪枝,再做无结构剪枝,反而优于单纯使用无结构剪枝。现有专家剪枝方法需$O(\frac{k^n}{\sqrt{n}})$次前向传播,难以扩展到现代大规模MoE。为此,我们提出一种复杂度仅为$O(1)$的可扩展方法,基于专家行为相似性,以贪心策略捕获联合剪枝效果。该方法极为高效:在480B参数、含128个专家的Snowflake Arctic模型上,仅用一台H100、两小时即可实现40%稀疏度下性能近乎无损,包括GSM8K等生成任务,优于当前最优的无结构剪枝方法。代码将公开。
原文摘要 · Abstract (English)
Mixture-of-experts (MoEs) have been adopted for reducing inference costs by sparsely activating experts in Large language models (LLMs). Despite this reduction, the massive number of experts in MoEs still makes them expensive to serve. In this paper, we study how to address this, by pruning MoEs. Among pruning methodologies, unstructured pruning has been known to achieve the highest performance for a given pruning ratio, compared to structured pruning, since the latter imposes constraints on the sparsification structure. This is intuitive, as the solution space of unstructured pruning subsumes that of structured pruning. However, our counterintuitive finding reveals that expert pruning, a form of structured pruning, can actually precede unstructured pruning to outperform unstructured-only pruning. As existing expert pruning, requiring $O(\frac{k^n}{\sqrt{n}})$ forward passes for $n$ experts, cannot scale for recent MoEs, we propose a scalable alternative with $O(1)$ complexity, yet outperforming the more expensive methods. The key idea is leveraging a latent structure between experts, based on behavior similarity, such that the greedy decision of whether to prune closely captures the joint pruning effect. Ours is highly effective -- for Snowflake Arctic, a 480B-sized MoE with 128 experts, our method needs only one H100 and two hours to achieve nearly no loss in performance with 40% sparsity, even in generative tasks such as GSM8K, where state-of-the-art unstructured pruning fails to. The code will be made publicly available.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。