はじめる

Quantizationは、元のデータの精度を維持しつつ、より少ないビット数でデータを表現することに焦点を当てています。これは通常、同じ情報をより少ないビット数で表現するためにデータ型を変換することを意味します。例えば、モデルの重みが32ビット浮動小数点数として保存されている場合、これを16ビット浮動小数点数にQuantizationすると、モデルサイズが半分になり、保存が容易になり、メモリ使用量が削減されます。精度が低くなることで、より少ないビット数での計算に要する時間が短くなるため、推論の高速化にもつながります。

Diffusers は、Flux などの大規模な拡散モデルをより利用しやすくするために、複数の量子化バックエンドをサポートしています。このガイドでは、事前学習済みまたは非量子化のチェックポイントからパイプラインを初期化する際に、PipelineQuantizationConfig クラスを使用してパイプラインを量子化する方法を示します。

パイプラインレベルの量子化

量子化設定にどの程度のカスタマイズを適用したいかによって、PipelineQuantizationConfigの使用方法は2つあります。

• 基本的なユースケースでは、quant_backend、quant_kwargs、components_to_quantize の引数を定義します。
• 詳細な量子化制御を行うには、quant_mappingを定義して、個々のモデルコンポーネントに対する量子化設定を指定します。

基本的な量子化

PipelineQuantizationConfig を以下のパラメータで初期化します。

• quant_backend は使用する量子化バックエンドを指定します。現在サポートされているバックエンドには、bitsandbytes_4bit、bitsandbytes_8bit、gguf、quanto、torchao が含まれます。
• quant_kwargs は、使用する量子化引数を指定します。

これらの quant_kwargs 引数は、バックエンドごとに異なります。各バックエンドの引数については、Quantization API のドキュメントを参照してください。

• components_to_quantize は、パイプラインのどのコンポーネントを量子化するかを指定します。通常、transformer などの計算負荷が最も大きいコンポーネントを量子化すべきです。FluxPipeline のように複数のテキストエンコーダーを持つパイプラインでは、テキストエンコーダーも量子化の対象として検討すべきです。以下の例では、FluxPipeline の T5 テキストエンコーダーを量子化し、CLIP モデルはそのままにしています。

  components_to_quantize は、複数のモデルを指定する場合はリストを、単一のモデルを指定する場合は文字列を受け取ります。

以下の例では、BitsAndBytesConfigから load_in_4bit、bnb_4bit_quant_type、bnb_4bit_compute_dtype という引数を指定して、bitsandbytes バックエンドを読み込みます。

import torch
from diffusers import DiffusionPipeline
from diffusers.quantizers import PipelineQuantizationConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_backend="bitsandbytes_4bit",
    quant_kwargs={"load_in_4bit": True, "bnb_4bit_quant_type": "nf4", "bnb_4bit_compute_dtype": torch.bfloat16},
    components_to_quantize=["transformer", "text_encoder_2"],
)

パイプラインを量子化するには、pipeline_quant_config を from_pretrained() に渡します。

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

image = pipe("photo of a cute dog").images[0]

高度な量子化

quant_mapping 引数により、パイプライン内の各コンポーネントの量子化方法をより詳細に指定できます。異なる量子化バックエンドを組み合わせることも可能です。

PipelineQuantizationConfig を初期化し、quant_mapping を渡します。quant_mapping を使用すると、transformer や text encoder など、パイプライン内の各コンポーネントに対する量子化オプションを指定できます。

以下の例では、transformerとテキストエンコーダーにそれぞれ2つの量子化バックエンド、QuantoConfigとtransformers.BitsAndBytesConfigを使用しています。

import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers.quantization_config import QuantoConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": QuantoConfig(weights_dtype="int8"),
        "text_encoder_2": TransformersBitsAndBytesConfig(
            load_in_4bit=True, compute_dtype=torch.bfloat16
        ),
    }
)

Transformersには別のbitsandbytesバックエンドがあります。Transformers由来のコンポーネントには、transformers.BitsAndBytesConfigをインポートして使用する必要があります。例えば、FluxPipeline内のtext_encoder_2はTransformersのT5EncoderModelであるため、diffusers.BitsAndBytesConfigではなくtransformers.BitsAndBytesConfigを使用する必要があります。

これらの個別のインポートを管理したくない場合、または各パイプラインコンポーネントの由来が不明な場合は、上記の基本的な量子化メソッドを使用してください。

import torch
from diffusers import DiffusionPipeline
from diffusers import BitsAndBytesConfig as DiffusersBitsAndBytesConfig
from diffusers.quantizers import PipelineQuantizationConfig
from transformers import BitsAndBytesConfig as TransformersBitsAndBytesConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={
        "transformer": DiffusersBitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16),
        "text_encoder_2": TransformersBitsAndBytesConfig(
            load_in_4bit=True, compute_dtype=torch.bfloat16
        ),
    }
)

パイプラインを量子化するには、pipeline_quant_config を from_pretrained() に渡します。

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

image = pipe("photo of a cute dog").images[0]

事前量子化済みチェックポイントの読み込み

一部の量子化バックエンドは、すでに量子化されて Hub にアップロードされたチェックポイントの読み込みをサポートしています。この場合、量子化設定はモデルの config.json に保存され、自動的に読み込まれます。そのため、PipelineQuantizationConfig を作成したり、quantization_config 引数を渡したりする必要はありません。

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

以下のバックエンドは、プリクオンタイズ済みチェックポイントの読み込みをそのままサポートしています。

リソース

量子化について詳しく学ぶには、以下のリソースをご覧ください。

• 量子化が初めての方は、DeepLearning.AI との協力により提供されている、初心者向けの以下のコースを参照することをお勧めします。

  • Hugging Faceによる量子化の基礎
  • 量子化の深掘り

• 新しい量子化手法の追加にご興味がある場合は、新しい量子化手法の貢献ガイドを参照してください。

• Transformers の量子化 概要 では、異なる量子化バックエンドの長所と短所について概説しています。

• Diffusersにおける量子化バックエンドの探索のブログ記事では、各量子化バックエンドの簡単な紹介、バックエンドの選択方法、および量子化と他のメモリ最適化の組み合わせについて解説しています。