Diffusers に Nunchaku 4 ビット推論を導入

記事一覧に戻る

大規模な拡散トランスフォーマーは、驚くほど美しい画像（さらには動画や音声のスニペット、そしてテキスト）を生成できます。しかし、最新のテキストから画像へのモデルを BF16 精度でロードするには、20〜30 GB の VRAM が必要になることが多く、これは一般的な消費者向け GPU では対応が難しいレベルです。

この問題に対する強力な解決策が量子化（Quantization）です。Diffusers はすでに bitsandbytes、GGUF、torchao、Quanto など複数の量子化バックエンドを統合しており、これらについては Diffusers における量子化バックエンドの探索 で詳しく解説しています。

これらのバックエンドの多くは「ウェイトのみ（weight-only）」方式です。これは、モデルの重みを低精度で保存し、計算時に高精度に展開（デクォンタイズ）する仕組みを指します。これによりメモリ使用量は大幅に削減されますが、通常は推論速度が向上するわけではなく、むしろわずかなレイテンシの増加を招くこともあります。

人気のある推論エンジン Nunchaku の背後にある量子化手法 SVDQuant は、これとは異なるアプローチを採用しています。これは主要なトランスフォーマー層を 4 ビット重みと活性化（W4A4）で実行し、メモリ使用量を削減しながらノイズ除去ループの速度も向上させます。

詳細は後述しますが、これまでこれらのチェックポイントを使用するには、専用の推論ライブラリが必要でした。

現在のDiffusersでは、`kernels` パッケージのおかげでローカルでの CUDA コンパイルが不要となり、Nunchaku チェックポイントの読み込みは from_pretrained() を呼び出すだけで完了します。また、サードパーティ製の diffuse-compressor ツールキットを使えば、独自のアーキテクチャを量子化して、通常のDiffusersリポジトリとして公開することも可能です。

目次

• Nunchaku Lite の始め方

• 背景：SVDQuant と Nunchaku

• Nunchaku Lite の紹介

• Diffusers でのネイティブ読み込み

• さらなる高速化とメモリ削減

• ベンチマーク結果

• 独自のモデルを量子化する

• すぐに使えるチェックポイント

• 結論

• 謝辞

Nunchaku Lite の始め方

まずは必要なパッケージをインストールしましょう。最新の Diffusers と、Hugging Face 公式の kernels パッケージが必要です。

pip install -U diffusers transformers accelerate kernels bitsandbytes

次に、他のDiffusersモデルと同様に、事前量子化済みのパイプラインを読み込みます。

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

カスタムパイプラインクラスを用意したり、別個の推論エンジンを使ったりする必要はありません。ローカルでのコンパイルも不要です。NVFP4 カーネルは、初めて使用した際に Nunchaku Lite カーネルページ から Hugging Face Hub 経由で自動的にダウンロードされます。

このチェックポイントは、Nunchaku NVFP4 トランスフォーマーと bitsandbytes NF4 テキストエンコーダーを組み合わせたもので、RTX 5090 で 1024x1024 の画像を約 1.7 秒で生成できます。ピークメモリ使用量は約 12 GB です。一方、BF16 パイプラインでは約 24 GB を必要とします。

Nunchaku Lite チェックポイントの詳細なフォーマットについては、公式の Diffusers ドキュメント をご覧ください。

NVFP4 チェックポイントは NVIDIA Blackwell GPU（RTX 50 シリーズ、RTX PRO 6000、B200）が必要です。それ以前の世代では INT4 バリアントを使用してください。詳細は以下の ハードウェアサポート テーブルをご確認ください。

背景：SVDQuant と Nunchaku

SVDQuantは、Nunchakuの背後にある量子化手法であり、そのリファレンスCUDA推論エンジンです。拡散トランスフォーマーでは重みと活性化値の両方に大きな外れ値が含まれるため、標準的な4ビット量子化は困難です。SVDQuantはこの課題に対し、活性化の外れ値を重みに移動させることで解決します。具体的には、各重み行列の中で最も扱いが難しい部分を小さな16ビット低ランクブランチで表現し、残りの部分（残差）を4ビットに量子化します。Nunchakuは、4ビット経路と低ランクブランチの両方に対する融合カーネルを実装することで、この処理を高速化しています。

