LLM-Pruner

On the Structural Pruning of Large Language Models

🦙 🦙 🦙 🦙 🦙 Compress your LLMs to any size! 🦙 🦙 🦙 🦙 🦙

Introduction

LLM-Pruner: On the Structural Pruning of Large Language Models [arXiv]
Xinyin Ma, Gongfan Fang, Xinchao Wang
National University of Singapore

Why LLM-Pruner

Task-agnostic compression: The compressed LLM should retains its original ability as a multi-task solver.
Less training corpus: In this work, we use only 50k publicly available samples (alpaca) to post-train the LLM.
Efficient compression: 3 minutes for pruning and 3 hours for post-training. (You can make it longer)
Automatic structural pruning: Pruning new LLMs with minimal human effort (In progress).

Supported LLMs:

TODO List:

Code for the Official LLaMA-7B
Code for ChatGLM
A tutorial for pruning new LLMs.
Scaling up the finetuning with large-scale corpus.
Support .from_pretrained() for loading the mode.

Quick Start

Installation

pip install -r requirement.txt

Minimal Example

bash script/llama_prune.sh

This script would compress the LLaMA-7B model with ～20% parameters pruned. All the pre-trained models and the dataset would be automatically downloaded, so you do not need to manually download the resource. When running this script for the first time, it will require some time to download the model and the dataset.

Step-by-step Instructions

It takes three steps to prune an LLM:

Discovery Stage: Discover the complicated inter-dependency in LLMs and find the minimally-removable unit, group.
Estimation Stage: Estimate the contribution of each group to the overall performance of the model and decide which group to prune.
Recover Stage: Fast post-training to recover model performance.

After pruning and post-training, we follow lm-evaluation-harness for evaluation.

1. Pruning (Discovery Stage + Estimation Stage)

🦙 LLaMA-7B pruning with ~20% parameters pruned:

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

Arguments:

Pruning Strategy: Choose between block-wise, channel-wise, or layer-wise pruning using the respective command options: {--block_wise}, {--channel_wise}, {--layer_wise --layer NUMBER_OF_LAYERS}. For block-wise pruning, specify the start and end layers to be pruned. Channel-wise pruning does not require extra arguments. For layer pruning, use --layer NUMBER_OF_LAYERS to specify the desired number of layers to be kept after pruning.
Importance Criterion: Select from l1, l2, random, or taylor using the --pruner_type argument. For the taylor pruner, choose one of the following options: vectorize, param_second, param_first, param_mix. By default, param_mix is used, which combines approximated second-order hessian and first-order gradient. If using l1, l2, or random, no extra arguments are required.
Pruning Ratio: Specifies the pruning ratio of groups. It differs from the pruning rate of parameters, as groups are removed as the minimal units.
Device and Eval_device: Pruning and evaluation can be performed on different devices. Taylor-based methods require backward computation during pruning, which may require significant GPU RAM. Our implementation uses the CPU for importance estimation (also support GPU, simply use --device cuda). eval_device is used to test the pruned model.

🦙 Vicuna Pruning

If you want to try Vicuna, please specify the argument --base_model to the path to vicuna weight. Please follow https://github.com/lm-sys/FastChat to get Vicuna weights.

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

🦙 ChatGLM Pruning

Comming Soon...

2. Post-Training (Recover Stage)

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

Make sure to replace PATH_TO_PRUNE_MODEL with the path to the pruned model in step 1, and replace PATH_TO_SAVE_TUNE_MODEL with the desired location where you want to save the tuned model.

3. Generation

How to load pruned/pre-trained models:

For the pruned model, simply use the following command to load your model.

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

Due to the different configurations between modules in the pruned model, where certain layers may have larger width while others have undergone more pruning, it becomes impractical to load the model using the .from_pretrained() as provided by Hugging Face. Currently, we employ the torch.save to store the pruned model.

Since the pruned model has different configuration in each layer, like some layers might be wider but some layers have been pruned more, the model cannot be loaded with the .from_pretrained() in Hugging Face. Currently, we simply use the torch.save to save the pruned model and torch.load to load the pruned model.

Generation with Gradio Interface

We provide a simple script to geneate texts using pre-trained / pruned models / pruned models with post-training.

LLaMA-7B Pre-trained

python generate.py --model_type pretrain

Pruned Model without Post-Training

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

Pruned Model with Post-Training

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

The above instructions will deploy your LLMs locally.

4. Testing MACs, Params and Memory

Pre-trained

python test_speedup.py --model_type pretrain

Pruned Model

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

Zero-shot Evaluation

A brief quantitative results for LLaMA-7B:

The results for Vicuna-7B:

The results for ChatGLM-6B:

Statistics for pruned models:

More results can be found in the paper.

More Examples

Original Model	Output	Pruned Model	Output
LLaMa-7B	The universe is the entirety of space, time, matter, and energy that exists. It includes all matter and energy, whether visible or not, and in fact the whole universe is invisible. There are many theories regarding its origin and ultimate fate...	LLaMa-5.4B	The universe is the entirety of space, time, matter, and energy that exists. The laws that make up physics — electrons orbiting atoms, magnets attracting other objects, planets orbiting the Sun—have been consistent over ...
Vicuna-7B	10 steps to build an ios app: 1. Research and develop the idea. 2. Decide on the platform you want to use (iOS, Android, or both). 3. Create a plan and budget. 4. Hire a team of developers. 5. Design the user interface (UI). 6. Develop the user interface (UI). 7. Test and iterate on the app. 8. Submit the app to the app store. 9. Optimize the app for performance and user experience (UX). 10. Market and promote the app.	Vicuna-5.4B	10 steps to build an ios app: 1. Plan your app: decide what it should do, what it will look like, and what it will cost to build it. 2. Hire a developer: A good iOS developer will help you design the user interface, create the app, and make sure it’s bug-free. 3. Create a user interface: This is where the magic happens. Your app’s user interface will be designed and mocked up at this point...
ChatGLM-6B	Write a poetry about love and peace. In the midst of life's hustle and bustle, I find your love a soothing peace. It is a feeling that cannot be conveyed, A connection that cannot be distanced. With your smile, your eyes, your embrace, I feel a sense of belonging that only you can give. It is a love that lasts forever, A love that cannot be defeat...	ChatGLM-5.6B	Write a poetry about love and peace. Love and peace, two such wonderful things\ A pure and endless desire in my heart\ And both of them, I must seek for\ A long, long time, I know..\ Love, I know, is a feeling of being\ A perfect partner, in every sense\ And peace, I need it, so much, one day\ A long, long way, my heart will go..

Limitations

Although we only used 50K data and trained for three hours, more data would definitely be better. We are testing on this.
The current compressed model still has several issues, such as generating repetitive tokens or producing nonsensical sentences. We believe there is significant room for improvement in the quality of the compressed model.
There are still some models for which we cannot automatically identify the mapping of indexes after concatenation and view operations. Therefore, we need to perform additional manual operations.

Acknowledgement

Logo is generated by Stable Diffusion
The evaluation of the 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

Citation

If you find this project useful, please cite

@misc{ma2023llmpruner,
      title={LLM-Pruner: On the Structural Pruning of Large Language Models}, 
      author={Xinyin Ma and Gongfan Fang and Xinchao Wang},
      year={2023},
      eprint={2305.11627},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

@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}
}

SUSTechBruce / LLM-Pruner