Nunchaku Lite

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

Nunchaku Liteは、元のNunchaku推論エンジン、DeepCompressor量子化ライブラリ、およびSVDQuant論文を基盤としています。

CUDAカーネルをインストールします

kernels パッケージは、最適化された CUDA カーネルを提供します。このパッケージは自動的に読み込まれるため、最初にインストールしてください。

pip install -U kernels

Nunchaku Lite は、rootonchair/nunchaku-lite-kernels リポジトリからカーネルを読み込みます。このリポジトリの公開者は、Hub 上で信頼できるカーネル公開者ではありません。読み込み時に Hub からコードをダウンロードして実行するため、Diffusers では DIFFUSERS_TRUST_REMOTE_KERNELS=true を設定して明示的にオプトインする必要があります。詳細は リモートカーネルの信頼 を参照してください。

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

config.jsonから量子化設定を読み込む from_pretrained() を使用して、プリクオンタイズされたパイプラインを読み込みます。

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")

エクスポートされた state dict は、ターゲットの Diffusers モデルのアーキテクチャと完全に一致する必要があります。例えば、融合された QKV 投影で量子化されたチェックポイントは、個別の Q、K、V 投影モジュールを想定するモデル設定には読み込めません。

サポートされている量子化タイプ

Nunchaku Lite は、以下の量子化された線形レイヤーの形式をサポートしています。

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）

sm_90 H100やH200などのHopper GPUは、現在サポートされていません。

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

NunchakuLiteQuantizationConfig

config.json には 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つの手法のみを使用できるため、同じターゲットを両方に列挙しないでください。

{
  "_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エンジンは、モデル固有の融合実行パスによって大幅な高速化を実現しています。Q、K、Vの射影をRMSNormおよびRoPEと組み合わせ、MLPには融合されたGELUカーネルを使用します。一方、Nunchaku Liteは標準的なDiffusersモデルと汎用的な量子化線形層を使用するため、これらの融合は含まれていません。

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

これらの融合がない場合、GPU は projection、normalization、rotary-embedding、MLP の各操作を個別のカーネルとして起動するため、起動オーバーヘッドが発生します。torch.compile はこのオーバーヘッドの一部を軽減する可能性がありますが、元のエンジンの融合カーネルを必ずしも再現するわけではありません。対象のモデルとワークロードに対して、コンパイル済みパイプラインのベンチマークを実行してください。

torch.compile

Nunchaku Lite カーネルと量子化された線形層は、torch.compileと互換性があります。 パイプラインを読み込んだ後に量子化されたTransformerをコンパイルすると、推論が高速化されます。

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
• Nunchakuのインストール要件