compressed-tensors

compressed-tensors は safetensors ファイルを拡張し、圧縮テンソルデータ型をサポートすることで、dense、int 量子化（int8）、float 量子化（fp8）、および pack 量子化（int4 または int8 の重み量子化を int32 にパックしたもの）といったさまざまな量子化形式の保存と読み込みのための統一されたチェックポイント形式を提供します。

compressed-tensors は PEFT によるファインチューニングをサポートしており、以下の機能も備えています。

fp8、int4、int8の重みおよび活性化の精度。
テンソル、チャネル、グループ、ブロック、トークン向けの量子化スケールおよびゼロポイント戦略。
トークンごとの動的アクティベーション量子化（または任意の静的戦略）。
nn.Linearモジュールに限らず、任意のモジュールの量子化。
名前またはクラスを指定して、特定のモジュールを対象にサポートします。

最新安定版（推奨）を取得するには、PyPIからcompressed-tensorsをインストールしてください。最新機能を利用するには、ソースからインストールしてください。

PyPI

ソースコード

Hugging Face Hub で互換性のあるモデルを検索するには、compressed-tensors の tag を使用してください。

事前量子化済みモデルは直接ロードできます。モデルを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 の 2 つの引数によって決まります。

設定	読み込み後の重み	実行
default	左圧縮	compressed-tensors がレイヤーを管理し、最初の forward pass でモデルをデコンプレスします
`dequantize=True`	モデルのdtype(例:BF16)にdequantizeされます	通常の稠密行列積であり、モデルはそのdtypeでファインチューニングや保存を行うことができます
`use_optimized_inference=True`	量子化されたまま	カーネルが実行されるスキームを持つレイヤー、現在は W8A8 fp8; 推論専用

FP8カーネルの高速化

use_optimized_inference=True を渡すことで、FP8 compressed-tensors モデルを FP8 のまま保持し、matmul をハードウェアアクセラレーション付きの FP8 カーネル (torch.nn.functional.scaled_mm、これは torch._scaled_mm_v2 にディスパッチされます。古い torch バージョンでは torch._scaled_mm にフォールバックします) 経由で実行します。これにより、重みを BF16 にデクアンタイズし直す必要がなくなります。推論全体で重みを FP8 のままにしておくことで、メモリ使用量を削減し、計算を高速化できます。これは推論専用なので、ファインチューニングを行う場合は無効にしてください。

デバイス	カーネル	注記
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 silicon (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 を指定して有効にし、モデルの config が 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 が最初の forward pass の際にそれらをデコンプレスします。モデルのファインチューニングや、元の精度（例: BF16）での保存を行いたい場合は、代わりに読み込み時にデクアンタイズするために dequantize=True を渡します。

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 への重みおよびアクティベーションの量子化を含みます。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 に引き渡されます。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 は最初の forward pass でモデル全体をデコンプレスします。一方、dequantize=True を指定すると、読み込み中にデコンプレスが行われるため、dense な重みを得るために forward pass を実行する必要はありません。

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 設定グループの対象となるレイヤーは、行方向の matmul カーネルが期待するレイアウトで fp8 重みとそのスケールを保持する CompressedTensorsFP8Linear に置き換えられます。これらの重みは、フォワードパスを含むすべての処理において 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