AutoRound

AutoRound — это продвинутый инструментарий для квантования. Он обеспечивает высокую точность при сверхнизких разрядностях (2–4 бита) с минимальной настройкой, используя спуск по знаковому градиенту и обеспечивая широкую совместимость с аппаратным обеспечением. Подробнее о наших работах SignRoundV1 и SignRoundV2.

Установите auto-round (версия ≥ 0.13.0):

pip install "auto-round>=0.13.0"

Чтобы использовать ядро Marlin для более быстрого вывода на CUDA, установите gptqmodel:

pip install "gptqmodel>=5.8.0"

Загрузка квантованной модели

Загрузите предварительно квантованную модель AutoRound, передав AutoRoundConfig в from_pretrained(). Метод работает с любой моделью, которая загружается через Accelerate и содержит слои torch.nn.Linear.

Вы можете использовать PipelineQuantizationConfig для квантования конкретных компонентов конвейера:

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")

Или загрузите компонент квантованной модели напрямую:

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 в Diffusers поддерживает только загрузку предварительно квантованных моделей. Чтобы выполнить квантование модели с нуля, используйте CLI или Python API AutoRound напрямую, а затем загрузите результат с помощью Diffusers.

torch.compile

AutoRound совместим с torch.compile для более быстрого вывода. Вы можете скомпилировать квантованный трансформер (DiT) для лучшей производительности:

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)

Бэкенды

AutoRound поддерживает несколько бэкендов для инференса квантованных моделей с Weight-only квантованием. Бэкенд определяет, какое ядро выполняет деквантование во время прямого прохода. Установите параметр backend в AutoRoundConfig, чтобы выбрать один из них:

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")

Сохранение и загрузка

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")

Поддерживаемые схемы квантования

AutoRound поддерживает несколько схем:

• 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(поддерживаются все форматы Q_K,Q_0,Q*_1, предоставляемые llamacpp)
• NVFP4(экспериментальная функция, рекомендуется экспорт в формат llm_compressor с параметрами data_type nvfp4,act_data_type nvfp4,static_global_scale,group_size 16)
• MXFP4(Исследовательская функция, нет реального ядра, Стандартный MXFP4, data_type mxfp,act_data_type mxfp,bits 4, act_bits 4, group_size 32)
• MXINT4(Исследовательская функция, нет реального ядра, Стандартный MXINT4, data_type mxint,act_data_type mxint,bits 4, act_bits 4, group_size 32)
• MXFP4_RCEIL(Исследовательская функция, реального ядра нет, вариант NVIDIA, data_type mxfp, act_data_type mxfp_rceil, bits 4, act_bits 4, group_size 32)
• MXFP8(Исследовательская функция, нет реального ядра, data_type mxfp,act_data_type mxfp_rceil,group_size 32)
• FPW8A16(Исследовательская функция, нет реального ядра, data_type fp8,group_size 0->per tensor )
• FP8_STATIC(Исследовательская функция, нет реального ядра, data_type:fp8,act_data_type:fp8,group_size -1 ->по каналам, act_group_size=0->по тензору)

Кроме того, вы можете изменить group_size, bits, sym и многие другие параметры конфигурации, хотя для них, возможно, не существует реальных ядер.

Ресурсы

• Модели AutoRound с предварительным квантованием на Hub