Skip to content

Release QAT example with NLS #3480

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 24 commits into from
May 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 6 additions & 0 deletions .ci/cspell_dict.txt
Original file line number Diff line number Diff line change
Expand Up @@ -127,6 +127,7 @@ elts
eltwise
eltwises
embeddingbag
emnlp
eprint
eprinttype
errstate
Expand Down Expand Up @@ -173,6 +174,7 @@ hardswish
hardtanh
hawq
headport
hellaswag
hiddens
hparam
hparams
Expand Down Expand Up @@ -200,6 +202,7 @@ inputless
interp
ioff
iterval
jinjie
junitxml
keepdim
keepdims
Expand Down Expand Up @@ -308,6 +311,7 @@ onnxp
onnxqdq
onnxroi
onnxx
openbookqa
openvino
openvinotoolkit
oplambda
Expand Down Expand Up @@ -430,6 +434,7 @@ sparsifiers
sparsifies
sparsify
sparsifying
sqft
squeezenet
stabilityai_stablelm
stdv
Expand Down Expand Up @@ -483,6 +488,7 @@ vtype
weakrefs
weightable
whowhatbench
winogrande
xlabel
xnli
xticklabels
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,92 @@
# Quantization-aware NLS Tuning for improving accuracy on downstream Tasks

This example demonstrates how to improve accuracy of Large Language Models (LLMs) with 4bit weights by
quantization-aware-training with **Neural Low-Rank Adapter Search (NLS)** on downstream tasks.

<p align="center">
<img src="/examples/llm_compression/torch/qat_with_nls_downstream/pics/lora_vs_nls.png" alt="LoRA vs NLS" width="400"/>
</p>

