LLM-Pruner

关于大语言模型 (LLM) 的结构化剪枝 (Structural Pruning)

:llama: :llama: :llama: :llama: :llama: 将您的 LLM 压缩至任意规模！:llama: :llama: :llama: :llama: :llama:

简介

LLM-Pruner: On the Structural Pruning of Large Language Models [arXiv]
Xinyin Ma, Gongfan Fang, Xinchao Wang
新加坡国立大学 (National University of Singapore)

为什么选择 LLM-Pruner

• 任务无关的压缩 (Task-agnostic compression)：压缩后的 LLM 应保留其作为多任务求解器的原始能力。
• 较少的训练语料 (Less training corpus)：在本工作中，我们仅使用 5 万条公开可用的样本 (alpaca) 对 LLM 进行后训练 (Post-training)。
• 高效压缩 (Efficient compression)：剪枝耗时 3 分钟，后训练耗时 3 小时。（您可以根据需要延长）
• 自动结构化剪枝 (Automatic structural pruning)：以最少的人工干预剪枝新的 LLM。（进行中）。

支持的 LLM：

• Llama-3.1
• Llama-3
• Llama-2
• LLaMA
• BLOOM
• Vicuna
• Baichuan
• TinyLlama

更新日志：

• 2024 年 7 月 27 日：:rocket: 支持 GQA (群查询注意力)! 现在 LLM-Pruner 可以在 Llama3 和 Llama 3.1 上运行。我们仍在测试新 LLM（Llama3, Llama3.1, Gemma）的剪枝结果，您可以在 此处 找到剪枝结果。
• 2023 年 8 月 30 日：LLM-Pruner 现在支持 BLOOM :cherry_blossom:
• 2023 年 8 月 14 日：代码 和 结果 现已支持使用大规模语料进行微调 (Finetuning)。微调后的 LLaMA-5.4B 模型平均准确率达到 62.36%，非常接近原始 LLaMA-7B (63.25%)。
• 2023 年 7 月 19 日：:fire: LLM-Pruner 现在支持 Llama-2-7b 和 Llama-2-13b（Huggingface 版本）
• 2023 年 7 月 18 日：:rocket: 支持 Baichuan，一款双语 LLM。
• 2023 年 5 月 20 日：:tada: 代码和预印本论文发布！

待办事项 (TODO List)：

• 针对新 LLM 剪枝的教程。
• 支持使用 .from_pretrained() 加载模型。

联系我们：

加入我们的微信群进行交流：

• 微信群 群组-2 (>200/500), 群组-1 (500/500, 已满).

目录

• 快速开始
• 逐步说明
• 零样本评估 (Zero-shot Evaluation)
• 更多示例
• 版本信息
• 局限性
• 致谢
• 引用

快速开始

安装

pip install -r requirement.txt

最小示例

bash script/llama_prune.sh

此脚本将压缩 LLaMA-7B 模型，剪去约 20% 的参数。所有预训练模型和数据集都将自动下载，因此您无需手动下载资源。首次运行此脚本时，需要一些时间来下载模型和数据集。

逐步说明

剪枝一个 LLM 需要三个步骤：

• 发现阶段 (Discovery Stage)：发现 LLM 中复杂的相互依赖关系，并找到可移除的最小单元，即组 (group)。
• 估计阶段 (Estimation Stage)：估计每个组对模型整体性能的贡献，并决定剪枝哪个组。
• 恢复阶段 (Recover Stage)：通过快速后训练恢复模型性能。

剪枝和后训练完成后，我们遵循 lm-evaluation-harness 进行评估。

1. 剪枝（发现阶段 + 估计阶段）

:llama: LLaMA/Llama-2 模型剪枝，约剪去 20% 参数：

python hf_prune.py --pruning_ratio 0.25 \
      --block_wise \
      --block_mlp_layer_start 4 --block_mlp_layer_end 30 \
      --block_attention_layer_start 4 --block_attention_layer_end 30 \
      --pruner_type taylor \
      --test_after_train \
      --device cpu  --eval_device cuda \
      --save_ckpt_log_name llama_prune 

参数说明：

