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 模型。工作流程包括:
- 檢查 (Inspection):通用掃描器識別 SVDQ W4A4 目標(重複的 Transformer 區塊)與 AWQ W4A16 目標(調製線性層)。
- 量化 (Quantization):執行 SVDQuant 以建立量化檢查點(支援
int4或nvfp4)。 - 封裝 (Packaging):將量化後的 Transformer 與基礎流水線結合,並在
transformer/config.json中建立nunchaku_lite配置。 - 驗證 (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 檢查點。