LLM-Pruner简介
LLM-Pruner是一个用于结构化剪枝大型语言模型(LLM)的开源框架,由新加坡国立大学的研究人员开发。它的主要目标是在保持模型性能的同时,显著减小模型规模,提高模型的推理效率。
LLM-Pruner的主要特点包括:
- 任务无关压缩:压缩后的LLM仍然保持原有的多任务解决能力
- 少量训练数据:只需使用5万个公开样本即可完成模型的后训练
- 高效压缩:剪枝仅需3分钟,后训练仅需3小时
- 自动结构化剪枝:可以用最少的人工干预剪枝新的LLM
目前LLM-Pruner支持的模型包括Llama 3/3.1、Llama 2、LLaMA、BLOOM、Vicuna、Baichuan、TinyLlama等。
快速上手
要开始使用LLM-Pruner,只需按以下步骤操作:
- 安装依赖:
pip install -r requirement.txt
- 运行示例脚本:
bash script/llama_prune.sh
这个脚本会将LLaMA-7B模型压缩约20%的参数。所有预训练模型和数据集都会自动下载。
使用步骤
使用LLM-Pruner压缩模型主要包括三个阶段:
- 发现阶段:发现LLM中复杂的相互依赖关系,找到最小可移除单元"group"
- 估计阶段:估计每个group对模型整体性能的贡献,决定剪枝哪些group
- 恢复阶段:快速后训练以恢复模型性能
1. 剪枝
以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
2. 后训练
使用Alpaca数据集进行后训练:
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
3. 评估
使用lm-evaluation-harness进行评估:
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
相关资源
- GitHub仓库: https://github.com/horseee/LLM-Pruner
- 论文: https://arxiv.org/abs/2305.11627
- 模型权重: https://huggingface.co/collections/meta-llama/llama-31-669fc079a0c406a149a5738f
LLM-Pruner为压缩大型语言模型提供了一个简单高效的开源解决方案。欢迎感兴趣的读者尝试使用,为项目贡献代码,或在自己的研究中引用相关工作。如有任何问题,也可以通过GitHub Issues与作者团队交流。
