# compressed-tensors

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

compressed-tensors は [PEFT](https://huggingface.co/docs/peft) によるファインチューニングをサポートしており、以下の機能も備えています。

- fp8、int4、int8の重みおよび活性化の精度。
- [テンソル、チャネル、グループ、ブロック、トークン](https://github.com/neuralmagic/compressed-tensors/blob/83b2e7a969d70606421a76b9a3d112646077c8de/src/compressed_tensors/quantization/quant_args.py#L43-L52)向けの量子化スケールおよびゼロポイント戦略。
- トークンごとの動的アクティベーション量子化（または任意の静的戦略）。
- [nn.Linear](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html)モジュールに限らず、任意のモジュールの量子化。
- 名前またはクラスを指定して、特定のモジュールを対象にサポートします。

最新安定版（推奨）を取得するには、[PyPI](https://pypi.org/project/compressed-tensors)からcompressed-tensorsをインストールしてください。最新機能を利用するには、ソースからインストールしてください。

```bash
pip install compressed-tensors
```

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

Hugging Face Hub で互換性のあるモデルを検索するには、compressed-tensors の [tag](https://huggingface.co/models?other=compressed-tensors) を使用してください。

事前量子化済みモデルは直接ロードできます。モデルをcompressed-tensors形式に量子化するには、[llm-compressor](https://github.com/vllm-project/llm-compressor)を参照してください。また、モデルを独立して作成し、compressed-tensors設定でシリアライズすることもできます。

```python
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](/docs/transformers/main/ja/main_classes/quantization#transformers.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](https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_mm.html)、これは `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](https://huggingface.co/RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8-dynamic) |
| テンソルごとの静的 | [RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8](https://huggingface.co/RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8) |

### 事前量子化済みFP8モデルの読み込み

FP8カーネルはオプトイン方式です。`use_optimized_inference=True`を指定して要求すると、モデルのconfigがFP8量子化を指定しており、サポートされているGPUが利用可能な場合に使用されます。

```python
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` を渡します。

```python
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](https://huggingface.co/nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf/blob/main/config.json)の`config.json`ファイルから抜粋したものです。

圧縮中および圧縮後の柔軟な表現を可能にするため、多くのエントリが存在しますが、読み込みと推論のエントリは、いくつかの主要なエントリに焦点を当てることで簡素化できます。

```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](https://huggingface.co/nm-testing/Meta-Llama-3.1-8B-Instruct-FP8-hf?show_file_info=model.safetensors.index.json)を使用して、すべての[nn.Linear](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html)モジュールの量子化された重み、入力スケール、および重みスケールを確認できます。

| テンソル | 形状 | 精度 |
| ------- | ----- | --------- |
|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](https://pytorch.org/docs/stable/generated/torch.nn.Linear.html) のインスタンスのままなので、`print` にはそれが表示されます。最近の compressed-tensors のバージョンでは、それらを `CompressedLinear` サブクラスでラップしなくなりました。`ignore` に列挙されたモジュール（`lm_head` など）は変更されません。

`dequantize=False`（デフォルト）の場合、読み込みが完了しても重みは圧縮されたままです。compressed-tensors は最初の forward pass でモデル全体をデコンプレスします。一方、`dequantize=True` を指定すると、読み込み中にデコンプレスが行われるため、dense な重みを得るために forward pass を実行する必要はありません。

```python
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 のまま保持されます。

```python
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
```