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

要約

大規模言語モデル (LLM) の専門家をまばらにアクティブにすることで推論コストを削減するために、専門家混合 (MoE) が採用されています。
この削減にもかかわらず、教育省には膨大な数の専門家が存在するため、サービス提供に依然として費用がかかります。
この論文では、MoE を枝刈りすることでこれに対処する方法を研究します。
枝刈り手法の中でも、構造化枝刈りはスパース化構造に制約を課すため、非構造化枝刈りは構造化枝刈りに比べて、特定の枝刈り率で最高のパフォーマンスを達成することが知られています。
非構造化枝刈りの解空間は構造化枝刈りの解空間を包含するため、これは直感的です。
しかし、私たちの直観に反する発見は、構造化剪定の一形式であるエキスパート剪定が実際には非構造化剪定に先行して、非構造化のみの剪定を上回るパフォーマンスを発揮できることを明らかにしました。
既存のエキスパートプルーニングでは、$n$ エキスパートに $O(\frac{k^n}{\sqrt{n}})$ フォワードパスが必要で、最近の MoE に合わせて拡張できないため、$O(1) を使用したスケーラブルな代替案を提案します。
$ 複雑ですが、より高価な方法よりも優れたパフォーマンスを発揮します。
重要なアイデアは、行動の類似性に基づいて専門家間の潜在的な構造を活用し、枝刈りをするかどうかの貪欲な決定によって共同枝刈り効果を厳密に捉えることができるようにすることです。
私たちの手法は非常に効果的です。128 人の専門家を擁する 480B 規模の MoE である Snowflake Arctic では、GSM8K などの生成タスクであっても、40% のスパース性でほぼパフォーマンスの損失を達成するために、私たちの手法は 1 つの H100 と 2 時間しか必要としません。
最先端の非構造化枝刈りは失敗します。
コードは公開されます。

要約(オリジナル)

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.

arxiv情報

提供元, 利用サービス

arxiv.jp, Google