*Nunchakuは、低ランク降下投影を量子化カーネルと融合させ、低ランク上昇投影を4ビット計算カーネルと融合させることで、16ビットブランチにおけるメモリアクセスのオーバーヘッドを排除しています。図出典：SVDQuant論文。

Nunchaku Lite の紹介*

オリジナルの Nunchaku エンジン は、QKV 射影や GELU/MLP カーネルを統合した実行パスといった、モデル固有の最適化により高速化を実現しています。これらの最適化は各アーキテクチャのモジュール構成とチェックポイント形式に依存しているため、新しいモデルファミリーへの対応には通常、モデルごとの統合作業が必要となります。

Nunchaku Liteは、Diffusersにおける新しい統合パスです。これにより、Diffusersはカスタムパイプラインや別個の推論エンジンなしで、Nunchaku形式のチェックポイントを読み込めるようになります。仕組みとしては、チェックポイントをロードする前に、既存のDiffusersモデル内の関連するnn.Linearモジュールを、ランタイムでSVDQ/AWQ対応の線形層にパッチ適用します。CUDAカーネルはkernelsパッケージを通じてHubから取得されます。ここでは2つのカーネルファミリーが使用されます。

• svdq_w4a4: SVDQuantによる低ランク補正を備えた、4ビット重みと活性化値を使用する層です。この層は、計算のほとんどが集中するトランスフォーマーのアテンションやMLP投影に用いられ、INT4およびNVFP4の両バリアントで利用可能です。

• awq_w4a16: 適応型正規化やモジュレーション投影（FLUX の adanorm_single / adanorm_zero や Qwen-Image のモジュレーション層など）に使用される、4 ビット重みと 16 ビット活性化の組み合わせです。これらの層はメモリアクセスがボトルネックとなり、精度への感度も高いため、AWQ を適用することでメモリやストレージを節約しつつ、必要な精度を維持できます。

アーキテクチャ固有の融合カーネルやモジュールがない場合、Nunchaku Lite はオリジナルの Nunchaku エンジンほどの高速化は達成できません。しかし、最小限の実装でも VRAM 削減効果 を維持したまま、約 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 の世代とチェックポイントの精度に応じて異なるカーネルバリアントを使用します：

Volta および Hopper アーキテクチャの GPU は、現在 4 ビットカーネルに対応していません。量子化器はロード時に GPU の CUDA コア能力を検証し、誤った出力を生成するのではなく明確なエラーを返します。

より高速でメモリ効率の高い運用へ

Nunchaku Lite は、Diffusers が提供する他のメモリ最適化や速度向上機能と組み合わせることができます。

torch.compile. トランスフォーマーをコンパイルすることで、エンドツーエンドの速度向上率が 1.35 倍から 1.8 倍に改善されます：

pipe.transformer.compile(fullgraph=True)

# or compile_repeated_blocks() for faster compilation

pipe.transformer.compile_repeated_blocks(fullgraph=True)

量子化されたテキストエンコーダー。 トランスフォーマーだけが大きなメモリを消費するわけではありません。T5 や Qwen3 などのテキストエンコーダーも、単体で数ギガバイトの容量を占有することがあります。bitsandbytes の NF4 を用いてテキストエンコーダーをさらに量子化すると、ベンチマークではピーク VRAM が約 22% 削減されます。

オフローディング。 パイプラインをより小さな GPU に収めたい場合、enable_model_cpu_offload() や enable_sequential_cpu_offload() といった Diffusers のオフローディングヘルパーは通常どおり機能します。

ベンチマーク結果

以下の数値はすべて、NVIDIA RTX PRO 6000 (Blackwell) を使用し、1024x1024 の解像度で測定したものです。モデルには rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder を使用しました。

エンドツーエンドのレイテンシとメモリ使用量

上記の結果が示す通り、Nunchaku はピーク VRAM を最大 50% 削減しながらも、レイテンシを約 30% 短縮しています。残りのオーバーヘッドは主に追加のカーネル起動に起因しますが、torch.compile を活用することでこれを緩和でき、フルパイプラインの実行時間を 1.68 秒まで短縮し、BF16 ベースラインと比較して 1.8 倍高速化できます。

画像品質

