跳转到内容
4ALL API 文档
主站 新闻 控制台

# 将 Nunchaku 4 位扩散推理引入 Diffusers

· Hugging Face
教程模型卡

大型扩散 Transformer 现在可以通过 Nunchaku Lite 在 Diffusers 中原生加载

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

这些后端大多是仅权重量化（weight-only）。这意味着它们会以低精度存储权重，并在计算时再反量化回高精度。这样可以显著降低显存占用，但通常并不会让推理更快，甚至还可能带来一点额外延迟。

SV DQuant 是热门推理引擎 Nunchaku 背后的量化方法，它采取了不同的思路。它使用 4 位权重和激活（W4A4）运行主要的 Transformer 层，在降低内存占用的同时也加快了去噪循环。下面会详细介绍其原理，但在此之前，使用这些 checkpoint 一直都需要单独的推理库。

在当前的 Diffusers 中，加载 Nunchaku checkpoint 只需调用 from_pretrained() 即可，无需本地编译 CUDA，这要归功于 kernels 包。此外，配套的 diffuse-compressor 工具包还允许你自行量化新的架构，并将其像普通 Diffusers 仓库一样发布。

Nunchaku Lite 图像质量与性能对比

目录

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 kernels 会在首次使用时通过 Nunchaku Lite kernels 页面 从 Hub 下载。这个 checkpoint 将 Nunchaku NVFP4 Transformer 与 bitsandbytes NF4 文本编码器配对，在 RTX 5090 上生成一张 1024x1024 图像大约需要 1.7 秒，峰值显存占用约 12 GB；相比之下，BF16 pipeline 约为 24 GB。关于 Nunchaku Lite checkpoint 格式的更多细节，可以查看 官方 Diffusers 文档。

NVFP4 checkpoint 需要 NVIDIA Blackwell GPU（RTX 50 系列、RTX PRO 6000、B200）。对于更早一代的 GPU，请使用 INT4 变体。有关详情请参见下方的 硬件支持 表格。

背景：SVDQuant 与 Nunchaku

SVDQuant 是 Nunchaku 背后的量化方法，后者是其参考 CUDA 推理引擎。对于扩散 Transformer 来说，标准的 4 位量化很难处理，因为权重和激活中都存在较大的离群值。SVDQuant 的做法是将激活中的离群值转移到权重中，用一个小型的 16 位低秩分支来表示每个权重矩阵中最难处理的部分，并将其余残差量化为 4 位。Nunchaku 通过针对 4 位路径和低秩分支的融合 kernel 使这一过程更快。

Nunchaku kernel 融合：低秩下投影与输入量化融合，低秩上投影与 4 位矩阵乘融合
Nunchaku 将低秩下投影与量化 kernel 融合，并将低秩上投影与 4 位计算 kernel 融合，从而消除了 16 位分支的内存访问开销。图源来自 SVDQuant 论文。

介绍 Nunchaku Lite

原始的 Nunchaku 引擎 很大一部分速度优势来自面向特定模型的融合执行路径，例如融合的 QKV 投影和融合的 GELU/MLP kernels。这些优化与每种架构的模块布局和 checkpoint 格式紧密相关，因此支持新的模型家族通常需要针对该模型进行集成工作。

Nunchaku Lite 是 Diffusers 中新的集成方式。借助它，Diffusers 可以加载 Nunchaku 风格的 checkpoint，而无需自定义 pipeline 或独立推理引擎。在底层，Nunchaku Lite 会在加载 checkpoint 之前，用运行时 SVDQ/AWQ 线性层替换标准 Diffusers 模型中相关的 nn.Linear 模块。CUDA kernels 则通过 kernels 包从 Hub 获取。这里使用两类 kernel：

  • svdq_w4a4：4 位权重和激活，并结合 SVDQuant 低秩修正。该层用于 Transformer 的注意力和 MLP 投影，这些部分几乎占据了全部计算量，提供 INT4 和 NVFP4 两种变体。

  • awq_w4a16：4 位权重、16 位激活，用于自适应归一化和调制投影，例如 FLUX 的 adanorm_single / adanorm_zero 或 Qwen-Image 的调制层。这些层是内存受限且对精度敏感的，因此 AWQ 很适合在节省内存和空间的同时尽量保留精度。

代价是，如果没有面向架构的融合 kernels 和模块，Nunchaku Lite 无法达到原始 Nunchaku 引擎那样的加速效果。不过，这个基础实现仍然能带来约 30% 的速度提升，同时保持相同级别的 显存降低。

Diffusers 中的原生加载

如果你在 Diffusers 中使用过 bitsandbytes 或 torchao，那么机械流程……