WRP : Weight Recover Prune for Structured Sparsity (2024)

Abstract: As the scale of Large Language Models (LLMs) increases, it is necessary to compress the models to reduce the substantial demand on computational resources. Network pruning significantly reduces the model size by converting the weight matrix from dense to sparse data format. Current methodologies advocate for one-shot pruning to avoid the expense of retraining, ensuring the maintenance of model performance under conditions of 50%-60% unstructured pruning. Nevertheless, matrices characterized by this level of sparsity could not be treated as sparse matrices, because the indices would incur significant costs. To mitigate this problem, NVIDIA introduced the 2:4 structured sparsity. However, we observe a notable decline in model performance when adopting 2:4 structured sparsity due to group constraints. In this paper, we introduce the Weight Recover Prune (WRP) approach. By recovering a minimal set of critical weights, WRP aims to enhance model performance while maintaining the efficiency of the compression. Our evaluation of the WRP method on the LLAMA2 and OPT models shows that it outperforms other 2:4 pattern one-shot pruning methods. Meanwhile, WRP can guarantee that the size of the pruned model is about 60% of the dense model. Our code is available at: https://github.com/TanZhendong/WRP.

## introduction nowadays, many large language models (llms) have been developed, based on the transformer architecture @xcite @xcite . these models have demonstrated astonishing capabilities across a variety of tasks. however, the deployment of llms, characterized by their billions of parameters, demands substantial hardware resources. for instance, the llama2-70b model, with a size of 129gb, necessitates at least two a100-80gb gpus for inference. to mitigate the extensive resource requirements for model deployment, pruning and quantization algorithms emerge as two prevalent strategies. existing quantization algorithms could compress llms to 4 bits without retraining @xcite @xcite , which could significantly reduce the size of the models. network pruning is a model compression approach orthogonal to quantization algorithms. based on the granularity of the pruning algorithm, it is principally categorized into unstructured pruning and structured pruning. unstructured pruning offers higher flexibility and typically results in less precision loss. it converts dense matrices into sparse matrices by setting certain values in the weight matrix to zero, thereby achieving model compression and acceleration. considering the significant training overhead of llms, some pruning algorithms pursue one-shot pruning-that is, they avoid retraining to recover accuracy @xcite . such pruning methods have demonstrated minimal accuracy loss but generally could not achieve high levels of sparsity, with an optimal sparsity between 50%-60%. as for the compression of sparse matrices, taking the compressed sparse row (csr) data format as an example, as shown in figure 1 , a sparsity level of over 70% is usually required to realize compression benefits. in this context, matrices couldn't be treated as sparse for compression and computation if they do not meet this level of sparsity. nvidia has introduced a solution known as structured sparsity, or 2:4 sparsity @xcite . this pattern mandates that within every group of 4 values, at most 2 values can be retained. firstly, this leads to a 50% degree of sparsity, which is beneficial for the performance of models following one-shot pruning. moreover, by compressing the indices, it ensures efficient compression under 50% sparsity. additionally, this approach could achieve a 2x math throughput increase on the nvidia ampere gpu architecture. as a result, this pattern and compressed format are more hardware-friendly than unstructured sparsity at 50%. however, due to its grouping constraints, there is a notable decline in accuracy compared to unstructured pruning at 50%. for instance, when applying wanda pruning @xcite to the llama-7b model, the resulting perplexity on the wikitext dataset under unstructured 50% and 2:4 pattern conditions are 7.26 and 11.53, respectively. this indicates that there is still room for improvement in performance of 2:4 pattern. in this work, we observe a performance gap between the 2:4 pattern and unstructured 50% pruning. consequently, our primary objective is to enhance the model performance of the one-shot 2:4 pattern pruning, while ensuring the model compression is achieved. in section 3.1, we observed that some crucial weights might be incorrectly pruned in the 2:4 pattern pruning. based on this phenomenon, we propose the weight recover prune (wrp) approach, which aims to improve model accuracy by restoring a minor portion of critical weights, while for the majority of the matrix is still adopting a 2:4 pattern. to safeguard the compression efficacy, we partition the weight matrix into two separately stored entities: one is the 2:4 pattern and the other is high sparsity matrix, as shown in figure 1 . this approach allows us to achieve a balance between model size and performance, addressing the challenge inherent in the structured 2:4 sparsity pattern. the main contributions of this work are: structured 2:4 sparsity. nvidia proposed the 2:4 sparsity pattern, which typically requires retraining to recover model accuracy @xcite . channel permutations could be utilized to enhance model accuracy and alleviate the limitations of group constraints @xcite . for the 2:4 sparse matrix format, cusparselt provides the corresponding operators (nvidia, 2020). py-torch has also implemented support for 2:4 sparsity @xcite . ## conclusions in this work, we propose the weight recover prune (wrp) methodology for achieving structured sparsity in llms. observing the notable performance gap between the 2:4 pruning pattern and the unstructured 50% pruning, our wrp technique enhances the model performance associated with the 2:4 pattern by recovering a minimal set of crucial weights, thereby ensuring the efficiency of model