Pruning Large Language Models with Semi-Structural Adaptive Sparse Training

TL;DR

The paper tackles the challenge of deploying large language models at scale by pruning with minimal performance loss. It introduces Adaptive Sparse Trainer (AST), which jointly learns masks and weights under an annealed SR-STE regime, complemented by knowledge distillation and Sparse Low-Rank Boosting (SLoRB) to maintain expressivity. Empirical results on LLaMA2-7B show that a $2:4$ N:M sparsity yields perplexity and zero-shot accuracy gaps as low as $0.6\%$ and $1.16\%$, using under $0.4\%$ of pretraining tokens, with additional gains from AWQ quantization. Overall, AST demonstrates the practicality of semi-structured sparse LLMs, offering a scalable path to highly compressed models when integrated with quantization techniques.

Abstract

The remarkable success of Large Language Models (LLMs) relies heavily on their substantial scale, which poses significant challenges during model deployment in terms of latency and memory consumption. Recently, numerous studies have attempted to compress LLMs using one-shot pruning methods. However, these methods often suffer from considerable performance degradation on complex language understanding tasks, raising concerns about the feasibility of pruning in LLMs. To address this issue, we propose Adaptive Sparse Trainer (AST), a novel and efficient retraining framework tailored for semi-structured sparse models. AST enables models to learn optimal masks during the weight update process without incurring additional computational overhead. Furthermore, we demonstrate that incorporating knowledge distillation significantly improves retraining efficiency and enhances model performance under fixed computational constraints. Additionally, a supplementary set of well-initialized parameters is integrated to further augment the model's efficacy. AST achieves state-of-the-art performance with minimal training cost. When applied to the LLaMA2-7B model, AST reduces the perplexity and zero-shot accuracy gap between dense and 2:4 semi-structured sparse models to 0.6 and 1.16%, respectively, utilizing less than 0.4% of the pretraining tokens and GPU hours. Our work demonstrates the feasibility of deploying semi-structured sparse LLMs and offers a promising alternative for achieving highly compressed models when combined with existing quantization techniques.

Figures (4)

• Figure 1: (Left) In the naive training baseline, the mask remains constant during training, which can result in suboptimal performance. (Right) Adaptive sparse training enables both mask and weight learning through a scheduled decay term. AST also utilizes distillation and SLoRB parameters to speed up convergence and improve performance.
• Figure 2: (Upper) The Wikitext perplexity curve for retraining GPT2 with and without knowledge distillation. (Lower) The training loss curve for pretraining and retraining.
• Figure 3: Visualization of SLoRB Initialization: Consider a weight matrix of size 3 by 8, with $k=4$. Weight $S_{ij}$ is broadcasted within group $G_{ij}$ after multiplication with the projection matrix.
• Figure 4: (Upper) Flip rate for static and Annealing SR-STE during the training of the GPT2 model. (Lower) Initial flip rate for static and Annealing SR-STE during the training of the GPT2 model.