過去的幾個月，我們目睹了使用基於 transformer 模型作為擴散模型的主幹網路來進行高解析度文生圖 (text-to-image，T2I) 的趨勢。和一開始的許多擴散模型普遍使用 UNet 架構不同，這些模型使用 transformer 架構作為擴散過程的主模型。由於 transformer 的性質，這些主幹網路表現出了良好的可擴充套件性，模型引數量可從 0.6B 擴充套件至 8B。
隨著模型越變越大，記憶體需求也隨之增加。對擴散模型而言，這個問題愈加嚴重，因為擴散流水線通常由多個模型串成: 文字編碼器、擴散主幹模型和影像解碼器。此外，最新的擴散流水線通常使用多個文字編碼器 - 如: Stable Diffusion 3 有 3 個文字編碼器。使用 FP16 精度對 SD3 進行推理需要 18.765GB 的 GPU 視訊記憶體。
這麼高的記憶體要求使得很難將這些模型執行在消費級 GPU 上，因而減緩了技術採納速度並使針對這些模型的實驗變得更加困難。本文，我們展示瞭如何使用 Diffusers 庫中的 Quanto 量化工具指令碼來提高基於 transformer 的擴散流水線的記憶體效率。
基礎知識
你可參考 這篇文章 以獲取 Quanto 的詳細介紹。簡單來說，Quanto 是一個基於 PyTorch 的量化工具包。它是 Hugging Face Optimum 的一部分，Optimum 提供了一套硬體感知的最佳化工具。
模型量化是 LLM 從業者必備的工具，但在擴散模型中並不算常用。Quanto 可以幫助彌補這一差距，其可以在幾乎不傷害生成質量的情況下節省記憶體。
我們基於 H100 GPU 配置進行基準測試，軟體環境如下:
- CUDA 12.2
- PyTorch 2.4.0
- Diffusers (從原始碼安裝，參考 此提交)
- Quanto (從原始碼安裝，參考 此提交)
除非另有說明，我們預設使用 FP16 進行計算。我們不對 VAE 進行量化以防止數值不穩定問題。你可於 此處 找到我們的基準測試程式碼。
截至本文撰寫時，以下基於 transformer 的擴散模型流水線可用於 Diffusers 中的文生圖任務:
- PixArt-Alpha 及 PixArt-Sigma
- Stable Diffusion 3
- Hunyuan DiT
- Lumina
- Aura Flow
另外還有一個基於 transformer 的文生影片流水線: Latte。
為簡化起見，我們的研究僅限於以下三個流水線: PixArt-Sigma、Stable Diffusion 3 以及 Aura Flow。下表顯示了它們各自的擴散主幹網路的引數量:
| 模型 | Checkpoint | 引數量（Billion） |
|---|---|---|
| PixArt | https://huggingface.co/PixArt-alpha/PixArt-Sigma-XL-2-1024-MS | 0.611 |
| Stable Diffusion 3 | https://huggingface.co/stabilityai/stable-diffusion-3-medium-diffusers | 2.028 |
| Aura Flow | https://huggingface.co/fal/AuraFlow/ | 6.843 |
用 Quanto 量化 DiffusionPipeline
使用 Quanto 量化模型非常簡單。
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(pipeline.transformer, weights=qfloat8)
freeze(pipeline.transformer)
我們對需量化的模組呼叫 quantize() ，以指定我們要量化的部分。上例中，我們僅量化引數，保持啟用不變，量化資料型別為 FP8。最後，呼叫 freeze() 以用量化引數替換原始引數。
然後，我們就可以如常呼叫這個 pipeline 了:
image = pipeline("ghibli style, a fantasy landscape with castles").images[0]
| FP16 | 將 transformer 擴散主幹網路量化為 FP8 |
|---|---|
 |
 |
