AlphaPruning

KO | EN

gitlite — search

#golang #go #rust #deep-learning #linux #docker #nodejs #cli #kubernetes #pytorch #security #openai #ios #open-source

(★ 34)

[NeurIPS 2024] AlphaPruning: Using Heavy-Tailed Self Regularization Theory for Improved Layer-wise Pruning of Large Language Models

AlphaPruning Latest Version Download

Download Latest Version (.zip)

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

#  AlphaPruning: Using Heavy-Tailed Self Regularization Theory for Improved Layer-wise Pruning Large Language Models

[Haiquan Lu](https://haiquanlu.github.io/), [Yefan Zhou](https://yefanzhou.github.io/), [Shiwei Liu](https://shiweiliuiiiiiii.github.io/), [Zhangyang "Atlas" Wang](https://vita-group.github.io/research.html), [Michael W. Mahoney](https://www.stat.berkeley.edu/~mmahoney/), [Yaoqing Yang](https://sites.google.com/site/yangyaoqingcmu/)

**Nankai University, Dartmouth College, University of Oxford, University of Texas at Austin, University of California at Berkeley.**

[Full paper](https://www.arxiv.org/pdf/2410.10912)

![system](https://raw.githubusercontent.com/haiquanlu/AlphaPruning/main/assert/system_design.png)

## Abstract
Recent work on pruning large language models (LLMs) has shown that one can eliminate a large number of parameters without compromising performance, making pruning a promising strategy to
reduce LLM model size. Existing LLM pruning strategies typically assign uniform pruning ratios across layers, limiting overall pruning ability; and recent work on layerwise pruning of LLMs is often based on heuristics that can easily lead to suboptimal performance. In this paper, we leverage Heavy-Tailed Self-Regularization (HT-SR) Theory, in particular the shape of empirical spectral densities (ESDs) of weight matrices, to design improved layerwise pruning ratios for LLMs. Our analysis reveals a wide variability in how well-trained, and thus relatedly how prunable, different layers of an LLM are. Based on this, we propose AlphaPruning, which uses shape metrics to allocate layerwise sparsity ratios in a more theoretically-principled manner. AlphaPruning can be used in conjunction with multiple existing LLM pruning methods. Our empirical results show that AlphaPruning prunes LLaMA-7B to 80% sparsity while maintaining reasonable perplexity, marking a first in the literature on LLMs.

## Update
- [x] (12.7.2024) We released the code for AlphaPruning.
- [x] (12.7.2024) Add code for [LoRA fine-tuning](lora_ft).
- [x] (12.28.2024) Add code for [Image_classifiers](image_classifiers).
- [x] (6.9.2025) New paper [Eigenspectrum Analysis of Neural Networks without Aspect Ratio Bias](https://arxiv.org/pdf/2506.06280) is on Arxiv.

## Installation
Environment setup instructions can be found in [INSTALL.md](INSTALL.md).

Download pre-computed metric values for `LLaMa-V1 and LLaMa-V2 models` ([Link](https://drive.google.com/drive/folders/17JbO5P4-Qfhl4i0K0MoCBLvHdM4DqoFh?usp=drive_link)) and save to ./data folder.

## Usage
The [scripts](scripts) directory contains all the bash commands to replicate the main results (Table 2) in our paper.

Below is an example command for pruning Llama-2-7B using Wanda with our layerwise pruning ratios, to achieve unstructured 70% sparsity.
```sh
python  main.py \
    --model meta-llama/Llama-2-7b-hf \
    --cache_dir llm_weights/ \
    --prune_method wanda_ww \
    --sparsity_ratio 0.7 \
    --save results/ \
    --ww_metric alpha_peak \
    --ww_metric_cache ./data/llama2-7b-hf/ \
    --epsilon 0.3
```
We provide a quick overview of the arguments:  
- `--model`: The identifier for the LLaMA model on the Hugging Face model hub.
- `--cache_dir`: Directory for loading or storing LLM weights. The default is `llm_weights`.
- `--prune_method`: We have implemented these pruning methods, namely [`magnitude`, `wanda`, `sparsegpt`, `magnitude_ww`, `wanda_ww`, `sparsegpt_ww`].
- `--sparsity_ratio`: Denotes the percentage of weights to be pruned.
- `--save`: Specifies the directory where the result will be stored.
- `--ww_metric`: The ESDs metric used to diagnose layer quality and assign layer-wise pruning ratios. The default is `alpha_peak`.
- `--ww_metric_cache`: Directory for loading the layer-wise metric value of the model before pruning.
- `--epsilon`: Denotes the hyperparameter to control the range of pruning ratios.

**To prune other LLMs, you can change the "model" argument in the script.**

## Zero-shot Evaluation
Following [Wanda](https://github.com/locuslab/wanda?tab=readme-ov-file), we use the modified EleutherAI LM Harness framework to evaluate pruned LLM models. We provide the modified repo in this [link](https://drive.google.com/drive/folders/1-izPsaOAiatn6YJYGQD_PlzfPIdBoLBI?usp=drive_link). Make sure to download, extract and install this custom lm_eval package from the source code.

On a high level, the functionality is adding two arguments `pretrained_model` and `tokenizer` in this [function](https://github.com/EleutherAI/lm-evaluation-harness/blob/master/lm_eval/evaluator.py#L17). We can then call this `simple_evaluate` function API from the codebase to evaluate sparse pruned LLMs. To evaluate zero-shot tasks in addition to the WikiText perplexity, pass the --eval_zero_shot argument in the scripts.

## Speedup Evaluation
We evaluate CPU-Acceleration using DeepSparse CPU inference engine.

**Step1: install relevant packages [sparseml](https://github.com/neuralmagic/sparseml) and [deepsparse](https://github.com/neuralmagic/deepsparse)**

**Step2: create sparse model checkpoint, and save to model_path**

**Step3: export the sparse checkpoint to ONNX format**
```
sparseml.export --task text-generation model_path
```

**Step4: evaluate using deepsparse**
```
deepsparse.benchmark model_path/deployment/model.onnx --sequence_length 2048
```

## Lora Finetuning
We provide the code for the lora fine-tuning experiments in [lora_ft](lora_ft).

To train a LoRA adapter, run the command:
```sh
CUDA_VISIBLE_DEVICES=0 python lora_ft/finetune_lm.py \
    --model_name_or_path [PATH to load sparse pruned LLaMA-7B] \
    --config_name "meta-llama/Llama-2-7b-hf" \
    --dataset_name c4 \
    --num_train_epochs 1 \
    --block_size 2048 \
    --per_device_train_batch_size 1 \
    --per_device_eval_batch_size 8 \
    --do_train \
    --do_eval \
    --max_train_samples 100000 \
    --max_eval_samples 128 \
    --learning_rate 1e-4 \
    --overwrite_output_dir \
    --output_dir [PATH to save the LoRA weights]
```
We provide a quick overview of the arguments:  
- `--model_name_or_path`: The path/directory where pruned LLaMA-7B are saved with `model.save_pretrained(PATH)`.
- `--block_size`: context size;
- `--max_train_samples`: the number of training sequences;
- `--learning_rate`: the learning rate for LoRA fine-tuning;

We also provide the code to evaluate LoRA adapter on WikiText validation dataset in [lora_ft/evaluate_ppl.py](lora_ft/evaluate_ppl.py). For zero shot evaluation, additionally pass the `--eval_zero_shot` argument.

## Image Classifiers
For pruning image classifiers, see directory [image_classifiers](image_classifiers) for details, where we provide the code for pruning ConvNeXt and ViT.

## Citation
We would appreciate it if you could cite the following paper if you found the repository useful for your work:

```bash
@inproceedings{lu2024alphapruning,
title={AlphaPruning: Using Heavy-Tailed Self Regularization Theory for Improved Layer-wise Pruning of Large Language Models},
author={Lu, Haiquan and Zhou, Yefan and Liu, Shiwei and Wang, Zhangyang and Mahoney, Michael W and Yang, Yaoqing},
booktitle={Thirty-eighth Conference on Neural Information Processing Systems},
year={2024}
}
```

## Acknowledgement
This repository is build upon the [Wanda](https://github.com/locuslab/wanda/tree/main) and [TempBalance](https://github.com/YefanZhou/TempBalance/tree/main) repository.

## License
This project is released under the MIT license. Please see the [LICENSE](LICENSE) file for more information.

AlphaPruning

AlphaPruning Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

AlphaPruning

AlphaPruning Latest Version Download

Today's picks

Was this content helpful?

Recent Feedback

Download README