通过权重和激活进行LLM剪枝

Wanda（通过权重和激活进行剪枝）的官方PyTorch实现，如论文所述：

一种简单有效的大型语言模型剪枝方法
孙明杰*、刘庄*、Anna Bair、J. Zico Kolter （*表示同等贡献）
卡内基梅隆大学、Meta AI研究与博世人工智能中心
论文 - 项目页面

@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}
}

---

与仅基于权重大小的幅度剪枝不同，我们的剪枝方法Wanda通过权重大小与输入激活范数的乘积，按输出维度进行权重剪枝。

更新

• (9.22.2023) 增加对LLaMA-2的支持
• (9.22.2023) 增加重现论文中OBS权重更新消融实验的代码
• (10.6.2023) 增加对消融实验中权重更新分析的新支持 可自由尝试！
• (10.6.2023) 增加对零样本评估的支持
• (10.20.2023) 增加对OPT模型剪枝的代码
• (10.23.2023) 增加LoRA微调的代码

设置

安装说明请参见INSTALL.md

使用

scripts目录包含所有复制论文主结果（表2）的bash命令。

以下是以Wanda剪枝LLaMA-7B实现无结构50%稀疏度的示例命令：

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/ 

我们提供参数的简要说明：

• --model: Hugging Face模型库中LLaMA模型的标识符。
• --cache_dir: 加载或存储LLM权重的目录。默认为llm_weights。
• --prune_method: 实现了三种剪枝方法，即[magnitude, wanda, sparsegpt].
• --sparsity_ratio: 表示要剪枝的权重百分比。
• --sparsity_type: 指定稀疏度类型 [unstructured, 2:4, 4:8].
• --use_variant: 是否使用Wanda变体，默认为False。
• --save: 指定结果存储的目录。

对于结构化N:M稀疏度，将参数--sparsity_type设置为"2:4"或"4:8"。以下是一个示例命令：

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/ 

剪枝LLaMA-2

对于LLaMA-2模型，将--model替换为meta-llama/Llama-2-7b-hf（以7b为例）：

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结果：（截至9.22.2023，LLaMA-2-34b尚未发布）

稀疏度	ppl	llama2-7b	llama2-13b	llama2-70b
-	dense	5.12	4.57	3.12
无结构50%	magnitude	14.89	6.37	4.98
无结构50%	sparsegpt	6.51	5.63	3.98
无结构50%	wanda	6.42	5.56	3.98
4:8	magnitude	16.48	6.76	5.58
4:8	sparsegpt	8.12	6.60	4.59
4:8	wanda	7.97	6.55	4.47
2:4	magnitude	54.59	8.33	6.33
2:4	sparsegpt	10.17	8.32	5.40
2:4	wanda	11.02	8.27	5.16

OBS权重更新消融实验

为重现权重更新分析，我们提供了该消融实验的实现。所有命令可在此脚本中找到。

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 

此处ablate_{mag/wanda}_{seq/iter}表示我们使用幅度剪枝或wanda在每一层获得剪枝掩码，然后每128个输入通道应用权重更新过程（顺序式或迭代式）。详情请参见论文第5节。

零样本评估

为评估零样本任务，我们修改了EleutherAI LM Harness框架，使其能够评估剪枝后的LLM模型。我们提供修改后的仓库在此链接。请确保从源代码下载、解压并安装此自定义lm_eval包。

为保证可复现性，我们使用了主分支上的commit df3da98。所有任务均在任务版本0上评估，除了BoolQ任务，其任务版本为1。

总体而言，我们提供的功能是在此函数中新增了两个参数pretrained_model和tokenizer。我们可以通过代码库调用此simple_evaluate函数API来评估稀疏剪枝的LLM。若要评估除WikiText外的零样本任务，请传递--eval_zero_shot参数。

加速评估

通过计算每种方法在剪枝过程中所花费的累计时间（每层）来评估剪枝速度，不包含前向传播的时间。

对于结构稀疏性的推理加速，读者可参考这篇博客文章，其中结构稀疏性在 PyTorch >= 2.1 中得到支持。您可以通过此处切换 CUTLASS 或 CuSPARSELt 内核。

最后，关于图像分类器的剪枝，请参见目录 image_classifiers 获取详细信息。

致谢

本仓库基于 SparseGPT 仓库构建。

许可证

本项目采用 MIT 许可证。更多详情请参见 LICENSE 文件。

问题

欢迎通过问题或邮件与我们讨论论文/代码！
mingjies@cs.cmu.edu
liuzhuangthu@gmail.com

Wanda 快速上手指南

环境准备

• 系统要求：Python 3.8+，PyTorch 2.0+
• 前置依赖：

  pip install torch torchvision torchaudio
  pip install transformers==4.35.0
  pip install -i https://pypi.tuna.tsinghua.edu.cn/simple sentencepiece
  

安装步骤

1. 克隆仓库：

   git clone https://github.com/locuslab/wanda
   cd wanda
   

2. 安装依赖（参考 INSTALL.md）：

   pip install -r requirements.txt
   

基本使用

最简单示例：LLaMA-7B 无结构50%剪枝

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/

支持模型

• LLaMA-2 使用：

  --model meta-llama/Llama-2-7b-hf
  

• OPT 模型剪枝支持
• LoRA 微调支持（需额外配置）

结构化剪枝示例（2:4）

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/