SDNQ

SDNQ (SD.Next Quantization) es un backend de cuantización sin entrenamiento que admite una amplia gama de dtypes de pesos (int8 hasta 2 bits, FP8 y otros formatos de coma flotante de bajo bit) con corrección SVD opcional, rotación de Hadamard y matmul cuantizado INT8/FP8 para una inferencia más rápida. Funciona en CUDA, ROCm, XPU, MPS y CPU.

Instale la biblioteca sdnq para usar este backend.

pip install sdnq

Cargar un punto de control precuantizado

La configuración de cuantización de un checkpoint SDNQ precuantizado se guarda en su config.json, por lo que se carga como cualquier otro modelo, sin necesidad del argumento quantization_config. El único requisito es tener sdnq instalado.

import torch
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained(
    "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

image = pipe("a cat holding a sign that says hello").images[0]
image.save("output.png")

Establezca DIFFUSERS_SDNQ_TRANSFORMERS=1 (antes de importar Diffusers) cuando cargue un modelo de Transformers cuantizado con SDNQ de forma independiente, como un codificador de texto cuantizado o un componente LLM autónomo en una pipeline modular. Esto permite que sdnq se registre con Transformers para que el modelo cuantizado se cargue correctamente. No es necesario cuando se carga un transformer de diffusers cuantizado con SDNQ junto con él, ya que la carga de ese transformer ya importa sdnq y lo registra. Está desactivado de forma predeterminada para que las pipelines que no usan SDNQ nunca importen sdnq.

Cuantizar un modelo sobre la marcha

Pase un SDNQConfig a from_pretrained para cuantizar un modelo durante la carga. Use PipelineQuantizationConfig para cuantizar componentes específicos de la pipeline.

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, SDNQConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": SDNQConfig(weights_dtype="uint4", use_svd=True)}
)
pipe = DiffusionPipeline.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

O cuantice directamente un componente individual del modelo.

import torch
from diffusers import ZImageTransformer2DModel, SDNQConfig

transformer = ZImageTransformer2DModel.from_pretrained(
    "Tongyi-MAI/Z-Image-Turbo",
    subfolder="transformer",
    quantization_config=SDNQConfig(weights_dtype="int8"),
    dtype=torch.bfloat16,
)

Todos los argumentos se reenvían a sdnq.SDNQConfig. Consulte la documentación de SDNQ para obtener la lista completa de dtypes y opciones admitidas.

Elegir qué capas cuantizar

Algunas capas son más sensibles a la cuantización que otras, y omitir las sensibles mejora la calidad con un coste de memoria mínimo. SDNQ gestiona esto de forma nativa: incluye listas de omisión integradas (una genérica y listas específicas por modelo para muchas arquitecturas) que se aplican automáticamente, de modo que las capas sensibles habituales permanecen sin cuantizar sin necesidad de configuración. Para un control más fino, también puede seleccionar capas dinámicamente, midiendo el error de cuantización por capa y omitiendo las que superan un umbral de pérdida. Consulte la guía de cuantización dinámica de SDNQ para más detalles.

Inferencia más rápida con matmul cuantizado

Por defecto, SDNQ descuantiza los pesos al tipo de datos de cálculo antes de cada matmul. Si se habilita el matmul cuantizado, el matmul se ejecuta directamente en un tipo de datos de baja precisión (INT8/FP8/FP16), lo cual es más rápido en hardware con una instalación funcional de Triton (CUDA, ROCm o XPU).

Actívelo en un modelo cargado con apply_sdnq_options_to_model, que lee la configuración de cuantización existente del modelo y vuelve a aplicar la opción in situ. Esto también funciona para puntos de control precuantizados, donde no se pasa una SDNQConfig.

import torch
from diffusers import DiffusionPipeline
from sdnq.common import use_torch_compile as triton_is_available
from sdnq.loader import apply_sdnq_options_to_model

pipe = DiffusionPipeline.from_pretrained(
    "Disty0/Z-Image-Turbo-SDNQ-uint4-svd-r32",
    dtype=torch.bfloat16,
).to("cuda")  # or "mps", "xpu", "cpu"

if triton_is_available and (torch.cuda.is_available() or torch.xpu.is_available()):
    pipe.transformer = apply_sdnq_options_to_model(pipe.transformer, use_quantized_matmul=True)
    pipe.text_encoder = apply_sdnq_options_to_model(pipe.text_encoder, use_quantized_matmul=True)

torch.compile

Los modelos cuantizados SDNQ son compatibles con torch.compile para obtener aún más velocidad, incluso en combinación con group offloading.

transformer.compile(fullgraph=True)

Guardar un modelo cuantizado

Los modelos cuantizados se pueden serializar con save_pretrained y recargarlos sin una quantization_config.

transformer.save_pretrained("z-image-turbo-sdnq-int8")

Recursos

Update on GitHub