IntactKV: Improving Large Language Model Quantization by Keeping Pivot Tokens Intact

This repository contains the PyTorch implementation of IntactKV: Improving Large Language Model Quantization by Keeping Pivot Tokens Intact.

IntactKV is a simple and orthogonal method to enhance the quantized LLMs. It can be feasibly combined with various existing quantization approaches (e.g., AWQ, OmniQuant, GPTQ, QuaRot) with no inference overhead on various LLMs (LLaMA, Vicuna, OPT, Mistral e.t.c.). IntactKV is built on a valuable observation that pivot tokens do exist in current LLMs with massive values and highly concentrated attention scores, and thay are critical to the performance of quantized LLMs. A concurrent work Massive Activations also discovers such tokens and provides more detailed studies on this phenomenon.

Preparations

Installation

conda create -n intactkv python=3.10 -y
conda activate intactkv
pip install -r requirements.txt

Data Preparation

Download datasets in ./datasets.

Calibration set or PPL evaluation

MMLU evaluation

Commonsense QA evaluation

Model Preparation

Download models in ./modelzoo.

Weight-only Quantization

Model Quantization

GPTQ Quantize model with AutoGPTQ. The quantized model will be available in ./modelzoo/autogptq.

# w3g128 quantization of Vicuna-v1.5-7B on GPU 0
bash ./scripts/quantization/auto_gptq.sh vicuna-v1.5 vicuna-v1.5-7b 3 128 0

AWQ Download pre-computed AWQ parameters from AWQ modelzoo, or reproduce with the following script. The search results will be saved in ./modelzoo/llm-awq.

# w3g128 quantization of Vicuna-v1.5-7B on GPU 0
bash ./scripts/quantization/llm_awq.sh vicuna-v1.5 7b 3 128 0

IntactKV_[B]

Evaluation

IntactKV_[B] can be directly integrated with various quantization methods (e.g., AWQ, GPTQ, RTN) without training or inference overhead, and can be evaluated on PPL, MMLU, and QA tasks, where [BOS] token is prepended to the inputs.

PPL

# evaluate w3g128 AWQ-quantized LLaMA-2-7B model on GPU0, port 29500
bash ./scripts/eval/eval.sh llama-2 7b awq 3 16 ppl 29500 0

MMLU

# evaluate w3g128 AWQ-quantized Vicuna-v1.5-7B model on GPU0, port 29500
bash ./scripts/eval/eval.sh vicuna-v1.5 7b awq 3 16 mmlu 29500 0

Commonsense QA

# evaluate w3g128 AWQ-quantized Vicuna-v1.5-7B model on GPU0, port 29500
bash ./scripts/eval/eval.sh vicuna-v1.5 7b awq 3 16 qa 29500 0

IntactKV_[P]

IntactKV as Trainable Parameters

IntactKV can be optionally calibrated on a calibration set of size 128 to compensate for the quantization error.

# calibrate IntactKV of w3g128 AWQ-quantized Vicuna-v1.5-7B model on GPU0
bash ./scripts/train/train.sh vicuna-v1.5 7b awq 3 0

Evaluation

MT-bench

1. Generate answers to MT-bench with the following script.

   # generate answers to MT-bench for w3g128 AWQ-quantized Vicuna-v1.5-7B model on GPU0
   bash scripts/eval/gen_mtbench_answer.sh vicuna-v1.5 7b awq 3 0

2. Score the answers with GPT4 using LLM Judge. Reference answer of gpt-4-0125-preview can be found in ./fastchat/data/mt_bench/reference_answer.

Weight and Activation Quantization

We integrate IntactKV with a SOTA INT4 weight and activation quantization method QuaRot, which uses hadamard transformation to alleviate outlier issues. Run the following script to obtain PPL evaluation results.

# LLaMA-2-7B model on GPU0
bash ./scripts/eval/quarot.sh llama-2 7b 0

KV Cache Quantization

We implement a simple asymmetric per-head dynamic quantization strategy for KV cache. Run the following scripts to obtain PPL/MMLU/QA evaluation results.

IntactKV is also available in another SOTA KV cache only quantization method KVQuant, and can be evaluated with KVQuant's official code.

PPL

# evaluate w3g128kv4 AWQ-quantized LLaMA-2-7B model on GPU0, port 29500
bash ./scripts/eval/eval.sh llama-2 7b awq 3 4 ppl 29500 0

MMLU

# evaluate w3g128kv4 AWQ-quantized Vicuna-v1.5-7B model on GPU0, port 29500
bash ./scripts/eval/eval.sh vicuna-v1.5 7b awq 3 4 mmlu 29500 0

Commonsense QA

# evaluate w3g128kv4 AWQ-quantized Vicuna-v1.5-7B model on GPU0, port 29500
bash ./scripts/eval/eval.sh vicuna-v1.5 7b awq 3 4 qa 29500 0

Visualizations

Visualizations of pivot tokens Visualize output activations and corresponding attention maps of LLMs. Output PDFs will be saved in ./outputs/visualizations.

# LLaMA-2-7B model on GPU0
bash ./scripts/visualization/plot_act.sh llama-2 7b 0

Plot quantization loss w.r.t. IntactKV size Plot the line chart in Figure 2, which strongly demonstrates the importance of pivot tokens.

# LLaMA-2-7B model on GPU0
bash ./scripts/visualization/motivation.sh llama-2 7b 0

Results

Table 2. INT3-group128 weight-only quantization results of Vicuna models on 5-shot MMLU tasks.

Table 3. INT3-group128 weight-only quantization results of Vicuna models on 0-shot QA tasks.

Reference

If you find IntactKV helpful, please cite our paper:

@inproceedings{liu2024intactkv,
  title={IntactKV: Improving Large Language Model Quantization by Keeping Pivot Tokens Intact},
  author={Liu, Ruikang and Bai, Haoli and Lin, Haokun and Li, Yuening and Gao, Han and Xu, Zhengzhuo and Hou, Lu and Yao, Jun and Yuan, Chun},
  booktitle={Findings of the Association for Computational Linguistics: ACL 2024},
  year={2024}
}