# Nunchaku Lite

Nunchaku Lite は、Diffusers で事前量子化されたチェックポイントをロードするための量子化バックエンドです。互換性のあるチェックポイントは [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor) で作成します。transformer を量子化してエクスポートし、Diffusers パイプラインとしてパッケージングします。

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/kernels/rootonchair/nunchaku-lite-kernels) リポジトリからカーネルを読み込みます。このリポジトリの公開者は、Hub 上で信頼できるカーネル公開者ではありません。読み込み時に Hub からコードをダウンロードして実行するため、Diffusers では `DIFFUSERS_TRUST_REMOTE_KERNELS=true` を設定して明示的にオプトインする必要があります。詳細は [リモートカーネルの信頼](../optimization/attention_backends#trusting-remote-kernels) を参照してください。

## 量子化されたパイプラインを読み込む

`config.json`から量子化設定を読み込む [from_pretrained()](/docs/diffusers/main/ja/api/pipelines/overview#diffusers.DiffusionPipeline.from_pretrained) を使用して、事前量子化されたパイプラインを読み込みます。

```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]
> `sm_90` H100やH200などのHopper GPUは、現在サポートされていません。

`nvfp4` チェックポイントには、Blackwell 以降の NVIDIA GPU が必要です。Blackwell GPU を使用する場合、PyTorch >= 2.7 と CUDA >= 12.8 を使用してください。`int4` チェックポイントには、Turing 以降の NVIDIA GPU が必要です。

| メソッド | 精度 | グループサイズ | 注記 |
|---|---:|---:|---|
| `svdq_w4a4` | `nvfp4` | 16 | SVDQ低ランク補正を用いたNVFP4ランタイムカーネルを使用します。 |
| `svdq_w4a4` | `int4` | 64 | SVDQ低ランク補正を用いたINT4 W4A4カーネルを使用します。 |
| `awq_w4a16` | `int4` | 64 | INT4の重み専用AWQスタイルのカーネルを使用します。 |

## NunchakuLiteQuantizationConfig

`config.json` には [NunchakuLiteQuantizationConfig](/docs/diffusers/main/ja/api/quantization#diffusers.NunchakuLiteQuantizationConfig) を含める必要があります。これにより、実行時の `compute_dtype` と、各 Nunchaku Lite 量子化手法のターゲットモジュールが定義されます。

- `compute_dtype`: 量子化モジュール内の浮動小数点バッファのランタイム 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` の下に、量子化したい各モジュールを列挙してください。モジュールは1つの手法のみを使用できるため、同じターゲットを両方に列挙しないでください。

```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モデルと汎用的な量子化線形層を使用するため、これらの融合は含まれていません。

以下の測定結果は、Flux Schnell を RTX 5090 で使用して、元の Nunchaku エンジンにおけるこれらの融合によるステップごとのレイテンシへの影響を示しています。

| 最適化 | OFF時のレイテンシ | ON時のレイテンシ | 高速化 |
| --- | --- | --- | --- |
| QKV+norm+RoPEの融合（統合） | 1.057 s | 0.834 s | **1.27×** |
| — Q/K/V を 1 つの matmul にまとめる、単独 | 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 は projection、normalization、rotary-embedding、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 パイプライン（RTX PRO 6000 上で 2.271 秒 → 1.675 秒）と比較して 1.8 倍高速に動作します。

## リソース

- [diffuse-compressor](https://github.com/rootonchair/diffuse-compressor)
- [Nunchakuのインストール要件](https://nunchaku.tech/docs/nunchaku/installation/installation.html)