compressed-tensors

compressed-tensors 扩展了 safetensors 文件，支持压缩张量数据类型，提供统一的检查点格式，用于存储和加载各种量化和稀疏格式，例如密集、整数量化（int8）、浮点量化（fp8）以及打包量化（int4 或 int8 权重量化打包成 int32）。

compressed-tensors 支持使用 PEFT 进行微调，并且还包含以下功能。

fp8、int4、int8 权重和激活精度。
针对张量、通道、组、块、token 的量化尺度和零点策略。
动态的每 token 激活量化（或任何静态策略）。
权重稀疏性（非结构化或半结构化，如 2:4）可以与量化结合，实现极致压缩。
量化任意模块，而不仅仅是 nn.Linear 模块。
通过名称或类进行目标模块支持。

从 PyPI 安装 compressed-tensors 以获取最新的稳定版本（推荐），或者从源代码安装以获取最新功能。

PyPI

源代码

使用 compressed-tensors 标签在 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-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 进行权重和激活量化。lm_head 模块未进行量化，如 ignore 键所示。

要更详细地查看模型权重，请使用模型卡上的 safetensors 查看器，查看所有 nn.Linear 模块的量化权重、输入尺度和权重尺度。

张量（Tensors）	形状	精度
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 模型时，量化配置中指定的所有 nn.Linear 模块都将被 CompressedLinear 模块替换，这些模块负责管理压缩权重和推理的前向传播。lm_head 模块仍然保留为未量化的 nn.Linear 模块。

from transformers import AutoModelForCausalLM

ct_model = AutoModelForCausalLM.from_pretrained("nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf")
print(ct_model)
"""
LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 4096)
    (layers): ModuleList(
      (0-31): 32 x LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): CompressedLinear(
            in_features=4096, out_features=4096, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (k_proj): CompressedLinear(
            in_features=4096, out_features=1024, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (v_proj): CompressedLinear(
            in_features=4096, out_features=1024, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (o_proj): CompressedLinear(
            in_features=4096, out_features=4096, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (rotary_emb): LlamaRotaryEmbedding()
        )
        (mlp): LlamaMLP(
          (gate_proj): CompressedLinear(
            in_features=4096, out_features=14336, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (up_proj): CompressedLinear(
            in_features=4096, out_features=14336, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (down_proj): CompressedLinear(
            in_features=14336, out_features=4096, bias=False
            (input_observer): MovingAverageMinMaxObserver()
            (weight_observer): MovingAverageMinMaxObserver()
          )
          (act_fn): SiLU()
        )
        (input_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
        (post_attention_layernorm): LlamaRMSNorm((4096,), eps=1e-05)
      )
    )
    (norm): LlamaRMSNorm((4096,), eps=1e-05)
    (rotary_emb): LlamaRotaryEmbedding()
  )
  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)
)
"""

在 GitHub 上更新

Transformers

compressed-tensors

模型检查点