AQLM

Official PyTorch implementation for Extreme Compression of Large Language Models via Additive Quantization

[2025.04] Released aqlm v1.1.7. Added support for arbitrary 8-dimensional codebooks on GPU, improved accuracy for 1-bit models, e.g. ISTA-DASLab/Llama-2-7b-AQLM-1Bit-1x8-hf at ~1 bit achieves WikiText 2 PPL 7.85. To quantize your own models this way, use num_codebooks=1, nbits_per_codebook=256 as per the tutorial below.

[2024.11] PV-tuning was accepted to NeurIPS'2024 for oral presentation!

[2024.05] AQLM was accepted to ICML'2024! If you're attending, meet us around this poster.

[2024.06] We released a new paper that extends AQLM with new finetuning algorithm called PV-tuning. We're also releasing PV-tuned AQLM models in this collection

[2024.08] We have merged the PV-Tuning branch into the main branch. To reproduce results with old finetuning (before Aug 21), use commit 559a366.

Inference

Demo

Learn how to run the prequantized models using this Google Colab examples:

| Basic AQLM

generation | Streaming with

GPU/CPU | Inference with CUDA

graphs (3x speedup) | Fine-tuning

with PEFT | Serving with

vLLM | |:-----------:|:-------:|:---------------:|:----------:|:--------:| | | | | | |

Browser demo (Rust/WASM)

If you want to try AQLM+PV inference on CPU directly in your browser, check out aqlm-rs:

• Live demo: galqiwi.github.io/aqlm-rs/about.html
• Source code: galqiwi/demo-aqlm-rs

Models

This repository is currently designed to work with models of LLaMA, Mistral and Mixtral families. The models reported below use full model fine-tuning as described in appendix A, with cross-entropy objective with teacher logits.

We provide a number of prequantized AQLM models without PV-Tuning (scroll down for PV-Tuned models):

You can also download AQLM models tuned via PV-tuning:

Note that models with "g16" in their scheme require aqlm inference library v1.1.6 or newer:

pip install aqlm[gpu,cpu]>=1.1.6

Above perplexity is evaluated on 4k context length for Llama 2 models and 8k for Mistral/Mixtral and Llama 3. Please also note that token-level perplexity can only be compared within the same model family, but should not be compared between models that use different vocabularies. While Mistral has a lower perplexity than Llama 3 8B but this does not mean that Mistral is better: Llama's perplexity is computed on a much larger dictionary and has higher per-token perplexity because of that.

For more evaluation results and detailed explanations, please see our papers: Egiazarian et al. (2024) for pure AQLM and Malinovskii et al. (2024) for PV-Tuned models.

Inference kernels

AQLM quantization setpus vary mainly on the number of codebooks used as well as the codebook sizes in bits. The most popular setups, as well as inference kernels they support are:

Installation

To run the models, one would have to install an inference library:

pip install aqlm[gpu,cpu]

, specifying either gpu, cpu or both based on one's inference setting.

Then, one can use the familiar .from_pretrained method provided by the transformers library:

from transformers import AutoModelForCausalLM

quantized_model = AutoModelForCausalLM.from_pretrained(
    "ISTA-DASLab/Llama-2-7b-AQLM-2Bit-1x16-hf",
    trust_remote_code=True, torch_dtype="auto"
).cuda()

Notice that torch_dtype should be set to either torch.float16 or "auto" on GPU and torch.float32 on CPU. After that, the model can be used exactly the same as one would use and unquantized model.

Quantization

Dependencies

Install packages from requirements.txt:

pip install -r requirements.txt

Loading / caching datasets and tokenizer

The script