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

Ray

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

安装依赖:

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