• 基础模型：从 LLaMA 或 Llama-2 中选择基础模型，并将 pretrained_model_name_or_path 传递给 --base_model。模型名称用于 AutoModel.from_pretrained 加载预训练的大语言模型（LLM）。例如，如果你想使用 130 亿参数的 llama-2，则将 meta-llama/Llama-2-13b-hf 传递给 --base_model。
• 剪枝策略：使用相应的命令选项选择块级（block-wise）、通道级（channel-wise）或层级（layer-wise）剪枝：{--block_wise}，{--channel_wise}，{--layer_wise --layer NUMBER_OF_LAYERS}。对于块级剪枝，指定要剪枝的起始和结束层。通道级剪枝不需要额外参数。对于层级剪枝，使用 --layer NUMBER_OF_LAYERS 指定剪枝后保留的层数。
• 重要性准则：使用 --pruner_type 参数从 l1, l2, random 或 taylor 中选择。对于泰勒剪枝器（Taylor pruner），选择以下选项之一：vectorize, param_second, param_first, param_mix。默认使用 param_mix，它结合了近似二阶海森矩阵（Hessian）和一阶梯度（Gradient）。如果使用 l1, l2 或 random，则不需要额外参数。
• 剪枝比例：指定组的剪枝比例。它与参数剪枝率不同，因为组是最小的移除单位。
• 设备 和 评估设备：剪枝和评估可以在不同的设备上执行。基于泰勒的方法在剪枝期间需要反向计算，这可能需要大量的 GPU 显存。我们的实现使用 CPU 进行重要性估计（也支持 GPU，只需使用 --device cuda）。eval_device 用于测试剪枝后的模型。

:llama: Vicuna 剪枝

python hf_prune.py --pruning_ratio 0.25 \
      --block_wise \
      --block_mlp_layer_start 4 --block_mlp_layer_end 30 \
      --block_attention_layer_start 4 --block_attention_layer_end 30 \
      --pruner_type taylor \
      --test_after_train \
      --device cpu  --eval_device cuda \
      --save_ckpt_log_name llama_prune \
      --base_model PATH_TO_VICUNA_WEIGHTS

:llama: Baichuan 剪枝

:llama: Llama3/Llama3.1 剪枝

python llama3.py --pruning_ratio 0.25 \
                 --device cuda --eval_device cuda \
                 --base_model meta-llama/Meta-Llama-3-8B-Instruct \
                 --block_wise --block_mlp_layer_start 4 --block_mlp_layer_end 30 \
                 --block_attention_layer_start 4 --block_attention_layer_end 30 \
                 --save_ckpt_log_name llama3_prune \
                 --pruner_type taylor --taylor param_first \
                 --max_seq_len 2048 \
                 --test_after_train --test_before_train --save_model 

2. 后训练（恢复阶段）

• 使用 Alpaca 进行训练，包含 50,000 个样本。以下是单 GPU 训练的示例：

CUDA_VISIBLE_DEVICES=X python post_training.py --prune_model prune_log/PATH_TO_PRUNE_MODEL/pytorch_model.bin \
      --data_path yahma/alpaca-cleaned \
      --lora_r 8 \
      --num_epochs 2 \ 
      --learning_rate 1e-4 \ 
      --batch_size 64\
      --output_dir tune_log/PATH_TO_SAVE_TUNE_MODEL \ 
      --wandb_project llama_tune

请确保将 PATH_TO_PRUNE_MODEL 替换为步骤 1 中剪枝模型的路径，并将 PATH_TO_SAVE_TUNE_MODEL 替换为你希望保存微调模型的位置。

提示：在 float16 下训练 LLaMA-2 不推荐，已知会产生 nan；因此，模型应在 bfloat16 下训练。

• 使用 MBZUAI/LaMini-instruction 进行训练，包含 2.59M 个样本。以下是使用多 GPU 进行训练的示例：

deepspeed --include=localhost:1,2,3,4 post_training.py \
      --prune_model prune_log/PATH_TO_PRUNE_MODEL/pytorch_model.bin \
      --data_path MBZUAI/LaMini-instruction  \
      --lora_r 8 \
      --num_epochs 3  \
      --output_dir tune_log/PATH_TO_SAVE_TUNE_MODEL \
      --extra_val_dataset wikitext2,ptb \
      --wandb_project llmpruner_lamini_tune \
      --learning_rate 5e-5 \
      --cache_dataset

3. 生成

如何加载剪枝/预训练模型：

对于剪枝模型，只需使用以下命令加载你的模型。

  pruned_dict = torch.load(YOUR_CHECKPOINT_PATH, map_location='cpu')
  tokenizer, model = pruned_dict['tokenizer'], pruned_dict['model']

由于剪枝模型中模块之间的配置不同，某些层可能具有更大的宽度，而其他层则经过了更多的剪枝，因此无法像 Hugging Face 提供的那样使用 .from_pretrained() 加载模型。目前，我们采用 torch.save 来存储剪枝模型。

由于剪枝模型在每一层的配置不同，比如某些层可能更宽，但某些层被剪枝得更多，模型无法通过 Hugging Face 的 .from_pretrained() 加载。目前，我们简单地使用 torch.save 保存剪枝模型，并使用 torch.load 加载剪枝模型。

