arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2409.06211cs.LGcs.CL

STUN:面向可扩展 MoE 剪枝的结构化后非结构化剪枝

STUN: Structured-Then-Unstructured Pruning for Scalable MoE Pruning

Snowflake AI Research(Snowflake AI研究院);Seoul National University(首尔国立大学)

Jaeseong Lee, seung-won hwang, Aurick Qiao, Daniel F Campos, Zhewei Yao, Yuxiong He

AI总结:

本文提出 STUN 方法,通过先利用专家行为相似性进行 O(1) 复杂度的结构化专家剪枝、再进行非结构化剪枝,在 480B 规模 MoE 上以极低计算成本实现 40% 稀疏度下几乎无损的性能。

摘要:

专家混合(Mixture-of-experts, MoEs)已被用于通过在大语言模型(LLMs)中稀疏激活专家来降低推理成本。尽管有这种降低,MoEs 中数量庞大的专家仍然使其服务成本高昂。在本文中,我们研究如何通过剪枝 MoEs 来解决这一问题。在剪枝方法中,已知在给定剪枝比例下,非结构化剪枝相比结构化剪枝能够取得最高性能,因为后者对稀疏化结构施加了约束。这是直观的,因为非结构化剪枝的解空间包含结构化剪枝的解空间。然而,我们反直觉的发现表明,专家剪枝(一种结构化剪枝形式)实际上可以先于非结构化剪枝进行,从而优于仅使用非结构化剪枝的方法。由于现有的专家剪枝对 n 个专家需要 O(k^n/√n) 次前向传播,无法扩展到近期的 MoEs,我们提出了一种具有 O(1) 复杂度的可扩展替代方案,同时性能优于更昂贵的方法。其关键思想是利用专家之间基于行为相似性的潜在结构,使得是否剪枝的贪心决策能够紧密捕捉联合剪枝效应。我们的方法非常有效——对于 Snowflake Arctic,一个具有 128 个专家的 480B 规模 MoE,我们的方法仅需一块 H100 和两小时,即可在 40% 稀疏度下实现几乎无损的性能,即使在 GSM8K 等生成式任务中也是如此,而最先进的非结构化剪枝在这些任务中会失败。代码将公开发布。

英文摘要:

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.

机构、总结、中文摘要均由AI分析而来,可能存在错误,请以论文正文为准。

↑