*シードと設定を同一にした場合の、BF16 と 4-bit の出力比較。

独自のモデルを量子化する

Diffusers における Nunchaku Lite のサポートはアーキテクチャに依存せず、diffuse-compressor ツールキットが Diffusers モデル向けの SVDQuant ワークフロー（キャリブレーション、量子化、パッケージ化、公開）をエンドツーエンドで提供します。

以下では、FLUX.2 Klein 4B の量子化を例に主要な手順を解説します。具体的には、モデルの調査、トランスフォーマー部分のキャリブレーションと量子化、結果を Diffusers パイプラインとしてパッケージ化、そして検証と Hub へのアップロードです。詳細なフラグの説明については、完全なチュートリアルをご覧ください。

1. 量子化する対象の調査

汎用的なスキャナがモデルを走査し、どの部分をターゲットにするかを判断します。繰り返しトランスフォーマーブロックスタック内の互換性のある線形層は SVDQ W4A4 の対象となり、認識されたモデュレーション線形層は AWQ W4A16 の対象となります。それ以外の部分はそのまま密（dense）な状態を維持されます。

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 の場合、期待される結果は SVDQ ターゲットが 100 個、AWQ ターゲットが 3 つ、そして重複や欠落がない 6 つの密な外線形層（dense outer linears）です。

2. 量子化の実行

以下のコマンドは Transformer に 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 パイプラインをパッケージ化する

コンバータは量子化されたトランスフォーマーとベースパイプラインの他のコンポーネントを結合し、コンパクトな nunchaku_lite 設定を transformer/config.json に書き出します。また、必要に応じてテキストエンコーダーも NF4 形式に変換できます。

python examples/convert_nunchaku_lite_diffusers.py \
  --checkpoint outputs/checkpoints/svdq-int4_r32-flux-2-klein-4b.safetensors \
  --model-id black-forest-labs/FLUX.2-klein-4B \
  --bnb4-text-encoder text_encoder \
  --compute-dtype bfloat16 \
  --output-dir outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder

4. Hub へのロード、検証、プッシュ

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "outputs/diffusers/FLUX.2-klein-4B-nunchaku-lite-int4-bnb4-text-encoder",
    device_map="cuda",
)
image = pipe(
    "A glass robot in a greenhouse, cinematic lighting",
    num_inference_steps=4, guidance_scale=1.0,
    generator=torch.Generator("cuda").manual_seed(12345),
).images[0]

出力結果に問題がなければ、pipe.push_to_hub("your-name/your-model-nunchaku-lite-int4") を実行してください。これで他のユーザーも、上記と同じ from_pretrained() パターンでモデルをロードできるようになります。

構造的書き換えによるモデルの量子化

一般的な手順では、アーキテクチャが構造的な書き換えなしに量子化可能であることを前提としています。さらに高速化を図るため、元の Nunchaku エンジンは Diffusers のレイヤー群を融合モジュールとして書き換えます。一方、一般的な手順では、Q、K、V 投影を一つのモジュールに結合したり、融合された投影を複数のモジュールに分割したりといった変更を自動的に推論することはできません。

FLUX.1-dev の QKV 投影は具体的な例です。Diffusers では、3 つの別モジュールを定義しています。

self.to_q = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_k = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)
self.to_v = torch.nn.Linear(query_dim, self.inner_dim, bias=bias)

Nunchaku の FLUX モジュールは、これら 3 つの層を結合し、量子化された `to_qkv` モジュールとして実装しています。

to_qkv = fuse_linears([other.to_q, other.to_k, other.to_v])
self.to_qkv = SVDQW4A4Linear.from_linear(to_qkv, **kwargs)

このグループ化されたモジュールが必要とされるのは、Nunchaku の融合演算子が QKV 投影、Q/K 正規化、回転埋め込みをまとめて処理するためです。一方、デフォルトの Diffusers パス ではこれらが個別に実行されます。

query = attn.to_q(hidden_states)
key = attn.to_k(hidden_states)
value = attn.to_v(hidden_states)

query = query.unflatten(-1, (attn.heads, -1))
key = key.unflatten(-1, (attn.heads, -1))
value = value.unflatten(-1, (attn.heads, -1))

query = attn.norm_q(query)
key = attn.norm_k(key)