我們注意到使用 FP8 可以節省視訊記憶體，且幾乎不影響生成質量; 我們也看到量化模型的延遲稍有變長:
| Batch Size | 量化 | 記憶體 (GB) | 延遲 (秒) |
|---|---|---|---|
| 1 | 無 | 12.086 | 1.200 |
| 1 | FP8 | 11.547 | 1.540 |
| 4 | 無 | 12.087 | 4.482 |
| 4 | FP8 | 11.548 | 5.109 |
我們可以用相同的方式量化文字編碼器:
quantize(pipeline.text_encoder, weights=qfloat8)
freeze(pipeline.text_encoder)
文字編碼器也是一個 transformer 模型，我們也可以對其進行量化。同時量化文字編碼器和擴散主幹網路可以帶來更大的視訊記憶體節省:
| Batch Size | 量化 | 是否量化文字編碼器 | 視訊記憶體 (GB) | 延遲 (秒) |
|---|---|---|---|---|
| 1 | FP8 | 否 | 11.547 | 1.540 |
| 1 | FP8 | 是 | 5.363 | 1.601 |
| 4 | FP8 | 否 | 11.548 | 5.109 |
| 4 | FP8 | 是 | 5.364 | 5.141 |
量化文字編碼器後生成質量與之前的情況非常相似:
上述攻略通用嗎？
將文字編碼器與擴散主幹網路一起量化普遍適用於我們嘗試的很多模型。但 Stable Diffusion 3 是個特例，因為它使用了三個不同的文字編碼器。我們發現 _ 第二個 _ 文字編碼器量化效果不佳，因此我們推薦以下替代方案:
- 僅量化第一個文字編碼器 (
CLIPTextModelWithProjection) 或 - 僅量化第三個文字編碼器 (
T5EncoderModel) 或 - 同時量化第一個和第三個文字編碼器
下表給出了各文字編碼器量化方案的預期記憶體節省情況 (擴散 transformer 在所有情況下均被量化):
| Batch Size | 量化 | 量化文字編碼器 1 | 量化文字編碼器 2 | 量化文字編碼器 3 | 視訊記憶體 (GB) | 延遲 (秒) |
|---|---|---|---|---|---|---|
| 1 | FP8 | 1 | 1 | 1 | 8.200 | 2.858 |
| 1 ✅ | FP8 | 0 | 0 | 1 | 8.294 | 2.781 |
| 1 | FP8 | 1 | 1 | 0 | 14.384 | 2.833 |
| 1 | FP8 | 0 | 1 | 0 | 14.475 | 2.818 |
| 1 ✅ | FP8 | 1 | 0 | 0 | 14.384 | 2.730 |
| 1 | FP8 | 0 | 1 | 1 | 8.325 | 2.875 |
| 1 ✅ | FP8 | 1 | 0 | 1 | 8.204 | 2.789 |
| 1 | 無 | - | - | - | 16.403 | 2.118 |
| 量化文字編碼器: 1 | 量化文字編碼器: 3 | 量化文字編碼器: 1 和 3 |
|---|---|---|
 |
 |
 |
其他發現
在 H100 上 bfloat16 通常表現更好
對於支援 bfloat16 的 GPU 架構 (如 H100 或 4090)，使用 bfloat16 速度更快。下表列出了在我們的 H100 參考硬體上測得的 PixArt 的一些數字: Batch Size 精度 量化 視訊記憶體 (GB) 延遲 (秒) 是否量化文字編碼器
| Batch Size | 精度 | 量化 | 視訊記憶體（GB） | 延遲（秒） | 是否量化文字編碼器 |
|---|---|---|---|---|---|
| 1 | FP16 | INT8 | 5.363 | 1.538 | 是 |
| 1 | BF16 | INT8 | 5.364 | 1.454 | 是 |
| 1 | FP16 | FP8 | 5.363 | 1.601 | 是 |
| 1 | BF16 | FP8 | 5.363 | 1.495 | 是 |
qint8 的前途
我們發現使用 qint8 (而非 qfloat8 ) 進行量化，推理延遲通常更好。當我們對注意力 QKV 投影進行水平融合 (在 Diffusers 中呼叫 fuse_qkv_projections() ) 時，效果會更加明顯，因為水平融合會增大 int8 運算元的計算維度從而實現更大的加速。我們基於 PixArt 測得了以下資料以證明我們的發現:
| Batch Size | 量化 | 視訊記憶體 (GB) | 延遲 (秒) | 是否量化文字編碼器 | QKV 融合 |
|---|---|---|---|---|---|
| 1 | INT8 | 5.363 | 1.538 | 是 | 否 |
| 1 | INT8 | 5.536 | 1.504 | 是 | 是 |
| 4 | INT8 | 5.365 | 5.129 | 是 | 否 |
| 4 | INT8 | 5.538 | 4.989 | 是 | 是 |
INT4 咋樣？
在使用 bfloat16 時，我們還嘗試了 qint4 。目前我們僅支援 H100 上的 bfloat16 的 qint4 量化，其他情況尚未支援。透過 qint4 ，我們期望看到記憶體消耗進一步降低，但代價是推理延遲變長。延遲增加的原因是硬體尚不支援 int4 計算 - 因此權重使用 4 位，但計算仍然以 bfloat16 完成。下表展示了 PixArt-Sigma 的結果:
| Batch Size | 是否量化文字編碼器 | 視訊記憶體 (GB) | 延遲 (秒) |
|---|---|---|---|
| 1 | 否 | 9.380 | 7.431 |
| 1 | 是 | 3.058 | 7.604 |
但請注意，由於 INT4 量化比較激進，最終結果可能會受到影響。所以，一般對於基於 transformer 的模型，我們通常不量化最後一個投影層。在 Quanto 中，我們做法如下:
quantize(pipeline.transformer, weights=qint4, exclude="proj_out")
freeze(pipeline.transformer)
"proj_out" 對應於 pipeline.transformer 的最後一層。下表列出了各種設定的結果:
| 量化文字編碼器: 否 , 不量化的層: 無 | 量化文字編碼器: 否 , 不量化的層: "proj_out" | 量化文字編碼器: 是 , 不量化的層: 無 | 量化文字編碼器: 是 , 不量化的層: "proj_out" |
|---|---|---|---|
 |
 |
 |
 |
