compressed-tensors

compressed-tensors extends safetensors files to compressed tensor data types to provide a unified checkpoint format for storing and loading various quantization formats such as dense, int-quantized (int8), float-quantized (fp8), and pack-quantized (int4 or int8 weight-quantized packed into int32).

compressed-tensors supports fine-tuning with PEFT and includes the following features as well.

• fp8, int4, int8 weight and activation precisions.
• Quantization scales and zero-points strategies for tensor, channel, group, block, token.
• Dynamic per-token activation quantization (or any static strategy).
• Quantization of arbitrary modules, not just nn.Linear modules.
• Targeted support for specific modules by name or class.

Install compressed-tensors from PyPI to get the latest stable release (recommended) or install it from source to get the latest features.

pip install compressed-tensors

git clone https://github.com/neuralmagic/compressed-tensors
cd compressed-tensors
pip install -e .

Search using the compressed-tensors tag to find a compatible model on the Hugging Face Hub.

Pre-quantized models can be loaded directly. To quantize a model into the compressed-tensors format, see llm-compressor. Alternatively, models can be created independently and serialized with a compressed-tensors config.

from transformers import AutoModelForCausalLM

ct_model = AutoModelForCausalLM.from_pretrained("nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf", device_map="auto")

# measure memory usage
mem_params = sum([param.nelement()*param.element_size() for param in ct_model.parameters()])
print(f"{mem_params/2**30:.4f} GB")
# 8.4575 GB

Loading modes

A compressed-tensors checkpoint stores its weights compressed (fp8, or packed int4/int8). How they are executed is up to two [CompressedTensorsConfig] arguments.

FP8 kernel acceleration

Pass use_optimized_inference=True to keep an FP8 compressed-tensors model in FP8 and run its matmuls through hardware-accelerated FP8 kernels (torch.nn.functional.scaled_mm, which dispatches to torch._scaled_mm_v2; older torch versions fall back to torch._scaled_mm), instead of dequantizing the weights back to BF16. Keeping weights in FP8 throughout inference lowers memory usage and speeds up computation. This is inference only, so leave it off to fine-tune.

The FP8 kernel path supports these quantization layouts.

Loading a pre-quantized FP8 model

The FP8 kernels are opt-in: ask for them with use_optimized_inference=True, and they are used when the model's config specifies FP8 quantization and a supported GPU is available.

from transformers import AutoModelForCausalLM, AutoTokenizer, CompressedTensorsConfig

model = AutoModelForCausalLM.from_pretrained(
    "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic",
    quantization_config=CompressedTensorsConfig(use_optimized_inference=True),
    device_map="auto",
)
tokenizer = AutoTokenizer.from_pretrained("RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic")
inputs = tokenizer("Hello, how are you?", return_tensors="pt").to(model.device)
outputs = model.generate(**inputs, max_new_tokens=20)
print(tokenizer.decode(outputs[0], skip_special_tokens=True))

Dequantizing at load time

Without use_optimized_inference=True, the model takes the regular compressed-tensors route: the weights are left compressed and compressed-tensors decompresses them on the first forward pass. Pass dequantize=True to dequantize them during loading instead, which is what you want to fine-tune the model or save it in its original precision (e.g. BF16).

from transformers import AutoModelForCausalLM, CompressedTensorsConfig

model = AutoModelForCausalLM.from_pretrained(
    "RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic",
    quantization_config=CompressedTensorsConfig(dequantize=True),
    device_map="auto",
)

Model checkpoint

Compressed-tensor models are defined through its configuration entry. The following example is taken from the nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf config.json file.

There are a lot of entries to allow for flexible expression both during and after compression, but the entries for loading and inference can be simplified to focus on just a few key entries.

"quantization_config": {
  "config_groups": {
    "group_0": {
      "input_activations": {
        "num_bits": 8,
        "strategy": "tensor",
        "type": "float"
      },
      "targets": ["Linear"],
      "weights": {
        "num_bits": 8,
        "strategy": "tensor",
        "type": "float"
      }
    }
  },
  "format": "naive-quantized",
  "ignore": ["lm_head"],
  "quant_method": "compressed-tensors",
  "quantization_status": "frozen"
},

The config file specifies the quantization of a config group (group_0), which includes weight and activation quantization to fp8 with a static per-tensor strategy. The lm_head module is unquantized as shown in the ignore key.

For a more detailed look at the model weights, use the safetensors viewer on the model card to see the quantized weights, input scale, and weight scale for all nn.Linear modules.

When loading a compressed-tensors model with the [~quantizers.HFQuantizer] integration, the targeted modules are handed over to compressed-tensors: it attaches the resolved quantization_scheme, sets quantization_status, registers the parameters the checkpoint stores (weight in fp8, plus weight_scale and, for a static strategy, input_scale) and installs its own forward pass over them. They stay nn.Linear instances, so that is what print shows — recent compressed-tensors versions no longer wrap them in a CompressedLinear subclass. Modules listed under ignore, such as lm_head, are left untouched.

With dequantize=False (the default), the weights are still compressed once loading is over, and compressed-tensors decompresses the whole model on the first forward pass. dequantize=True does it during loading instead, so no forward pass is needed to get dense weights.

import torch
from transformers import AutoModelForCausalLM, CompressedTensorsConfig

model_id = "nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf"

ct_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=CompressedTensorsConfig(dequantize=False),
    device_map="auto",
)
q_proj = ct_model.model.layers[0].self_attn.q_proj
print(q_proj, q_proj.quantization_status)
# Linear(in_features=4096, out_features=4096, bias=False) QuantizationStatus.COMPRESSED
# ^ compressed-tensors module: fp8 weight, weight_scale, and its own forward

ct_model(input_ids=torch.tensor([[0, 1, 2]], device=ct_model.device))
print(q_proj, q_proj.quantization_status)
# Linear(in_features=4096, out_features=4096, bias=False) QuantizationStatus.DECOMPRESSED
# ^ weight is BF16 now, decompressed by that forward pass

ct_model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=CompressedTensorsConfig(dequantize=True),
    device_map="auto",
)
print(ct_model.model.layers[0].self_attn.q_proj)
# Linear(in_features=4096, out_features=4096, bias=False)      weight: BF16

With use_optimized_inference=True, the layers covered by an fp8 config group are replaced by CompressedTensorsFP8Linear, which holds the fp8 weight and its scale in the layout its row-wise matmul kernel expects. Those weights stay in fp8, forward passes included.

from transformers import AutoModelForCausalLM, CompressedTensorsConfig

ct_model = AutoModelForCausalLM.from_pretrained(
    "nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf",
    quantization_config=CompressedTensorsConfig(use_optimized_inference=True),
    device_map="auto",
)
print(ct_model.model.layers[0].self_attn.q_proj)
# CompressedTensorsFP8Linear(in_features=4096, out_features=4096, bias=False)      weight: F8_E4M3