StructPrune: Structured Global Pruning asymptotics with $\mathcal{O}(\sqrt{N})$ GPU Memory

TL;DR

This work tackles the memory bottleneck of global structured pruning for billion-parameter LLMs by introducing STRUPRUNE, an ADMM-based divide-and-conquer framework that achieves global structured pruning with memory scaling of $O(\sqrt{N})$. It develops a closed-form layer-wise sparsity solution and an energy-based softmax allocation to adaptively allocate pruning across layers, enabling efficient pruning while preserving cross-layer dependencies. The approach is validated on OPT models, showing perplexity comparable to global structured pruning while dramatically reducing GPU memory requirements, thus enabling practical deployment at scale. By integrating structured masks with memory-efficient optimization and a lightweight importance score (Wanda), STRUPRUNE offers hardware-friendly pruning that maintains model fidelity under substantial sparsity.

Abstract

Pruning is critical for scaling large language models (LLMs). Global pruning achieves strong performance but requires $\mathcal{O}(N)$ memory, which is infeasible for billion-parameter models. Local pruning reduces GPU memory usage to that of a single layer by pruning layers independently, but it neglects inter-layer dependencies and often leads to suboptimal performance in high-sparsity regimes. Unlike unstructured pruning, structured pruning produces regular sparsity patterns that align well with GPU kernels and library optimizations, making it more hardware-efficient. However, structured pruning typically relies on global pruning, since structured patterns are more prone to severe performance degradation under local optimization. To jointly achieve structured pruning and the memory efficiency of local pruning, we propose a divide-and-conquer strategy that decomposes the global pruning problem into coordinated subproblems across different modules, each of which fits within limited GPU memory. Building on this idea, we design \textbf{STRUPRUNE}, an ADMM-based framework that integrates structured sparsity into the pruning process, combining the memory efficiency of local pruning with the hardware compatibility of structured methods. We derive a closed-form analytical solution for structured pruning masks that provides an explicit rule for layer-wise sparsity allocation, and further develop an energy-based asymptotic framework yielding a softmax-form allocation scheme that simplifies optimization while adapting to heterogeneous layer importance. Experiments demonstrate that STRUPRUNE matches the perplexity of global structured pruning while reducing memory cost from $\mathcal{O}(N)$ to $\mathcal{O}(\sqrt{N})$, enabling practical deployment at the billion-parameter scale.

Figures (1)

• Figure 1: Depth (number of layers) and width (parameters per layer) grow roughly proportionally as model size increases, for both OPT and LLaMA-2.

Theorems & Definitions (3)

• Lemma 3.1: Layer-wise Sparsity and Mask Construction
• Lemma 3.1: Layer-wise Sparsity and Mask Construction
• Lemma 3.2: Asymptotic Layer-wise Sparsity Allocation