Hugging Face 将 Nunchaku 4-bit Diffusion 推理集成至 Diffusers
Hugging Face 将 Nunchaku 4-bit Diffusion 推理集成至 Diffusers
Hugging Face 已将 Nunchaku Lite 集成到 Diffusers 库中，允许用户使用 4-bit 权重和激活（W4A4）运行大型扩散 Transformer。与 BF16 基准相比，这种集成将 VRAM 需求降低了高达 50%，并将去噪延迟提高了约 30%，使得在消费级 GPU 上进行高性能图像生成成为可能。
SVDQuant 与 Nunchaku 引擎
Nunchaku Lite 基于 SVDQuant，这是一种旨在处理扩散 Transformer 的权重和激活中存在的巨大离群值的量化方法。与标准的 4-bit 量化不同，SVDQuant 通过将激活离群值移至权重中，并使用小的 16-bit 低秩分支来表示每个权重矩阵中最困难的部分，从而隔离激活离群值。剩余的残差随后被量化为 4-bit。
参考 Nunchaku CUDA 推理引擎使用融合算子（fused kernels）优化了这一过程，这些算子将低秩下投影与量化算子结合，并将低秩上投影与 4-bit 计算算子结合，有效地消除了 16-bit 分支的内存访问开销。
Nunchaku Lite 在 Diffusers 中的集成
Nunchaku Lite 在 Diffusers 中提供了一条流线型的集成路径，允许通过 from_pretrained() 加载 Nunchaku 风格的检查点，而无需单独的推理引擎或本地 CUDA 编译。它通过使用运行时 SVDQ/AWQ 线性层对 nn.Linear 模块进行修补来实现这一点。
Nunchaku Lite 采用了两个主要的算子家族：
svdq_w4a4：用于 Transformer 注意力和 MLP 投影，这是大部分计算发生的地方。它提供 INT4 和 NVFP4 变体。awq_w4a16：用于内存受限且对精度敏感的层，如自适应归一化和调制投影（例如 FLUXadanorm_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 基准 | 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：提供带有 NF4 文本编码器的 INT4 和 NVFP4 变体。
- Krea 2 Turbo：提供 NVFP4 版本。
- lite-infer：一系列额外的 Nunchaku Lite 检查点。