将 Nunchaku 4 位扩散推理集成到 Diffusers

来源：https://huggingface.co/blog/nunchaku-diffusers 返回文章列表 (https://huggingface.co/blog)

Pham Hong Vinh 的头像 (https://huggingface.co/rootonchair)

Sayak Paul 的头像 (https://huggingface.co/sayakpaul)

• 目录 (https://huggingface.co/blog/nunchaku-diffusers#table-of-contents)
• 快速上手 Nunchaku Lite (https://huggingface.co/blog/nunchaku-diffusers#getting-started-with-nunchaku-lite)
• 背景：SVDQuant 与 Nunchaku (https://huggingface.co/blog/nunchaku-diffusers#background-svdquant-and-nunchaku)
• 介绍 Nunchaku Lite (https://huggingface.co/blog/nunchaku-diffusers#introducing-nunchaku-lite)
• Diffusers 中的原生加载 (https://huggingface.co/blog/nunchaku-diffusers#native-loading-in-diffusers)
• 硬件支持 (https://huggingface.co/blog/nunchaku-diffusers#hardware-support)
• 提升速度与降低内存 (https://huggingface.co/blog/nunchaku-diffusers#getting-more-speed-and-lower-memory)
• 基准测试 (https://huggingface.co/blog/nunchaku-diffusers#benchmarks)
  • 端到端延迟与内存 (https://huggingface.co/blog/nunchaku-diffusers#end-to-end-latency-and-memory)
  • 图像质量 (https://huggingface.co/blog/nunchaku-diffusers#image-quality)
• 量化你自己的模型 (https://huggingface.co/blog/nunchaku-diffusers#quantizing-your-own-model)
  • 1. 检查哪些部分会被量化 (https://huggingface.co/blog/nunchaku-diffusers#1-inspect-what-will-be-quantized)
  • 2. 执行量化 (https://huggingface.co/blog/nunchaku-diffusers#2-run-quantization)
  • 3. 打包成 Diffusers pipeline (https://huggingface.co/blog/nunchaku-diffusers#3-package-a-diffusers-pipeline)
  • 4. 加载、验证并推送到 Hub (https://huggingface.co/blog/nunchaku-diffusers#4-load-verify-and-push-to-the-hub)
  • 量化经过结构重写的模型 (https://huggingface.co/blog/nunchaku-diffusers#quantizing-models-with-structural-rewrites)
• 现成可用的检查点 (https://huggingface.co/blog/nunchaku-diffusers#ready-to-use-checkpoints)
• 结论 (https://huggingface.co/blog/nunchaku-diffusers#conclusion)
• 致谢 (https://huggingface.co/blog/nunchaku-diffusers#acknowledgements)

大型扩散变换器能够生成令人惊艳的图像（甚至视频、音频片段，以及现在的文本），但加载一个现代文生图模型的 BF16 精度版本通常需要 20-30 GB 的显存，这让大多数消费级 GPU 难以企及。量化是解决这个问题的有力方案，Diffusers 已经集成了多个量化后端，例如 bitsandbytes、GGUF、torchao 和 Quanto，我们在《探索 Diffusers 中的量化后端》(https://huggingface.co/blog/diffusers-quantization) 中介绍过这些。

这些后端中的大多数是仅权重量化的。这意味着它们以低精度存储权重，在计算时再将它们反量化为高精度。这大大减少了内存使用，但通常不会加快推理速度，有时甚至会带来微小的延迟开销。

SVDQuant (https://arxiv.org/abs/2411.05007) 是流行推理引擎 Nunchaku (https://github.com/nunchaku-tech/nunchaku) 背后的量化方法，它采用了不同的思路。该方法使用 4 位权重和激活（W4A4）运行主要的变换器层，在减少内存的同时还加速了去噪循环。具体细节将在下文介绍，但在此之前，使用这些检查点需要一个单独的推理库。

而现在的 Diffusers 中，加载一个 Nunchaku 检查点只需要调用 from_pretrained()，无需本地 CUDA 编译，这要归功于 kernels (https://github.com/huggingface/kernels) 包。此外，配套的 diffuse-compressor (https://github.com/rootonchair/diffuse-compressor) 工具包让你能够自己量化新的架构，并将它们发布为普通的 Diffusers 仓库。

Nunchaku Lite 图像质量与性能对比

目录

• 快速上手 Nunchaku Lite (https://huggingface.co/blog/nunchaku-diffusers#getting-started-with-nunchaku-lite)
• 背景：SVDQuant 与 Nunchaku (https://huggingface.co/blog/nunchaku-diffusers#background-svdquant-and-nunchaku)
• 介绍 Nunchaku Lite (https://huggingface.co/blog/nunchaku-diffusers#introducing-nunchaku-lite)
• Diffusers 中的原生加载 (https://huggingface.co/blog/nunchaku-diffusers#native-loading-in-diffusers)
• 提升速度与降低内存 (https://huggingface.co/blog/nunchaku-diffusers#getting-more-speed-and-lower-memory)
• 基准测试 (https://huggingface.co/blog/nunchaku-diffusers#benchmarks)
• 量化你自己的模型 (https://huggingface.co/blog/nunchaku-diffusers#quantizing-your-own-model)
• 现成可用的检查点 (https://huggingface.co/blog/nunchaku-diffusers#ready-to-use-checkpoints)
• 结论 (https://huggingface.co/blog/nunchaku-diffusers#conclusion)
• 致谢 (https://huggingface.co/blog/nunchaku-diffusers#acknowledgements)

快速上手 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")

BF16 和 Nunchaku Lite 对红狐狸提示词的输出

无需自定义 pipeline 类或单独的推理引擎，也不需要本地编译任何东西。NVFP4 kernel 会在首次使用时通过Nunchaku Lite kernel 页面 (https://huggingface.co/kernels/rootonchair/nunchaku-lite-kernels)从 Hub 下载。该检查点将 Nunchaku NVFP4 变换器与 bitsandbytes NF4 文本编码器配对，在 RTX 5090 上生成 1024x1024 图像大约需要 1.7 秒，峰值内存约 12 GB，而 BF16 pipeline 则需要约 24 GB。更多关于 Nunchaku Lite 检查点格式的细节，请参见官方 Diffusers 文档 (https://huggingface.co/docs/diffusers/main/en/quantization/nunchaku)。

NVFP4 检查点需要 NVIDIA Blackwell GPU（RTX 50 系列、RTX PRO 6000、B200）。对于早期代际，请使用 INT4 变体。详见下文硬件支持表。

背景：SVDQuant 与 Nunchaku

SVDQuant 是 Nunchaku 推理引擎的量化方法。标准的 4 位量化对于扩散变换器来说较为困难，因为权重和激活都存在大量离群值。SVDQuant 通过将激活离群值转移到权重中来解决这个问题，它使用一个小的 16 位低秩分支来表示每个权重矩阵中最难的部分，并将剩余部分量化为 4 位。Nunchaku 通过融合后的 kernel 来加速低秩分支和 4 位路径。

Nunchaku kernel 融合：低秩下投影与输入量化融合，低秩上投影与 4 位矩阵乘法融合

Nunchaku 将低秩下投影与量化 kernel 融合，低秩上投影与 4 位计算 kernel 融合，消除了 16 位分支的内存访问开销。图来自SVDQuant 论文 (https://arxiv.org/abs/2411.05007)。

介绍 Nunchaku Lite

原始的Nunchaku 引擎 (https://github.com/nunchaku-ai/nunchaku) 其速度主要来源于模型特定的融合执行路径 (https://huggingface.co/blog/nunchaku-diffusers#quantizing-models-with-structural-rewrites)，例如融合的 QKV 投影和融合的 GELU/MLP kernel。这些优化与每个架构的模块布局和检查点格式紧密绑定，因此支持新的模型家族通常需要针对特定模型的集成工作。

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

• svdq_w4a4：4 位权重和激活，带有 SVDQuant 低秩校正。该层用于变换器的注意力层和 MLP 投影，这些地方几乎消耗了所有计算量，有 INT4 和 NVFP4 两种变体。
• awq_w4a16：4 位权重与 16 位激活，用于自适应归一化和调制投影，例如 FLUX 的 adanorm_single / adanorm_zero 或 Qwen-Image 的调制层。这些层是内存密集且对精度敏感的，因此 AWQ 是保持精度的同时节省内存和空间的好选择。

其代价是，由于没有架构特定的融合 kernel 和模块，Nunchaku Lite 无法达到原始 Nunchaku 引擎的速度提升。然而，这个基础实现仍然提供了大约30% 的加速，同时保持了相同水平的显存减少。

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

Volta 和 Hopper GPU 目前不受 4 位 kernel 支持。量化器会在加载时验证 GPU 的 CUDA 能力，并给出明确的错误信息，而不是产生错误输出。

提升速度与降低内存

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

torch.compile。编译变换器可以将端到端加速从 1.35 倍提升到 1.8 倍：

pipe.transformer.compile(fullgraph=True)

# 或者使用 compile_repeated_blocks() 以加快编译速度

pipe.transformer.compile_repeated_blocks(fullgraph=True)

量化文本编码器。变换器并不是唯一占用大量内存的组件。像 T5 或 Qwen3 这样的文本编码器本身就可能占用几个 GB。使用 bitsandbytes NF4 进一步量化文本编码器，在我们的基准测试中减少了大约 22% 的峰值显存。

卸载。如果需要将 pipeline 放置到更小的 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 (https://huggingface.co/rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder) 检查点。

端到端延迟与内存

如上所示，Nunchaku 将峰值显存减少了多达 50%，同时还将延迟改善了大约 30%。剩余的开销主要来自额外的 kernel 启动，而 torch.compile 可以缓解这个问题，使完整 pipeline 降至 1.68 秒，比 BF16 基线快 1.8 倍。

图像质量

质量对比网格，BF16 与 4 位输出，使用相同的种子和设置。

量化你自己的模型

Diffusers 中的 Nunchaku Lite 支持与架构无关，而 diffuse-compressor (https://github.com/rootonchair/diffuse-compressor) 工具包为 Diffusers 模型提供了一个端到端的 SVDQuant 工作流：校准、量化、打包和发布。

下面，我们以量化 FLUX.2 Klein 4B 为例进行演示。它涵盖了主要步骤：检查模型、校准并量化变换器、将结果打包为 Diffusers pipeline，然后验证并推送到 Hub。完整教程 (https://github.com/rootonchair/diffuse-compressor/blob/main/docs/quantize_new_hf_model.md) 详细介绍了每个标志。

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 个稠密的外部线性层，没有遗漏的模式或重复的名称。

2. 执行量化

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