Diffusers 原生支持 Nunchaku 4-bit 扩散推理

大型扩散 Transformer 可以生成令人惊叹的图像（甚至视频、音频片段，如今还有文本），但以 BF16 精度加载现代文本到图像模型通常需要 20-30 GB 的显存，这让大多数消费级 GPU 无法运行这些模型。量化是解决这一问题的有力方案，Diffusers 已经集成了多种量化后端，如 bitsandbytes、GGUF、torchao 和 Quanto，我们在 Exploring Quantization Backends in Diffusers 中对此进行了介绍。

这些后端大多是仅权重量化。这意味着它们以低精度存储权重，并在计算时将其反量化回高精度。这显著减少了内存占用，但通常不会加快推理速度，甚至可能带来少量延迟开销。

SVDQuant 是流行的 Nunchaku 推理引擎背后的量化方法，它采用了不同的思路。它以 4 位权重和激活（W4A4）运行主要的 Transformer 层，在减少内存的同时加速去噪循环。具体细节见下文，但在此之前，使用这些检查点需要一个单独的推理库。

在当前的 Diffusers 中，加载 Nunchaku 检查点只需调用 from_pretrained()，得益于 kernels 包，无需本地 CUDA 编译。此外，配套的 diffuse-compressor 工具包让你可以自行量化新架构，并将其作为常规 Diffusers 仓库发布。

目录

• Nunchaku Lite 入门
• 背景：SVDQuant 与 Nunchaku
• 介绍 Nunchaku Lite
• 在 Diffusers 中原生加载
• 获得更高速度和更低内存
• 基准测试
• 量化你自己的模型
• 即用型检查点
• 结论
• 致谢

Nunchaku Lite 入门

首先，安装依赖项。你需要最新版本的 Diffusers 和 Hugging Face kernels 包：

pip install -U diffusers transformers accelerate kernels bitsandbytes

然后像加载其他任何 Diffusers 模型一样加载预量化的 pipeline：

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")

无需自定义 pipeline 类或单独的推理引擎，也无需在本地编译任何内容。NVFP4 内核在首次使用时通过 Nunchaku Lite kernels 页面从 Hub 下载。此检查点将 Nunchaku NVFP4 Transformer 与 bitsandbytes NF4 文本编码器配对，在 RTX 5090 上约 1.7 秒内生成 1024x1024 图像，峰值内存使用约 12 GB，而 BF16 pipeline 约为 24 GB。你可以在官方 Diffusers 文档中找到有关 Nunchaku Lite 检查点格式的更多详细信息。

NVFP4 检查点需要 NVIDIA Blackwell GPU（RTX 50 系列、RTX PRO 6000、B200）。对于更早的世代，请使用 INT4 变体。详情请参阅下方的硬件支持表。

背景：SVDQuant 与 Nunchaku

SVDQuant 是 Nunchaku 背后的量化方法，也是其参考 CUDA 推理引擎。标准 4 位量化对扩散 Transformer 来说很困难，因为权重和激活中都包含大量离群值。SVDQuant 通过将激活离群值移入权重来处理这一问题，用一个小型 16 位低秩分支表示每个权重矩阵中最难的部分，并将剩余的残差量化为 4 位。Nunchaku 通过为 4 位路径和低秩分支提供融合内核，使这一过程变得快速。

推出 Nunchaku Lite

原始 Nunchaku 引擎 的很大一部分速度来自 模型特定的融合执行路径，例如融合的 QKV 投影和融合的 GELU/MLP 内核。这些优化与每种架构的模块布局和检查点格式绑定，因此支持新的模型系列通常需要针对特定模型进行集成工作。

Nunchaku Lite 是 Diffusers 中新的集成路径。借助它，Diffusers 无需自定义 pipeline 或单独的推理引擎即可加载 Nunchaku 风格的检查点。在底层，Nunchaku Lite 会在加载检查点之前，用运行时 SVDQ/AWQ 线性层替换标准 Diffusers 模型中的相关 nn.Linear 模块。CUDA 内核通过 kernels 包从 Hub 获取。使用了两类内核：

• svdq_w4a4：4 位权重和激活，并带有 SVDQuant 低秩校正。该层用于 Transformer 的注意力和 MLP 投影，几乎所有计算都花在这里，并提供 INT4 和 NVFP4 变体。
• awq_w4a16：4 位权重和 16 位激活，用于自适应归一化和调制投影，例如 FLUX adanorm_single / adanorm_zero 或 Qwen-Image 调制层。这些层受内存带宽限制且对精度敏感，因此 AWQ 非常适合在节省内存和空间的同时保持精度。

