AutoRound

AutoRound es un kit de herramientas avanzado de cuantización. Logra una alta precisión en anchos de bits ultrabajos (2-4 bits) con un ajuste mínimo, aprovechando el descenso de gradiente de signo y ofreciendo una amplia compatibilidad con el hardware. Consulte nuestros artículos SignRoundV1 y SignRoundV2 para obtener más detalles.

Instala auto-round (versión ≥ 0.13.0):

pip install "auto-round>=0.13.0"

Para usar el kernel Marlin para una inferencia CUDA más rápida, instale gptqmodel:

pip install "gptqmodel>=5.8.0"

Cargar un modelo cuantizado

Cargue un modelo AutoRound pre-cuantizado pasando AutoRoundConfig a from_pretrained(). El método funciona con cualquier modelo que se cargue mediante Accelerate y que tenga capas torch.nn.Linear.

Puede usar PipelineQuantizationConfig para cuantizar componentes específicos de un pipeline:

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, AutoRoundConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": AutoRoundConfig(backend="auto")}
)
pipe = DiffusionPipeline.from_pretrained(
    "INCModel/Z-Image-W4A16-AutoRound",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda",  # or "mps", "xpu", "cpu"
)

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

O cargue directamente un componente de modelo cuantizado:

import torch
from diffusers import ZImageTransformer2DModel, ZImagePipeline, AutoRoundConfig

model_id = "INCModel/Z-Image-W4A16-AutoRound"

quantization_config = AutoRoundConfig(backend="auto")
transformer = ZImageTransformer2DModel.from_pretrained(
    model_id,
    subfolder="transformer",
    quantization_config=quantization_config,
    dtype=torch.bfloat16,
    device_map="cuda",  # or "mps", "xpu", "cpu"
)

pipe = ZImagePipeline.from_pretrained(
    model_id,
    transformer=transformer,
    dtype=torch.bfloat16,
    device_map="cuda",
)

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

AutoRound en Diffusers solo admite la carga de modelos precuantizados. Para cuantizar un modelo desde cero, use directamente la CLI o la API de Python de AutoRound y luego cargue el resultado con Diffusers.

torch.compile

AutoRound es compatible con torch.compile para una inferencia más rápida. Puedes compilar el transformer (DiT) cuantizado para un mejor rendimiento:

import torch
from diffusers import DiffusionPipeline, PipelineQuantizationConfig, AutoRoundConfig

pipeline_quant_config = PipelineQuantizationConfig(
    quant_mapping={"transformer": AutoRoundConfig(backend="auto")}
)
pipe = DiffusionPipeline.from_pretrained(
    "INCModel/Z-Image-W4A16-AutoRound",
    quantization_config=pipeline_quant_config,
    dtype=torch.bfloat16,
    device_map="cuda",  # or "mps", "xpu", "cpu"
)

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

Backends

AutoRound admite varios backends de inferencia para modelos cuantizados solo en pesos. El backend controla qué kernel gestiona la decuantización durante la pasada hacia adelante. Establece el parámetro backend en AutoRoundConfig para elegir uno:

from diffusers import AutoRoundConfig

# Auto-select (default)
config = AutoRoundConfig()

# Explicit Triton backend for CUDA
config = AutoRoundConfig(backend="tritonv2")

# Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0)
config = AutoRoundConfig(backend="marlin")

# ExllamaV2 backend for good CUDA performance (requires gptqmodel>=5.8.0)
config = AutoRoundConfig(backend="exllamav2")

# PyTorch backend for CPU/CUDA inference
config = AutoRoundConfig(backend="torch")

Guardar y cargar

from auto_round import AutoRound

autoround = AutoRound(
    "Tongyi-MAI/Z-Image",
    scheme="W4A16",  # W4G128 symmetric
    enable_torch_compile=True,
    num_inference_steps=3,
    guidance_scale=7.5,
    dataset="coco2014",
)
autoround.quantize_and_save("Z-Image-W4A16-AutoRound")

import torch
from diffusers import ZImageTransformer2DModel, ZImagePipeline

model_id = "INCModel/Z-Image-W4A16-AutoRound"

# The inference backend will be automatically selected.
pipe = ZImagePipeline.from_pretrained(
    model_id,
    dtype=torch.bfloat16,
    device_map="cuda",  # or "mps", "xpu", "cpu"
)

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

Esquemas de cuantización admitidos

AutoRound admite varios esquemas:

• W4A16(bits:4,group_size:128,sym:True,act_bits:16)
• W8A16(bits:8,group_size:128,sym:True,act_bits:16)
• W3A16(bits:3,group_size:128,sym:True,act_bits:16)
• W2A16(bits:2,group_size:128,sym:True,act_bits:16)
• GGUF:Q4_K_M(se admiten todos los Q_K,Q_0,Q*_1 proporcionados por llamacpp)
• NVFP4(Función experimental, se recomienda exportar al formato llm_compressor con data_type nvfp4,act_data_type nvfp4,static_global_scale,group_size 16)
• MXFP4(Función de investigación, sin kernel real, MXFP4 estándar, data_type mxfp,act_data_type mxfp,bits 4, act_bits 4, group_size 32)
• MXINT4(Función de investigación, sin kernel real, MXINT4 estándar, data_type mxint,act_data_type mxint,bits 4, act_bits 4, group_size 32)
• MXFP4_RCEIL(Función de investigación, sin kernel real, variante de NVIDIA, data_type mxfp, act_data_type mxfp_rceil, bits 4, act_bits 4, group_size 32)
• MXFP8(Función de investigación, sin kernel real, data_type mxfp,act_data_type mxfp_rceil,group_size 32)
• FPW8A16(Función de investigación, sin kernel real, data_type fp8,group_size 0->per tensor )
• FP8_STATIC(Función de investigación, sin kernel real, data_type:fp8,act_data_type:fp8,group_size -1 ->por canal, act_group_size=0->por tensor)

Además, puedes modificar group_size, bits, sym y muchas otras configuraciones que desees, aunque es posible que no existan kernels reales.

Recursos

• Modelos AutoRound pre-cuantizados en el Hub