# Nunchaku Lite

Nunchaku Lite 是用于在 Diffusers 中加载预量化检查点的量化后端。请使用 [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor) 创建兼容的检查点。它会量化并导出 transformer，然后将其打包为 Diffusers pipeline。

Nunchaku Lite 基于原始的 [Nunchaku](https://github.com/nunchaku-ai/nunchaku) 推理引擎、[DeepCompressor](https://github.com/nunchaku-ai/deepcompressor) 量化库以及 [SVDQuant 论文](https://arxiv.org/abs/2411.05007) 构建。

## 安装 CUDA 内核

kernels 包提供经过优化的 CUDA 内核，它们会自动加载。请先安装该包。

```bash
pip install -U kernels
```

Nunchaku Lite 会从 [`rootonchair/nunchaku-lite-kernels`](https://huggingface.co/rootonchair/nunchaku-lite-kernels) 仓库加载其内核，而该仓库的发布者并非 Hub 上受信任的内核发布者。加载该内核会从 Hub 下载并执行代码，因此 Diffusers 要求您通过设置 `DIFFUSERS_TRUST_REMOTE_KERNELS=true` 来明确选择加入。有关详细信息，请参阅 [信任远程内核](../optimization/attention_backends#trusting-remote-kernels)。

## 加载量化流水线

使用 [from_pretrained()](/docs/diffusers/v0.43.0/zh/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained) 加载预量化流水线，该函数会从 `config.json` 中读取量化配置。

```python
import torch
from diffusers import DiffusionPipeline

model_id = "rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4"

pipe = DiffusionPipeline.from_pretrained(
    model_id, dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

prompt = "A modern red armchair in a quiet studio, soft window light, realistic product photography"
image = pipe(
    prompt=prompt,
    height=1024,
    width=1024,
    num_inference_steps=8,
    guidance_scale=1.0,
).images[0]
image.save("ernie-image-turbo-nunchaku-lite.png")
```

> [!NOTE]
> 导出的 state dict 必须与目标 Diffusers 模型架构完全匹配。例如，使用融合 QKV 投影进行量化的检查点无法加载到期望使用独立 Q、K 和 V 投影模块的模型配置中。

## 支持的量化类型

Nunchaku Lite 支持以下量化线性层格式。

> [!TIP]
> 在 Blackwell GPU 上使用 `nvfp4`。在 Blackwell 上运行 `int4` 检查点可能比 `nvfp4` 更慢。

目前，CUDA 内核支持以下 NVIDIA GPU 架构：

- `sm_75`（Turing，例如 RTX 2080）
- `sm_80`（Ampere，例如 A100）
- `sm_86`（Ampere，例如 RTX 3090 和 RTX A6000）
- `sm_89`（例如 RTX 4090 的 Ada）
- `sm_120`（Blackwell，例如 RTX 5090）

> [!NOTE]
> Hopper 架构的 GPU（例如 `sm_90` H100 和 H200）目前尚不支持。

`nvfp4` 检查点需要 Blackwell 或更新版本的 NVIDIA GPU。在 Blackwell GPU 上，请使用 PyTorch >= 2.7 和 CUDA >= 12.8。`int4` 检查点需要 Turing 或更新版本的 NVIDIA GPU。

| 方法 | 精度 | 组大小 | 备注 |
|---|---:|---:|---|
| `svdq_w4a4` | `nvfp4` | 16 | 使用 NVFP4 运行时内核，并结合 SVDQ 低秩校正。 |
| `svdq_w4a4` | `int4` | 64 | 使用 INT4 W4A4 内核，并结合 SVDQ 低秩校正。 |
| `awq_w4a16` | `int4` | 64 | 使用 INT4 仅权重的 AWQ 风格内核。 |

## NunchakuLiteQuantizationConfig

`config.json` 文件必须包含一个 [NunchakuLiteQuantizationConfig](/docs/diffusers/v0.43.0/zh/api/quantization#diffusers.NunchakuLiteQuantizationConfig)。它定义了运行时 `compute_dtype` 以及每种 Nunchaku Lite 量化方法的目标模块。

- `compute_dtype`：量化模块中浮点缓冲区的运行时数据类型，通常为 `torch.bfloat16`。
- `svdq_w4a4`：包含 `precision`、`group_size`、`rank` 和 `targets` 的 SVDQ W4A4 目标配置。
- `awq_w4a16`：包含 `precision`、`group_size` 和 `targets` 的 AWQ W4A16 目标配置。

`targets` 中的每个条目都必须指向一个线性层。Diffusers 会将每个 `svdq_w4a4` 目标替换为 SVDQ W4A4 层，并将每个 `awq_w4a16` 目标替换为 AWQ W4A16 层。下面的示例展示了使用缩短的目标列表时的预期形状。

在 `svdq_w4a4` 或 `awq_w4a16` 下列出每个需要量化的模块。一个模块只能使用一种方法，因此请勿将同一目标同时列在两者之下。

```json
{
  "_class_name": "ErnieImageTransformer2DModel",
  "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"]
    },
    "awq_w4a16": {
      "precision": "int4",
      "group_size": 64,
      "targets": ["final_linear"]
    }
  }
}
```

## 融合内核

原始的 [Nunchaku](https://github.com/nunchaku-ai/nunchaku) 引擎的大部分速度优势来自于针对特定模型的融合执行路径。它将 Q、K 和 V 投影与 RMSNorm 及 RoPE 相结合，并在 MLP 中使用融合 GELU 内核。而 Nunchaku Lite 则使用标准的 Diffusers 模型和通用的量化线性层，因此不包含这些融合操作。

以下测量结果展示了在原始 Nunchaku 引擎中，这些融合操作对每步延迟的影响，测量环境为在 RTX 5090 上运行 Flux Schnell。

| 优化 | 关闭时的延迟 | 开启后的延迟 | 加速 |
| --- | --- | --- | --- |
| QKV+norm+RoPE 融合（组合） | 1.057 s | 0.834 s | **1.27×** |
| — 将 Q/K/V 合并为单个矩阵乘法，单独 | 1.059 s | 1.034 s | 1.025× |
| — 在分组的基础上，实现了 RMSNorm+RoPE 后处理融合 | 1.034 s | 0.836 s | 1.236× |
| GELU-MLP 融合 | 0.766 s | 0.722 s | **1.06×** |

如果没有这些融合操作，GPU 会将投影、归一化、旋转嵌入和 MLP 操作作为独立的内核启动，从而增加启动开销。`torch.compile` 可能会减少部分开销，但它不一定能复现原始引擎的融合内核。请针对您的模型和工作负载对编译后的流水线进行基准测试。

## torch.compile

Nunchaku Lite 内核和量化线性层与 [`torch.compile`](../optimization/fp16#torchcompile) 兼容。在加载流水线后对量化 Transformer 进行编译，以加快推理速度。

```python
pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=True)
```

编译后的 Nunchaku Lite NVFP4 流水线比原始 BF16 流水线快 1.8 倍（在 RTX PRO 6000 上，从 2.271 秒降至 1.675 秒）。

## 资源

- [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor)
- [Nunchaku 安装要求](https://nunchaku.tech/docs/nunchaku/installation/installation.html)