Hugging Face Diffusers 量化後端

Hugging Face Diffusers 現在整合了多種量化後端，使大型擴散模型（例如 FLUX.1-dev）更易於使用，透過降低記憶體與計算需求。這些後端讓使用者在不大幅犧牲影像品質的情況下縮小模型，8 位元量化常能產生與高精度 BF16 模型幾乎無法區分的結果。

FLUX.1-dev 模型架構與記憶體基準

要了解量化的影響，需要先確定 FLUX.1-dev 模型在 BF16 精度下的記憶體基準，約需 31.447 GB 記憶體。該模型由四個主要組件構成：

  • Transformer (MMDiT)： 核心的生成式多模態擴散 Transformer，需 23.8 GB。
  • Text Encoder 2 (T5)： 用於細緻的理解與文字渲染，需 9.52 GB。
  • Text Encoder 1 (CLIP)： 用於初步的文字理解，需 246 MB。
  • VAE： 在像素與潛在空間之間轉換圖像，需 168 MB。

量化工作主要針對 transformer 與 T5 文字編碼器，以達成最大的記憶體節省。

支援的量化後端

bitsandbytes (BnB)

bitsandbytes 提供 8 位元與 4 位元量化。對於 FLUX.1-dev，4 位元量化（使用 NF4）將載入後的記憶體降低至 12.584 GB，峰值記憶體至 17.281 GB，同時在 NVIDIA H100 80GB GPU 上維持 12 秒的推論時間。8 位元量化則會增加記憶體使用（載入 19.273 GB / 峰值 24.432 GB）且推論時間變為 27 秒。

torchao

torchao 是一個原生於 PyTorch 的架構最佳化函式庫。它支援多種僅權重格式：

  • int4_weight_only： 記憶體佔用最低（載入 10.635 GB / 峰值 14.654 GB），但推論速度最慢（109 秒）。
  • int8_weight_only： 記憶體適中（載入 17.020 GB / 峰值 21.482 GB），且推論快速（15 秒）。
  • float8_weight_only： 記憶體與速度與 int8 相似（載入 17.016 GB / 峰值 21.488 GB；15 秒）。

Quanto

透過 optimum 函式庫整合，Quanto 支援多種精度：

  • INT4： 載入 12.254 GB / 峰值記憶體 16.139 GB；推論時間 109 秒。
  • INT8： 載入 17.330 GB / 峰值記憶體 21.814 GB；推論時間 15 秒。
  • FP8： 載入 16.395 GB / 峰值記憶體 20.898 GB；推論時間 16 秒。

GGUF

Diffusers 支援 GGUF 檔案格式，允許透過 from_single_file 使用來自 llama.cpp 社群的預先量化模型。對 FLUX.1-dev 的基準測試顯示：

  • Q2_k： 載入 13.264 GB / 峰值記憶體 17.752 GB；推論時間 26 秒。
  • Q4_1： 載入 16.838 GB / 峰值記憶體 21.326 GB；推論時間 23 秒。
  • Q8_0： 載入 21.502 GB / 峰值記憶體 25.973 GB；推論時間 15 秒。

FP8 分層轉換

enable_layerwise_casting 是一種記憶體最佳化方式，將權重以 FP8 (e4m3) 儲存，但在計算時動態轉換為較高的計算精度（例如 BF16）。此方式將載入後的記憶體降低至 23.682 GB，峰值記憶體至 28.451 GB，推論時間為 13 秒。

結合量化與記憶體最佳化

量化後端可與其他 Diffusers 的最佳化技術結合，以進一步降低 VRAM 使用量：

  • Model CPU Offloading (enable_model_cpu_offload)： 在 CPU 與 GPU 之間移動整個組件。將 BnB 4 位元與此方法結合，可將峰值記憶體降低至 12.383 GB。
  • Group Offloading (enable_group_offload)： 將內部層組移至 CPU。將 FP8 分層轉換與群組卸載結合，可將載入後的記憶體降低至 9.264 GB，峰值記憶體至 14.232 GB。
  • torch.compile： 透過 PyTorch 2.x 加速執行。雖然不會降低記憶體，但能顯著加快推論。例如，torchao int4_weight_only 的推論時間在編譯後從 109 秒降至 6 秒（儘管編譯約需 285 秒）。

後端選擇指南

根據硬體與效能目標，建議使用不同的後端：

  • NVIDIA 使用者（最簡單）： 使用 bitsandbytes 4/8 位元。
  • 推論速度： 使用 torchao、GGUF，或結合 torch.compile() 的 bitsandbytes。
  • 硬體彈性（CPU/MPS）或 FP8： 使用 Quanto。
  • Hopper/Ada 架構： 使用 FP8 分層轉換。
  • 現有 GGUF 模型： 透過 from_single_file 載入 GGUF。

Sources