Buckets:

mishig's picture
|
download
raw
7.23 kB

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 desde 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.

Método Precisión Tamaño del grupo Notas
svdq_w4a4 nvfp4 16 Utiliza kernels de tiempo de ejecución NVFP4 con corrección de bajo rango SVDQ.
svdq_w4a4 int4 64 Utiliza núcleos INT4 W4A4 con corrección de bajo rango SVDQ.
awq_w4a16 int4 64 Utiliza núcleos de solo pesos INT4 de estilo AWQ.

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.

Optimización Latencia con él DESACTIVADO Latencia con él ACTIVADO Aceleración
Fusión QKV+norm+RoPE (combinada) 1.057 s 0.834 s 1.27×
— agrupar Q/K/V en una sola matmul, por separado 1.059 s 1.034 s 1.025×
— Fusión del epílogo de RMSNorm+RoPE, además del agrupamiento 1.034 s 0.836 s 1.236×
Fusión GELU-MLP 0.766 s 0.722 s 1.06×

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

Xet Storage Details

Size:
7.23 kB
·
Xet hash:
17ef19035b59ec524ff3c30d981114411da5a0907625b0b565b2ef1ea2cf5b38

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.