代价是，如果没有针对特定架构的融合内核和模块，Nunchaku Lite 无法达到原始 Nunchaku 引擎的加速效果。不过，这个精简实现仍能带来约 30% 的加速，同时保持相同程度的 显存降低。

在 Diffusers 中原生加载

如果你在 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": [
            "layers.0.self_attention.to_q",
            "layers.0.self_attention.to_k",
            "..."
        ]
    },
    "awq_w4a16": {
        "precision": "int4",
        "group_size": 64,
        "targets": [
            "adaLN_modulation.1",
            "..."
        ]
    }
}

该配置告诉 Diffusers 哪些模块被量化、它们使用哪种方案，以及要实例化哪个 Nunchaku Lite 运行时层（SVDQW4A4Linear 或 AWQW4A16Linear）。

由于量化模型保留了稠密模型完全相同的模块结构，下游的一切（调度器、LoRA 加载钩子、卸载、torch.compile）看到的都是普通的 Diffusers 模型。

硬件支持

Nunchaku Lite 根据 GPU 代际和检查点精度使用不同的内核变体：

Volta 和 Hopper GPU 目前不受 4-bit 内核支持。量化器会在加载时验证 GPU 的 CUDA 计算能力，并抛出明确的错误，而不是产生错误的输出。

获得更高速度和更低内存占用

Nunchaku Lite 可以与其他 Diffusers 内存和速度优化结合使用。

torch.compile。 编译 transformer 可将端到端加速从 1.35 倍提升至 1.8 倍：

pipe.transformer.compile(fullgraph=True)

# or compile_repeated_blocks() for faster compilation

pipe.transformer.compile_repeated_blocks(fullgraph=True)

量化文本编码器。 transformer 并不是唯一占用大量内存的组件。诸如 T5 或 Qwen3 之类的文本编码器本身就可能占用数 GB 内存。在我们的基准测试中，使用 bitsandbytes NF4 进一步量化文本编码器可将峰值 VRAM 降低约 22%。

卸载。 如果你需要将流水线适配到更小的 GPU 上，Diffusers 的卸载辅助工具（如 enable_model_cpu_offload() 和 enable_sequential_cpu_offload()）可照常使用。

基准测试

以下所有数据均在 NVIDIA RTX PRO 6000（Blackwell）上以 1024x1024 分辨率、使用 rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder 测得。

端到端延迟与内存

如上所示，Nunchaku 可将峰值 VRAM 降低多达 50%，同时仍将延迟改善约 30%。剩余的开销主要来自额外的内核启动，torch.compile 可以缓解这一问题，使完整流水线降至 1.68 s，即比 BF16 基线快 1.8 倍。

图像质量

量化你自己的模型

Diffusers 中对 Nunchaku Lite 的支持与架构无关，diffuse-compressor 工具包为 Diffusers 模型提供了端到端的 SVDQuant 工作流：校准、量化、打包和发布。

下面，我们以量化 FLUX.2 Klein 4B 为例进行演示。它涵盖了主要步骤：检查模型、校准并量化 transformer、将结果打包为 Diffusers 流水线，然后验证并推送到 Hub。完整教程详细介绍了每一个标志。

1. 检查将要量化的内容

通用扫描器会遍历模型并决定目标：重复的 transformer 块堆栈内兼容的线性层会成为 SVDQ W4A4 目标，识别出的调制线性层会成为 AWQ W4A16 目标，其余所有内容保持稠密。

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
  --precision int4 --rank 32 --inspect-config

在量化之前务必阅读此报告。对于 FLUX.2 Klein 4B，预期结果是 100 个 SVDQ 目标、3 个 AWQ 目标和 6 个稠密外部线性层，没有缺失的模式或重复的名称。

2. 运行量化

以下命令在 transformer 上运行 SVDQuant，并将量化后的检查点写入 outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors：

python examples/text_to_image/quantize_hf.py black-forest-labs/FLUX.2-klein-4B \
  --precision int4 \
  --output outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors

将 --precision int4 替换为 nvfp4 以构建 Blackwell 原生权重。

3. 打包 Diffusers 流水线

转换器将量化后的 transformer 与基础流水线的其他组件组合在一起，将紧凑的 nunchaku_lite 配置写入 transformer/config.json，并可选择将文本编码器转换为 NF4：

python examples/convert_nunchaku_lite_diffusers.py \
  --checkpoint outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors \
  --model-id black-forest-labs/FLUX.2-klein-4B \
  --bnb4-text-encoder text_encoder \
  --compute-dtype bfloat16 \
  --output-dir outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder

