壓縮張量

compressed-tensors 擴充套件了 safetensors 檔案，支援壓縮張量資料型別，為儲存和載入各種量化和稀疏格式（如稠密、整數量化（int8）、浮點量化（fp8）和打包量化（int4 或 int8 權重打包到 int32））提供統一的檢查點格式。

compressed-tensors 支援使用 PEFT 進行微調，並且還包括以下功能。

fp8、int4、int8 權重和啟用精度。
量化比例和零點策略，適用於張量、通道、組、塊、令牌。
動態逐令牌啟用量化（或任何靜態策略）。
權重稀疏性（非結構化或半結構化如 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

模型檢查點

壓縮張量模型透過其配置條目定義。以下示例取自 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

壓縮張量

模型檢查點