Transformers documentation
compressed-tensors
compressed-tensors
compressed-tensors 扩展了 safetensors 文件，以支持压缩张量数据类型，从而为存储和加载各种量化格式（如 dense、int 量化（int8）、float 量化（fp8）以及 pack 量化（将 int4 或 int8 权重量化打包为 int32））提供统一的检查点格式。
compressed-tensors 支持使用 PEFT 进行微调，并包含以下功能。
- fp8、int4、int8 权重和激活精度。
- 针对 tensor, channel, group, block, token 的量化缩放因子和零点策略。
- 动态逐 token 激活量化（或任何静态策略）。
- 对任意模块进行量化，而不仅仅是 nn.Linear 模块。
- 按名称或类对特定模块进行针对性支持。
从 PyPI 安装 compressed-tensors 以获取最新的稳定版本（推荐），或从源代码安装以获取最新功能。
pip install compressed-tensors
git clone https://github.com/neuralmagic/compressed-tensors
cd compressed-tensors
pip install -e .使用压缩张量 标签 在 Hugging Face Hub 上搜索兼容的模型。
预量化模型可以直接加载。若要将模型量化为 compressed-tensors 格式，请参阅 llm-compressor。或者，可以独立创建模型，并使用 compressed-tensors 配置进行序列化。
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加载模式
compressed-tensors 检查点以压缩形式（fp8，或打包的 int4/int8）存储其权重。其执行方式取决于 CompressedTensorsConfig 中的两个参数。
| 配置 | 加载后的权重 | 执行 |
|---|---|---|
| default | 以 compressed-tensors 格式保存：每个序列化张量（weight_packed、weight_scale、weight_zero_point、……）均按存储方式加载 | compressed-tensors 负责管理各层，并在首次前向传播时解压模型 |
dequantize=True | 在加载到模型数据类型（例如 BF16）时进行解压缩 | 一个普通的稠密模型，没有保留任何quantization_scheme，因此 w8a8 激活值也不会被量化；可以以该 dtype 对其进行微调或保存 |
use_optimized_inference=True | 保持量化 | 其方案中包含内核执行的层，目前为 W8A8 fp8；仅用于推理 |
默认情况下，save_pretrained() 会将压缩张量原样写回，因此加载的模型会以相同的 compressed-tensors 格式保存。
FP8 内核加速
传入 use_optimized_inference=True 可让 FP8 compressed-tensors 模型保持 FP8 格式，并通过硬件加速的 FP8 内核（torch.nn.functional.scaled_mm，它会分发到 torch._scaled_mm_v2；较旧版本的 torch 会回退到 torch._scaled_mm）执行矩阵乘法，而不是将权重反量化回 BF16。在整个推理过程中保持权重为 FP8 可以降低内存占用并加快计算速度。此选项仅适用于推理，因此在进行微调时应将其关闭。
| 设备 | Kernel | 备注 |
|---|---|---|
| Intel XPU | torch.nn.functional.scaled_mm | 所有支持 FP8 的 XPU 设备 |
| NVIDIA CUDA (SM89+) | torch.nn.functional.scaled_mm | Ada Lovelace（L4、L40）、Hopper（H100）、Blackwell 及更新架构 |
| Apple 芯片（MPS） | torch.nn.functional.scaled_mm | 需要 torch >= 2.15 |
| CPU / CUDA SM80 (A100) / MPS 搭配 torch < 2.15 | 回退 | use_optimized_inference=True 会被忽略，模型将以反量化后的状态运行 |
FP8 内核路径支持这些量化布局。
| 策略 | 示例模型 |
|---|---|
| 按通道动态 | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic |
| 按张量静态 | RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8 |
加载预量化的 FP8 模型
FP8 内核是可选启用的：通过设置 use_optimized_inference=True 来请求使用它们，当模型的配置指定了 FP8 量化且存在受支持的 GPU 时，这些内核将被使用。
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))在加载时进行反量化
如果不设置 use_optimized_inference=True，模型将走常规的 compressed-tensors 路径：权重保持压缩状态，由 compressed-tensors 在第一次前向传播时进行解压缩。如果传入 dequantize=True，则会在加载阶段完成解压缩，这在需要对模型进行微调或以原始精度（例如 BF16）保存模型时是推荐的做法。最终得到的是一个普通的稠密模型：由于不再附加 compressed-tensors 方案，激活量化（例如 w8a8）也会被移除。
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",
)模型检查点
Compressed-tensor 模型通过其配置项进行定义。以下示例取自 nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf 的 config.json 文件。
为了在压缩期间和压缩之后都能实现灵活的表达，配置项有很多，但加载和推理相关的配置项可以简化，只需关注几个关键项即可。
"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"
},配置文件指定了配置组（group_0）的量化方式，其中包括将权重和激活值量化为 fp8，并采用静态的 per-tensor 策略。如 ignore 键所示，lm_head 模块未进行量化。
如需更详细地查看模型权重，请使用模型卡片上的 safetensors viewer 来查看所有 nn.Linear 模块的量化权重、输入缩放和权重缩放。
| 张量 | 形状 | 精度 |
|---|---|---|
| model.layers.0.input_layernorm.weight | [4 096] | BF16 |
| model.layers.0.mlp.down_proj.input_scale | [1] | BF16 |
| model.layers.0.mlp.down_proj.weight | [4 096, 14 336] | F8_E4M3 |
| model.layers.0.mlp.down_proj.weight_scale | [1] | BF16 |
| model.layers.0.mlp.gate_proj.input_scale | [1] | BF16 |
| model.layers.0.mlp.gate_proj.weight | [14 336, 4 096] | F8_E4M3 |
| model.layers.0.mlp.gate_proj.weight_scale | [1] | BF16 |
| model.layers.0.mlp.up_proj.input_scale | [1] | BF16 |
| model.layers.0.mlp.up_proj.weight | [14 336, 4 096] | F8_E4M3 |
| model.layers.0.mlp.up_proj.weight_scale | [1] | BF16 |
| model.layers.0.post_attention_layernorm.weight | [4 096] | BF16 |
| model.layers.0.self_attn.k_proj.input_scale | [1] | BF16 |
| model.layers.0.self_attn.k_proj.weight | [1 024, 4 096] | F8_E4M3 |
| model.layers.0.self_attn.k_proj.weight_scale | [1] | BF16 |
| model.layers.0.self_attn.o_proj.input_scale | [1] | BF16 |
| model.layers.0.self_attn.o_proj.weight | [4 096, 4 096] | F8_E4M3 |
| model.layers.0.self_attn.o_proj.weight_scale | [1] | BF16 |
| model.layers.0.self_attn.q_proj.input_scale | [1] | BF16 |
| model.layers.0.self_attn.q_proj.weight | [4 096, 4 096] | F8_E4M3 |
| model.layers.0.self_attn.q_proj.weight_scale | [1] | BF16 |
| model.layers.0.self_attn.v_proj.input_scale | [1] | BF16 |
| model.layers.0.self_attn.v_proj.weight | [1 024, 4 096] | F8_E4M3 |
| model.layers.0.self_attn.v_proj.weight_scale | [1] | BF16 |
使用 ~quantizers.HFQuantizer 集成加载 compressed-tensors 模型时，目标模块会被移交给 compressed-tensors：它会附加已解析的 quantization_scheme，设置 quantization_status，注册检查点中存储的参数（以 fp8 存储的 weight，以及 weight_scale，对于静态策略还包括 input_scale），并在这些参数上安装其自身的 forward pass。这些模块仍然是 nn.Linear 实例，因此 print 显示的就是它们——较新版本的 compressed-tensors 不再将它们包装在 CompressedLinear 子类中。列在 ignore 下的模块（例如 lm_head）将保持原样。
当 dequantize=False（默认值）时，加载完成后权重仍处于压缩状态，compressed-tensors 会在第一次前向传播时解压整个模型。dequantize=True 则在加载每个权重时立即对其进行解压，因此无需前向传播即可获得稠密权重，且这些模块是普通的稠密层，不包含 quantization_scheme。
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当 use_optimized_inference=True 时，被 fp8 配置组覆盖的层会被替换为 CompressedTensorsFP8Linear，该层以行式矩阵乘法内核所期望的布局保存 fp8 权重及其缩放因子。这些权重始终保持为 fp8 格式，前向传播也不例外。
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