使用 Gradio 界面进行生成

我们提供了一个简单的脚本，用于使用预训练/剪枝模型/经过后训练的剪枝模型生成文本。

• LLaMA-7B 预训练模型

python generate.py --model_type pretrain

• 未经过后训练的剪枝模型

python generate.py --model_type pruneLLM --ckpt <YOUR_MODEL_PATH_FOR_PRUNE_MODEL>

• 经过后训练的剪枝模型

python generate.py --model_type tune_prune_LLM --ckpt <YOUR_CKPT_PATH_FOR_PRUNE_MODEL> --lora_ckpt <YOUR_CKPT_PATH_FOR_LORA_WEIGHT>

上述指令将在本地部署您的大语言模型（LLM）。

4. 评估

为了评估剪枝模型（pruned model）的性能，我们遵循 lm-evaluation-harness（语言模型评估工具集） 来评估模型：

• 步骤 1：如果您只需要评估剪枝模型，请跳过此步骤并跳转到步骤 2。 此步骤是为了整理文件以满足 lm-evaluation-harness 的输入要求。来自后训练阶段的微调检查点 将按以下格式保存：

- PATH_TO_SAVE_TUNE_MODEL
  | - checkpoint-200
      | - pytorch_model.bin
      | - optimizer.pt
      ...
  | - checkpoint-400
  | - checkpoint-600
  ...
  | - adapter_config.bin
  | - adapter-config.json

通过以下命令整理文件：

cd PATH_TO_SAVE_TUNE_MODEL
export epoch=YOUR_EVALUATE_EPOCH
cp adapter_config.json checkpoint-$epoch/
mv checkpoint-$epoch/pytorch_model.bin checkpoint-$epoch/adapter_model.bin

如果您想评估 checkpoint-200，则通过 export epoch=200 将 epoch（轮次）设置为 200。

• 步骤 2：

export PYTHONPATH='.'
python lm-evaluation-harness/main.py --model hf-causal-experimental \
       --model_args checkpoint=PATH_TO_PRUNE_MODEL,peft=PATH_TO_SAVE_TUNE_MODEL,config_pretrained=PATH_OR_NAME_TO_BASE_MODEL \
       --tasks openbookqa,arc_easy,winogrande,hellaswag,arc_challenge,piqa,boolq \
       --device cuda:0 --no_cache \
       --output_path PATH_TO_SAVE_EVALUATION_LOG 

在此处，将 PATH_TO_PRUNE_MODEL 和 PATH_TO_SAVE_TUNE_MODEL 替换为您保存剪枝模型和微调模型的路径，而 PATH_OR_NAME_TO_BASE_MODEL 用于加载基础模型的配置文件。

[更新]：如果您想使用微调检查点评估剪枝模型，我们上传了一个脚本来简化评估过程。直接使用以下命令：

CUDA_VISIBLE_DEVICES=X bash scripts/evaluate.sh PATH_OR_NAME_TO_BASE_MODEL PATH_TO_SAVE_TUNE_MODEL  PATH_TO_PRUNE_MODEL EPOCHS_YOU_WANT_TO_EVALUATE

在命令中替换您模型的必要信息。最后一个参数用于在一个命令中迭代不同的 epoch，以便评估多个检查点。例如：

CUDA_VISIBLE_DEVICES=1 bash scripts/evaluate.sh decapoda-research/llama-7b-hf tune_log/llama_7B_hessian prune_log/llama_prune_7B 200 1000 2000

5. 测试 MACs（乘加运算次数）、Params（参数量）和内存

• 预训练

python test_speedup.py --model_type pretrain

• 剪枝模型

python test_speedup.py --model_type pruneLLM --ckpt <YOUR_MODEL_PATH_FOR_PRUNE_MODEL>

零样本 (Zero-shot) 评估

LLaMA-7B 的简要定量结果：

Vicuna-7B 的结果：

ChatGLM-6B 的结果：

剪枝模型的统计数据：

LLM-Pruner 使用 2.59M 样本的结果：

更多结果可在论文中找到。

更多示例

版本信息

由于本项目使用的模型和仓库版本发生变化，我们列出了一些已知的版本问题以及复现我们方法所需的具体版本：

1. lm-eval-harness：我们使用了 lm-evaluation-harness 的 此 commit，代码也包含在此仓库中。详细信息请查看 Issue #25。
2. LLaMA1-7B：我们在实验中使用的是 decapoda-research/llama-7b-hf 的检查点，目前该链接不可用。请考虑使用复制的版本，例如 baffo32/decapoda-research-llama-7B-hf。

局限性

