大型拡散 Transformer が Nunchaku Lite を介して Diffusers でネイティブにロード可能に

大型拡散 Transformer（Diffusion Transformer）は素晴らしい画像（さらには動画、オーディオクリップ、そして最近ではテキストも）を生成できますが、現代のテキスト画像生成モデルを BF16 精度でロードするには通常 20〜30 GB の VRAM が必要となり、ほとんどのコンシューマー向け GPU では手が届きません。量子化はこの問題に対する強力な解決策であり、Diffusers はすでに bitsandbytes、GGUF、torchao、Quanto など、さまざまな量子化バックエンドを統合しています。これについては Exploring Quantization Backends in Diffusers で紹介しました。

これらのバックエンドのほとんどは重みのみの量子化（weight-only）です。これは、重みを低精度で保存し、計算時に高精度に逆量子化することを意味します。これにより VRAM のフットプリントは大幅に削減されますが、通常は推論が高速化されるわけではなく、むしろわずかな追加遅延が発生することさえあります。

SVDQuant は人気のある推論エンジン Nunchaku の背後にある量子化手法であり、異なるアプローチを取っています。これは、主要な Transformer レイヤーを 4 ビットの重みとアクティベーション（W4A4）で実行し、メモリ使用量を削減すると同時にデノイジングループを高速化します。その仕組みについては後ほど詳しく説明しますが、これまではこれらのチェックポイントを使用するために個別の推論ライブラリが必要でした。

現在の Diffusers では、kernels パッケージのおかげで、ローカルでの CUDA コンパイルなしで、単に from_pretrained() を呼び出すだけで Nunchaku チェックポイントをロードできます。さらに、付属の diffuse-compressor ツールキットを使用すると、独自の新しいアーキテクチャを量子化し、通常の Diffusers リポジトリと同様に公開できます。

Nunchaku Lite の画質とパフォーマンスの比較

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

BF16 と Nunchaku Lite によるアカギツネのプロンプト出力の比較

カスタムパイプラインクラスや個別の推論エンジンは不要で、ローカルでコンパイルするものもありません。NVFP4 カーネルは、初回使用時に Nunchaku Lite kernels ページから Hub 経由でダウンロードされます。このチェックポイントは、Nunchaku NVFP4 Transformer と bitsandbytes NF4 テキストエンコーダーを組み合わせており、RTX 5090 上で 1024x1024 の画像を生成するのに約 1.7 秒かかり、ピーク VRAM 使用量は約 12 GB です。これに対し、BF16 パイプラインは約 24 GB を消費します。Nunchaku Lite チェックポイント形式のより詳細な情報については、公式 Diffusers ドキュメントを参照してください。

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

背景：SVDQuant と Nunchaku

SVDQuant は、リファレンス CUDA 推論エンジンである Nunchaku の背後にある量子化手法です。拡散 Transformer において、標準的な 4 ビット量子化は、重みとアクティベーションの両方に大きな外れ値（アウトライヤー）が存在するため困難です。SVDQuant のアプローチは、アクティベーションの外れ値を重みに移行し、各重み行列の最も処理が難しい部分を小さな 16 ビットの低ランク（low-rank）ブランチで表現し、残りの残差を 4 ビットに量子化することです。Nunchaku は、4 ビットパスと低ランクブランチ向けの融合（fused）カーネルによって、このプロセスを高速化します。

Nunchaku カーネル融合：低ランクダウンプロジェクションと入力量子化の融合、低ランクアッププロジェクションと 4 ビット行列積の融合
Nunchaku は、低ランクダウンプロジェクションを量子化カーネルと融合させ、低ランクアッププロジェクションを 4 ビット計算カーネルと融合させることで、16 ビットブランチのメモリ書き込み・読み出しオーバーヘッドを排除します。図は SVDQuant 論文より引用。

Nunchaku Lite の紹介

オリジナルの Nunchaku エンジンの速度面での優位性の大部分は、融合 QKV プロジェクションや融合 GELU/MLP カーネルなど、特定のモデルに向けた融合実行パスによるものです。これらの最適化は、各アーキテクチャのモジュールレイアウトやチェックポイント形式と密接に関連しているため、新しいモデルファミリーをサポートするには、通常、そのモデルに特化した統合プロセスが必要になります。

Nunchaku Lite は、Diffusers における新しい統合方法です。これにより、Diffusers はカスタムパイプラインや独立した推論エンジンを必要とせずに、Nunchaku スタイルのチェックポイントをロードできます。内部的には、Nunchaku Lite はチェックポイントをロードする前に、標準の Diffusers モデル内の関連する nn.Linear モジュールを実行時の SVDQ/AWQ 線形レイヤーに置き換えます。CUDA カーネルは、kernels パッケージを介して Hub から取得されます。ここでは 2 種類のカーネルが使用されます。

svdq_w4a4：

# Nunchaku 4ビット拡散

大型拡散 Transformer が Nunchaku Lite を介して Diffusers でネイティブにロード可能に

目次

Nunchaku Lite を始める

背景：SVDQuant と Nunchaku

Nunchaku Lite の紹介