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

TL;DR

This paper introduces a scalable structured pruning method for mixture-of-experts models that outperforms traditional unstructured pruning, significantly reducing computational costs while maintaining model performance.

Contribution

Proposes a novel scalable expert pruning approach leveraging latent structure, outperforming unstructured pruning in large language models.

Findings

Achieves 40% sparsity with minimal performance loss on Snowflake Arctic model.

Requires only one H100 GPU and two hours for pruning large models.

Outperforms state-of-the-art unstructured pruning in generative tasks like GSM8K.

Abstract

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.