Quanto

Quantoバックエンドは非推奨となり、バージョン1.0.0で削除されます。bitsandbytesやtorchaoなど、他のサポートされている量子化バックエンドへの切り替えをご検討ください。

Quantoは、Optimum向けのPyTorch量子化バックエンドです。多様性とシンプルさを念頭に設計されています：

すべての機能は eager モードで利用可能です（トレースできないモデルでも動作します）
量子化認識トレーニングをサポート
量子化モデルはtorch.compileと互換性があります
量子化モデルはデバイス非依存です（例: CUDA、XPU、MPS、CPU）。

Quantoバックエンドを使用するには、まずoptimum-quanto>=0.2.6とaccelerateをインストールする必要があります。

pip install optimum-quanto accelerate

QuantoConfig オブジェクトを from_pretrained() メソッドに渡すことで、モデルの量子化を行うことができます。Quanto ライブラリでは nn.Conv2d および nn.LayerNorm モジュールの量子化も可能ですが、現時点で Diffusers がサポートしているのは、モデルの nn.Linear レイヤーの重みの量子化のみです。以下のスニペットは、Quanto を用いて float8 量子化を適用する方法を示しています。

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

pipe = FluxPipeline.from_pretrained(model_id, transformer=transformer, dtype=dtype)
pipe.to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A cat holding a sign that says hello world"
image = pipe(
    prompt, num_inference_steps=50, guidance_scale=4.5, max_sequence_length=512
).images[0]
image.save("output.png")

特定のモジュールでのQuantizationのスキップ

QuantoConfig の modules_to_not_convert 引数を使用することで、特定のモジュールに対して量子化を適用しないようにすることができます。この引数に渡すモジュールが、state_dict 内のモジュールのキーと一致していることを確認してください。

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8", modules_to_not_convert=["proj_out"])
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)

Quanto バックエンドでの from_single_file の使用

QuantoConfig は ~FromOriginalModelMixin.from_single_file と互換性があります。

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

ckpt_path = "https://huggingface.co/black-forest-labs/FLUX.1-dev/blob/main/flux1-dev.safetensors"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_single_file(ckpt_path, quantization_config=quantization_config, dtype=torch.bfloat16)

量子化モデルの保存

Diffusers は、~ModelMixin.save_pretrained メソッドを使用して Quanto モデルのシリアライズをサポートしています。

Quanto ライブラリで直接量子化されたモデルと、Quanto をバックエンドとして使用して Diffusers で量子化されたモデルでは、シリアライズおよび読み込みの要件が異なります。現在、Quanto で直接量子化されたモデルを ~ModelMixin.from_pretrained を使用して Diffusers に読み込むことはできません。

import torch
from diffusers import FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="float8")
transformer = FluxTransformer2DModel.from_pretrained(
      model_id,
      subfolder="transformer",
      quantization_config=quantization_config,
      dtype=torch.bfloat16,
)
# save quantized model to reuse
transformer.save_pretrained("<your quantized model save path>")

# you can reload your quantized model with
model = FluxTransformer2DModel.from_pretrained("<your quantized model save path>")

Quanto と torch.compile を使用する場合

現在、Quantoバックエンドは、以下の量子化タイプに対してtorch.compileをサポートしています。

int8重み

import torch
from diffusers import FluxPipeline, FluxTransformer2DModel, QuantoConfig

model_id = "black-forest-labs/FLUX.1-dev"
quantization_config = QuantoConfig(weights_dtype="int8")
transformer = FluxTransformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
)
transformer = torch.compile(transformer, mode="max-autotune", fullgraph=True)

pipe = FluxPipeline.from_pretrained(
    model_id, transformer=transformer, dtype=dtype
)
pipe.to("cuda")  # or "mps", "xpu", "cpu"
images = pipe("A cat holding a sign that says hello").images[0]
images.save("flux-quanto-compile.png")

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

重み

float8
int8
int4
int2

Update on GitHub