UniSVQ & LC-QAT: 2-Bit LLM Quantization with Linear-Constrained Vector Quantization
This repository contains the models of two companion papers on 2-bit LLM quantization:
- UniSVQ (arXiv:2606.10520): A post-training quantization (PTQ) framework that unifies scalar and vector quantization. It parameterizes codewords as an affine transform of integer lattices, preserving compatibility with optimized integer kernels while retaining the flexibility of VQ.
- LC-QAT (arXiv:2606.10531): A data-efficient 2-bit quantization-aware training (QAT) framework. It introduces linear-constrained vector quantization (LCVQ), replacing discrete codebook lookup with a learned linear projection for end-to-end differentiable training.
Codes are available at Our Github page.
The model here is an instuction-tuned quantized model. We strengthen the model's performance on math, code, and complex tasks using approximately 8B tokens of conversational supervised data. The performance of our 2-bit model after 8B supervised fine-tuning on Qwen3 1.7B is shown below:
| Model | OpenbookQA | If | MMLU | GSM8K | MATH | HumanEval | BBH | Avg. |
|---|---|---|---|---|---|---|---|---|
| Qwen3 1.7B fp16 | 64.40 | 74.34 | 63.87 | 83.70 | 71.20 | 60.98 | 60.47 | 68.42 |
| BitNet 2B4T | 41.60 | 53.48 | 53.17 | 58.38 | 43.40 | 38.40 | 49.83 | 48.32 |
| 2bit | 66.00 | 71.94 | 58.98 | 70.36 | 42.20 | 46.95 | 52.58 | 58.43 |
Usage
1. Install Python Dependencies
pip install torch==2.9.0
pip install fast-hadamard-transform
2. Install HadaQuant CUDA Extension
HadaQuant CUDA Extension can be found at Our Github page. The HadaQuant package provides optimized CUDA kernels for 128-block Hadamard transforms and packed 2-bit dequantization:
cd hadaquant/csrc
pip install -e .
cd ../..
The kernels are compiled for SM 80/86/89/90 (A100, RTX 3090, RTX 4090, H100). Edit setup.py to add/remove GPU architectures as needed.
3: Inference Throughput Evaluation
After installing the HadaQuant CUDA extension:
#! /usr/bin/env bash
set -ue
export CUDA_VISIBLE_DEVICES=1 # Set visible GPUs, e.g., "0,1,2,3" for 4 GPUs or "0-7" for 8 GPUs
python -W ignore quantized_infer.py \
--model_path "YOUR_MODEL_PATH" \
--batch_size 1
Key Technical Details
Linear Codebook (LCVQ)
Traditional vector quantization uses a discrete codebook lookup, which requires a straight-through estimator (STE) to pass gradients. Our linear codebook (lib/codebook/index_codebook.py) replaces this with an orthogonal linear projection using a fixed orthogonal matrix. This achieves better performance than conventional 2-bit scalar methods while allowing end-to-end QAT training of all quantization weights.
Block-wise Hadamard Transform
The original method applies Hadamard transforms over the full input/output dimensions, which is expensive for inference. We constrain the transform to 128×128 blocks, enabling CUDA kernel fusion of the Hadamard transform with SU/SV scaling. This block-wise design also supports tensor parallelism for distributed training.
Differentiable Quantization (DGE)
During QAT training, we adopt the differentiable gradient estimator proposed by Ruizhe Wang et al. in Optimizing Large Language Model Training Using FP4 Quantization. We further introduce stochastic gradient masking to stabilize training, and enable end-to-end training of the codebook projection, scales (SU/SV), and latent weights.
Citation
If you use this code in your research, please cite:
@article{wang2026lcqat,
title = {LC-QAT: Data-Efficient 2-Bit QAT for LLMs via Linear-Constrained Vector Quantization},
author = {Wang, Haoyu and Yu, Xingyu and Zhao, Haiyan and Wang, Fengxiang and Han, Xu},
journal = {arXiv preprint arXiv:2606.10531},
year = {2026}
}
@article{wang2026unisvq,
title = {UniSVQ: 2-bit Unified Scalar-Vector Quantization},
author = {Wang, Haoyu and Zhao, Haiyan and Yu, Xingyu and Yao, Zhangyang and Han, Xu and Liu, Zhiyuan and Sun, Maosong},
journal = {arXiv preprint arXiv:2606.10520},
year = {2026}
}
- Downloads last month
- 498