為了恢復損失的影像質量，常見的做法是進行量化感知訓練，Quanto 也支援這種訓練。這項技術超出了本文的範圍，如果你有興趣，請隨時與我們聯絡！
本文的所有實驗結果都可以在 這裡 找到。
加個雞腿 - 在 Quanto 中儲存和載入 Diffusers 模型
以下程式碼可用於對 Diffusers 模型進行量化並儲存量化後的模型:
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")
此程式碼生成的 checkpoint 大小為 587MB ，而不是原本的 2.44GB。然後我們可以載入它:
from optimum.quanto import QuantizedPixArtTransformer2DModel
import torch
transformer = QuantizedPixArtTransformer2DModel.from_pretrained("pixart-sigma-fp8")
transformer.to(device="cuda", dtype=torch.float16)
最後，在 DiffusionPipeline 中使用它:
from diffusers import DiffusionPipeline
import torch
pipe = DiffusionPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
transformer=None,
torch_dtype=torch.float16,
).to("cuda")
pipe.transformer = transformer
prompt = "A small cactus with a happy face in the Sahara desert."
image = pipe(prompt).images[0]
將來，我們計劃支援在初始化流水線時直接傳入 transformer 就可以工作:
pipe = PixArtSigmaPipeline.from_pretrained(
"PixArt-alpha/PixArt-Sigma-XL-2-1024-MS",
- transformer=None,
+ transformer=transformer,
torch_dtype=torch.float16,
).to("cuda")
QuantizedPixArtTransformer2DModel 實現可參考 此處。如果你希望 Quanto 支援對更多的 Diffusers 模型進行儲存和載入，請在 此處 提出需求並 @sayakpaul 。
小訣竅
- 根據應用場景的不同，你可能希望對流水線中不同的模組使用不同型別的量化。例如，你可以對文字編碼器進行 FP8 量化，而對 transformer 擴散模型進行 INT8 量化。由於 Diffusers 和 Quanto 的靈活性，你可以輕鬆實現這類方案。
- 為了最佳化你的用例，你甚至可以將量化與 Diffuser 中的其他 記憶體最佳化技術 結合起來，如
enable_model_cpu_offload()。
總結
本文，我們展示瞭如何量化 Diffusers 中的 transformer 模型並最佳化其記憶體消耗。當我們同時對文字編碼器進行量化時，效果變得更加明顯。我們希望大家能將這些工作流應用到你的專案中並從中受益🤗。
感謝 Pedro Cuenca 對本文的細緻審閱。
英文原文: https://hf.co/blog/quanto-diffusers
原文作者: Sayak Paul，David Corvoysier
譯者: Matrix Yao (姚偉峰)，英特爾深度學習工程師，工作方向為 transformer-family 模型在各模態資料上的應用及大規模模型的訓練推理。