Hugging Face 將 Nunchaku 4-bit Diffusion 推論整合至 Diffusers

Hugging Face 將 Nunchaku 4-bit Diffusion 推論整合至 Diffusers

Hugging Face 已將 Nunchaku Lite 整合進 Diffusers 函式庫,讓使用者能夠使用 4-bit 權重與激活值 (W4A4) 執行大型擴散 Transformer。與 BF16 基線相比,此整合減少了高達 50% 的 VRAM 需求,並將去噪延遲縮短了約 30%,使得在消費級 GPU 上進行高性能圖像生成變得觸手可及。

SVDQuant 與 Nunchaku Engine

Nunchaku Lite 基於 SVDQuant,這是一種專為處理擴散 Transformer 中權重與激活值中出現的大量離群值 (outliers) 而設計的量化方法。與標準的 4-bit 量化不同,SVDQuant 透過將激活離群值移至權重中,並使用一個小的 16-bit 低階秩 (low-rank) 分支來表示每個權重矩陣中最困難的部分,從而隔離激活離群值。其餘的殘差則被量化為 4-bit。

參考實現的 Nunchaku CUDA 推論引擎使用融合內核 (fused kernels) 優化了此過程,將低階秩下投影與量化內核結合,並將低階秩上投影與 4-bit 計算內核結合,有效地消除了 16-bit 分支的記憶體存取開銷。

Nunchaku Lite 在 Diffusers 中的整合

Nunchaku Lite 在 Diffusers 中提供了一條流線型的整合路徑,允許透過 from_pretrained() 載入 Nunchaku 風格的檢查點 (checkpoints),而無需額外的推論引擎或本地 CUDA 編譯。它透過使用運行時 SVDQ/AWQ 線性層來修補 nn.Linear 模組來實現這一點。

Nunchaku Lite 採用了兩大核心內核家族:

  • svdq_w4a4:用於 Transformer 的 Attention 與 MLP 投影,這是大部分計算發生的地方。它提供 INT4 與 NVFP4 變體。
  • awq_w4a16:用於受記憶體限制且對精度敏感的層,例如自適應歸一化 (adaptive normalization) 與調製投影 (modulation projections)(例如 FLUX 的 adanorm_single / adanorm_zero),使用 4-bit 權重與 16-bit 激活值。

雖然 Nunchaku Lite 並未包含原始 Nunchaku 引擎中針對特定架構的融合執行路徑,但它仍能提供約 30% 的加速並達到相同的 VRAM 減少效果。

硬體支援與效能

對 Nunchaku Lite 內核的支援因 GPU 代際與精度而異:

方案 精度 支援的 GPU
svdq_w4a4 nvfp4 Blackwell (RTX 50 series, RTX PRO 6000, B200)
svdq_w4a4 int4 Turing / Ampere / Ada (RTX 30 & 40 series, A100, L40S)
awq_w4a16 int4 Turing / Ampere / Ada (RTX 30 & 40 series, A100, L40S)

目前尚不支援 Volta 與 Hopper GPU。

基準測試

在 NVIDIA RTX PRO 6000 (Blackwell) 上,使用 ERNIE-Image-Turbo 檢查點,於 1024x1024 解析度下測得:

配置 完整流水線 去噪迴圈 峰值 VRAM 加速比
BF16 baseline 3.00 s 2.86 s 31.1 GB 1.0x
Nunchaku Lite NVFP4 2.27 s 2.13 s 20.6 GB 1.35x
Nunchaku Lite NVFP4 + torch.compile 1.68 s 1.53 s 20.6 GB 1.8x
Nunchaku Lite NVFP4 + NF4 text encoder 2.29 s 2.13 s 16.0 GB 1.35x

將 Nunchaku Lite 與 torch.compile 結合可將端到端加速提升至 1.8x,而加入 bitsandbytes NF4 量化文本編碼器則能進一步將峰值 VRAM 降低至 16.0 GB。

使用 diffuse-compressor 量化自定義模型

Nunchaku Lite 與架構無關。diffuse-compressor 工具包允許使用者校準、量化、封裝並發布他們自己的 Diffusers 模型。工作流程包括:

  1. 檢查 (Inspection):通用掃描器識別 SVDQ W4A4 目標(重複的 Transformer 區塊)與 AWQ W4A16 目標(調製線性層)。
  2. 量化 (Quantization):執行 SVDQuant 以建立量化檢查點(支援 int4nvfp4)。
  3. 封裝 (Packaging):將量化後的 Transformer 與基礎流水線結合,並在 transformer/config.json 中建立 nunchaku_lite 配置。
  4. 驗證 (Verification):透過 DiffusionPipeline.from_pretrained() 載入模型,並將最終結果推送到 Hugging Face Hub。

結構重寫

為了獲得最大效能,某些模型需要進行結構重寫——例如將獨立的 Q、K、V 投影合併為單個 to_qkv 模組。雖然通用的 diffuse-compressor 路徑無法推斷這些變更,但可以透過 rootonchair/nunchaku-lite 提供的模型特定目標配置與運行時適配器來實現,以啟用融合操作。

現成可用的檢查點

Hub 上提供數個預量化的檢查點,包括:

  • ERNIE-Image-Turbo:提供 INT4 與 NVFP4 變體,並配備 NF4 文本編碼器。
  • Krea 2 Turbo:提供 NVFP4 版本。
  • lite-infer:一系列額外的 Nunchaku Lite 檢查點。

Sources