Skip to main content
`Ctrl`+`K`

# LLM Pruning

## Overview of LLM Pruning

Deployment of LLMs is constrained by their intensive computational demands.

To make LLMs more accessible and efficient for practical use, LLM pruning shows the ability to reduce the number of parameters or create a Small LLM from a larger one.

Through pruning, either by dropping layers (depth pruning) or dropping neurons and attention heads, and embedding channels (width pruning).

After pruning, the higher the pruning ratio, the higher the accuracy degradation will be. As a result, Pruning is often accompanied by some fine-tuning or retraining for accuracy recovery.

## Mainstream Methods Overview

Unstructured pruning: Concentrate on individual weights, set the individual elements in weight to 0. Meaning that sparsity of rate s% will eliminate s% of the entries in a weight matrix. But the efficiency of the pruned sparse network cannot be realized on general-purpose GPU hardware.

Semi-structured pruning: This approach enforces exactly N non-zero values in each block of M consecutive weights. Whether can get GPU acceleration depends on specific hardware architectures.

Structured pruning: The pruning methods can be categorized into two different grain levels.

- Channel-wise: Based on the weight importance, delete entire rows or columns of weights, providing a more hardware-friendly solution that reduces storage requirements and improves GPU inference performance.

- Layer-Wise: Delete the entire decode layer, which reduces the model parameters and model structure more straightforwardly. As a result, this can bring a more significant GPU inference performance improvement.

By viewing the above short introduction, even though several methods have been proposed by researchers. Considering the actual production environments, Unstructured pruning and Semi-structured pruning may have competitive results in accuracy. It may not decrease the amount of computation, especially may not bring GPU inference acceleration.

In order to get a real LLM model size reduction and forward computation acceleration. As a result, in the Quark Pruning tool, we mainly consider the structured pruning methods, which can directly reduce the model size and accelerate GPU inference.

## Quark Pruning Tool Feature

### Supported Methods Overview

- OSSCAR Pruner: A structured pruning method, this method needs to use a small amount of calibration data to perform layer-wise reconstruction and minimize the layer reconstruction error.

  - This method will reduce the intermediate dimension (output channel of the former FC and the input channel of the following FC layer ) in fully connected layers.

  - This method will not change LLM model structure.

  - After pruning, if the remaining channel number is power-of-2, it usually can bring GPU inference time acceleration. (e.g 4096 = 2**12).

- Depth-Wise Pruner: Based on the PPL influence, the consecutive decode layers that even after deletion have less influence on the final PPL will be regarded as having less influence on the LLM. These layers can be deleted.

  - Will delete consecutive decode layers based on pre-defined prune ratio.

  - Will directly reduce model size, as it directly deletes the decode layers. This way will improve GPU inference performance.

- For example, the following code block shows a typical `Llama` model structure.

  - For `OSSCAR` pruner: will decrease the output channel in `gate_proj` and `up_proj`, meanwhile decrease the input channel of `down_proj`.

  - For depth pruner: will directly delete the entire `LlamaDecoderLayer` in the model, for example, after deleting the 10 consecutive layers, this model will remain 70 (80-10) layers, which can directly improve GPU inference performance.

    LlamaForCausalLM(
      (model): LlamaModel(
        (embed_tokens): Embedding(128256, 8192)
        (layers): ModuleList(
          (0-79): 80 x LlamaDecoderLayer(
            (self_attn): LlamaAttention(
              (q_proj): Linear(in_features=8192, out_features=8192, bias=False)
              (k_proj): Linear(in_features=8192, out_features=1024, bias=False)
              (v_proj): Linear(in_features=8192, out_features=1024, bias=False)
              (o_proj): Linear(in_features=8192, out_features=8192, bias=False)
            )
            (mlp): LlamaMLP(
              (gate_proj): Linear(in_features=8192, out_features=28672, bias=False)
              (up_proj): Linear(in_features=8192, out_features=28672, bias=False)
              (down_proj): Linear(in_features=28672, out_features=8192, bias=False)
              (act_fn): SiLU()
            )
            (input_layernorm): LlamaRMSNorm((8192,), eps=1e-05)
            (post_attention_layernorm): LlamaRMSNorm((8192,), eps=1e-05)
          )
        )
        (norm): LlamaRMSNorm((8192,), eps=1e-05)
        (rotary_emb): LlamaRotaryEmbedding()
      )
      (lm_head): Linear(in_features=8192, out_features=128256, bias=False)
    )

