Small LLMs: Pruning vs Training from Scratch

Yufeng Xu¹, Taiming Lu¹, Jiachen Zhu², Mingjie Sun³, Kunjun Li¹, and Zhuang Liu¹

1 Princeton. 2 NYU. 3 CMU.

---

This is a Jax-based repo for LLM Prunning, It contains:

• the implementations of various LLM pruning methods of different granularity.
• pretraining and fine-tuning code for both GPU and TPU platforms.
• evaluation scripts for assessing model performance.

We gratefully acknowledge the generous support of the Google TPU Research Cloud (TRC), which provided the computational resources used to build this repository.

The repo is organized as follows:

├── pruning
│   ├── FLAP # including Wanda-sp and FLAP
│   ├── LLM-Pruner
│   ├── llmshearing # sheared llama
│   ├── minitron # including shortgpt
│   ├── SLEB # including sleb
│   ├── SliceGPT # including slicegpt
│   └── wanda # including sparsegpt and magnitude pruning
├── training
│   ├── fms_fsdp
│   └── maxtext
└── eval

where pruning is the collection of the pruning methods we experimented; training contains the LLM training frameworks we used, and we provided options for both TPU and GPU; eval contains JAX-compatible eval scripts we used to evaluate the pruned models.

---

Supported Features

Pruning Methods

• Minitron
• ShortGPT
• Wanda
• SliceGPT
• SparseGPT
• Magnitude
• Sheared Llama
• SLEB
• LLM Pruner
• FLAP

Training Frameworks

• FMS-FSDP
• MaxText

Evaluation

• accelerate lm-eval-harness for maxtext. (by 2-4x times!)

Get Started

Pruning

In order to reproduce the results of the different pruning methods, we need to set up separate environments for different methods. The installation and command guide can be found at pruning/<method>/README.md. Below is an overview:

Minitron

cd pruning/minitron
bash scripts/install.sh
bash scripts/prune_llama3.1-8b.sh # contains minitron depth and width for llama3.1-8b

ShortGPT

cd pruning/minitron
bash scripts/install.sh
bash scripts/prune_llama2-7b.sh 

Wanda, SparseGPT, Magnitude

cd pruning/wanda
bash scripts/install.sh
bash scripts/prune_llama3.1-8b.sh # contains wanda, sparsegpt, and magnitude for llama3.1-8b
bash scripts/prune_llama2-7b.sh
bash scripts/prune_llama-7b.sh

LLM-Pruner

cd pruning/LLM-Pruner
bash scripts/install.sh
bash scripts/prune_llama-7b.sh
bash scripts/prune_llama2-7b.sh
bash scripts/prune_llama3.1-8b.sh

Sheared Llama

cd pruning/llmshearing
bash scripts/install.sh

mkdir -p llmshearing/data/red_pajama && cd llmshearing/data/red_pajama
huggingface-cli download Zephyr271828/redpajama-for-prune --repo-type dataset --local-dir for_prune
cd -

bash scripts/hf2composer.sh
bash scripts/prune_llama2-2.7b.sh
bash scripts/prune_llama2-1.3b.sh
bash scripts/prune_llama2-370m.sh
bash scripts/composer2hf.sh

Training

GPU To train on GPUs, please refer to the guide of fms-fsdp for details.

TPU To train on TPUs, please refer to guide of MaxText for details.

Evaluation

GPU
For evaluation on GPUS, you may run the following evaluation script on your HF checkpoint:

cd training/fms_fsdp
bash scripts/install.sh

cd ../../eval
bash scripts/eval.sh

Note for LLM-Pruner and Wanda, they have specified a specific version of lm-eval to use, which is included in their respective directory, and the evaluation code is included in the pruning process.
For all other methods, you may eval with the script provided.

TPU Please refer to the guide of MaxText for details.

Reproduction Results

In this section, we show some of our results to verify that we can reproduce the results from the original pruning papers.

• For Minitron, the original papers did not report evaluation results after pruning and before retraining, so we attempt to reproduce the plots from LLM Pruning and Distillation in Practice: The Minitron Approach.
• For ShortGPT, although evaluation results are provided, yet we noticed there are inconsistencies between results in the table (also see ShortGPT: Layers in Large Language Models are More Redundant Than You Expect). Therefore, we choose to reproduce the block importance plot from the paper, which implies the correctness of our implementation.
• For all other methods, both official implementation and evaluation results are provided, so we simply provide comparison with the reported results in paper.

Minitron-Winogrande
Left: plot from the paper; Right: plot made by us.

Minitron-Wikitext

ShortGPT

Wanda

Llama-2-7b-hf:

SparseGPT

Llama-2-7b-hf:

Magnitude

Llama-2-7b-hf:

Sheared Llama

LLM-Pruner

Note: The results are obtained by running the exact pruning and evaluation scripts from the LLM-Pruner repo. Still, some results differ the reported results in the paper. My conjecture is that only 10 samples are randomly selected from the bookcorpus dataset for importance estimation, and this caused some randomness, even though we fixed the random seed.