LLM-Pruner
关于大型语言模型的结构剪枝
:llama: :llama: :llama: :llama: :llama: 压缩你的LLM到任意大小!:llama: :llama: :llama: :llama: :llama:
介绍
LLM-Pruner:关于大型语言模型的结构剪枝 [arXiv]
Xinyin Ma, Gongfan Fang, Xinchao Wang
新加坡国立大学
为什么选择 LLM-Pruner
- 任务无关的压缩:压缩后的LLM应保留其作为多任务解决者的原本能力。
- 更少的训练语料:在此工作中,我们仅使用50,000个公开可用的样本(alpaca)来对LLM进行后训练。
- 高效压缩:剪枝3分钟,后训练3小时。(你可以延长时间)
- 自动结构剪枝:以最少的人力剪枝新的LLM(正在进行中)。
受支持的LLMs:
更新:
- 2024年7月27日: :rocket: 支持GQA!现在LLM-Pruner可以用于Llama3和Llama 3.1。我们仍在测试新LLMs(Llama3,Llama3.1,Gemma)的剪枝结果,你可以在这里找到剪枝结果。
- 2023年8月30日: LLM-Pruner现在支持 BLOOM :cherry_blossom:
- 2023年8月14日: 代码 和 结果 现已可用于大规模语料库的微调。微调后的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 列表:
- 用于剪枝新LLMs的教程。
- 支持
.from_pretrained()
用于加载模型。
联系我们:
加入我们的Discord或微信交流群进行交流:
- Discord: 链接
- 微信交流群 Group-2, Group-1 (500/500, 已满)。
目录
快速开始
安装
pip install -r requirement.txt
最简示例
bash script/llama_prune.sh
该脚本将压缩LLaMA-7B模型并剪枝约20%的参数。所有预训练模型和数据集将被自动下载,因此你不需要手动下载资源。首次运行该脚本时,可能需要一些时间下载模型和数据集。
分步指南
修剪一个大型语言模型(LLM)需要三步:
- 发现阶段:发现LLM中复杂的相互依赖关系并找到最小的可移除单元,组。
- 估算阶段:估算每个组对模型整体性能的贡献,并决定修剪哪些组。
- 恢复阶段:通过快速后训练恢复模型性能。
在修剪和后训练之后,我们使用 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 --layer 数层}。对于区块级修剪,指定要修剪的起始和结束层。通道级修剪不需要额外参数。层修剪使用--layer number_of_layers
指定修剪后的层数。重要性标准
:使用--pruner_type
参数选择l1、l2、随机或taylor。对于taylor修剪器,选择以下选项之一:vectorize、param_second、param_first、param_mix。默认使用param_mix,它结合了近似二阶hessian和一阶梯度。如果使用l1、l2或随机方式,则不需要额外的参数。修剪比率
:指定组的修剪比率。它不同于参数的修剪率,因为组作为最小单位被移除。设备
和评估设备
:修剪和评估可以在不同的设备上进行。基于taylor的方法在修剪期间需要反向计算,可能需要大量显存。我们的实现使用CPU进行重要性估计(也支持GPU,只需使用--device cuda
)。评估设备用于测试修剪后的模型。
:llama: Vicuna修剪
详情:
如果你想尝试Vicuna,请将参数--base_model
指定为vicuna权重的路径。请按照https://github.com/lm-sys/FastChat获取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修剪
详情:
有关详细信息,请参阅 Example/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
确保用步骤1中的修剪模型路径替换PATH_TO_PRUNE_MODEL
,并用你希望保存调整后模型的位置替换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
存储修剪模型并使用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. 评估
为了评估修剪模型的性能,我们遵循lm-evaluation-harness进行模型评估:
- 第一步:如果你只需评估修剪模型,请跳过这一步并进入第二步。
这一步是安排文件以满足
lm-evaluation-harness
的输入要求。来自后训练步骤的tuned checkpoint将保存为以下格式:
- 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。
- 第二步:
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
将命令中的模型信息替换为您的实际信息。最后一个参数用于在一个命令中遍历不同的时代,如果您想评估多个检查点。例如:
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, 参数和内存
- 预训练模型
python test_speedup.py --model_type pretrain
- 剪枝模型
python test_speedup.py --model_type pruneLLM --ckpt <YOUR_MODEL_PATH_FOR_PRUNE_MODEL>
零样本评估
LLaMA-7B 的简要定量结果:
Vicuna-7B 的结果:
ChatGLM-6B 的结果:
剪枝模型的统计数据:
LLM-Pruner 使用 2.59M 样本的结果:
剪枝比 | 参数数量 | 内存 | 延迟 | 加速 | BoolQ | PIQA | HellaSwag | WinoGrande | ARC-e | ARC-c | OBQA | 平均结果 |
---|---|---|---|---|---|---|---|---|---|---|---|---|
LLaMA-7B | 6.74B | 12884.5MiB | 69.32s | 1x | 73.18 | 78.35 | 72.99 | 67.01 | 67.45 | 41.38 | 42.40 | 63.25 |
LLaMA-5.4B with Alpaca(50k) | 5.47B | 10488.4MiB | 58.55s | 1.18x | 64.62 | 77.20 | 68.80 | 63.14 | 64.31 | 36.77 | 39.80 | 59.23 |
LLaMA-5.4B with LaMini(2.59M) | 5.47B | 10488.4MiB | 58.55s | 1.18x | 76.57 | 77.37 | 66.60 | 65.82 | 70.62 | 40.70 | 38.80 | 62.36 |
更多结果可以在论文中找到。
更多示例
原始模型 | 输出 | 剪枝模型 | 输出 |
---|---|---|---|
LLaMa-7B | 宇宙是所有空间、时间、物质和能量的总和。 包括所有物质和能量,无论是否可见,实际上整个宇宙都是不可见的。关于它的起源和最终命运有许多理论... | LLaMa-5.4B | 宇宙是所有空间、时间、物质和能量的总和。 构成物理学的定律——电子绕原子旋转、磁铁吸引其他物体、行星绕太阳公转——一直很一致... |
Vicuna-7B | 构建 iOS 应用程序的 10 个步骤: 1. 研究和开发想法。2. 决定您要使用的平台(iOS、Android 或两者)。3. 制定计划和预算。4. 雇用开发团队。5. 设计用户界面(UI)。6. 开发用户界面(UI)。7. 测试和迭代应用程序。8. 将应用程序提交到应用商店。9. 优化应用程序的性能和用户体验(UX)。10. 推广和宣传应用程序。 | Vicuna-5.4B | 构建 iOS 应用程序的 10 个步骤: 1. 计划您的应用程序:决定它应该做什么,它的外观以及构建它的成本。2. 雇用开发人员:优秀的 iOS 开发人员将帮助您设计用户界面、创建应用程序并确保其无错误。3. 创建用户界面:这是魔法发生的地方。您的应用程序的用户界面将在这个阶段设计和模拟... |
ChatGLM-6B | 写一首关于爱与和平的诗。 在生活的喧嚣和混乱之中,我发现你的爱是一种安慰的和平。这是一种无法传达的感觉,一种无法远离的联结。你的微笑,你的眼睛,你的拥抱,给我一种只有你能给予的归属感。这是一种永恒的爱,一种无法被打败的爱... | ChatGLM-5.6B | 写一首关于爱与和平的诗。 爱与和平,两种如此美好的东西\ 我心中纯洁而无尽的渴望\ 而我必须追求它们\ 我知道这是一段长长的旅程..\ 爱,我知道,是一种存在的感觉\ 完美的伴侣,无论从哪个角度\ 和平,我如此需要,总有一天\ 我的心会走很长、很长的路.. |
版本信息
由于本项目所用模型和库版本的变化,我们列出了一些已知的版本问题和重现我们方法所需的具体版本:
- lm-eval-harness:我们使用了 lm-eval-harness 的这个提交,代码也包含在本仓库中。详情请查看 Issue #25。
- LLaMA1-7B:我们在实验中使用了 decapoda-research/llama-7b-hf 的检查点,目前该检查点不可用。请考虑使用备份版本,例如 baffo32/decapoda-research-llama-7B-hf。
限制
- 虽然我们仅使用了 50K 数据并训练了三个小时,但更多的数据肯定会更好。我们正在对此进行测试。
- 当前的压缩模型仍然存在一些问题,例如生成重复的标记或产生无意义的句子。我们认为在压缩模型的质量方面仍有很大的改进空间。
- 仍然有一些模型我们无法自动识别索引在连接和视图操作后的映射。因此,我们需要执行额外的手动操作。
致谢
- 标志由 Stable Diffusion 生成
- 大型语言模型的评估: 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}
}