Nunchaku 4-bit 量化推理已集成 Hugging Face Diffusers

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

SVDQuant 是流行的 Nunchaku 推理引擎背后的量化方法，采用了不同的方法。它使用 4 位权重和激活（W4A4）运行主变换器层，既减少内存又加快了去噪循环。下面会介绍详细内容，但到目前为止，使用这些检查点需要一个单独的推理库。

使用最新的 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 模型一样加载预量化的管道：

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 内核从 Hub 通过 Nunchaku Lite kernels 页面在第一次使用时下载。这个检查点将 Nunchaku NVFP4 变换器与 bitsandbytes NF4 文本编码器配对，在 RTX 5090 上生成 1024x1024 的图像仅需约 1.7 秒，峰值内存使用约 12 GB，相比 BF16 管道的约 24 GB。你可以在官方 Diffusers 文档中找到关于 Nunchaku Lite 检查点格式的更多细节。

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

背景：SVDQuant 和 Nunchaku

SVDQuant 是 Nunchaku 背后的量化方法，Nunchaku 是它的参考 CUDA 推理引擎。标准 4 位量化对于扩散变换器很困难，因为权重和激活都包含大的异常值。SVDQuant 通过将激活异常值移动到权重中来处理这个问题，用小的 16 位低秩分支表示每个权重矩阵最难的部分，并将剩余的残差量化为 4 位。Nunchaku 通过为 4 位路径和低秩分支的融合内核使这一步变得高效。

介绍 Nunchaku Lite

原始 Nunchaku 引擎从模型特定的融合执行路径（如融合 QKV 投影和融合 GELU/MLP 内核）获得大部分速度。这些优化与每个架构的模块布局和检查点格式相关，所以支持新的模型族通常需要模型特定的集成工作。

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

svdq_w4a4：4 位权重和激活，带有 SVDQuant 低秩修正。这个层用于变换器的注意力和 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 存储库。唯一特殊的部分是变换器 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 位内核。量化器在加载时验证 GPU 的 CUDA 能力，并提出明确的错误，而不是产生不正确的输出。

获得更快的速度和更低的内存

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

torch.compile。编译变换器将端到端加速从 1.35x 改进到 1.8x：

pipe.transformer.compile(fullgraph=True)

# or compile_repeated_blocks() for faster compilation

pipe.transformer.compile_repeated_blocks(fullgraph=True)

量化的文本编码器。变换器不是唯一有大量内存占用的组件。文本编码器如 T5 或 Qwen3 自己可以占用几 GB。用 bitsandbytes NF4 进一步量化文本编码器会在我们的基准测试中减少约 22% 的峰值 VRAM。

卸载。Diffusers 卸载助手如 enable_model_cpu_offload() 和 enable_sequential_cpu_offload() 在你需要将管道装入较小 GPU 时照常工作。

下面的所有数字都在 NVIDIA RTX PRO 6000（Blackwell）上以 1024x1024 分辨率使用 rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder 测量。

端到端延迟和内存

如上所示，Nunchaku 将峰值 VRAM 减少高达 50%，同时仍然改进延迟约 30%。剩余的开销在很大程度上来自额外的内核启动，torch.compile 可以缓解，将完整管道降低到 1.68 秒，或比 BF16 基线快 1.8 倍。

量化你自己的模型

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

下面，我们通过量化 FLUX.2 Klein 4B 作为示例进行演练。它涵盖了主要步骤：检查模型、校准和量化变换器、将结果打包为 Diffusers 管道，然后验证并推送到 Hub。完整教程涵盖每个标志的详细信息。

1. 检查将要被量化的内容

通用扫描器遍历模型并决定要定位的内容：重复变换器块堆栈中的兼容线性变成 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 个密集外层线性，没有缺失的模式或重复的名称。

以下命令对变换器运行 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

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

3. 打包一个 Diffusers 管道

转换器将量化变换器与基础管道的其他组件结合，将紧凑的 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 内核是在消费级硬件上高效运行扩散变换器的最有效方法之一，现在在 Diffusers 中原生支持。预量化的检查点通过 from_pretrained() 加载，diffuse-compressor 工具包使得无需等待引擎支持就能量化新架构成为可能。通过量化权重和激活，W4A4 路径降低了内存使用，同时改进了去噪延迟，保持图像质量接近原始 BF16。

如果你量化和发布了新的模型，我们很乐意听说。在 Hub 上分享，让我们知道！如果你对这个功能有任何疑问，欢迎加入我们的 Discord。

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

Diffusers Nunchaku 文档

集成 PR (huggingface/diffusers#14100)

SVDQuant 论文和 Nunchaku 引擎

之前的文章：Exploring Quantization Backends in Diffusers 和 Memory-efficient Diffusion Transformers with Quanto and Diffusers

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

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