Nunchaku Lite

Nunchaku Lite là một backend lượng tử hóa để nạp các checkpoint đã lượng tử hóa trước trong Diffusers. Tạo các checkpoint tương thích bằng diffuse-compressor. Công cụ này lượng tử hóa và xuất một transformer, sau đó đóng gói nó thành một pipeline Diffusers.

Nunchaku Lite được xây dựng dựa trên bộ máy suy luận Nunchaku ban đầu, thư viện lượng tử hóa DeepCompressor và bài báo SVDQuant.

Cài đặt các kernel CUDA

Gói kernels cung cấp các kernel CUDA đã được tối ưu hóa, được tải tự động. Hãy cài đặt nó trước tiên.

pip install -U kernels

Nunchaku Lite tải các kernel của nó từ kho rootonchair/nunchaku-lite-kernels, nơi nhà phát hành không phải là nhà phát hành kernel đáng tin cậy trên Hub. Việc tải nó sẽ tải xuống và thực thi mã từ Hub, do đó Diffusers yêu cầu bạn phải xác nhận rõ ràng bằng cách đặt DIFFUSERS_TRUST_REMOTE_KERNELS=true. Xem Trusting remote kernels để biết thêm chi tiết.

Tải pipeline đã lượng tử hóa

Tải pipeline đã lượng tử hóa trước bằng from_pretrained(), hàm này sẽ đọc cấu hình lượng tử hóa từ 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 được xuất ra phải khớp chính xác với kiến trúc của model Diffusers đích. Ví dụ, một checkpoint được lượng tử hóa với các phép chiếu QKV đã hợp nhất sẽ không thể nạp vào cấu hình model kỳ vọng các module chiếu Q, K và V riêng biệt.

Các kiểu lượng hóa được hỗ trợ

Nunchaku Lite hỗ trợ các định dạng lớp tuyến tính lượng tử hóa sau.

Sử dụng nvfp4 trên GPU Blackwell. Việc chạy các checkpoint int4 trên Blackwell có thể chậm hơn so với nvfp4.

Các kernel CUDA hiện hỗ trợ các kiến trúc GPU NVIDIA sau:

• sm_75 (Turing, ví dụ RTX 2080)
• sm_80 (Ampere, ví dụ A100)
• sm_86 (Ampere, ví dụ RTX 3090 và RTX A6000)
• sm_89 (Ada, ví dụ RTX 4090)
• sm_120 (Blackwell, ví dụ RTX 5090)

Các GPU Hopper, chẳng hạn như H100 và H200 sm_90, hiện chưa được hỗ trợ.

nvfp4 checkpoints yêu cầu GPU NVIDIA Blackwell hoặc mới hơn. Trên GPU Blackwell, hãy dùng PyTorch >= 2.7 với CUDA >= 12.8. int4 checkpoints yêu cầu GPU NVIDIA Turing hoặc mới hơn.

NunchakuLiteQuantizationConfig

Tệp config.json phải bao gồm một NunchakuLiteQuantizationConfig. Nó định nghĩa compute_dtype thời gian chạy và các module đích cho từng phương pháp lượng tử hóa Nunchaku Lite.

• compute_dtype: dtype thời gian chạy cho các bộ đệm số dấu phẩy động trong các module lượng tử hóa, thường là torch.bfloat16.
• svdq_w4a4: Cấu hình đích SVDQ W4A4 với precision, group_size, rank và targets.
• awq_w4a16: cấu hình đích AWQ W4A16 với precision, group_size và targets.

Mỗi mục trong targets phải trỏ đến một lớp tuyến tính. Diffusers thay thế mỗi mục tiêu svdq_w4a4 bằng một lớp SVDQ W4A4 và mỗi mục tiêu awq_w4a16 bằng một lớp AWQ W4A16. Ví dụ bên dưới cho thấy hình dạng mong đợi với các danh sách mục tiêu được rút ngắn.

Liệt kê từng module bạn muốn lượng tử hóa dưới svdq_w4a4 hoặc awq_w4a16. Một module chỉ có thể sử dụng một phương pháp, vì vậy đừng liệt kê cùng một mục tiêu dưới cả hai.

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

Các kernel đã hợp nhất

Động cơ Nunchaku gốc đạt được phần lớn tốc độ từ các đường dẫn thực thi hợp nhất (fused execution paths) dành riêng cho từng mô hình. Nó kết hợp các phép chiếu Q, K và V với RMSNorm và RoPE, đồng thời sử dụng kernel GELU hợp nhất cho MLP. Nunchaku Lite thay vào đó sử dụng mô hình Diffusers tiêu chuẩn với các lớp tuyến tính lượng tử hóa (quantized linear layers) tổng quát, do đó không bao gồm các phép hợp nhất này.

Các phép đo dưới đây cho thấy tác động của độ trễ mỗi bước do các phép hợp nhất này gây ra trong engine Nunchaku gốc, được đo với Flux Schnell trên RTX 5090.

Nếu không có các phép hợp nhất này, GPU sẽ khởi chạy các thao tác chiếu, chuẩn hóa, nhúng quay và MLP dưới dạng các kernel riêng biệt, điều này làm tăng chi phí khởi chạy. torch.compile có thể giảm một phần chi phí này, mặc dù nó không nhất thiết tái tạo các kernel đã hợp nhất của engine gốc. Hãy chạy benchmark cho pipeline đã biên dịch của bạn cho mô hình và khối lượng công việc cụ thể.

torch.compile

Nunchaku Lite kernels và các lớp tuyến tính lượng tử hóa tương thích với torch.compile. Biên dịch transformer đã lượng tử hóa sau khi tải pipeline để suy luận nhanh hơn.

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

Pipeline Nunchaku Lite NVFP4 đã được biên dịch chạy nhanh hơn 1,8 lần so với pipeline BF16 gốc (2,271s → 1,675s trên RTX PRO 6000).

Tài nguyên

• diffuse-compressor
• Yêu cầu cài đặt Nunchaku