Hugging Face Diffusers 量子化バックエンド

Hugging Face Diffusers は、FLUX.1-dev などの大規模拡散モデルのメモリと計算要件を削減し、より手軽に利用できるようにするため、複数の量子化バックエンドを統合しました。これらのバックエンドにより、画像品質を大幅に損なうことなくモデルを圧縮でき、8 ビット量子化は高精度 BF16 モデルとほぼ区別がつかない結果を得られることが多いです。

FLUX.1-dev モデルアーキテクチャとメモリベースライン

量子化の影響を評価するには、FLUX.1-dev モデルの BF16 精度でのメモリベースラインを把握する必要があります。このモデルは約 31.447 GB のメモリを必要とします。モデルは主に以下の 4 つのコンポーネントから構成されます。

  • Transformer (MMDiT): 生成的マルチモーダル拡散トランスフォーマーのコアで、23.8 GB 必要。
  • Text Encoder 2 (T5): 微細な理解とテキスト生成に使用され、9.52 GB 必要。
  • Text Encoder 1 (CLIP): 初期テキスト理解に使用され、246 MB 必要。
  • VAE: ピクセル空間と潜在空間の相互変換を行い、168 MB 必要。

量子化は主に Transformer と T5 テキストエンコーダを対象とし、最も大きなメモリ削減を狙います。

サポートされている量子化バックエンド

bitsandbytes (BnB)

bitsandbytes は 8 ビットと 4 ビットの量子化を提供します。FLUX.1-dev に対しては、4 ビット量子化（NF4 使用）によりロード後メモリが 12.584 GB、ピークメモリが 17.281 GB に削減され、NVIDIA H100 80GB GPU 上で推論時間は 12 秒です。8 ビット量子化はメモリ使用量が増加（ロード時 19.273 GB / ピーク 24.432 GB）し、推論は 27 秒に遅くなります。

torchao

torchao は PyTorch ネイティブのアーキテクチャ最適化ライブラリです。以下の weight‑only フォーマットをサポートします。

  • int4_weight_only: 最小メモリフットプリント（ロード時 10.635 GB / ピーク 14.654 GB）だが、最も遅い推論（109 秒）。
  • int8_weight_only: 中程度のメモリ（ロード時 17.020 GB / ピーク 21.482 GB）で高速推論（15 秒）。
  • float8_weight_only: int8 と同等のメモリと速度（ロード時 17.016 GB / ピーク 21.488 GB、15 秒）。

Quanto

optimum ライブラリ経由で統合されている Quanto は様々な精度をサポートします。

  • INT4: ロード時 12.254 GB / ピーク 16.139 GB、推論 109 秒。
  • INT8: ロード時 17.330 GB / ピーク 21.814 GB、推論 15 秒。
  • FP8: ロード時 16.395 GB / ピーク 20.898 GB、推論 16 秒。

GGUF

Diffusers は GGUF ファイル形式をサポートし、from_single_file を介して llama.cpp コミュニティの事前量子化モデルを利用できます。FLUX.1-dev のベンチマークは次の通りです。

  • Q2_k: ロード時 13.264 GB / ピーク 17.752 GB、推論 26 秒。
  • Q4_1: ロード時 16.838 GB / ピーク 21.326 GB、推論 23 秒。
  • Q8_0: ロード時 21.502 GB / ピーク 25.973 GB、推論 15 秒。

FP8 Layerwise Casting

enable_layerwise_casting は、重みを FP8（e4m3）で保存し、計算時に動的に高精度（例: BF16）へキャストするメモリ最適化手法です。これによりロード後メモリが 23.682 GB、ピークメモリが 28.451 GB に削減され、推論時間は 13 秒です。

量子化とメモリ最適化の組み合わせ

量子化バックエンドは、Diffusers の他の最適化手法と組み合わせて VRAM 使用量をさらに削減できます。

  • Model CPU Offloading (enable_model_cpu_offload): コンポーネント全体を CPU と GPU 間で転送します。BnB の 4 ビット量子化と併用すると、ピークメモリが 12.383 GB に低減します。
  • Group Offloading (enable_group_offload): 内部レイヤーのグループを CPU に移動します。FP8 layerwise casting と組み合わせると、ロード後メモリが 9.264 GB、ピークメモリが 14.232 GB に削減されます。
  • torch.compile: PyTorch 2.x のコンパイル機能で実行を高速化します。メモリは変わりませんが、推論が大幅に速くなります。例として、torchao の int4_weight_only 推論時間はコンパイル前の 109 秒からコンパイル後は 6 秒に短縮されます（コンパイル自体に約 285 秒かかります）。

バックエンド選択ガイド

ハードウェア構成やパフォーマンス目標に応じて、以下のバックエンドが推奨されます。

  • NVIDIA ユーザー（最も簡単）: bitsandbytes の 4/8 ビットを使用。
  • 推論速度重視: torchao、GGUF、または bitsandbytes と torch.compile() の組み合わせ。
  • ハードウェア柔軟性（CPU/MPS）または FP8 が必要: Quanto。
  • Hopper/Ada アーキテクチャ: FP8 Layerwise Casting。
  • 既存の GGUF モデルがある場合: from_single_file による GGUF ローディング。

Sources