if image_rotary_emb is not None:
    query = apply_rotary_emb(query, image_rotary_emb, sequence_dim=1)
    key = apply_rotary_emb(key, image_rotary_emb, sequence_dim=1)

Nunchaku パス では、グループ化された投影、正規化モジュール、回転埋め込みを 1 つの融合演算子に渡します。

qkv = fused_qkv_norm_rottary(
    hidden_states, attn.to_qkv, attn.norm_q, attn.norm_k, image_rotary_emb
)

これは、一般的なパスでは推論できない構造的な書き換えです。Diffusers には、to_q、to_k、to_v パラメータを持つ 3 つの宛先モジュールがあります。

一方、Nunchaku は 1 つのグループ化されたモジュールを to_qkv の下に配置しています。

モデル固有のターゲット設定またはアダプタでは、Q、K、V パラメータを出力次元に沿ってその順序で連結し、to_qkv に読み込む必要があると明記する必要があります。

構造的な書き換えは、量子化時にモデル固有のターゲット設定ファイルで指定され、チェックポイント読み込み時には小型のランタイムアダプターによって処理されます。

構造的に書き換えられたチェックポイントを生成するための具体的なターゲット設定例については、FLUX.2 Klein 4B 量子化スクリプト をご覧ください。また、グループ化された QKV テンソルの読み込みや融合された射影の分割、その他の融合演算を処理するために必要なランタイムアダプターは、rootonchair/nunchaku-lite で提供されています。

完全なワークフローについては、新しいモデルの追加 ガイドをご参照ください。

すぐに使えるチェックポイント

すぐに使い始めるには、以下のリポジトリをチェックしてください：

• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-int4-bnb4-text-encoder: bitsandbytes NF4 形式のテキストエンコーダーを搭載した INT4 版 ERNIE-Image-Turbo

• rootonchair/ERNIE-Image-Turbo-nunchaku-lite-nvfp4-bnb4-text-encoder: bitsandbytes NF4 形式のテキストエンコーダーを搭載した NVFP4 版 ERNIE-Image-Turbo

• OzzyGT/Krea_2_Turbo_nunchaku_lite_nvfp4: NVFP4 形式の Krea 2 Turbo チェックポイント

Hugging Face の lite-infer では、Nunchaku Lite のチェックポイントやコレクションがさらに追加されています。

結論

Nunchaku の SVDQuant カーネルは、コンシューマー向けハードウェア上で拡散トランスフォーマーを効率的に実行するための最も効果的な手法の一つです。現在では Diffusers でもネイティブサポートが実現されており、事前量子化されたチェックポイントは from_pretrained() でそのまま読み込めます。また、diffuse-compressor ツールキットを使えば、エンジン側のサポート待ちなしで新しいアーキテクチャを量子化することが可能です。

重みと活性化の両方を量子化する W4A4 パスでは、メモリ使用量を削減しつつ推論レイテンシも短縮できます。その結果、画像品質は BF16 版のオリジナルにほぼ匹敵するレベルを維持しています。

もしあなたが新しいモデルを量子化して公開した場合、ぜひその成果をお知らせください。Hugging Face Hub で共有し、お知らせいただければ幸いです。本機能についてご質問がある場合は、Discord へお気軽にご参加ください。

詳しくは以下のリソースをご覧ください：

• Diffusers Nunchaku ドキュメント
• 統合に関する PR (huggingface/diffusers#14100)
• SVDQuant 論文 と Nunchaku エンジン
• diffuse-compressor

前回の投稿：Diffusers における量子化バックエンドの探求 および Quanto と Diffusers を活用したメモリ効率の高い拡散トランスフォーマー

謝辞

Diffusers のメンテナーの皆様には、統合プロセスを通じてレビューとご指導をいただき、MIT HAN Lab / Nunchaku チームには SVDQuant の元となる研究を提供していただき、心より感謝申し上げます。また、本ブログ記事に対するフィードバックをいただいた Marc Sun 氏にもお礼申し上げます。

nunchaku-lite を実際に試していただき、貴重なご意見を寄せてくださった Álvaro Somoza 氏にも深く感謝いたします。

rootonchair は、この研究への支援と開発の多くが行われた環境を提供してくれた SilverAI 社にも厚く御礼を申し上げます。