LLM-Pruner入门指南:高效压缩大型语言模型的开源工具

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,只需按以下步骤操作:

1. 安装依赖:

pip install -r requirement.txt

2. 运行示例脚本:

bash script/llama_prune.sh

这个脚本会将LLaMA-7B模型压缩约20%的参数。所有预训练模型和数据集都会自动下载。

使用步骤

使用LLM-Pruner压缩模型主要包括三个阶段:

1. 发现阶段:发现LLM中复杂的相互依赖关系,找到最小可移除单元"group"
2. 估计阶段:估计每个group对模型整体性能的贡献,决定剪枝哪些group
3. 恢复阶段:快速后训练以恢复模型性能

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与作者团队交流。