Nunchaku Lite

Nunchaku Lite es un backend de cuantización para cargar puntos de control precuantizados en Diffusers. Crea puntos de control compatibles con diffuse-compressor. Cuantiza y exporta un transformer y luego lo empaqueta como un pipeline de Diffusers.

Nunchaku Lite se basa en el motor de inferencia original Nunchaku, la biblioteca de cuantización DeepCompressor y el artículo SVDQuant.

Instalar los núcleos de CUDA

El paquete kernels proporciona los kernels de CUDA optimizados, que se cargan automáticamente. Instálalo primero.

pip install -U kernels

Nunchaku Lite carga sus kernels desde el repositorio rootonchair/nunchaku-lite-kernels, cuyo editor no es un editor de kernels de confianza en el Hub. Al cargarlo, se descarga y ejecuta código desde el Hub, por lo que Diffusers requiere que aceptes explícitamente esta acción estableciendo DIFFUSERS_TRUST_REMOTE_KERNELS=true. Consulta Confianza en kernels remotos para más detalles.

Cargar un pipeline cuantizado

Cargue el pipeline precuantizado con from_pretrained(), que lee la configuración de cuantización de 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")

El state dict exportado debe coincidir exactamente con la arquitectura del modelo Diffusers de destino. Por ejemplo, un checkpoint cuantizado con proyecciones QKV fusionadas no se cargará en una configuración de modelo que espere módulos de proyección Q, K y V separados.

Tipos de cuantización admitidos

Nunchaku Lite admite los siguientes formatos de capas lineales cuantizadas.

Use nvfp4 en GPUs Blackwell. Ejecutar puntos de control int4 en Blackwell puede ser más lento que nvfp4.

Los núcleos de CUDA actualmente admiten las siguientes arquitecturas de GPU de NVIDIA:

• sm_75 (Turing, por ejemplo RTX 2080)
• sm_80 (Ampere, por ejemplo A100)
• sm_86 (Ampere, por ejemplo RTX 3090 y RTX A6000)
• sm_89 (Ada, por ejemplo RTX 4090)
• sm_120 (Blackwell, por ejemplo, RTX 5090)

Las GPU Hopper, como sm_90 H100 y H200, no están soportadas actualmente.

Los puntos de control nvfp4 requieren una GPU NVIDIA Blackwell o más reciente. En las GPUs Blackwell, use PyTorch >= 2.7 con CUDA >= 12.8. Los puntos de control int4 requieren una GPU NVIDIA Turing o más reciente.

NunchakuLiteQuantizationConfig

El archivo config.json debe incluir una NunchakuLiteQuantizationConfig. Define el compute_dtype de tiempo de ejecución y los módulos objetivo para cada método de cuantización de Nunchaku Lite.

• compute_dtype: dtype de ejecución para los búferes de punto flotante en los módulos cuantizados, típicamente torch.bfloat16.
• svdq_w4a4: configuración objetivo SVDQ W4A4 con precision, group_size, rank y targets.
• awq_w4a16: configuración objetivo AWQ W4A16 con precision, group_size y targets.

Cada entrada en targets debe apuntar a una capa lineal. Diffusers reemplaza cada destino svdq_w4a4 por una capa SVDQ W4A4 y cada destino awq_w4a16 por una capa AWQ W4A16. El ejemplo siguiente muestra la forma esperada con listas de destinos abreviadas.

Liste cada módulo que desea cuantizar bajo svdq_w4a4 o awq_w4a16. Un módulo solo puede usar un método, por lo que no liste el mismo destino bajo ambos.

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

Kernels fusionados

El motor original Nunchaku obtiene gran parte de su velocidad de las rutas de ejecución fusionadas específicas del modelo. Combina las proyecciones Q, K y V con RMSNorm y RoPE, y utiliza un kernel GELU fusionado para el MLP. Nunchaku Lite, en cambio, utiliza el modelo estándar de Diffusers con capas lineales cuantizadas genéricas, por lo que no incluye estas fusiones.

Las siguientes mediciones muestran el impacto en la latencia por paso de estas fusiones en el motor original de Nunchaku, medido con Flux Schnell en una RTX 5090.

Sin estas fusiones, la GPU lanza las operaciones de proyección, normalización, incrustación rotacional y MLP como núcleos separados, lo que añade una sobrecarga de lanzamiento. torch.compile puede reducir parte de esta sobrecarga, aunque no necesariamente reproduce los núcleos fusionados del motor original. Realice una prueba de rendimiento del pipeline compilado para su modelo y carga de trabajo.

torch.compile

Los kernels de Nunchaku Lite y las capas lineales cuantizadas son compatibles con torch.compile. Compila el transformer cuantizado después de cargar el pipeline para una inferencia más rápida.

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

El pipeline compilado de Nunchaku Lite NVFP4 se ejecuta 1,8 veces más rápido que el pipeline original BF16 (2,271 s → 1,675 s en una RTX PRO 6000).

Recursos

• diffuse-compressor
• Requisitos de instalación de Nunchaku