Nunchaku Lite en Diffusers: inferencia 4-bit para modelos de difusión

El problema: grandes modelos y memoria limitada

Los modelos de difusión de última generación producen imágenes impresionantes, pero su tamaño y precisión (por ejemplo BF16) suelen exigir entre 20 y 30 GB de VRAM. Ese requisito coloca a muchos GPU de consumo fuera del alcance para ejecutar pipelines text-to-image a alta resolución.

La cuantización ha sido una solución práctica: almacenando pesos en baja precisión se reduce notablemente el uso de memoria. Backends como bitsandbytes, GGUF, torchao y Quanto ya están integrados en Diffusers y ofrecen opciones de cuantización peso-solo. Sin embargo, la mayoría de estas soluciones de peso-solo descomprimen a alta precisión en tiempo de cómputo, lo que reduce memoria pero no siempre acelera la inferencia; en algunos casos añade latencia.

Qué es SVDQuant y por qué importa

SVDQuant es la técnica de cuantización detrás del motor Nunchaku. A diferencia de la cuantización tradicional de 4 bits, SVDQuant aborda un problema clave en transformadores de difusión: tanto pesos como activaciones a menudo contienen outliers que dañan la precisión si se cuantizan de forma directa.

La estrategia de SVDQuant consiste en mover los outliers de activación hacia los pesos, representando la parte más difícil de cada matriz de pesos con una rama de baja rango en 16 bits y cuantizando el residuo a 4 bits. Esto permite mantener precisión donde más se necesita y reducir drásticamente el tamaño del resto del modelo.

Nunchaku acelera esta técnica mediante kernels fusionados que eliminan el overhead de memoria de la rama de baja precisión, integrando las proyecciones de baja y alta precisión para minimizar accesos a memoria y latencia.

Nunchaku Lite: integración nativa en Diffusers

El Nunchaku original obtiene gran parte de su rendimiento de optimizaciones específicas de cada arquitectura: kernels fusionados para QKV, GELU/MLP y layouts de módulo particulares. Eso suele implicar trabajo de integración por modelo.

Nunchaku Lite propone un camino diferente: permite cargar checkpoints con estilo Nunchaku directamente desde Diffusers, sin necesidad de un motor de inferencia separado ni clases de pipeline personalizadas. Bajo el capó, Nunchaku Lite parchea en tiempo de ejecución los módulos nn.Linear relevantes de un modelo Diffusers con capas lineales SVDQ/AWQ antes de cargar el checkpoint. Los kernels CUDA necesarios se descargan desde el Hub a través del paquete kernels.

Se usan dos familias de kernels principales:

• svdq_w4a4: pesos y activaciones en 4 bits con corrección de bajo rango (SVDQuant). Se aplica en las proyecciones de atención y MLP, donde reside la mayor parte del cómputo. Disponible en variantes INT4 y NVFP4.
• awq_w4a16: pesos en 4 bits y activaciones en 16 bits; útil para proyecciones que son sensibles a la precisión (normalizaciones adaptativas y módulos de modulación), reduciendo el uso de memoria sin sacrificar la estabilidad numérica.

Sin las fusiones específicas del motor Nunchaku original, Nunchaku Lite no alcanza exactamente los mismos picos de velocidad, pero ofrece una implementación más universal y que aún consigue alrededor de un 30% de mejora en velocidad junto con una reducción significativa de VRAM.

Cómo empezar con Nunchaku Lite

Instalar los requisitos es sencillo. Diffusers y el paquete kernels son necesarios, junto con transformers, accelerate y bitsandbytes:

pip install -U diffusers transformers accelerate kernels bitsandbytes

Después, cargar un pipeline cuantizado funciona igual que con cualquier checkpoint en Diffusers. Por ejemplo, un pipeline pre-cuantisado ErnieImage:

import torch
from diffusers import ErnieImagePipeline

pipe = ErnieImagePipeline.from_pretrained(
  "lite-infer/ERNIE-Image-Turbo-nunchaku-lite-nvfp4_r32-bnb4-text-encoder",
  torch_dtype=torch.bfloat16,
).to("cuda")

image = pipe(
  prompt="A cinematic portrait of a red fox in a misty forest at sunrise, detailed fur, volumetric light",
  height=1024, width=1024, num_inference_steps=8, guidance_scale=1.0,
  generator=torch.Generator("cuda").manual_seed(42),
).images[0]

image.save("output.png")

No hay que compilar CUDA localmente ni usar un motor separado: las kernels NVFP4 se descargan desde el Hub la primera vez que se requieren.

Rendimiento y soporte de hardware

Un ejemplo práctico citado por los desarrolladores muestra que un checkpoint NVFP4 que combina un transformador Nunchaku NVFP4 con un text-encoder bitsandbytes NF4 puede generar una imagen 1024×1024 en aproximadamente 1.7 segundos en una RTX 5090, usando un pico de memoria cercano a 12 GB. Para el mismo pipeline en BF16, el uso de memoria rondaría los 24 GB.

Importante: los checkpoints NVFP4 requieren GPUs NVIDIA Blackwell (por ejemplo: RTX 50 series, RTX PRO 6000, B200). Para generaciones anteriores de GPU se recomienda usar las variantes INT4. Revisen la tabla de soporte de hardware del repositorio o la documentación de Diffusers para confirmar compatibilidades con sus GPUs.

Cuantizar y publicar sus propios modelos

Si su organización o equipo quiere producir checkpoints cuantizados, el toolkit companion diffuse-compressor permite cuantizar nuevas arquitecturas y publicarlas como repositorios regulares de Diffusers. Esto facilita distribuir checkpoints optimizados a usuarios finales sin que tengan que cambiar su flujo de trabajo: desde from_pretrained() se cargan como cualquier otro modelo.

Consideraciones prácticas para América Latina

En la región, muchas empresas y centros de investigación usan GPUs de generaciones anteriores o tarjetas de consumo con memoria limitada. Nunchaku Lite y las variantes INT4 ofrecen una vía para ejecutar modelos grandes sin invertir inmediatamente en hardware de última generación. Esto puede acelerar adopciones en startups, equipos de I+D y departamentos de datos que buscan prototipar capacidades generativas con recursos moderados.

Sin embargo, verifiquen compatibilidad con su flota de GPUs antes de migrar cargas; la versión NVFP4 exige hardware Blackwell, y la ganancia real depende de la arquitectura del modelo y del mix de kernels usados.

Conclusión

Nunchaku Lite acerca la inferencia de 4-bit al ecosistema Diffusers sin requerir compilaciones locales ni motores externos. Al aplicar SVDQuant y un set de kernels especializados, permite reducir VRAM y mejorar la latencia en cargas de difusión, manteniendo la experiencia de uso simple y compatible con los pipelines existentes.

Para equipos en América Latina y mercados con recursos limitados, esto representa una alternativa práctica para ejecutar modelos generativos avanzados con menor inversión en memoria GPU. Si desean llevar esta optimización a su catálogo de modelos, el diffuse-compressor facilita cuantizar y publicar checkpoints listos para usar.

Agradecimientos

Créditos al equipo de Hugging Face por desarrollar la integración y la documentación que hacen posible esta transición hacia inferencia más eficiente.