Diffusers 原生支持 Nunchaku 4-bit 推理

大型扩散 Transformer 能生成惊艳的图像（甚至视频、音频和文本），但以 BF16 精度加载现代文生图模型往往需要 20-30 GB 显存，远超多数消费级 GPU 的容量。量化是解决这一问题的有效手段，Diffusers 已集成 bitsandbytes、GGUF、torchao、Quanto 等多个量化后端。

这些后端大多仅对权重进行量化（weight-only）：存储时用低精度，计算时反量化回高精度。这显著降低了显存占用，但通常不会加速推理，甚至可能带来轻微延迟开销。

SVDQuant（背后的推理引擎 Nunchaku）采用不同路线——将主 Transformer 层以 4-bit 权重和激活（W4A4）运行，在降低显存的同时加速去噪循环。此前使用这些检查点需额外安装推理库。

现在，Diffusers 已原生支持 Nunchaku 检查点。加载方式与普通模型无异，仅需调用 from_pretrained()，且无需本地 CUDA 编译，因为内核通过 kernels 包分发。配套工具 diffuse-compressor 还支持用户自行量化新架构并发布为常规 Diffusers 仓库。

快速上手 Nunchaku Lite

安装依赖：

pip install -U diffusers transformers accelerate kernels bitsandbytes

加载预量化流水线：

import torch
from diffusers import ErnieImagePipeline

pipe = ErnieImagePipeline.from_pretrained(
    "lite-infer/ERNIE-Image-Turbo-nunchaku-lite-nvfp4_r32-bnb4-text-encoder",
    torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
    prompt="A cinematic portrait of a red fox in a misty forest at sunrise, detailed fur, volumetric light",
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(42),
).images[0]
image.save("output.png")

该检查点将 Nunchaku NVFP4 Transformer 与 bitsandbytes NF4 文本编码器配对，在 RTX 5090 上生成 1024×1024 图像约需 1.7 秒，峰值显存约 12 GB，而 BF16 流水线约为 24 GB。NVFP4 检查点需要 NVIDIA Blackwell GPU（RTX 50 系列、RTX PRO 6000、B200），早期 GPU 可使用 INT4 变体。

背景：SVDQuant 与 Nunchaku

标准 4-bit 量化对扩散 Transformer 困难，因为权重和激活均包含大量异常值。SVDQuant 将激活异常值转移到权重中，用小型 16-bit 低秩分支表示最难部分，其余残差量化为 4-bit。Nunchaku 通过融合内核加速 4-bit 路径和低秩分支。

Nunchaku Lite 介绍

原始 Nunchaku 引擎依赖模型特定的融合执行路径（如融合 QKV 投影、GELU/MLP 内核），新增模型支持需针对性集成。Nunchaku Lite 是 Diffusers 中的新集成路径，无需自定义流水线或独立推理引擎即可加载 Nunchaku 风格检查点。底层通过运行时 SVDQ/AWQ 线性层替换标准 nn.Linear 模块，CUDA 内核来自 Hub 的 kernels 包。

两种内核族：

• svdq_w4a4：4-bit 权重和激活 + SVDQuant 低秩校正，用于注意力与 MLP 投影。
• awq_w4a16：4-bit 权重 + 16-bit 激活，用于自适应归一化和调制投影（如 FLUX 的 adanorm 层）。

缺乏架构特定融合内核，Nunchaku Lite 速度提升不及原始引擎，但仍可实现约 30% 加速 和同等级别的 显存降低。

Diffusers 原生加载

使用过 bitsandbytes 或 torchao 的用户会感到熟悉。Nunchaku Lite 模型仓库是普通 Diffusers 仓库，仅在 Transformer 的 config.json 中包含 quantization_config 块，指明模块量化方式。

"quantization_config": {
    "quant_method": "nunchaku_lite",
    "compute_dtype": "bfloat16",
    "svdq_w4a4": { "precision": "nvfp4", "group_size": 16, "rank": 32, "targets": ["..."] },
    "awq_w4a16": { "precision": "int4", "group_size": 64, "targets": ["..."] }
}

由于保持与密集模型相同的模块结构，所有下游功能（调度器、LoRA 加载、卸载、torch.compile）均可正常工作。

硬件支持

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

Volta 和 Hopper 暂不支持。

更多加速与更低显存

• torch.compile：编译 Transformer 可将端到端加速从 1.35x 提升至 1.8x。
• 量化文本编码器：使用 bitsandbytes NF4 进一步量化 T5 或 Qwen3 等编码器，显存降低约 22%。
• 卸载：enable_model_cpu_offload() 等正常工作。

基准测试

在 NVIDIA RTX PRO 6000 (Blackwell) 上，1024×1024 分辨率，使用 INT4 检查点：

配置	完整流水线	去噪循环	峰值显存	加速比
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
+ torch.compile	1.68 s	1.53 s	20.6 GB	1.8x
+ NF4 文本编码器	2.29 s	2.13 s	16.0 GB	1.35x

显存最多降低 50%，延迟优化约 30%，torch.compile 可进一步优化。

量化自己的模型

diffuse-compressor 提供端到端 SVDQuant 工作流：校准、量化、打包、发布。以 FLUX.2 Klein 4B 为例：

1. 检查量化目标：python quantize_hf.py black-forest-labs/FLUX.2-klein-4B --precision int4 --rank 32 --inspect-config
2. 运行量化：python quantize_hf.py ... --precision int4
3. 打包为 Diffusers 流水线：python convert_nunchaku_lite_diffusers.py ...
4. 加载验证并推送：pipe.push_to_hub("...")

对于需要结构重写的模型（如 FLUX.1-dev 的 QKV 融合），需使用模型特定目标配置和运行时适配器，详情参见 Adding A New Model。

现成检查点

• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder
• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoder
• OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4
• lite-infer 更多集合

结语

Nunchaku 的 SVDQuant 内核是在消费级硬件上高效运行扩散 Transformer 的最有效方式之一，现已原生集成至 Diffusers。预量化检查点可直接用 from_pretrained() 加载，diffuse-compressor 工具让量化新架构无需等待引擎支持。通过同时量化权重和激活，W4A4 路径在降低显存的同时提升去噪延迟，图像质量接近 BF16 原始版本。

---

原文链接：Hugging Face
本文由前途科技编辑整理