# torchao

[torchao](https://github.com/pytorch/ao) は、PyTorch モデルの推論と学習向けに、量子化とスパース性に基づく高性能な dtype と最適化を提供します。[Accelerate](https://hf.co/docs/accelerate/index) による読み込みに対応し、`torch.nn.Linear` レイヤーを含むモデルであれば、モダリティを問わずサポートされています。

以下のコマンドで、Pytorch 2.5以上とtorchaoがインストールされていることを確認してください。

```bash
uv pip install -U torch torchao
```

各量子化 dtype は、[AOBaseConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.core.config.AOBaseConfig.html) クラスの個別のインスタンスとして利用可能です。これにより、利用可能な引数がより多く公開されるため、より柔軟な設定オプションが提供されます。

`AOBaseConfig` を [from_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained) の [TorchAoConfig](/docs/diffusers/main/ja/api/quantization#diffusers.TorchAoConfig) に渡します。量子化データ型として、[Int4WeightOnlyConfig](https://docs.pytorch.org/ao/main/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html) のようなものを指定します。

```py
import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int8WeightOnlyConfig(group_size=128, version=2))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)
```

`device_map="cuda"` を指定すると、各レイヤーがロードされる際に GPU 上で量子化されます。この方法は高速ですが、元の重みと量子化された重みの両方を一時的に GPU メモリに保持するため、追加の GPU メモリが必要です。モデルがすでに GPU メモリの大部分を使用している場合、メモリ不足エラーによりロードが失敗する可能性があります。その場合は、`device_map` を省略し、ロード後にパイプラインを GPU に移動してください。

```py
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    torch_dtype=torch.bfloat16,
)
pipeline.to("cuda")  # or "mps", "xpu", "cpu"
```

`device_map`を指定しない場合、DiffusersはCPU上でレイヤーを量子化します。これにより速度は低下しますが、量子化中の一時的なGPUメモリ使用量の急増を回避できます。GPUメモリ使用量をさらに削減するには、代わりに[enable_model_cpu_offload()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.enable_model_cpu_offload)を使用してください。また、テキストエンコーダーなどの追加コンポーネントも量子化できます。

## torch.compile

torchao は [torch.compile](../optimization/fp16#torchcompile) をサポートしており、1行のコードで推論を高速化できます。

```python
import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, TorchAoConfig
from torchao.quantization import Int4WeightOnlyConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": TorchAoConfig(Int4WeightOnlyConfig(group_size=128))}
)
pipeline = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-dev",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda"  # or "mps", "xpu", "cpu"
)

pipeline.transformer.compile(transformer, mode="max-autotune", fullgraph=True)
```

FluxおよびCogVideoXにおける推論速度とメモリ使用量のベンチマークについては、この[表](https://github.com/huggingface/diffusers/pull/10009#issue-2688781450)を参照してください。各種ハードウェアに関するさらなるベンチマークは、torchaoの[リポジトリ](https://github.com/pytorch/ao/tree/main/torchao/quantization#benchmarks)でも利用できます。

> [!TIP]
> torchao の FP8 後学習量子化スキームは、計算能力が 8.9 以上の GPU（RTX-4090、Hopper など）で有効です。画像や動画の生成において、FP8 は速度、メモリ、品質のバランスにおいて最も優れたトレードオフを提供することがよくあります。GPU が対応している場合は、FP8 と torch.compile の組み合わせを推奨します。

## サポートされている量子化タイプ

torchao は、int8、float3-float8、および uint1-uint7 における、重みのみの量子化と、重みおよび動的アクティベーションの量子化をサポートしています。

Weight-only quantization は、モデルの重みを特定の低ビットのデータ型で保存し、`bfloat16` などの高精度なデータ型で計算を実行します。これにより、モデルの重みによるメモリ要件は低減されますが、アクティベーションの計算に伴うメモリピークは維持されます。

Dynamic activation quantization は、モデルの重みを低ビットの dtype で保存すると同時に、アクティベーションをオンザフライで量子化することで、さらなるメモリを節約します。これにより、モデルの重みによるメモリ要件が低下するだけでなく、アクティベーションの計算に伴うメモリオーバーヘッドも低減されます。ただし、場合によっては品質とのトレードオフが生じる可能性があるため、異なるモデルを徹底的にテストすることをお勧めします。

利用可能な量子化手法について詳しく理解するには、[torchaoの公式ドキュメント](https://docs.pytorch.org/ao/stable/index.html)を参照してください。設定オプションの完全なリストは[こちら](https://docs.pytorch.org/ao/main/workflows/inference.html#inference-workflows)で確認できます。

以下に、よく使われる量子化設定の例を示します。

| **カテゴリ** | **設定クラス** |
|---|---|
| **整数量子化** | [`Int4WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int4WeightOnlyConfig.html), [`Int8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8WeightOnlyConfig.html), [`Int8DynamicActivationInt8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Int8DynamicActivationInt8WeightConfig.html) |
| **浮動小数点8ビット量子化** | [`Float8WeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8WeightOnlyConfig.html), [`Float8DynamicActivationFloat8WeightConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.Float8DynamicActivationFloat8WeightConfig.html) |
| **符号なし整数量子化** | [`IntxWeightOnlyConfig`](https://docs.pytorch.org/ao/stable/api_reference/generated/torchao.quantization.IntxWeightOnlyConfig.html) |

## 量子化モデルのシリアライズとデシリアライズ

指定された dtype で量子化モデルをシリアライズするには、まず希望する量子化 dtype でモデルを読み込み、[save_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.save_pretrained) メソッドを使用して保存します。

```python
import torch
from diffusers import AutoModel, TorchAoConfig
from torchao.quantization import Int8WeightOnlyConfig

quantization_config = TorchAoConfig(Int8WeightOnlyConfig())
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_int8wo", safe_serialization=False)
```

シリアライズされた量子化モデルを読み込むには、[from_pretrained()](/docs/diffusers/main/ja/api/models/overview#diffusers.ModelMixin.from_pretrained)メソッドを使用します。

```python
import torch
from diffusers import FluxPipeline, AutoModel

transformer = AutoModel.from_pretrained("/path/to/flux_int8wo", dtype=torch.bfloat16, use_safetensors=False)
pipe = FluxPipeline.from_pretrained("black-forest-labs/Flux.1-Dev", transformer=transformer, dtype=torch.bfloat16)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A cat holding a sign that says hello world"
image = pipe(prompt, num_inference_steps=30, guidance_scale=7.0).images[0]
image.save("output.png")
```

`torch<=2.6.0` を使用している場合、`uint4` の weight-only などの一部の量子化手法では、モデルの読み込みに際して直接読み込めず、`UnpicklingError` が発生する可能性があります。一方、保存時には正常に動作します。この問題を回避するには、state dict をモデルに手動で読み込む方法があります。ただし、これには `torch.load` で `weights_only=False` を指定する必要があるため、重みが信頼できるソースから取得された場合のみ実行してください。

```python
import torch
from accelerate import init_empty_weights
from diffusers import FluxPipeline, AutoModel, TorchAoConfig
from torchao.quantization import IntxWeightOnlyConfig

# Serialize the model
transformer = AutoModel.from_pretrained(
    "black-forest-labs/Flux.1-Dev",
    subfolder="transformer",
    quantization_config=TorchAoConfig(IntxWeightOnlyConfig(dtype=torch.uint4)),
    dtype=torch.bfloat16,
)
transformer.save_pretrained("/path/to/flux_uint4wo", safe_serialization=False, max_shard_size="50GB")
# ...

# Load the model
state_dict = torch.load("/path/to/flux_uint4wo/diffusion_pytorch_model.bin", weights_only=False, map_location="cpu")
with init_empty_weights():
    transformer = AutoModel.from_config("/path/to/flux_uint4wo/config.json")
transformer.load_state_dict(state_dict, strict=True, assign=True)
```

> [!TIP]
> 以下の例に示すように、[AutoModel](/docs/diffusers/main/ja/api/models/auto_model#diffusers.AutoModel) API は PyTorch >= 2.6 でサポートされています。

## リソース

- [TorchAO 量子化 API](https://docs.pytorch.org/ao/stable/index.html)
- [Diffusers-TorchAOの例](https://github.com/sayakpaul/diffusers-torchao)