Hugging FaceがDiffusersにNunchaku 4ビット拡散推論を統合

Hugging FaceがDiffusersにNunchaku 4ビット拡散推論を統合

Hugging FaceはNunchaku LiteをDiffusersライブラリに統合し、4ビットの重みとアクティベーション（W4A4）で大規模な拡散トランスフォーマーを実行できるようにしました。この統合により、VRAMの必要量が最大50％削減され、BF16ベースラインと比較してノイズ除去レイテンシが約30％改善され、消費者向けGPUでも高性能な画像生成が可能になります。

SVDQuantとNunchakuエンジン

Nunchaku LiteはSVDQuantに基づいており、拡散トランスフォーマーの重みとアクティベーションに見られる大きな外れ値を扱うために設計された量子化手法です。標準的な4ビット量子化とは異なり、SVDQuantはアクティベーションの外れ値を重みに移動させ、各重み行列の最も困難な部分を小さな16ビット低ランクブランチで表現します。残りの残差は4ビットに量子化されます。

参照実装のNunchaku CUDA推論エンジンは、低ランクのダウンドプロジェクションと量子化カーネル、低ランクのアッププロジェクションと4ビット計算カーネルを融合したカーネルを使用してこのプロセスを最適化し、16ビットブランチに対するメモリアクセスのオーバーヘッドを実質的に排除します。

DiffusersへのNunchaku Lite統合

Nunchaku LiteはDiffusers内でシンプルな統合パスを提供し、from_pretrained()を介してNunchakuスタイルのチェックポイントをロードできるようにします。別途推論エンジンやローカルCUDAコンパイルは不要で、nn.LinearモジュールをランタイムのSVDQ/AWQ線形層でパッチングすることで実現しています。

Nunchaku Liteは主に以下の2つのカーネルファミリーを使用します。

  • svdq_w4a4: トランスフォーマーの注意機構やMLP投影など、計算が集中する箇所で使用されます。INT4とNVFP4のバリエーションがあります。
  • awq_w4a16: メモリバウンドで精度に敏感な層（例：適応正規化や変調投影、FLUXのadanorm_single / adanorm_zeroなど）で使用され、4ビット重みと16ビットアクティベーションを利用します。

Nunchaku Liteは元のNunchakuエンジンのアーキテクチャ固有の融合実行パスは含みませんが、約30％の速度向上と同等のVRAM削減を実現します。

ハードウェアサポートとパフォーマンス

Nunchaku LiteカーネルのサポートはGPU世代と精度により異なります。

スキーム 精度 対応GPU
svdq_w4a4 nvfp4 Blackwell (RTX 50 series, RTX PRO 6000, B200)
svdq_w4a4 int4 Turing / Ampere / Ada (RTX 30 & 40 series, A100, L40S)
awq_w4a16 int4 Turing / Ampere / Ada (RTX 30 & 40 series, A100, L40S)

VoltaおよびHopper GPUは現在未対応です。

ベンチマーク

NVIDIA RTX PRO 6000（Blackwell）で1024×1024解像度、ERNIE-Image-Turboチェックポイントを使用して測定しました。

設定 フルパイプライン デノイズループ ピークVRAM 速度向上
BF16ベースライン 3.00 s 2.86 s 31.1 GB 1.0x
Nunchaku Lite NVFP4 2.27 s 2.13 s 20.6 GB 1.35x
Nunchaku Lite NVFP4 + torch.compile 1.68 s 1.53 s 20.6 GB 1.8x
Nunchaku Lite NVFP4 + NF4テキストエンコーダ 2.29 s 2.13 s 16.0 GB 1.35x

Nunchaku Liteにtorch.compileを組み合わせるとエンドツーエンドの速度向上が1.8倍に達し、bitsandbytesのNF4量子化テキストエンコーダを追加するとピークVRAMがさらに16.0 GBに削減されます。

diffuse-compressorでカスタムモデルを量子化

Nunchaku Liteはアーキテクチャに依存しません。diffuse-compressorツールキットを使えば、ユーザーは自分のDiffusersモデルをキャリブレーション、量子化、パッケージ化、公開できます。ワークフローは以下の通りです。

  1. Inspection（検査）: 汎用スキャナがSVDQ W4A4対象（繰り返しトランスフォーマーブロック）とAWQ W4A16対象（変調線形）を特定します。
  2. Quantization（量子化）: SVDQuantを実行して量子化チェックポイントを作成します（int4またはnvfp4に対応）。
  3. Packaging（パッケージ化）: 量子化されたトランスフォーマーをベースパイプラインと結合し、transformer/config.jsonにnunchaku_lite設定を作成します。
  4. Verification（検証）: DiffusionPipeline.from_pretrained()でモデルをロードし、最終結果をHugging Face Hubにプッシュします。

構造的リライト

最大性能を引き出すために、一部のモデルでは構造的リライトが必要です。たとえば、個別のQ、K、V投影を単一のto_qkvモジュールに統合するなどです。汎用のdiffuse-compressorパスではこれらの変更を自動推論できませんが、rootonchair/nunchaku-liteが提供するモデル固有のターゲット設定とランタイムアダプタを使用すれば、融合操作を有効にできます。

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

Hub上で利用可能な事前量子化チェックポイントがいくつかあります。

  • ERNIE-Image-Turbo: INT4とNVFP4のバリエーションがあり、NF4テキストエンコーダを使用可能。
  • Krea 2 Turbo: NVFP4で利用可能。
  • lite-infer: 追加のNunchaku Liteチェックポイントのコレクション。

Sources