torchao

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

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

uv pip install -U torch torchao

各量子化 dtype は、AOBaseConfig クラスの個別のインスタンスとして利用可能です。これにより、利用可能な引数がより多く公開されるため、より柔軟な設定オプションが提供されます。

AOBaseConfig を from_pretrained() の TorchAoConfig に渡します。量子化データ型として、Int4WeightOnlyConfig のようなものを指定します。

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 に移動してください。

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()を使用してください。また、テキストエンコーダーなどの追加コンポーネントも量子化できます。

torch.compile

torchao は torch.compile をサポートしており、1行のコードで推論を高速化できます。

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における推論速度とメモリ使用量のベンチマークについては、この表を参照してください。各種ハードウェアに関するさらなるベンチマークは、torchaoのリポジトリでも利用できます。

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の公式ドキュメントを参照してください。設定オプションの完全なリストはこちらで確認できます。

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

カテゴリ	設定クラス
整数量子化	`Int4WeightOnlyConfig`, `Int8WeightOnlyConfig`, `Int8DynamicActivationInt8WeightConfig`
浮動小数点8ビット量子化	`Float8WeightOnlyConfig`, `Float8DynamicActivationFloat8WeightConfig`
符号なし整数量子化	`IntxWeightOnlyConfig`

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

指定された dtype で量子化モデルをシリアライズするには、まず希望する量子化 dtype でモデルを読み込み、save_pretrained() メソッドを使用して保存します。

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()メソッドを使用します。

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 を指定する必要があるため、重みが信頼できるソースから取得された場合のみ実行してください。

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)

以下の例に示すように、AutoModel API は PyTorch >= 2.6 でサポートされています。

リソース

Update on GitHub