優化 Stable Diffusion XL (SDXL) 的推理速度與記憶體使用
Hugging Face 詳細介紹了一系列簡單的優化方法，旨在讓 Stable Diffusion XL (SDXL) 在推理時更具實用性，特別是在消費級 GPU 上。透過實施低精度、記憶體高效注意力機制以及各種卸載（offloading）技術，記憶體佔用量可以從 28GB 降低至 11.47GB，且推理延遲可以從 72.2 秒降低到大約 10.3 秒。
推理速度優化
對於擴散模型來說，優化速度至關重要，因為圖像生成的迭代性質通常需要多次運行才能達到理想結果。以下技術專注於降低延遲：
低精度 (fp16)
使用 float16 (fp16) 代替標準的 float32 (fp32) 可以將記憶體使用量減少一半並提高計算速度，因為 fp16 捕捉的浮點數範圍較窄，且受到現代 GPU 硬體的更好支援。在測試過的 SDXL pipeline 中，切換到 fp16 將記憶體使用量降低至 21.7GB，並將推理時間從 72.2 秒縮短至 14.8 秒。
記憶體高效注意力機制 (SDPA)
Transformer 中的注意力模組可能會造成記憶體瓶頸，因為記憶體需求會隨輸入序列長度呈平方級增長。PyTorch 2.0 引入了 Scaled Dot Product Attention (SDPA)，它提供了 Flash Attention 和記憶體高效注意力 (xFormers) 的融合實作。在 🤗Diffusers 中，PyTorch ≥ 2.0 預設啟用 SDPA，在保持 21.7GB 記憶體佔用的同時，進一步將推理時間縮短至 11.4 秒。
使用 torch.compile 進行 JIT 編譯
PyTorch 2.0 的 torch.compile API 允許將 PyTorch 程式碼即時（JIT）編譯為優化後的內核（kernels）。使用 torch.compile（使用 mode="reduce-overhead"）封裝 SDXL UNet，可進一步將推理時間提升至 10.2 秒。請注意，第一次編譯運行會較慢，但隨後的調用速度會顯著加快。
減少模型記憶體佔用
由於 SDXL 的規模大約是先前 Stable Diffusion 模型的 3 倍（具有 3.5B 參數的 UNet），將其放入 VRAM 是主要的挑戰。幾種技術可以減少記憶體佔用：
CPU 卸載 (CPU Offloading)
模型卸載會在 pipeline 的組件不需要在 GPU 上主動運行時，將其移至 CPU：
- Model CPU Offloading：將 UNet 加載到 GPU 記憶體中，同時將文本編碼器（text encoders）和 VAE 保留在 CPU 上。這將記憶體使用量降低至 20.2GB。
- Sequential CPU Offloading：將個別 UNet 子模組的權重卸載到 CPU，僅在進行前向傳播（forward pass）前才將其加載到 GPU。這將記憶體降低至 19.9GB，但會顯著增加延遲至 67 秒。
VAE Slicing
變分自編碼器 (VAE) 將潛在空間（latents）解碼為圖像，這個過程的記憶體使用量會隨批次大小（batch size）而增加。VAE slicing 將輸入張量拆分為較小的切片，並分多個步驟進行解碼。這項優化可以在不顯著影響延遲的情況下，將記憶體使用量降低至 15.4GB。
快取計算結果
SDXL 利用兩個文本編碼器從提示詞（prompts）計算嵌入（embeddings）。由於這些嵌入在整個反向擴散過程中保持不變，因此可以預先計算並快取。一旦生成了嵌入，文本編碼器和分詞器（tokenizers）就可以從 GPU 記憶體中移除，結合 SDPA 和 fp16 後，記憶體佔用量為 21.9GB。
Tiny Autoencoder (TAESD)
使用蒸餾版本（distilled version）取代標準 VAE（例如 madebyollin 的 Tiny Autoencoder，約 10MB），可將記憶體使用量降低至 15.6GB 並減少推理延遲。然而，Tiny Autoencoder 可能會忽略細微細節，主要建議用於圖像預覽。
效能總結
在 A100 GPU (40 GB) 上針對每個提示詞生成 4 張圖像進行的測試顯示了記憶體與延遲之間的權衡：
| 技術 | 記憶體 (GB) | 推理延遲 (ms) |
|---|---|---|
| 未優化的 pipeline | 28.09 | 72200.5 |
| fp16 | 21.72 | 14800.9 |
| fp16 + SDPA (預設) | 21.72 | 11413.0 |
預設 + torch.compile |
21.73 | 10296.7 |
| 預設 + model CPU offload | 20.21 | 16082.2 |
| 預設 + sequential CPU offload | 19.91 | 67034.0 |
| 預設 + VAE slicing | 15.40 | 11232.2 |
| 預設 + VAE slicing + sequential CPU offload | 11.47 | 66869.2 |
| 預設 + 預先計算的文本嵌入 | 21.85 | 11909.0 |
| 預設 + Tiny Autoencoder | 15.48 | 10449.7 |