Quanto 量子化で Transformer ディフュージョンパイプラインのメモリ削減
TL;DR
Hugging Face Quanto を用いた量子化により、Transformer ベースのディフュージョンパイプライン（例：PixArt‑Sigma、Stable Diffusion 3、Aura Flow）の GPU メモリ使用量が約 12 GB から最小で約 5 GB に削減され、レイテンシへの影響は僅かで品質の低下はほとんどありません。
はじめに – ディフュージョン Transformer におけるメモリの重要性
Transformer バックボーンは、高解像度テキスト‑ツー‑イメージ ディフュージョンモデルの主流アーキテクチャとなっており、パラメータ数は 0.6 B から 8 B にスケールしています。モデルが大きくなるほど GPU メモリ消費は劇的に増加し、FP16 での Stable Diffusion 3 のフル推論は 18.8 GB を占有します。このメモリの壁は一般消費者向け GPU での導入を制限し、迅速な実験を妨げます。本稿では、Diffusers ライブラリに統合された Quanto の量子化ユーティリティが、視覚品質を保ちつつメモリフットプリントを大幅に削減できることを示します。
Quanto を用いた量子化の基本
Quanto は Hugging Face Optimum 内にある PyTorch ベースの量子化ツールキットです。FP8、INT8、INT4 などの低精度フォーマットへのウェイトのみの量子化をサポートし、任意の Diffusers モジュールに適用できます。
from optimum.quanto import freeze, qfloat8, quantize
from diffusers import PixArtSigmaPipeline
import torch
pipeline = PixArtSigmaPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", torch_dtype=torch.float16
).to("cuda")
# Quantize the diffusion transformer to FP8
quantize(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)
同じ quantize/freeze 呼び出しは、テキストエンコーダや他のサブモジュールでも機能します。
メモリとレイテンシの結果 – FP8 ウェイト量子化
ディフュージョン Transformer のみ
| バッチサイズ | メモリ (GB) | レイテンシ (秒) |
|---|---|---|
| 1 | 11.55 (FP8) vs 12.09 (FP16) | 1.54 vs 1.20 |
| 4 | 11.55 (FP8) vs 12.09 (FP16) | 5.11 vs 4.48 |
FP8 ウェイトはメモリを約 0.5 GB 削減し、レイテンシは僅かに増加します。
テキストエンコーダの量子化を追加
| バッチサイズ | テキストエンコーダを量子化? | メモリ (GB) | レイテンシ (秒) |
|---|---|---|---|
| 1 | No | 11.55 | 1.54 |
| 1 | Yes | 5.36 | 1.60 |
| 4 | No | 11.55 | 5.11 |
| 4 | Yes | 5.36 | 5.14 |
ディフュージョン Transformer とテキストエンコーダの両方を量子化することで、メモリ使用量が半減し、レイテンシはほぼ変わりません。
複数モデルへの汎用性
著者らは 3 つのパイプラインを評価しました：
- PixArt‑Sigma（0.61 B パラメータ）
- Stable Diffusion 3（ミディアム）（2.03 B パラメータ、テキストエンコーダ 3 つ）
- Aura Flow（6.84 B パラメータ）
PixArt‑Sigma と Aura Flow では、テキストエンコーダを量子化すると常に大幅なメモリ削減が得られました。Stable Diffusion 3 では、2 番目のテキストエンコーダ（中間の CLIP バリアント）を量子化すると品質が低下するため、選択的な量子化が必要です。推奨戦略は次のとおりです：
- 最初の CLIP エンコーダのみを量子化、または
- 3 番目の T5 エンコーダのみを量子化、または
- 最初と 3 番目のエンコーダの両方を量子化。
SD‑3（バッチ 1、ディフュージョン Transformer は常に FP8）の代表的な表では、メモリは 8.20 GB（3 つのエンコーダすべてを量子化）から 16.40 GB（量子化なし）までの範囲です。
追加の発見
H100 における bfloat16 と fp16 の比較
NVIDIA H100 GPU では、bfloat16 と INT8 または FP8 ウェイトを組み合わせることでレイテンシが改善されます：
| 精度 | 量子化 | メモリ (GB) | レイテンシ (秒) |
|---|---|---|---|
| FP16 | INT8 | 5.363 | 1.538 |
| BF16 | INT8 | 5.364 | 1.454 |
| FP16 | FP8 | 5.363 | 1.601 |
| BF16 | FP8 | 5.363 | 1.495 |
INT8（qint8）と融合された QKV プロジェクション
INT8 ウェイトは FP8 より高速で、特にアテンションの QKV プロジェクションを融合 (fuse_qkv_projections()) した場合に顕著です。PixArt‑Sigma（バッチ 1）では、レイテンシが 1.538 s（INT8、非融合）から 1.504 s（INT8、融合）に低下しました。
INT4（qint4）による積極的な圧縮
H100 上で bfloat16 と組み合わせると、INT4 はメモリを劇的に削減します（例：PixArt‑Sigma は 9.38 GB から 3.06 GB へ）。しかし計算は依然として bfloat16 で行われるためレイテンシは上昇し（≈7.6 秒）、品質低下が顕著です。著者らは、品質劣化を抑えるために最終プロジェクション層（proj_out）を量子化から除外することを推奨しています。
量子化された Diffusers モデルの保存とロード
Quanto は、永続化および再ロード可能なモデルクラスを提供します：
from diffusers import PixArtTransformer2DModel
from optimum.quanto import QuantizedPixArtTransformer2DModel, qfloat8
model = PixArtTransformer2DModel.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS", subfolder="transformer"
)
qmodel = QuantizedPixArtTransformer2DModel.quantize(model, weights=qfloat8)
qmodel.save_pretrained("pixart-sigma-fp8") # 587 MB checkpoint
ロードも同様のパターンで行え、量子化された Transformer を DiffusionPipeline に組み込むことができます。
開発者向け実践的なヒント
- モジュールごとに量子化タイプを組み合わせて使用（例：テキストエンコーダは FP8、ディフュージョン Transformer は INT8）し、メモリと速度のバランスを取ります。
- Quanto の量子化と Diffusers の既存メモリ削減ユーティリティ（例：
enable_model_cpu_offload()）を組み合わせて、さらなる削減を実現します。 - INT4 を使用する際は、画像の忠実度を保つために必ず最終プロジェクション層（
exclude="proj_out"）を除外してください。
結論
Hugging Face Quanto を用いた Transformer ベースのディフュージョンパイプラインの量子化により、GPU メモリ要件が最大 70 %（約 12 GB から約 5 GB）削減され、レイテンシは低く、視覚品質もほぼ維持されます。この手法は複数の最先端モデルで有効で、生成されるチェックポイントは劇的に小さくなります（例：587 MB 対 2.44 GB）。開発者は既存の Diffusers パイプラインに容易に組み込め、他のメモリ最適化戦略と併用することも可能です。
謝辞
本稿の詳細なレビューを行ってくれた Pedro Cuenca に感謝します。