Buckets:

hf-doc-build
/

doc

Files

xet

hf-doc-build/doc / diffusers /main /ja /quantization /torchao.md

mishig

8 days ago

preview code

download

raw

10.8 kB

# 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)

Xet Storage Details

Size:: 10.8 kB
Xet hash:: 7c05d481366aff13f8cd4cc00a98aeb7b2304d9296fbac06e87c248c8304ef36

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	# 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)