Home

Awesome

Pruning LLMs by Weights and Activations

Official PyTorch implementation of Wanda (Pruning by Weights and activations), as presented in our paper:

A Simple and Effective Pruning Approach for Large Language Models </br> Mingjie Sun*, Zhuang Liu*, Anna Bair, J. Zico Kolter (* indicates equal contribution) <br> Carnegie Mellon University, Meta AI Research and Bosch Center for AI <br> Paper - Project page

@article{sun2023wanda,
  title={A Simple and Effective Pruning Approach for Large Language Models}, 
  author={Sun, Mingjie and Liu, Zhuang and Bair, Anna and Kolter, J. Zico},
  year={2023},
  journal={arXiv preprint arXiv:2306.11695}
}

<p align="center"> <img src="https://user-images.githubusercontent.com/20168304/273351964-53c3807e-3453-49c5-b855-b620b1026466.png" width=100% height=100% class="center"> </p>

Compared to magnitude pruning which removes weights solely based on their magnitudes, our pruning approach Wanda removes weights on a per-output basis, by the product of weight magnitudes and input activation norms.

Update

Setup

Installation instructions can be found in INSTALL.md.

Usage

The scripts directory contains all the bash commands to replicate the main results (Table 2) in our paper.

Below is an example command for pruning LLaMA-7B with Wanda, to achieve unstructured 50% sparsity.

python main.py \
    --model decapoda-research/llama-7b-hf \
    --prune_method wanda \
    --sparsity_ratio 0.5 \
    --sparsity_type unstructured \
    --save out/llama_7b/unstructured/wanda/ 

We provide a quick overview of the arguments:

For structured N:M sparsity, set the argument --sparsity_type to "2:4" or "4:8". An illustrative command is provided below:

python main.py \
    --model decapoda-research/llama-7b-hf \
    --prune_method wanda \
    --sparsity_ratio 0.5 \
    --sparsity_type 2:4 \
    --save out/llama_7b/2-4/wanda/ 

Pruning LLaMA-2

For LLaMA-2 models, replace --model with meta-llama/Llama-2-7b-hf (take 7b as an example):

python main.py \
    --model meta-llama/Llama-2-7b-hf \
    --prune_method wanda \
    --sparsity_ratio 0.5 \
    --sparsity_type unstructured \
    --save out/llama2_7b/unstructured/wanda/

LLaMA-2 results: (LLaMA-2-34b is not released as of 9.22.2023)

sparsitypplllama2-7bllama2-13bllama2-70b
-dense5.124.573.12
unstructured 50%magnitude14.896.374.98
unstructured 50%sparsegpt6.515.633.98
unstructured 50%wanda6.425.563.98
4:8magnitude16.486.765.58
4:8sparsegpt8.126.604.59
4:8wanda7.976.554.47
2:4magnitude54.598.336.33
2:4sparsegpt10.178.325.40
2:4wanda11.028.275.16

Ablation on OBS weight update

To reproduce the analysis on weight update, we provide our implementation for this ablation. All commands can be found in this script.

for method in ablate_mag_seq ablate_wanda_seq ablate_mag_iter ablate_wanda_iter 
do 
CUDA_VISIBLE_DEVICES=0 python main.py \
  --model decapoda-research/llama-7b-hf \
  --sparsity_ratio 0.5 \
  --sparsity_type unstructured \
  --prune_method ${method} \
  --save out/llama_7b_ablation/unstructured/
done 

Here ablate_{mag/wanda}_{seq/iter} means that we use magnitude pruning or wanda to obtain the pruned mask at each layer, then apply weight update procedure with either a sequential style or an iterative style every 128 input channels. For details, please see Section 5 of our paper.

Zero-Shot Evaluation

For evaluating zero-shot tasks, we modify the EleutherAI LM Harness framework so that it could evaluate pruned LLM models. We provide the modified repo in this link. Make sure to download, extract and install this custom lm_eval package from the source code.

For reproducibility, we used commit df3da98 on the main branch. All tasks were evaluated on task version of 0 except for BoolQ, where the task version is 1.

On a high level, the functionality we provide is adding two arguments pretrained_model and tokenizer in this function. We can then call this simple_evaluate function API from our codebase to evaluate sparse pruned LLMs. To evaluate zero-shot tasks in addition to the WikiText perplexity, pass the --eval_zero_shot argument.

Speedup Evaluation

The pruning speed for each method is evaluated by the cumulated time spent on pruning (for each layer), without the forward passes.

For inference speedup with structured sparsity, we refer the reader to this blog post, where structured sparsity is supported by PyTorch >= 2.1. You can switch between the CUTLASS or CuSPARSELt kernel here.

Last, for pruning image classifiers, see directory image_classifiers for details.

Acknowledgement

This repository is build upon the SparseGPT repository.

License

This project is released under the MIT license. Please see the LICENSE file for more information.

Questions

Feel free to discuss papers/code with us through issues/emails!

mingjies at cs.cmu.edu
liuzhuangthu at gmail.com