使用 Quanto 和 Diffusers 實現記憶體高效的 Diffusion Transformer

在過去幾個月中，我們見證了基於 Transformer 的擴散模型骨幹網路在高解析度文字到影像 (T2I) 生成領域的興起。這些模型使用 Transformer 架構作為擴散過程的基礎模組，取代了許多早期擴散模型中普遍使用的 UNet 架構。得益於 Transformer 的特性，這些骨幹網路表現出良好的可擴充套件性，模型引數從 6 億到 80 億不等。

隨著模型越來越大，記憶體需求也隨之增加。這個問題更加嚴峻，因為一個擴散 pipeline 通常由幾個元件組成：一個文字編碼器、一個擴散模型骨幹和一個影像解碼器。此外，現代擴散 pipeline 使用多個文字編碼器——例如，Stable Diffusion 3 就有三個。使用 FP16 精度執行 SD3 推理需要 18.765 GB 的 GPU 記憶體。

這些高記憶體需求使得在消費級 GPU 上使用這些模型變得困難，從而減緩了它們的普及和實驗程序。在本文中，我們將展示如何利用 Diffusers 庫中 Quanto 的量化工具來提高基於 Transformer 的擴散 pipeline 的記憶體效率。

目錄

• 預備知識
• 使用 Quanto 量化 DiffusionPipeline
• 觀察的普適性
• 其他發現
  • bfloat16 作為主要計算資料型別
  • qint8 的前景
  • INT4
• 附贈 - 在 Quanto 中儲存和載入 Diffusers 模型
• 技巧
• 結論

預備知識

有關 Quanto 的詳細介紹，請參閱這篇文章。簡而言之，Quanto 是一個基於 PyTorch 構建的量化工具包。它是 Hugging Face Optimum 的一部分，這是一套用於硬體最佳化的工具。

模型量化是 LLM 從業者中一種流行的工具，但在擴散模型中卻不那麼常見。Quanto 可以幫助彌補這一差距，以極少或無質量下降為代價，實現記憶體節省。

為了進行基準測試，我們使用 H100 GPU，環境如下：

• CUDA 12.2
• PyTorch 2.4.0
• Diffusers (從此 commit 安裝)
• Quanto (從此 commit 安裝)

除非另有說明，我們預設使用 FP16 進行計算。我們選擇不對 VAE 進行量化，以防止出現數值不穩定問題。我們的基準測試程式碼可以在這裡找到。

在撰寫本文時，我們在 Diffusers 中有以下基於 Transformer 的用於文字到影像生成的擴散 pipeline：

• PixArt-Alpha 和 PixArt-Sigma
• Stable Diffusion 3
• Hunyuan DiT
• Lumina
• Aura Flow

我們還有 Latte，一個基於 Transformer 的文字到影片生成 pipeline。

為簡潔起見，我們的研究僅限於以下三個：PixArt-Sigma、Stable Diffusion 3 和 Aura Flow。下表顯示了它們擴散模型骨幹的引數數量：

使用 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	FP8 的 Diffusion Transformer

我們注意到，使用 FP8 時，記憶體消耗減少，延遲略有增加，但質量幾乎沒有下降：

我們可以用同樣的方法量化文字編碼器：

quantize(pipeline.text_encoder, weights=qfloat8)
freeze(pipeline.text_encoder)

文字編碼器也是一個 Transformer 模型，我們也可以對其進行量化。同時量化文字編碼器和擴散模型骨幹可以帶來更大的記憶體改進：

量化文字編碼器產生的結果與之前的情況非常相似：

觀察的普適性

將文字編碼器與擴散模型骨幹一起量化通常適用於我們嘗試過的模型。Stable Diffusion 3 是一個特例，因為它使用了三個不同的文字編碼器。我們發現量化*第二個*文字編碼器效果不佳，因此我們建議採用以下替代方案：

• 只量化第一個文字編碼器 (CLIPTextModelWithProjection) 或
• 只量化第三個文字編碼器 (T5EncoderModel) 或
• 量化第一個和第三個文字編碼器

下表給出了各種文字編碼器量化組合（所有情況下擴散 Transformer 都被量化）的預期記憶體節省情況：

量化的文字編碼器：1	量化的文字編碼器：3	量化的文字編碼器：1 和 3

其他發現

在 H100 上 bfloat16 通常更好

對於支援的 GPU 架構，如 H100 或 4090，使用 bfloat16 可能會更快。下表展示了在我們的 H100 參考硬體上測量的 PixArt 的一些資料：

qint8 的前景

我們發現，使用 `qint8`（而不是 `qfloat8`）進行量化在推理延遲方面通常更好。當我們水平融合注意力 QKV 投影（在 Diffusers 中呼叫 `fuse_qkv_projections()`）時，這種效果會更加明顯，從而加厚了 int8 核心的維度以加速計算。我們下面為 PixArt 提供了一些證據：

INT4 效果如何？

我們還額外試驗了在使用 `bfloat16` 時採用 `qint4`。這僅適用於 H100 上的 `bfloat16`，因為其他配置尚不支援。使用 `qint4`，我們可以預期在記憶體消耗方面看到更多改進，但代價是推理延遲增加。延遲增加是預料之中的，因為沒有原生硬體支援 int4 計算——權重使用 4 位傳輸，但計算仍然以 `bfloat16` 完成。下表顯示了我們對 PixArt-Sigma 的結果：

然而，請注意，由於 INT4 的激進離散化，最終結果可能會受到影響。這就是為什麼，對於基於 Transformer 的模型，我們通常將最終的投影層排除在量化之外。在 Quanto 中，我們這樣做：

quantize(pipeline.transformer, weights=qint4, exclude="proj_out")
freeze(pipeline.transformer)

"proj_out" 對應於 pipeline.transformer 中的最後一層。下表展示了各種設定下的結果：

量化 TE：否，層排除：無	量化 TE：否，層排除："proj_out"	量化 TE：是，層排除：無	量化 TE：是，層排除："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")

生成的模型權重大小為 *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]

未來，我們期望在初始化 pipeline 時可以直接傳遞 `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 模型以便於儲存和載入，請在此處提交一個 issue 並提及 `@sayakpaul`。

技巧

• 根據您的需求，您可能希望對不同的 pipeline 模組應用不同型別的量化。例如，您可以對文字編碼器使用 FP8，但對擴散 Transformer 使用 INT8。得益於 Diffusers 和 Quanto 的靈活性，這可以無縫完成。
• 為了最佳化您的用例，您甚至可以將量化與 Diffusers 中的其他記憶體最佳化技術相結合，例如 `enable_model_cpu_offload()`。

結論

在本文中，我們展示瞭如何量化 Diffusers 中的 Transformer 模型並最佳化其記憶體消耗。當我們額外量化混合中的文字編碼器時，量化的效果變得更加明顯。我們希望您能將一些工作流程應用到您的專案中並從中受益 🤗。

感謝 Pedro Cuenca 對本文的詳盡審閱。