大型擴散 Transformer 模型可以創造出令人驚嘆的圖像(甚至是影片、音訊片段,現在還有文字),但以 BF16 精度載入現代文字轉圖像模型通常需要 20-30 GB 的顯示記憶體,這使得這些模型超出了大多數消費級 GPU 的能力範圍。
量化是解決這個問題的強大方案,而 Diffusers 已經整合了多種量化後端,例如 bitsandbytes、GGUF、torchao 和 Quanto,我們在《探索 Diffusers 中的量化後端》一文中曾介紹過。
這些後端大多是僅對權重進行量化。這表示它們以低精度儲存權重,並在計算時將其反量化回高精度。這顯著減少了記憶體使用量,但通常不會加快推論速度,甚至可能增加微小的延遲開銷。
SVDQuant 是流行推論引擎 Nunchaku 背後的量化方法,它採用了不同的方法。它使用 4 位元權重和激活值 (W4A4) 運行主要的 Transformer 層,在減少記憶體的同時也加速了去噪迴圈。詳細資訊將在下方介紹,但在此之前,使用這些檢查點需要一個單獨的推論函式庫。
透過目前的 Diffusers,載入 Nunchaku 檢查點就像呼叫 from_pretrained() 一樣簡單,由於核心套件的支援,無需進行本地 CUDA 編譯。此外,配套的 diffuse-compressor 工具包讓您可以自行量化新的架構,並將其發布為常規的 Diffusers 儲存庫。
首先,您需要安裝所需的套件,包含最新版本的 Diffusers 和 Hugging Face 核心套件。接著,您可以像載入任何其他 Diffusers 模型一樣,載入預量化的管線。例如,您可以載入一個預量化的 ErnieImagePipeline 模型,並使用它來生成圖像。
無需自訂管線類別或單獨的推論引擎,也無需在本地編譯任何東西。NVFP4 核心會在首次使用時從 Hub 透過 Nunchaku Lite 核心頁面下載。這個檢查點將 Nunchaku NVFP4 Transformer 與 bitsandbytes NF4 文字編碼器配對,在 RTX 5090 上生成 1024x1024 圖像大約需要 1.7 秒,峰值記憶體使用量約為 12 GB,而 BF16 管線則約為 24 GB。
您可以在 Diffusers 官方文件中找到更多關於 Nunchaku Lite 檢查點格式的詳細資訊。
NVFP4 檢查點需要 NVIDIA Blackwell GPU(RTX 50 系列、RTX PRO 6000、B200)。對於早期世代的 GPU,請使用 INT4 變體。有關詳細資訊,請參閱下面的硬體支援表。
SVDQuant 是 Nunchaku 及其參考 CUDA 推論引擎背後的量化方法。對於擴散 Transformer 模型來說,標準的 4 位元量化很困難,因為權重和激活值都包含大量異常值。SVDQuant 透過將激活值異常值移至權重中來處理這個問題,用一個小的 16 位元低秩分支表示每個權重矩陣中最困難的部分,並將剩餘的殘差量化為 4 位元。Nunchaku 透過用於 4 位元路徑和低秩分支的融合核心使其快速運行。
Nunchaku 將低秩向下投影與量化核心融合,並將低秩向上投影與 4 位元計算核心融合,消除了 16 位元分支的記憶體存取開銷。
原始的 Nunchaku 引擎大部分速度來自於模型特定的融合執行路徑,例如融合的 QKV 投影和融合的 GELU/MLP 核心。這些優化與每個架構的模組佈局和檢查點格式相關聯,因此支援新的模型家族通常需要模型特定的整合工作。
Nunchaku Lite 是 Diffusers 中的新整合路徑。透過它,Diffusers 可以載入 Nunchaku 風格的檢查點,而無需自訂管線或單獨的推論引擎。在底層,Nunchaku Lite 會在載入檢查點之前,用執行時 SVDQ/AWQ 線性層修補標準 Diffusers 模型的相關 nn.Linear 模組。CUDA 核心來自 Hub 透過核心套件提供。使用了兩種核心家族:
svdq_w4a4:具有 SVDQuant 低秩校正的 4 位元權重和激活值。這個層用於 Transformer 的注意力與 MLP 投影,幾乎所有的計算都在這裡進行,並提供 INT4 和 NVFP4 變體。
awq_w4a16:具有 16 位元激活值的 4 位元權重,用於自適應正規化和調變投影,例如 FLUX adanorm_single / adanorm_zero 或 Qwen-Image 調變層。這些層受記憶體限制且對精度敏感,這使得 AWQ 成為在節省記憶體和空間的同時保持精度的良好選擇。
權衡是,在沒有架構特定融合核心和模組的情況下,Nunchaku Lite 無法達到原始 Nunchaku 引擎的速度提升。然而,這個精簡的實作仍然提供了約 30% 的速度提升,同時保持了相同的顯示記憶體減少水平。
如果您在 Diffusers 中使用過 bitsandbytes 或 torchao,其機制會讓您感到熟悉。Nunchaku Lite 模型儲存庫是一個普通的 Diffusers 儲存庫。唯一特殊的部分是 Transformer 的 config.json 中的 quantization_config 區塊。
這個配置告訴 Diffusers 哪些模組被量化、它們使用哪種方案,以及要實例化哪個 Nunchaku Lite 執行時層(SVDQW4A4Linear 或 AWQW4A16Linear)。
由於量化模型保留了密集模型的確切模組結構,因此所有下游功能(排程器、LoRA 載入掛鉤、卸載、torch.compile)都將其視為正常的 Diffusers 模型。
Nunchaku Lite 根據 GPU 世代和檢查點精度使用不同的核心變體。例如,svdq_w4a4 的 nvfp4 精度支援 Blackwell GPU,而 int4 精度則支援 Turing / Ampere / Ada 系列 GPU。awq_w4a16 的 int4 精度也支援 Turing / Ampere / Ada 系列 GPU。
Volta 和 Hopper GPU 目前不支援 4 位元核心。量化器會在載入時驗證 GPU 的 CUDA 能力,並在產生不正確輸出之前引發明確的錯誤。
Nunchaku Lite 可以與其他 Diffusers 記憶體和速度優化結合使用。編譯 Transformer 可以將端到端的速度提升從 1.35 倍提高到 1.8 倍。此外,進一步量化文字編碼器(例如使用 bitsandbytes NF4)可以將峰值顯示記憶體減少約 22%。
Diffusers 的卸載輔助功能,例如 enable_model_cpu_offload() 和 enable_sequential_cpu_offload(),如果需要將管線適應到較小的 GPU 上,也可以照常使用。
所有以下數據均在 NVIDIA RTX PRO 6000 (Blackwell) 上以 1024x1024 解析度使用 rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder 測量。
如上所示,Nunchaku 將峰值顯示記憶體減少了高達 50%,同時仍將延遲提高了約 30%。剩餘的開銷主要來自額外的核心啟動,torch.compile 可以緩解這些開銷,使整個管線的運行時間降至 1.68 秒,比 BF16 基準快 1.8 倍。
Diffusers 中對 Nunchaku Lite 的支援是架構無關的,而 diffuse-compressor 工具包為 Diffusers 模型提供了端到端的 SVDQuant 工作流程:校準、量化、打包和發布。下面,我們以量化 FLUX.2 Klein 4B 為例進行說明。
它涵蓋了主要步驟:檢查模型、校準和量化 Transformer、將結果打包為 Diffusers 管線,然後驗證並推送到 Hub。
通用掃描器會遍歷模型並決定要量化的目標:重複的 Transformer 區塊堆疊中相容的線性層成為 SVDQ W4A4 目標,識別出的調變線性層成為 AWQ W4A16 目標,而其他所有層則保持密集。在量化之前務必閱讀此報告。對於 FLUX.2 Klein 4B,預期結果是 100 個 SVDQ 目標、3 個 AWQ 目標和 6 個密集外部線性層,沒有遺漏模式或重複名稱。
執行 SVDQuant 於 Transformer 並將量化後的檢查點寫入指定路徑。您可以將 --precision int4 替換為 nvfp4 以建立 Blackwell 原生權重。
轉換器將量化後的 Transformer 與基礎管線的其他組件結合,將緊湊的 nunchaku_lite 配置寫入 transformer/config.json,並可選擇將文字編碼器轉換為 NF4。一旦輸出看起來良好,就可以將模型推送到 Hub。