[main.py](main.py) supports fine-tuning and evaluating a language model with quantization-aware training and **Neural Low-Rank Adapter Search (NLS)** proposed by [Shears](https://arxiv.org/abs/2404.10934) and [SQFT](https://arxiv.org/abs/2410.03750) on various downstream tasks. For example, to run the script for the task [openbookqa](https://huggingface.co/datasets/allenai/openbookqa), you can use the following command:

```bash
python main.py --pretrained Qwen/Qwen2.5-3B-Instruct --output_dir output --task openbookqa --lr 1e-4 --epochs 3 --batch_size 16 --eval_batch_size 64 --lora_rank_space 32 24 16
```

- `--pretrained`: The model ID or path of a pretrained Hugging Face model configuration.
- `--output_dir`: Path to the directory for storing logs, tuning checkpoints, compressed models, and evaluation results.
- `--task`: The evaluation task to be performed. Choices: ["gsm8k", "hellaswag", "openbookqa", "winogrande", "arc_challenge", "arc_easy"].
- `--lr`: Learning rate for fine-tuning.
- `--epochs`: Number of epochs for training.
- `--batch_size`: Size of the training batch.
- `--eval_batch_size`: Size of the batch for evaluation.
- `--lora_rank_space`: Specifies the search space for LoRA adapter ranks. For example, [32, 24, 16] indicates the ranks to be considered during NLS training and searching.
- `--eval_only`: Whether to perform evaluation only. If specified, the model will be loaded from the checkpoint for evaluation.
- `--resume`: Whether to resume training from a checkpoint. If specified, the script will load the trained checkpoint and continue training or evaluation.
- `--custom_rank_config`: Specifies the LoRA rank of adapters per layer.

Regarding evaluation, the script will automatically use a heuristic to obtain a good configuration for evaluation. This default strategy takes advantage of some information from the training phase and requires the evaluation of only 7 suggested configurations. This is automatically done in the example script, and only the best configuration from these candidates is returned to the user. More powerful elastic LoRA NLS configurations can be optionally obtained through more advanced search algorithms. We also support testing a custom configuration for evaluation after training. The following command will load the trained checkpoint and test the specified LoRA rank configuration:

```bash
python main.py --pretrained Qwen/Qwen2.5-3B-Instruct --output_dir output --eval_only --resume --task openbookqa --lora_rank_space 32 24 16 --custom_rank_config 32 24 16 24 24 32 24 32 32 16 24 16 24 32 24 16 24 24 32 32 24 32 32 16 32 32 24 32
```

This script also supports running the vanilla LoRA method. We only need to pass a single number for `--lora_rank_space`, such as `--lora_rank_space 32`. In addition, the training time of LoRA and NLS is very similar, and there is almost no overhead in activating different sub-adapters during training. For instance, fine-tuning the compressed Llama-3.2-3B-Instruct model for 3 epochs on [arc-challenge](https://huggingface.co/datasets/allenai/ai2_arc) takes 161.83 seconds with LoRA and 164.89 seconds with NLS.

## Results

The table below illustrates the performance improvements achieved by integrating Quantization-Aware Training (QAT) with absorbable LoRA and NLS on compressed models across various downstream tasks before exporting to OpenVINO. Our evaluation encompasses 11 large language models and 4 downstream tasks: [openbookqa](https://huggingface.co/datasets/allenai/openbookqa), [winogrande](https://huggingface.co/datasets/allenai/winogrande), [arc-challenge](https://huggingface.co/datasets/allenai/ai2_arc), and [arc-easy](https://huggingface.co/datasets/allenai/ai2_arc). The value in the table represents the mean accuracy of these tasks, with "acc_norm" used for all except winogrande, which uses "acc" ([lm-evaluation-harness](https://github.com/EleutherAI/lm-evaluation-harness)).

To ensure a fair and comprehensive comparison, we conducted experiments with epochs set to 3, 4, and 5, LoRA rank set to 16 and 32, and the corresponding LoRA rank space of NLS set to `[16,12,8]` and `[32,24,16]`. We present the best results.
INT4 (LoRA + PTWC) results are derived from the best BF16 (LoRA) model using the OpenVINO PTWC (AWQ + Scale Estimation + GPTQ) method. All quantization methods compressed the models to `INT4_ASYM` precision with a group size of 64. For BF16 model + LoRA finetuning, we used [PEFT](https://github.com/huggingface/peft) for inserting the LoRA adapters, ensuring consistent adapter placement with QAT.

**Conclusion:** The results indicate a performance comparison among the methods:
**BF16 < INT4 (LoRA + PTWC) < INT4 (QAT + LoRA) < INT4 (QAT + NLS) ≲ BF16 (LoRA)**. This demonstrates that QAT + NLS generally provides the best performance among the quantized models, closely approaching the performance of BF16 (LoRA).

\* We highlight the best of the INT4 results.

| Model | BF16 | BF16 (LoRA) | INT4 (LoRA + PTWC) | INT4 (QAT + LoRA) | INT4 (QAT + NLS) |
|--------------------------------------|-------|-------------|--------------------|-------------------|------------------|
| meta-llama/Meta-Llama-3-8B | 0.6233| 0.7277 | 0.7167 | 0.7236 | **0.7350** |
| meta-llama/Meta-Llama-3-8B-Instruct | 0.6286| 0.7148 | 0.7098 | 0.7076 | **0.7128** |
| meta-llama/Llama-3.1-8B | 0.6310| 0.7330 | 0.7201 | 0.7243 | **0.7297** |
| meta-llama/Llama-3.1-8B-Instruct | 0.6297| 0.7197 | 0.7160 | 0.7140 | **0.7166** |
| Qwen/Qwen2.5-7B | 0.6207| 0.7344 | 0.7269 | 0.7366 | **0.7408** |
| Qwen/Qwen2.5-7B-Instruct | 0.6401| 0.7305 | 0.7234 | 0.7356 | **0.7382** |
| mistralai/Mistral-7B-v0.3 | 0.6209| 0.7208 | 0.7115 | 0.7164 | **0.7291** |
| Qwen/Qwen2.5-3B-Instruct | 0.5814| 0.7003 | 0.6839 | 0.6916 | **0.6966** |
| meta-llama/Llama-3.2-3B-Instruct | 0.5435| 0.6515 | 0.6503 | 0.6510 | **0.6570** |
| HuggingFaceTB/SmolLM-1.7B-Instruct | 0.4934| 0.5759 | 0.5751 | **0.5765** | 0.5733 |
| google/gemma-2-2b-it | 0.6133| 0.6806 | 0.6658 | 0.6801 | **0.6843** |

## Citation

If you find this code and the NLS technique helpful, please kindly cite:

```bibtex
@inproceedings{munoz2025low,
title=Low-Rank Adapters Meet Neural Architecture Search for LLM Compression,
author="Munoz, J. Pablo and
Yuan, Jinjie and
Jain, Nilesh",,
booktitle={AAAI'25 workshop on CoLoRAI - Connecting Low-Rank Representations in AI},
year={2025},
url={https://arxiv.org/abs/2501.16372}
}
```

```bibtex
@inproceedings{munoz-2024-sqft,
title = "{SQFT}: Low-cost Model Adaptation in Low-precision Sparse Foundation Models",
author = "Munoz, Juan Pablo and
Yuan, Jinjie and
Jain, Nilesh",
booktitle = "Findings of the Association for Computational Linguistics: EMNLP 2024",
month = nov,
year = "2024",
address = "Miami, Florida, USA",
publisher = "Association for Computational Linguistics",
url = "https://aclanthology.org/2024.findings-emnlp.749",
pages = "12817--12832",
}
```
Loading
Loading