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/eva