AutoRound

AutoRound est un kit d’outils de quantification avancé. Il atteint une précision élevée à très faible nombre de bits (2 à 4 bits) avec un réglage minimal, en exploitant la descente de gradient signé et en offrant une large compatibilité matérielle. Consultez nos articles SignRoundV1 et SignRoundV2 pour plus de détails.

Installer auto-round(version ≥ 0.13.0) :

pip install "auto-round>=0.13.0"

Pour utiliser le noyau Marlin pour une inférence CUDA plus rapide, installez gptqmodel :

pip install "gptqmodel>=5.8.0"

Charger un modèle quantifié

Chargez un modèle AutoRound pré-quantifié en passant AutoRoundConfig à from_pretrained(). Cette méthode fonctionne avec tout modèle qui se charge via Accelerate et qui possède des couches torch.nn.Linear.

Vous pouvez utiliser PipelineQuantizationConfig pour quantifier des composants spécifiques d'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")

Ou chargez directement un composant de modèle quantifié :

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 dans Diffusers ne prend en charge que le chargement de modèles pré-quantifiés. Pour quantifier un modèle à partir de zéro, utilisez directement le CLI ou l’API Python d’AutoRound, puis chargez le résultat avec Diffusers.

torch.compile

AutoRound est compatible avec torch.compile pour une inférence plus rapide. Vous pouvez compiler le transformateur quantifié (DiT) pour de meilleures performances :

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 prend en charge plusieurs backends d'inférence pour les modèles quantifiés en Weight-only. Le backend détermine quel noyau gère la déquantification lors du passage avant. Définissez le paramètre backend dans AutoRoundConfig pour en choisir un :

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

Enregistrer et charger

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

Schémas de quantification pris en charge

AutoRound prend en charge plusieurs schémas :

• 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(tous les Q_K,Q_0,Q*_1 fournis par llamacpp sont pris en charge)
• NVFP4(Fonctionnalité expérimentale, il est recommandé d’exporter au format llm_compressor avec data_type nvfp4,act_data_type nvfp4,static_global_scale,group_size 16)
• MXFP4(Fonction de recherche, pas de noyau réel, MXFP4 standard, data_type mxfp,act_data_type mxfp,bits 4, act_bits 4, group_size 32)
• MXINT4(Fonction de recherche, pas de noyau réel, Standard MXINT4, data_type mxint,act_data_type mxint,bits 4, act_bits 4, group_size 32)
• MXFP4_RCEIL(Fonction de recherche, aucun kernel réel, variante de NVIDIA, data_type mxfp,act_data_type mxfp_rceil,bits 4, act_bits 4, group_size 32)
• MXFP8(Fonction de recherche, pas de noyau réel, data_type mxfp,act_data_type mxfp_rceil,group_size 32)
• FPW8A16(Fonction de recherche, pas de noyau réel, data_type fp8,group_size 0->par tenseur )
• FP8_STATIC(Fonction de recherche, pas de noyau réel, data_type:fp8,act_data_type:fp8,group_size -1 ->par canal, act_group_size=0->par tenseur)

Par ailleurs, vous pouvez modifier group_size, bits, sym et de nombreuses autres configurations, bien qu'il n'existe peut-être pas de vrais noyaux.

Ressources

• Modèles AutoRound pré-quantifiés sur le Hub