Nunchaku Lite

Nunchaku Lite adalah backend kuantisasi untuk memuat checkpoint prekuantisasi di Diffusers. Buat checkpoint yang kompatibel dengan diffuse-compressor. Alat ini mengkuantisasi dan mengekspor transformer, lalu mengemasnya sebagai pipeline Diffusers.

Nunchaku Lite dibangun di atas mesin inferensi Nunchaku asli, pustaka kuantisasi DeepCompressor, dan makalah SVDQuant.

Install kernel CUDA

Paket kernels menyediakan kernel CUDA yang dioptimalkan, yang dimuat secara otomatis. Instal paket ini terlebih dahulu.

pip install -U kernels

Nunchaku Lite memuat kernelnya dari repositori rootonchair/nunchaku-lite-kernels, yang penerbitnya bukan penerbit kernel yang tepercaya di Hub. Memuatnya akan mengunduh dan mengeksekusi kode dari Hub, sehingga Diffusers mengharuskan Anda untuk secara eksplisit mengaktifkannya dengan mengatur DIFFUSERS_TRUST_REMOTE_KERNELS=true. Lihat Mempercayai kernel jarak jauh untuk detailnya.

Muat pipeline terkuantisasi

Muat pipeline prequantized dengan from_pretrained(), yang membaca konfigurasi kuantisasi dari config.json.

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 yang diekspor harus cocok persis dengan arsitektur model Diffusers target. Sebagai contoh, checkpoint yang dikuantisasi dengan proyeksi QKV yang digabungkan tidak akan dapat dimuat ke dalam konfigurasi model yang mengharapkan modul proyeksi Q, K, dan V yang terpisah.

Jenis kuantisasi yang didukung

Nunchaku Lite mendukung format lapisan linear terkuantisasi berikut.

Gunakan nvfp4 pada GPU Blackwell. Menjalankan checkpoint int4 di Blackwell bisa lebih lambat daripada nvfp4.

Kernel CUDA saat ini mendukung arsitektur GPU NVIDIA berikut:

• sm_75 (Turing, misalnya RTX 2080)
• sm_80 (Ampere, misalnya A100)
• sm_86 (Ampere, misalnya RTX 3090 dan RTX A6000)
• sm_89 (misalnya Ada, seperti RTX 4090)
• sm_120 (Blackwell, misalnya RTX 5090)

GPU Hopper, seperti sm_90 H100 dan H200, saat ini belum didukung.

Checkpoint nvfp4 memerlukan GPU NVIDIA Blackwell atau yang lebih baru. Pada GPU Blackwell, gunakan PyTorch >= 2.7 dengan CUDA >= 12.8. Checkpoint int4 memerlukan GPU NVIDIA Turing atau yang lebih baru.

NunchakuLiteQuantizationConfig

File config.json harus menyertakan NunchakuLiteQuantizationConfig. Konfigurasi ini mendefinisikan compute_dtype saat runtime serta modul target untuk setiap metode kuantisasi Nunchaku Lite.

• compute_dtype: dtype runtime untuk buffer bilangan pecahan dalam modul terkuantisasi, biasanya torch.bfloat16.
• svdq_w4a4: konfigurasi target SVDQ W4A4 dengan precision, group_size, rank, dan targets.
• awq_w4a16: konfigurasi target AWQ W4A16 dengan precision, group_size, dan targets.

Setiap entri dalam targets harus menunjuk ke lapisan linear. Diffusers mengganti setiap target svdq_w4a4 dengan lapisan SVDQ W4A4 dan setiap target awq_w4a16 dengan lapisan AWQ W4A16. Contoh di bawah ini menunjukkan bentuk yang diharapkan dengan daftar target yang dipendekkan.

Daftarkan setiap modul yang ingin Anda kuantisasi di bawah svdq_w4a4 atau awq_w4a16. Sebuah modul hanya dapat menggunakan satu metode, jadi jangan mendaftarkan target yang sama di bawah keduanya.

{
  "_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"]
    }
  }
}

Kernel yang digabungkan

Mesin Nunchaku asli mendapatkan sebagian besar kecepatannya dari jalur eksekusi terpadu yang spesifik untuk model. Mesin ini menggabungkan proyeksi Q, K, dan V dengan RMSNorm dan RoPE, serta menggunakan kernel GELU terpadu untuk MLP. Nunchaku Lite, di sisi lain, menggunakan model Diffusers standar dengan lapisan linear terkuantisasi generik, sehingga tidak menyertakan fusi-fusi tersebut.

Pengukuran berikut menunjukkan dampak latensi per langkah dari fusi-fusi ini di mesin Nunchaku asli, diukur dengan Flux Schnell pada RTX 5090.

Tanpa fusi-fusi ini, GPU meluncurkan operasi proyeksi, normalisasi, rotary-embedding, dan MLP sebagai kernel terpisah, yang menambah overhead peluncuran. torch.compile dapat mengurangi sebagian overhead ini, meskipun tidak selalu mereproduksi kernel terfusi dari engine asli. Lakukan benchmark pada pipeline terkompilasi untuk model dan beban kerja Anda.

torch.compile

Kernel Nunchaku Lite dan lapisan linear terkuantisasi kompatibel dengan torch.compile. Kompilasi transformer terkuantisasi setelah memuat pipeline untuk inferensi yang lebih cepat.

pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=True)

Pipeline Nunchaku Lite NVFP4 yang dikompilasi berjalan 1,8x lebih cepat daripada pipeline BF16 asli (2,271s → 1,675s pada RTX PRO 6000).

Sumber daya

• diffuse-compressor
• Peryaratan instalasi Nunchaku