## API Usage

### Prepare LLM & Tokenizer

    model = AutoModelForCausalLM.from_pretrained(ckpt_path, device_map="auto", torch_dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(ckpt_path, padding_side="left")

### Init Prune Config & Calibration/Test Dataset

    from quark.torch.pruning.config import Config
    from quark.torch.pruning.config import OSSCARConfig, LayerImportancePruneConfig
    # Init Prune config
    # Using OSSCAR
    algo_config = OSSCARConfig.from_dict(json.load(algo_config_file))
    # or Using depth-wise prune
    algo_config = LayerImportancePruneConfig.from_dict((json.load(algo_config_file))

    pruning_config = Config(algo_config=pruning_algo_config)

    # Init the Calibration/Test dataset (Fake code)
    def get_wikitext_dataset(data_dir, tokenizer):
        testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")
        testenc = tokenizer("\n\n".join(testdata["text"]), return_tensors="pt")
        return testenc

    validation_data = get_wikitext_dataset(data_dir, tokenizer)

### Init Pruner & Perform Pruning

    from quark.torch import ModelPruner
    model_pruner = ModelPruner(pruning_config)
    model = model_pruner.pruning_model(model, validation_data)

### Evaluate the Metric & Save Pruned Model

    # evaluate
    eval_model(args, model, main_device) # evaluate the PPL on pruned model
    # save to target path
    model.save_pretrained(save_dir, safe_serialization=True)
    tokenizer.save_pretrained(save_dir)

## Prune Results

### OSSCAR:

Model Name | Model Size | Pruning Rate | Pruned Model Size | Before Pruning PPL On Wiki2 | After Pruning PPL On Wiki2
--- | --- | --- | --- | --- | ---
mistralai/Mixtral-8x7B-Instruct-v0.1 | 46.7B | 9.4838% | 42.2B | 4.1370 | 5.1195
CohereForAI/c4ai-command-r-08-2024 | 32.3B | 7.4025% | 29.9B | 4.5081 | 6.3794
Qwen/Qwen2.5-14B-Instruct | 14.8B | 7.0284% | 13.7B | 5.6986 | 7.5994
meta-llama/Meta-Llama-3-8B | 8.0B | 6.8945% | 7.5B | 6.1382 | 8.0755
meta-llama/Llama-2-7b-hf | 6.7B | 6.7224% | 6.2B | 5.4721 | 6.2462
facebook/opt-6.7b | 6.7B | 7.5651% | 6.2B | 10.8602 | 11.8958
THUDM/chatglm3-6b | 6.2B | 7.7590% | 5.6B | 29.9560 | 36.0010
microsoft/Phi-3.5-mini-instruct | 3.8B | 5.9274% | 3.6B | 6.1959 | 7.8074

### Depth-Wise Prune:

Model | PPL (Original) | PPL (After prune) | Prune Ratio | Other Info
--- | --- | --- | --- | ---
llama-2-7b-chat-hf | 6.9419 | 8.139 | 9.01% | delete 3 layer [11-13] total 32
llama-3.1-405B | 1.85624 | 2.8429 | 10.3% | delete 13 layer [18-30] total 126
Llama-2-70b-chat-h | 4.6460 | 5.5305 | 11.16% | delete 9 layer [27-35] total 80
Qwen2.5-14B-Instruct | 5.7010 | 7.0681 | 9.32% | delete 5 layers [22-26] total 48
Mixtral-8x7B-Instruct-v0.1 | 4.1378 | 5.0338 | 10% | delete 3 layer [12-14] total 32
facebook/opt-6.7b | 10.8605 | 14.4321 | 12.1% | delete 4 layer [21-24] total 32
deepseek-moe-16b-chat | 7.3593 | 8.94327 | 10.76% | delete 3 layer [11-13] total 28

## Example Code

All the example code can be found under `/examples/torch/language_modeling/llm_pruning`

## Other:

Reference: OSSCAR: One-Shot Structured Pruning in Vision and Language Models with Combinatorial Optimization