4. 加载、验证并推送到 Hub

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder",
    device_map="cuda",
)
image = pipe(
    "A glass robot in a greenhouse, cinematic lighting",
    num_inference_steps=4, guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(12345),
).images[0]

一旦输出看起来不错，运行 pipe.push_to_hub("your-name/your-model-nunchaku-lite-int4")。其他用户随后即可使用上面所示的相同 from_pretrained() 模式加载它。

量化具有结构重写的模型

注意，通用路径假设架构可以在不进行结构重写的情况下被量化。为了进一步加速，原始 Nunchaku 引擎会将成组的 Diffusers 层重写为融合模块。通用路径无法自行推断这些更改，例如将独立的 Q、K、V 投影合并为一个模块，或将一个融合投影拆分为多个模块。

FLUX.1-dev 的 QKV 投影就是一个具体例子。Diffusers 定义了三个独立的模块：

self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)

Nunchaku FLUX 模块将这些层合并为一个量化的 to_qkv 模块：

to_qkv = fuse_linears([other.to_q, other.to_k, other.to_v])
self.to_qkv = SVDQW4A4Linear.from_linear(to_qkv, **kwargs)

之所以需要这个分组模块，是因为 Nunchaku 的融合算子会同时消费 QKV 投影、Q/K 归一化和旋转嵌入。相比之下，默认的 Diffusers 路径会分别执行它们：

query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)

query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))

query = attn.norm_q(query)
key = attn.norm_k(key)

if image_rotary_emb is not None:
    query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
    key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)

Nunchaku 路径将分组投影、归一化模块和旋转嵌入提供给一个融合算子：

qkv = fused_qkv_norm_rottary(
    hidden_states, attn.to_qkv, attn.norm_q, attn.norm_k, image_rotary_emb
)

这就是通用路径无法推断的结构重写。Diffusers 有三个目标模块，分别带有 to_q、to_k 和 to_v 参数前缀，而 Nunchaku 在 to_qkv 下有一个分组模块。特定于模型的目标配置或适配器必须说明 Q、K 和 V 参数应按该顺序沿输出维度拼接，并加载到 to_qkv 中。

像这样的结构重写由特定于模型的目标配置在量化期间描述，并在加载检查点时由一个小型运行时适配器处理。 FLUX.2 Klein 4B 量化脚本提供了一个具体的目标配置示例，用于生成经过结构重写的检查点，而 rootonchair/nunchaku-lite 提供了加载分组 QKV 张量、拆分融合投影以及其他融合操作所需的运行时适配器。 有关完整工作流程，你可以查看添加新模型指南。

可直接使用的检查点

要立即开始，请查看以下仓库：

• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder：INT4 ERNIE-Image-Turbo，搭配 bitsandbytes NF4 文本编码器
• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoder：NVFP4 ERNIE-Image-Turbo，搭配 bitsandbytes NF4 文本编码器
• OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4：NVFP4 Krea 2 Turbo 检查点
• lite-infer：更多 Nunchaku Lite 检查点和集合

结论

Nunchaku 的 SVDQuant 内核是在消费级硬件上高效运行扩散 Transformer 的最有效方法之一，而且它们现在已在 Diffusers 中获得原生支持。预量化检查点可通过 from_pretrained() 加载，而 diffuse-compressor 工具包使得无需等待引擎支持即可量化新架构。通过同时量化权重和激活，W4A4 路径在降低内存使用的同时改善了去噪延迟，使图像质量接近 BF16 原始版本。

如果你量化并发布了新模型，我们很乐意了解。请在 Hub 上分享并告诉我们！如果你对此功能有任何疑问，欢迎加入我们的 Discord。

要了解更多信息，请查看以下资源：

• Diffusers Nunchaku 文档
• 集成 PR（huggingface/diffusers#14100）
• SVDQuant 论文和 Nunchaku 引擎
• diffuse-compressor
• 之前的文章：探索 Diffusers 中的量化后端 和 使用 Quanto 和 Diffusers 实现内存高效的 Diffusion Transformers

致谢

感谢 Diffusers 维护者在整个集成过程中提供的评审和指导，感谢 MIT HAN Lab / Nunchaku 团队所做的原始 SVDQuant 工作。感谢 Marc Sun 对这篇博客文章提供的反馈。感谢 Álvaro Somoza 试用 nunchaku-lite 并提供反馈。

rootonchair 还要感谢 SilverAI 对这项工作的支持，并提供了大部分开发工作得以进行的环境。