Quanto 量化減少 Transformer 擴散管線的記憶體使用

TL;DR

Quantization with Hugging Face Quanto reduces the GPU memory needed for Transformer‑based diffusion pipelines (e.g., PixArt‑Sigma, Stable Diffusion 3, Aura Flow) from ~12 GB to as low as ~5 GB with only minor latency impact and negligible quality loss.

介紹 – 為何記憶體對於擴散 Transformer 重要

Transformer 骨幹已成為高解析度文字轉圖像擴散模型的主流架構，參數規模從 0.6 B 擴展至 8 B。較大的模型會顯著提升 GPU 記憶體消耗；完整的 Stable Diffusion 3 在 FP16 下的推理佔用 18.8 GB。這個記憶體門檻限制了消費者級 GPU 的採用，也阻礙了快速實驗。本文展示了 Quanto 的量化工具（已整合於 Diffusers 函式庫）如何在保留視覺品質的同時，大幅縮減記憶體佔用。

Quanto 量化基礎

Quanto 是一個基於 PyTorch 的量化工具包，內嵌於 Hugging Face Optimum。它支援僅對權重進行量化至多種低精度格式（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) 延遲 (s)
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) 延遲 (s)
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 與文字編碼器可將記憶體使用減半，同時延遲基本保持不變。

各模型的通用性

作者評估了三條管線：

  • PixArt‑Sigma（0.61 B 參數）
  • Stable Diffusion 3 (medium)（2.03 B 參數，三個文字編碼器）
  • Aura Flow（6.84 B 參數）

對於 PixArt‑Sigma 與 Aura Flow，量化文字編碼器總是能帶來大量記憶體節省。Stable Diffusion 3 需要選擇性量化，因為第二個文字編碼器（中間的 CLIP 變體）在量化後會降低品質。建議的策略有：

  • 僅量化第一個 CLIP 編碼器，或
  • 僅量化第三個 T5 編碼器，或
  • 同時量化第一與第三個編碼器。

一個針對 SD‑3（批次 1，擴散 Transformer 永遠使用 FP8）的示例表格顯示，記憶體範圍從 8.20 GB（全部三個編碼器皆量化）到 16.40 GB（未量化）。

其他發現

bfloat16 與 fp16 在 H100 上的比較

在 NVIDIA H100 GPU 上，使用 bfloat16 搭配 INT8 或 FP8 權重可提升延遲表現：

精度 量化方式 記憶體 (GB) 延遲 (s)
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 且融合 QKV）。

INT4（qint4）用於激進壓縮

在 H100 上與 bfloat16 結合使用時，INT4 可大幅降低記憶體（例如 PixArt‑Sigma 從 9.38 GB 降至 3.06 GB），但延遲上升（≈7.6 s），因為計算仍在 bfloat16 中執行。品質損失明顯；作者建議在量化時排除最後的投影層（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 對本文的詳細審閱。

Sources