• 尽管我们仅使用了 5 万条数据并训练了三个小时，但更多的数据肯定会更好。我们正在对此进行测试。
• 当前的压缩模型仍然存在一些问题，例如生成重复的 token(词元) 或产生无意义的句子。我们相信压缩模型的质量还有很大的提升空间。
• 仍然有一些模型，我们无法在 concatenation(拼接) 和 view(视图) 操作后自动识别索引的映射关系。因此，我们需要执行额外的手动操作。

致谢

• Logo 由 Stable Diffusion 生成
• 大语言模型 (LLM) 评估： lm-evaluation-harness
• LLaMA: https://github.com/facebookresearch/llama
• Vicuna: https://github.com/lm-sys/FastChat
• Peft: https://github.com/huggingface/peft
• Alpaca-lora: https://github.com/tloen/alpaca-lora

引用

如果您觉得本项目有用，请引用

@inproceedings{ma2023llmpruner,
  title={LLM-Pruner: On the Structural Pruning of Large Language Models},
  author={Xinyin Ma and Gongfan Fang and Xinchao Wang},
  booktitle={Advances in Neural Information Processing Systems},
  year={2023},
}

@article{fang2023depgraph,
  title={DepGraph: Towards Any Structural Pruning},
  author={Fang, Gongfan and Ma, Xinyin and Song, Mingli and Mi, Michael Bi and Wang, Xinchao},
  journal={The IEEE/CVF Conference on Computer Vision and Pattern Recognition},
  year={2023}
}

LLM-Pruner 快速上手指南

LLM-Pruner 是一款专注于大语言模型（LLM）结构剪枝的开源工具，支持将主流 LLM 压缩至任意大小，同时保留其多任务处理能力。支持 Llama、Vicuna、Baichuan、BLOOM 等模型架构。

环境准备

• 系统要求: Linux / Windows / macOS
• Python 版本: 兼容标准 Python 环境
• 深度学习框架: PyTorch >= v1.7.1
• 硬件建议:
  • 剪枝与后训练阶段建议使用 GPU 以加速计算。
  • 使用 Taylor 剪枝策略时，需确保显存充足（也可在 CPU 上进行重要性估计，但速度较慢）。
• 网络提示: 首次运行会自动下载预训练模型和数据集，国内用户建议配置 HuggingFace 镜像以提升下载速度。

安装步骤

1. 克隆项目代码

   git clone https://github.com/horseee/LLM-Pruner.git
   cd LLM-Pruner
   

2. 安装依赖

   pip install -r requirement.txt
   

基本使用

1. 一键剪枝（最小示例）

项目提供了一个自动化脚本，可自动下载资源并压缩 LLaMA-7B 模型（约剪枝 20% 参数）。

bash script/llama_prune.sh

注：首次运行需要时间下载模型和数据集，请耐心等待。

2. 自定义剪枝流程

若需更精细控制，可使用命令行直接调用 hf_prune.py。以下示例演示了基于 LLaMA/Llama-2 的块级剪枝（Block-wise Pruning）：

python hf_prune.py --pruning_ratio 0.25 \
      --block_wise \
      --block_mlp_layer_start 4 --block_mlp_layer_end 30 \
      --block_attention_layer_start 4 --block_attention_layer_end 30 \
      --pruner_type taylor \
      --test_after_train \
      --device cpu  --eval_device cuda \
      --save_ckpt_log_name llama_prune

关键参数说明：

• --base_model: 指定基础模型路径（如 meta-llama/Llama-2-13b-hf）。
• --pruning_ratio: 剪枝比例（按组移除，非参数百分比）。
• --pruner_type: 重要性准则（可选 l1, l2, random, taylor）。
• --device: 剪枝计算设备（Taylor 方法建议 GPU，CPU 亦可）。

3. 后训练恢复（Post-Training）

剪枝后需使用少量数据（如 Alpaca 50k 样本）进行微调以恢复性能。

单卡训练示例：

CUDA_VISIBLE_DEVICES=X python post_training.py --prune_model prune_log/PATH_TO_PRUNE_MODEL/pytorch_model.bin \
      --data_path yahma/alpaca-cleaned \
      --lora_r 8 \
      --num_epochs 2 \ 
      --learning_rate 1e-4 \ 
      --batch_size 64 \
      --output_dir tune_log/PATH_TO_SAVE_TUNE_MODEL \ 
      --wandb_project llama_tune

4. 加载模型

由于剪枝后的模型结构发生变化，无法直接使用 HuggingFace 的 .from_pretrained()，请使用以下方式加载：

pruned_dict = torch.load(YOUR_CHECKPOINT_PATH, map_location='cpu')
tokenizer, model = pruned_dict['tokenizer'], pruned_dict['model']