Bringing Nunchaku 4-bit Diffusion Inference to Diffusers

将 Nunchaku 4 位扩散推理引入 Diffusers

返回文章列表

[-1

]

[0

2026 年 7 月 23 日发布

GitHub 上的更新

点赞

16

[

• +10

Pham Hong Vinh

rootonchair

关注

[1

guest

Sayak Paul

sayakpaul

大型扩散变压器可以生成令人惊叹的图像（甚至视频、音频片段和文本），但以 BF16 精度加载现代文本到图像模型通常需要 20-30 GB 的 VRAM，这使得这些模型对大多数消费级 GPU 来说难以使用。量化是解决这个问题的强大方法，Diffusers 已经集成了多种量化后端，如 bitsandbytes、GGUF、torchao 和 Quanto，我们在《探索 Diffusers 中的量化后端》中对此进行了介绍。

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

流行的 Nunchaku 推理引擎背后的量化方法 SVDQuant 采用了不同的方法。它使用 4 位权重和激活（W4A4）运行主要的变压器层，在减少内存的同时加速去噪循环。以下部分将详细介绍这些内容，但此前使用这些检查点需要单独的推理库。

通过当前的 Diffusers，加载 Nunchaku 检查点只需调用 from_pretrained()，无需本地 CUDA 编译，这得益于 kernels 包。此外，配套的 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 模型一样加载预量化管道：

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

无需自定义管道类或单独的推理引擎，也无需本地编译任何内容。首次使用时，NVFP4 内核会通过 Nunchaku Lite 内核页面从 Hub 下载。此检查点配对了 Nunchaku NVFP4 变压器和 bitsandbytes NF4 文本编码器，在 RTX 5090 上生成 1024x1024 图像大约需要 1.7 秒，峰值内存使用量约为 12 GB，而 BF16 管道约为 24 GB。有关 Nunchaku Lite 检查点格式的更多细节，请参阅官方 Diffusers 文档。

NVFP4检查点需要NVIDIA Blackwell GPU（RTX 50系列、RTX PRO 6000、B200）。对于早期版本，使用INT4变体。有关详细信息，请参阅下方的硬件支持表。

背景：SVDQuant和Nunchaku

SVDQuant是Nunchaku背后采用的量化方法，以及其参考CUDA推理引擎。由于扩散变压器的权重和激活值中包含大量异常值，标准4位量化存在困难。SVDQuant通过将激活值的异常值转移到权重中，用小的16位低秩分支表示每个权重矩阵中最难处理的部分，并将剩余的残差量化为4位来解决这个问题。Nunchaku通过融合4位路径和低秩分支的内核来加速这一过程。

Nunchaku将低秩下投影与量化内核融合，将低秩上投影与4位计算内核融合，消除了16位分支的内存访问开销。图源自SVDQuant论文。

引入Nunchaku Lite

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

Nunchaku Lite是Diffusers中的新集成路径。通过它，Diffusers可以在不使用自定义管道或单独推理引擎的情况下加载Nunchaku风格的检查点。在底层，Nunchaku Lite在检查点加载前，通过运行时SVDQ/AWQ线性层修补标准Diffusers模型中相关nn.Linear模块。CUDA内核通过kernels包从Hub获取。使用两种内核家族：

• svdq_w4a4：使用SVDQuant低秩校正的4位权重和激活值。该层用于变压器的注意力和MLP投影，这些部分几乎消耗了全部计算资源，提供INT4和NVFP4变体。
• awq_w4a16：4位权重与16位激活值，用于自适应归一化和调制投影，如FLUX adanorm_single/adanorm_zero或Qwen-Image调制层。这些层受内存和精度限制，AWQ在保持精度的同时节省内存和空间，是理想选择。

权衡在于，由于缺乏特定架构的融合内核和模块，Nunchaku Lite无法达到原始Nunchaku引擎的加速效果。不过，这种基础实现仍能提供约30%的加速，同时保持相同的VRAM减少水平。

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 世代和检查点精度使用不同的内核变体：

方案

精度

支持的 GPU

svdq_w4a4

nvfp4

Blackwell（RTX 50 系列、RTX PRO 6000、B200）

int4

Turing / Ampere / Ada（RTX 30 & 40 系列、A100、L40S）

awq_w4a16

当前 Volta 和 Hopper GPU 不支持 4-bit 内核。量化器在加载时会验证 GPU 的 CUDA 能力，并显示明确错误而不是生成错误输出。

获取更高速度和更低内存

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

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

pipe.transformer.
compile
(fullgraph=
True
)
# 或使用 compile_repeated_blocks() 以加快编译速度
pipe.transformer.compile_repeated_blocks(fullgraph=
True
)

量化文本编码器。Transformer 并不是唯一占用大量内存的组件。T5 或 Qwen3 等文本编码器本身可能占用数 GB 内存。使用 bitsandbytes NF4 进一步量化文本编码器，可在我们的基准测试中减少约 22% 的峰值 VRAM。

卸载。如果需要将管道适配到更小的 GPU，Diffusers 的卸载辅助工具如 enable_model_cpu_offload() 和 enable_sequential_cpu_offload() 会按常规方式工作。

基准测试

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

端到端延迟和内存

配置

完整管道

去噪循环

峰值 VRAM

加速比

BF16 基线

3.00 秒

2.86 秒

31.1 GB

1.0x

Nunchaku Lite NVFP4

2.27 秒

2.13 秒

20.6 GB

1.35x

Nunchaku Lite NVFP4 +

torch.compile

1.68 秒

1.53 秒

1.8x

Nunchaku Lite NVFP4 + NF4 文本编码器

2.29 秒

16.0 GB

如上所示，Nunchaku 最多可减少 50% 的峰值 VRAM，同时仍能将延迟降低约 30%。剩余的开销主要来自额外的内核启动，torch.compile 可缓解此问题，使完整管道降至 1.68 秒，比 BF16 基线快 1.8 倍。

图像质量

使用相同种子和设置的 BF16 与 4-bit 输出对比。

量化你自己的模型

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

下面以量化 FLUX.2 Klein 4B 为例，逐步演示过程。它涵盖了主要步骤：检查模型、校准并量化 Transformer、将结果打包为 Diffusers 管道，然后验证并推送到 Hub。完整教程会详细说明每个参数。

1. 检查将被量化的部分

通用扫描器遍历模型并决定目标：重复 Transformer-block 栈中的兼容线性层会成为 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 : 使用 bitsandbytes NF4 文本编码器的 INT4 ERNIE-Image-Turbo 模型
• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoder : 使用 bitsandbytes NF4 文本编码器的 NVFP4 ERNIE-Image-Turbo 模型
• OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4 : NVFP4 Krea 2 Turbo 检查点
• lite-infer : 更多 Nunchaku Lite 检查点和集合

结论

Nunchaku 的 SVDQuant 内核是目前在消费级硬件上高效运行扩散变换器最有效的方式之一，现在已原生支持于 Diffusers。预量化检查点可通过 from_pretrained() 加载，diffuse-compressor 工具包使得在无需等待引擎支持的情况下即可量化新架构。通过同时量化权重和激活值，W4A4 路径在降低内存使用的同时提升了去噪延迟，同时保持图像质量接近原始 BF16 模型。

如果你量化并发布了新模型，我们非常期待你的分享。请在 Hub 上发布并告诉我们！如果你对这个功能有任何疑问，欢迎加入我们的 Discord 社区。

了解更多内容，请查看以下资源：

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

致谢

感谢 Diffusers 维护者在整个集成过程中的审阅和指导，感谢 MIT HAN 实验室/Nunchaku 团队的原始 SVDQuant 工作。感谢 Marc Sun 对博文的反馈建议，感谢 Álvaro Somoza 测试 nunchaku-lite 并提供反馈。

rootonchair 也感谢 SilverAI 对这项工作的支持，以及提供使大部分开发得以进行的环境。

本文提到的模型 1

更多来自我们博客的文章

guide

diffusers

quantization

（LoRA）在消费级硬件上微调 FLUX.1-dev

• +1

107

2025年6月19日

探索 Diffusers 中的量化后端

45

2025年5月21日

社区

编辑

预览

/

上传图片、音频和视频，可通过将文件拖拽到文本输入框、粘贴或

点击此处

。

点击或粘贴此处以上传图片

评论

· 注册或登录以发表评论

• +4