From 42d4fdcdf8b6a898ef89ee605d929181dc7c185c Mon Sep 17 00:00:00 2001 From: Xin He Date: Fri, 10 Apr 2026 15:19:31 +0000 Subject: [PATCH 01/14] support auto_round Signed-off-by: Xin He --- src/diffusers/__init__.py | 21 +++ src/diffusers/quantizers/auto.py | 4 + .../quantizers/autoround/__init__.py | 1 + .../autoround/autoround_quantizer.py | 130 ++++++++++++++++++ .../quantizers/quantization_config.py | 73 ++++++++++ src/diffusers/utils/__init__.py | 1 + .../utils/dummy_auto_round_objects.py | 17 +++ src/diffusers/utils/import_utils.py | 5 + 8 files changed, 252 insertions(+) create mode 100644 src/diffusers/quantizers/autoround/__init__.py create mode 100644 src/diffusers/quantizers/autoround/autoround_quantizer.py create mode 100644 src/diffusers/utils/dummy_auto_round_objects.py diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index e9441ef71a31..25373be022ef 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -7,6 +7,7 @@ OptionalDependencyNotAvailable, _LazyModule, is_accelerate_available, + is_auto_round_available, is_bitsandbytes_available, is_flax_available, is_gguf_available, @@ -122,6 +123,18 @@ else: _import_structure["quantizers.quantization_config"].append("NVIDIAModelOptConfig") +try: + if not is_auto_round_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_auto_round_objects + + _import_structure["utils.dummy_auto_round_objects"] = [ + name for name in dir(dummy_auto_round_objects) if not name.startswith("_") + ] +else: + _import_structure["quantizers.quantization_config"].append("AutoRoundConfig") + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() @@ -935,6 +948,14 @@ else: from .quantizers.quantization_config import NVIDIAModelOptConfig + try: + if not is_auto_round_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_auto_round_objects import * + else: + from .quantizers.quantization_config import AutoRoundConfig + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 6cd24c459c9d..f61e4d8225eb 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -18,10 +18,12 @@ import warnings +from .autoround import AutoRoundQuantizer from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer from .gguf import GGUFQuantizer from .modelopt import NVIDIAModelOptQuantizer from .quantization_config import ( + AutoRoundConfig, BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, @@ -41,6 +43,7 @@ "quanto": QuantoQuantizer, "torchao": TorchAoHfQuantizer, "modelopt": NVIDIAModelOptQuantizer, + "auto-round": AutoRoundQuantizer, } AUTO_QUANTIZATION_CONFIG_MAPPING = { @@ -50,6 +53,7 @@ "quanto": QuantoConfig, "torchao": TorchAoConfig, "modelopt": NVIDIAModelOptConfig, + "auto-round": AutoRoundConfig, } diff --git a/src/diffusers/quantizers/autoround/__init__.py b/src/diffusers/quantizers/autoround/__init__.py new file mode 100644 index 000000000000..2fe2083d4a5f --- /dev/null +++ b/src/diffusers/quantizers/autoround/__init__.py @@ -0,0 +1 @@ +from .autoround_quantizer import AutoRoundQuantizer diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py new file mode 100644 index 000000000000..f6f46048bc66 --- /dev/null +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -0,0 +1,130 @@ +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import TYPE_CHECKING + +from ...utils import ( + is_auto_round_available, + is_torch_available, + logging, +) +from ..base import DiffusersQuantizer + + +if TYPE_CHECKING: + from ...models.modeling_utils import ModelMixin + + +if is_torch_available(): + import torch + +logger = logging.get_logger(__name__) + + +class AutoRoundQuantizer(DiffusersQuantizer): + r""" + Diffusers Quantizer for AutoRound (https://github.com/intel/auto-round). + + AutoRound is a weight-only quantization method that uses sign gradient descent to jointly optimize + rounding values and min-max ranges for weights. It supports W4A16 (4-bit weight, 16-bit activation) + quantization for efficient inference. + + This quantizer only supports loading pre-quantized AutoRound models. On-the-fly quantization + (calibration) is not supported through this interface. + """ + + # AutoRound requires data calibration — we only support loading pre-quantized checkpoints. + requires_calibration = True + required_packages = ["auto_round"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + def validate_environment(self, *args, **kwargs): + """ + Validates that the auto-round library (>= 0.5) is installed and captures the device_map + for later use during model conversion. + """ + self.device_map = kwargs.get("device_map", None) + if not is_auto_round_available(): + raise ImportError( + "Loading an AutoRound quantized model requires the auto-round library " + "(`pip install 'auto-round>=0.5'`)" + ) + + def _process_model_before_weight_loading( + self, + model: "ModelMixin", + device_map, + keep_in_fp32_modules: list[str] = [], + **kwargs, + ): + """ + Replaces target nn.Linear layers with AutoRound's quantized QuantLinear layers before + weights are loaded from the checkpoint. + + Uses `auto_round.inference.convert_model.convert_hf_model` which: + - Inspects the model architecture and the quantization config (bits, group_size, sym, backend). + - Replaces eligible nn.Linear modules with the appropriate QuantLinear variant + (the packed-weight layer that stores qweight, scales, qzeros). + - Returns the converted model and a set of used backend names. + + `infer_target_device` resolves the device_map into a single target device string + that AutoRound uses to select the correct kernel backend (e.g. "cuda", "cpu"). + """ + from auto_round.inference.convert_model import convert_hf_model, infer_target_device + + if self.pre_quantized: + target_device = infer_target_device(self.device_map) + model, used_backends = convert_hf_model(model, target_device) + self.used_backends = used_backends + + model.config.quantization_config = self.quantization_config + + def _process_model_after_weight_loading(self, model, **kwargs): + """ + Finalizes the model after all quantized weights (qweight, scales, qzeros, etc.) have + been loaded into the QuantLinear layers. + + Uses `auto_round.inference.convert_model.post_init` which: + - Performs backend-specific finalization (e.g. repacking weights into the kernel's + expected memory layout, moving buffers to the correct device). + - Freezes quantized parameters (requires_grad=False). + - Prepares the model for inference. + + Raises ValueError if the model is not pre-quantized, since AutoRound does not support + on-the-fly quantization through this loading path. + """ + if self.pre_quantized: + from auto_round.inference.convert_model import post_init + + post_init(model, self.used_backends) + else: + raise ValueError( + "AutoRound quantizer in diffusers only supports pre-quantized models. " + "Please provide a model that has already been quantized with AutoRound." + ) + return model + + @property + def is_trainable(self) -> bool: + """AutoRound W4A16 pre-quantized models do not support training.""" + return False + + @property + def is_serializable(self): + """AutoRound quantized models can be serialized (the quantization config may be + updated by the backend, e.g. for GPTQ/AWQ-compatible formats).""" + return True + diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index c3d829fde8cf..2cdfe75cb30b 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -48,6 +48,7 @@ class QuantizationMethod(str, Enum): TORCHAO = "torchao" QUANTO = "quanto" MODELOPT = "modelopt" + AUTOROUND = "auto-round" @dataclass @@ -749,3 +750,75 @@ def get_config_from_quant_type(self) -> dict[str, Any]: ) return BASE_CONFIG + + +@dataclass +class AutoRoundConfig(QuantizationConfigMixin): + """Configuration class for AutoRound quantization. + + AutoRound is a weight-only quantization algorithm that uses sign gradient descent to + jointly optimize weight rounding and min-max values. This config targets the W4A16 + (4-bit weights, 16-bit activations) setting. + + Reference: https://github.com/intel/auto-round + + Args: + bits (`int`, *optional*, defaults to `4`): + The number of bits to quantize weights to. For W4A16 this should be 4. + group_size (`int`, *optional*, defaults to `128`): + The group size for weight quantization. Weights in each group share the same + scale and zero-point. Common choices: 32, 64, 128, -1 (per-channel). + sym (`bool`, *optional*, defaults to `False`): + Whether to use symmetric quantization (zero-point fixed at 0) or asymmetric + quantization (zero-point is learned). Asymmetric is generally more accurate. + modules_to_not_convert (`list[str]` or `None`, *optional*, defaults to `None`): + List of module name patterns that should NOT be quantized. Useful for keeping + certain sensitive layers (e.g. the final output projection) in full precision. + backend (`str`, *optional*, defaults to `"auto"`): + The backend kernel to use for quantized inference. Options may include: + "auto" (automatically select), "gptq" (GPTQ-compatible kernels), etc. + The backend determines which QuantLinear packing format is used. + kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments forwarded to AutoRound (e.g. `iters`, `seqlen`, + `batch_size`, `lr`, `minmax_lr` for calibration when quantizing from scratch). + """ + + def __init__( + self, + bits: int = 4, + group_size: int = 128, + sym: bool = False, + modules_to_not_convert: list[str] | None = None, + backend: str = "auto", + **kwargs, + ) -> None: + self.quant_method = QuantizationMethod.AUTOROUND + self.bits = bits + self.group_size = group_size + self.sym = sym + self.modules_to_not_convert = modules_to_not_convert + self.backend = backend + for k, v in kwargs.items(): + setattr(self, k, v) + + def to_dict(self) -> dict: + """Serialize the config to a JSON-compatible dict. + + Output: A dict containing all config fields. The `quant_method` is stored as + its string value so it can be round-tripped through JSON. + """ + output = super().to_dict() + output["quant_method"] = output["quant_method"].value + return output + + @classmethod + def from_dict(cls, config_dict: dict, return_unused_kwargs: bool = False, **kwargs): + """Instantiate an AutoRoundConfig from a dictionary. + + Input: config_dict with keys like bits, group_size, sym, etc. + Output: An AutoRoundConfig instance (and optionally unused kwargs). + """ + # Filter out keys that are not constructor parameters + # (e.g. quant_method is set automatically) + config_dict = {k: v for k, v in config_dict.items() if k != "quant_method"} + return super().from_dict(config_dict, return_unused_kwargs=return_unused_kwargs, **kwargs) diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 23d7ac7c6c2d..f73cb77ba7c8 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -85,6 +85,7 @@ is_hpu_available, is_inflect_available, is_invisible_watermark_available, + is_auto_round_available, is_kernels_available, is_kernels_version, is_kornia_available, diff --git a/src/diffusers/utils/dummy_auto_round_objects.py b/src/diffusers/utils/dummy_auto_round_objects.py new file mode 100644 index 000000000000..be7a6b8403cb --- /dev/null +++ b/src/diffusers/utils/dummy_auto_round_objects.py @@ -0,0 +1,17 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class AutoRoundConfig(metaclass=DummyObject): + _backends = ["auto_round"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["auto_round"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["auto_round"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["auto_round"]) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index 551fa358a28d..0c7b9e89672e 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -230,6 +230,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _aiter_available, _aiter_version = _is_package_available("aiter", get_dist_name=True) _kornia_available, _kornia_version = _is_package_available("kornia") _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True) +_auto_round_available, _auto_round_version = _is_package_available("auto_round") _av_available, _av_version = _is_package_available("av") @@ -373,6 +374,10 @@ def is_nvidia_modelopt_available(): return _nvidia_modelopt_available +def is_auto_round_available(): + return _auto_round_available + + def is_timm_available(): return _timm_available From e1714a9f72eaecd398834f63a217c01f6432bd16 Mon Sep 17 00:00:00 2001 From: Xin He Date: Fri, 10 Apr 2026 15:40:26 +0000 Subject: [PATCH 02/14] add document and unit tests Signed-off-by: Xin He --- docs/source/en/quantization/autoround.md | 156 +++++++ .../quantizers/quantization_config.py | 9 +- src/diffusers/utils/testing_utils.py | 14 + tests/quantization/auto-round/__init__.py | 0 .../quantization/auto-round/test_autoround.py | 384 ++++++++++++++++++ 5 files changed, 560 insertions(+), 3 deletions(-) create mode 100644 docs/source/en/quantization/autoround.md create mode 100644 tests/quantization/auto-round/__init__.py create mode 100644 tests/quantization/auto-round/test_autoround.py diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md new file mode 100644 index 000000000000..20f5d9a9be92 --- /dev/null +++ b/docs/source/en/quantization/autoround.md @@ -0,0 +1,156 @@ + + +# AutoRound + +[AutoRound](https://github.com/intel/auto-round) is a weight-only quantization algorithm that uses **S**ign **G**radient **D**escent to jointly optimize rounding values and min-max ranges for weights. It targets the W4A16 configuration (4-bit weights, 16-bit activations), reducing model memory footprint while preserving inference accuracy. + +> **Paper:** [AutoRound: Optimize Weight Rounding via Signed Gradient Descent for the Quantization of LLMs](https://huggingface.co/papers/2309.05516) + +Before you begin, make sure you have `auto-round` installed (version ≥ 0.13.0): + +```bash +pip install "auto-round>=0.13.0" +``` + +For best CUDA inference performance with the Marlin kernel, also install `gptqmodel`: + +```bash +pip install "gptqmodel>=5.8.0" +``` + +## Quickstart + +Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelMixin.from_pretrained`]. This works for any model in any modality, as long as it supports loading with [Accelerate](https://hf.co/docs/accelerate/index) and contains `torch.nn.Linear` layers. + +```python +import torch +from diffusers import AutoModel, FluxPipeline, AutoRoundConfig + +model_id = "your-org/flux-autoround-w4g128" + +quantization_config = AutoRoundConfig(bits=4, group_size=128, sym=False) +transformer = AutoModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=quantization_config, + torch_dtype=torch.float16, +) +pipe = FluxPipeline.from_pretrained( + model_id, + transformer=transformer, + torch_dtype=torch.float16, +) +pipe.to("cuda") + +image = pipe("A cat holding a sign that says hello world").images[0] +image.save("output.png") +``` + +> **Note:** AutoRound in Diffusers only supports loading **pre-quantized** models. To quantize a model from scratch, use the [AutoRound CLI or Python API](https://github.com/intel/auto-round) directly, then load the result with Diffusers. + +## Inference backends + +AutoRound supports multiple inference backends. The backend determines which kernel is used for dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: + +| Backend | Value | Device | Requirements | Notes | +|---------|-------|--------|--------------|-------| +| **Auto** | `"auto"` | Any | — | Default. Automatically selects the best available backend. | +| **PyTorch** | `"auto_round:torch_zp"` | CPU / CUDA | — | Pure PyTorch implementation. Broadest compatibility. | +| **Triton** | `"auto_round:tritonv2_zp"` | CUDA | `triton` | Triton-based kernel for GPU inference. | +| **Marlin** | `"gptqmodel:marlin_zp"` | CUDA | `gptqmodel>=5.8.0` | Best CUDA performance via the Marlin kernel. | + +### Example: specifying a backend + +```python +from diffusers import AutoRoundConfig + +# Auto-select (default) +config = AutoRoundConfig(bits=4, group_size=128) + +# Explicit Triton backend for CUDA +config = AutoRoundConfig(bits=4, group_size=128, backend="auto_round:tritonv2_zp") + +# Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) +config = AutoRoundConfig(bits=4, group_size=128, backend="gptqmodel:marlin_zp") + +# PyTorch backend for CPU inference +config = AutoRoundConfig(bits=4, group_size=128, backend="auto_round:torch_zp") +``` + +## AutoRoundConfig + +The [`AutoRoundConfig`] class accepts the following parameters: + +- `bits` (`int`, defaults to `4`): Number of bits for weight quantization. Use `4` for W4A16. +- `group_size` (`int`, defaults to `128`): Group size for quantization. Weights in each group share the same scale and zero-point. Common values: `32`, `64`, `128`, or `-1` (per-channel). +- `sym` (`bool`, defaults to `False`): Whether to use symmetric quantization (`True`) or asymmetric quantization (`False`). Asymmetric is generally more accurate. +- `modules_to_not_convert` (`list[str]` or `None`, defaults to `None`): List of module name patterns to exclude from quantization. Use this to keep sensitive layers in full precision. +- `backend` (`str`, defaults to `"auto"`): The inference backend kernel. See the [Inference backends](#inference-backends) table above. + +## Supported quantization configurations + +AutoRound focuses on weight-only quantization. The primary configuration is W4A16 (4-bit weights, 16-bit activations), with flexibility in group size and symmetry: + +| Configuration | `bits` | `group_size` | `sym` | Description | +|--------------|--------|-------------|-------|-------------| +| W4G128 asymmetric | `4` | `128` | `False` | Default. Good balance of accuracy and compression. | +| W4G128 symmetric | `4` | `128` | `True` | Slightly faster dequantization, marginal accuracy loss. | +| W4G32 asymmetric | `4` | `32` | `False` | Finer granularity, better accuracy, slightly more metadata overhead. | + +## Serializing and deserializing quantized models + +AutoRound quantized models can be saved and reloaded using the standard [`~ModelMixin.save_pretrained`] and [`~ModelMixin.from_pretrained`] methods. + +### Save + +```python +import torch +from diffusers import AutoModel, AutoRoundConfig + +model_id = "your-org/flux-autoround-w4g128" +quantization_config = AutoRoundConfig(bits=4, group_size=128, sym=False) +model = AutoModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=quantization_config, + torch_dtype=torch.float16, +) +model.save_pretrained("path/to/saved_model") +``` + +### Load + +```python +import torch +from diffusers import AutoModel, FluxPipeline + +transformer = AutoModel.from_pretrained( + "path/to/saved_model", + torch_dtype=torch.float16, +) +pipe = FluxPipeline.from_pretrained( + "your-org/flux-autoround-w4g128", + transformer=transformer, + torch_dtype=torch.float16, +) +pipe.to("cuda") + +image = pipe("A beautiful sunset over the ocean").images[0] +image.save("output.png") +``` + +## Resources + +- [AutoRound GitHub repository](https://github.com/intel/auto-round) +- [AutoRound paper (arXiv:2309.05516)](https://arxiv.org/abs/2309.05516) +- [AutoRound Hugging Face integration (Transformers)](https://huggingface.co/docs/transformers/quantization/autoround) +- [Pre-quantized AutoRound models on the Hub](https://huggingface.co/models?search=autoround) diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 2cdfe75cb30b..53797ff4dd08 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -775,9 +775,12 @@ class AutoRoundConfig(QuantizationConfigMixin): List of module name patterns that should NOT be quantized. Useful for keeping certain sensitive layers (e.g. the final output projection) in full precision. backend (`str`, *optional*, defaults to `"auto"`): - The backend kernel to use for quantized inference. Options may include: - "auto" (automatically select), "gptq" (GPTQ-compatible kernels), etc. - The backend determines which QuantLinear packing format is used. + The backend kernel to use for quantized inference. Available backends: + - `"auto"`: Automatically select the best available backend for the current device. + - `"auto_round:torch_zp"`: Pure PyTorch kernel — works on CPU and CUDA. + - `"auto_round:tritonv2_zp"`: Triton-based kernel — requires CUDA. + - `"gptqmodel:marlin_zp"`: Marlin kernel via GPTQModel — requires CUDA and + `gptqmodel>=5.8.0`. Offers the best CUDA inference performance. kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments forwarded to AutoRound (e.g. `iters`, `seqlen`, `batch_size`, `lr`, `minmax_lr` for calibration when quantizing from scratch). diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index 619a37034949..2ea9d039e4f0 100644 --- a/src/diffusers/utils/testing_utils.py +++ b/src/diffusers/utils/testing_utils.py @@ -33,6 +33,7 @@ from .import_utils import ( BACKENDS_MAPPING, is_accelerate_available, + is_auto_round_available, is_bitsandbytes_available, is_compel_available, is_flax_available, @@ -654,6 +655,19 @@ def decorator(test_case): return decorator +def require_auto_round_version_greater_or_equal(auto_round_version): + def decorator(test_case): + correct_auto_round_version = is_auto_round_available() and version.parse( + version.parse(importlib.metadata.version("auto_round")).base_version + ) >= version.parse(auto_round_version) + return unittest.skipUnless( + correct_auto_round_version, + f"Test requires auto-round with version greater than {auto_round_version}.", + )(test_case) + + return decorator + + def require_kernels_version_greater_or_equal(kernels_version): def decorator(test_case): correct_kernels_version = is_kernels_available() and version.parse( diff --git a/tests/quantization/auto-round/__init__.py b/tests/quantization/auto-round/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/quantization/auto-round/test_autoround.py b/tests/quantization/auto-round/test_autoround.py new file mode 100644 index 000000000000..160e328861d1 --- /dev/null +++ b/tests/quantization/auto-round/test_autoround.py @@ -0,0 +1,384 @@ +import gc +import tempfile +import unittest + +from diffusers import AutoRoundConfig, FluxTransformer2DModel, FluxPipeline +from diffusers.utils import is_auto_round_available, is_torch_available +from diffusers.utils.testing_utils import ( + backend_empty_cache, + backend_reset_peak_memory_stats, + enable_full_determinism, + nightly, + numpy_cosine_similarity_distance, + require_accelerate, + require_big_accelerator, + require_auto_round_version_greater_or_equal, + require_torch_cuda_compatibility, + torch_device, +) + + +if is_torch_available(): + import torch + + from ..utils import get_memory_consumption_stat + + +def _is_gptqmodel_available(min_version="5.8.0"): + """Check if gptqmodel is installed with a minimum version.""" + try: + import importlib.metadata + + from packaging import version + + gptqmodel_version = importlib.metadata.version("gptqmodel") + return version.parse(gptqmodel_version) >= version.parse(min_version) + except importlib.metadata.PackageNotFoundError: + return False + + +enable_full_determinism() + + +@nightly +@require_big_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundBaseTesterMixin: + """Base test mixin for AutoRound quantized models. + + AutoRound is a weight-only quantization method (W4A16). It supports multiple inference + backends depending on the hardware: + - CPU: `auto_round:torch_zp` backend + - CUDA: `auto_round:tritonv2_zp` backend + - CUDA + GPTQModel>=5.8.0: `gptqmodel:marlin_zp` backend (best performance) + + When `backend="auto"`, AutoRound selects the best available backend automatically. + + Key differences from ModelOpt tests: + - Only pre-quantized model loading is supported (no on-the-fly quantization). + - `is_trainable` returns False, so no LoRA training test. + - No `test_dtype_assignment` (AutoRound doesn't restrict dtype changes). + - `requires_calibration = True` means we always load pre-quantized checkpoints. + """ + + # TODO: Replace with a real tiny AutoRound-quantized checkpoint on the Hub. + # This should be a small model that has been quantized with AutoRound and uploaded + # in the standard format (qweight, scales, qzeros, g_idx). + model_id = "hf-internal-testing/tiny-flux-pipe-autoround-w4g128" + model_cls = FluxTransformer2DModel + pipeline_cls = FluxPipeline + torch_dtype = torch.float16 + expected_memory_reduction = 0.0 + modules_to_not_convert = "" + _test_torch_compile = False + + def setUp(self): + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + gc.collect() + + def tearDown(self): + backend_reset_peak_memory_stats(torch_device) + backend_empty_cache(torch_device) + gc.collect() + + def get_dummy_init_kwargs(self): + """Returns the default AutoRoundConfig kwargs for W4A16 quantization. + + Subclasses override this to specify backend, group_size, sym, etc. + """ + return { + "bits": 4, + "group_size": 128, + "sym": False, + } + + def get_dummy_model_init_kwargs(self): + """Returns kwargs for model_cls.from_pretrained() with AutoRound quantization.""" + return { + "pretrained_model_name_or_path": self.model_id, + "torch_dtype": self.torch_dtype, + "quantization_config": AutoRoundConfig(**self.get_dummy_init_kwargs()), + "subfolder": "transformer", + } + + def get_dummy_inputs(self): + """Creates dummy inputs matching the model's expected forward signature. + + TODO: Adjust input shapes to match the tiny test checkpoint's config. + """ + batch_size = 1 + seq_len = 16 + height = width = 32 + num_latent_channels = 4 + caption_channels = 8 + + torch.manual_seed(0) + hidden_states = torch.randn((batch_size, num_latent_channels, height, width)).to( + torch_device, dtype=self.torch_dtype + ) + encoder_hidden_states = torch.randn((batch_size, seq_len, caption_channels)).to( + torch_device, dtype=self.torch_dtype + ) + timestep = torch.tensor([1.0]).to(torch_device, dtype=self.torch_dtype).expand(batch_size) + + return { + "hidden_states": hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "timestep": timestep, + } + + def test_autoround_layers(self): + """Verify that eligible nn.Linear layers have been replaced with AutoRound QuantLinear layers. + + After loading a pre-quantized model, all target linear layers should have been + converted to AutoRound's QuantLinear by `convert_hf_model`. + """ + from auto_round.inference.convert_model import QuantLinear + + model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) + has_quantized_layer = False + for name, module in model.named_modules(): + if isinstance(module, QuantLinear): + has_quantized_layer = True + assert has_quantized_layer, "No QuantLinear layers found in the model — quantization may have failed." + + def test_autoround_memory_usage(self): + """Compare peak memory between unquantized and AutoRound-quantized model. + + The quantized model should use significantly less memory due to 4-bit weight packing. + `expected_memory_reduction` defines the minimum ratio (unquantized / quantized). + """ + inputs = self.get_dummy_inputs() + inputs = { + k: v.to(device=torch_device, dtype=self.torch_dtype) for k, v in inputs.items() if not isinstance(v, bool) + } + + unquantized_model = self.model_cls.from_pretrained( + self.model_id, torch_dtype=self.torch_dtype, subfolder="transformer" + ) + unquantized_model.to(torch_device) + unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs) + + quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) + quantized_model.to(torch_device) + quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs) + + assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction + + def test_modules_to_not_convert(self): + """Verify that modules listed in `modules_to_not_convert` remain as standard nn.Linear.""" + from auto_round.inference.convert_model import QuantLinear + + init_kwargs = self.get_dummy_model_init_kwargs() + quantization_config_kwargs = self.get_dummy_init_kwargs() + quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert}) + quantization_config = AutoRoundConfig(**quantization_config_kwargs) + init_kwargs.update({"quantization_config": quantization_config}) + + model = self.model_cls.from_pretrained(**init_kwargs) + model.to(torch_device) + + for name, module in model.named_modules(): + if name in self.modules_to_not_convert: + assert not isinstance( + module, QuantLinear + ), f"Module '{name}' should NOT have been quantized but is a QuantLinear." + + def test_serialization(self): + """Test round-trip save and load of an AutoRound quantized model.""" + model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) + inputs = self.get_dummy_inputs() + + model.to(torch_device) + with torch.no_grad(): + model_output = model(**inputs) + + with tempfile.TemporaryDirectory() as tmp_dir: + model.save_pretrained(tmp_dir) + saved_model = self.model_cls.from_pretrained( + tmp_dir, + torch_dtype=self.torch_dtype, + ) + + saved_model.to(torch_device) + with torch.no_grad(): + saved_model_output = saved_model(**inputs) + + assert torch.allclose(model_output.sample, saved_model_output.sample, rtol=1e-5, atol=1e-5) + + def test_torch_compile(self): + """Test that the quantized model works with torch.compile.""" + if not self._test_torch_compile: + return + + model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) + compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True, dynamic=False) + + model.to(torch_device) + with torch.no_grad(): + model_output = model(**self.get_dummy_inputs()).sample + + compiled_model.to(torch_device) + with torch.no_grad(): + compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample + + model_output = model_output.detach().float().cpu().numpy() + compiled_model_output = compiled_model_output.detach().float().cpu().numpy() + + max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten()) + assert max_diff < 1e-3 + + def test_model_cpu_offload(self): + """Test that the quantized model works with pipeline CPU offload.""" + init_kwargs = self.get_dummy_init_kwargs() + transformer = self.model_cls.from_pretrained( + self.model_id, + quantization_config=AutoRoundConfig(**init_kwargs), + subfolder="transformer", + torch_dtype=self.torch_dtype, + ) + pipe = self.pipeline_cls.from_pretrained(self.model_id, transformer=transformer, torch_dtype=self.torch_dtype) + pipe.enable_model_cpu_offload(device=torch_device) + _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) + + +# ============================================================================ +# Backend: auto (auto-select best available backend) +# ============================================================================ + + +class AutoRoundW4G128AsymAutoBackendTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, asymmetric, backend='auto' (default — auto-selects best backend).""" + + expected_memory_reduction = 0.55 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": False, + "backend": "auto", + } + + +class AutoRoundW4G128SymAutoBackendTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, symmetric, backend='auto'.""" + + expected_memory_reduction = 0.55 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": True, + "backend": "auto", + } + + +class AutoRoundW4G32AsymAutoBackendTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=32, asymmetric, backend='auto' (finer granularity).""" + + expected_memory_reduction = 0.50 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 32, + "sym": False, + "backend": "auto", + } + + +# ============================================================================ +# Backend: auto_round:tritonv2_zp (CUDA, Triton-based kernel) +# ============================================================================ + + +@require_torch_cuda_compatibility(7.0) +class AutoRoundW4G128AsymTritonTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, asymmetric, backend='auto_round:tritonv2_zp' (CUDA Triton kernel).""" + + expected_memory_reduction = 0.55 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": False, + "backend": "auto_round:tritonv2_zp", + } + + +@require_torch_cuda_compatibility(7.0) +class AutoRoundW4G128SymTritonTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, symmetric, backend='auto_round:tritonv2_zp'.""" + + expected_memory_reduction = 0.55 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": True, + "backend": "auto_round:tritonv2_zp", + } + + +# ============================================================================ +# Backend: gptqmodel:marlin_zp (CUDA, requires GPTQModel>=5.8.0, best perf) +# ============================================================================ + + +@unittest.skipUnless(_is_gptqmodel_available("5.8.0"), "Test requires gptqmodel>=5.8.0") +@require_torch_cuda_compatibility(8.0) +class AutoRoundW4G128AsymMarlinTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, asymmetric, backend='gptqmodel:marlin_zp' (best CUDA performance).""" + + _test_torch_compile = True + expected_memory_reduction = 0.55 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": False, + "backend": "gptqmodel:marlin_zp", + } + + +@unittest.skipUnless(_is_gptqmodel_available("5.8.0"), "Test requires gptqmodel>=5.8.0") +@require_torch_cuda_compatibility(8.0) +class AutoRoundW4G128SymMarlinTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, symmetric, backend='gptqmodel:marlin_zp'.""" + + _test_torch_compile = True + expected_memory_reduction = 0.55 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": True, + "backend": "gptqmodel:marlin_zp", + } + + +# ============================================================================ +# Backend: auto_round:torch_zp (CPU, pure PyTorch kernel) +# ============================================================================ + + +class AutoRoundW4G128AsymTorchCPUTest(AutoRoundBaseTesterMixin, unittest.TestCase): + """W4A16, group_size=128, asymmetric, backend='auto_round:torch_zp' (CPU).""" + + expected_memory_reduction = 0.50 + + def get_dummy_init_kwargs(self): + return { + "bits": 4, + "group_size": 128, + "sym": False, + "backend": "auto_round:torch_zp", + } From c0daf156d51610840cffc60997fdba86854f4574 Mon Sep 17 00:00:00 2001 From: Xin He Date: Thu, 23 Apr 2026 13:34:33 +0000 Subject: [PATCH 03/14] fix CI Signed-off-by: Xin He --- .../{auto-round => auto_round}/__init__.py | 0 .../test_autoround.py | 102 ++++++++++-------- 2 files changed, 56 insertions(+), 46 deletions(-) rename tests/quantization/{auto-round => auto_round}/__init__.py (100%) rename tests/quantization/{auto-round => auto_round}/test_autoround.py (78%) diff --git a/tests/quantization/auto-round/__init__.py b/tests/quantization/auto_round/__init__.py similarity index 100% rename from tests/quantization/auto-round/__init__.py rename to tests/quantization/auto_round/__init__.py diff --git a/tests/quantization/auto-round/test_autoround.py b/tests/quantization/auto_round/test_autoround.py similarity index 78% rename from tests/quantization/auto-round/test_autoround.py rename to tests/quantization/auto_round/test_autoround.py index 160e328861d1..18b4b8b14ffa 100644 --- a/tests/quantization/auto-round/test_autoround.py +++ b/tests/quantization/auto_round/test_autoround.py @@ -2,7 +2,7 @@ import tempfile import unittest -from diffusers import AutoRoundConfig, FluxTransformer2DModel, FluxPipeline +from diffusers import AutoRoundConfig, ZImageTransformer2DModel, ZImagePipeline from diffusers.utils import is_auto_round_available, is_torch_available from diffusers.utils.testing_utils import ( backend_empty_cache, @@ -65,10 +65,10 @@ class AutoRoundBaseTesterMixin: # TODO: Replace with a real tiny AutoRound-quantized checkpoint on the Hub. # This should be a small model that has been quantized with AutoRound and uploaded # in the standard format (qweight, scales, qzeros, g_idx). - model_id = "hf-internal-testing/tiny-flux-pipe-autoround-w4g128" - model_cls = FluxTransformer2DModel - pipeline_cls = FluxPipeline - torch_dtype = torch.float16 + model_id = "INCModel/Z-Image-tiny-for-testing-W4A16-AutoRound" + model_cls = ZImageTransformer2DModel + pipeline_cls = ZImagePipeline + torch_dtype = torch.bfloat16 expected_memory_reduction = 0.0 modules_to_not_convert = "" _test_torch_compile = False @@ -104,45 +104,35 @@ def get_dummy_model_init_kwargs(self): } def get_dummy_inputs(self): - """Creates dummy inputs matching the model's expected forward signature. + """Creates dummy inputs matching ZImageTransformer2DModel.forward() signature. - TODO: Adjust input shapes to match the tiny test checkpoint's config. + ZImageTransformer2DModel expects: + - x: list of (C, F, H, W) tensors, one per batch item + - t: 1-D timestep tensor of shape (batch_size,) + - cap_feats: list of (seq_len, cap_feat_dim) tensors, one per batch item + + Dimensions are chosen to match the tiny test checkpoint + (in_channels=16, cap_feat_dim=512, patch_size=2, f_patch_size=1). """ batch_size = 1 - seq_len = 16 - height = width = 32 - num_latent_channels = 4 - caption_channels = 8 + in_channels = 16 # matches tiny model config + cap_feat_dim = 512 # matches tiny model config + height = width = 8 # must be divisible by patch_size=2 + frames = 1 # must be divisible by f_patch_size=1 + seq_len = 16 # caption token count (will be padded to multiple of 32) torch.manual_seed(0) - hidden_states = torch.randn((batch_size, num_latent_channels, height, width)).to( - torch_device, dtype=self.torch_dtype - ) - encoder_hidden_states = torch.randn((batch_size, seq_len, caption_channels)).to( - torch_device, dtype=self.torch_dtype - ) - timestep = torch.tensor([1.0]).to(torch_device, dtype=self.torch_dtype).expand(batch_size) - - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "timestep": timestep, - } - - def test_autoround_layers(self): - """Verify that eligible nn.Linear layers have been replaced with AutoRound QuantLinear layers. - - After loading a pre-quantized model, all target linear layers should have been - converted to AutoRound's QuantLinear by `convert_hf_model`. - """ - from auto_round.inference.convert_model import QuantLinear - - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - has_quantized_layer = False - for name, module in model.named_modules(): - if isinstance(module, QuantLinear): - has_quantized_layer = True - assert has_quantized_layer, "No QuantLinear layers found in the model — quantization may have failed." + x = [ + torch.randn((in_channels, frames, height, width)).to(torch_device, dtype=self.torch_dtype) + for _ in range(batch_size) + ] + cap_feats = [ + torch.randn((seq_len, cap_feat_dim)).to(torch_device, dtype=self.torch_dtype) + for _ in range(batch_size) + ] + t = torch.tensor([0.5] * batch_size).to(torch_device, dtype=self.torch_dtype) + + return {"x": x, "cap_feats": cap_feats, "t": t} def test_autoround_memory_usage(self): """Compare peak memory between unquantized and AutoRound-quantized model. @@ -151,8 +141,13 @@ def test_autoround_memory_usage(self): `expected_memory_reduction` defines the minimum ratio (unquantized / quantized). """ inputs = self.get_dummy_inputs() + # x and cap_feats are lists of tensors; move each element individually. inputs = { - k: v.to(device=torch_device, dtype=self.torch_dtype) for k, v in inputs.items() if not isinstance(v, bool) + k: [t.to(device=torch_device, dtype=self.torch_dtype) for t in v] + if isinstance(v, list) + else v.to(device=torch_device, dtype=self.torch_dtype) + for k, v in inputs.items() + if not isinstance(v, bool) } unquantized_model = self.model_cls.from_pretrained( @@ -169,7 +164,7 @@ def test_autoround_memory_usage(self): def test_modules_to_not_convert(self): """Verify that modules listed in `modules_to_not_convert` remain as standard nn.Linear.""" - from auto_round.inference.convert_model import QuantLinear + from auto_round.inference.convert_model import dynamic_import_inference_linear init_kwargs = self.get_dummy_model_init_kwargs() quantization_config_kwargs = self.get_dummy_init_kwargs() @@ -180,11 +175,23 @@ def test_modules_to_not_convert(self): model = self.model_cls.from_pretrained(**init_kwargs) model.to(torch_device) + # Resolve the actual backend used after model loading. + # When backend='auto', AutoRound selects the best available backend + # per-layer during convert_hf_model(); 'auto' itself is not a valid + # argument to dynamic_import_inference_linear. + used_backends = getattr(model.hf_quantizer, "used_backends", []) + resolved_backend = ( + used_backends[0] + if used_backends + else quantization_config_kwargs.get("backend", "auto_round:torch_zp") + ) + quant_linear_cls = dynamic_import_inference_linear(resolved_backend, quantization_config_kwargs) + for name, module in model.named_modules(): if name in self.modules_to_not_convert: assert not isinstance( - module, QuantLinear - ), f"Module '{name}' should NOT have been quantized but is a QuantLinear." + module, quant_linear_cls + ), f"Module '{name}' should NOT have been quantized but is a {quant_linear_cls}." def test_serialization(self): """Test round-trip save and load of an AutoRound quantized model.""" @@ -206,7 +213,9 @@ def test_serialization(self): with torch.no_grad(): saved_model_output = saved_model(**inputs) - assert torch.allclose(model_output.sample, saved_model_output.sample, rtol=1e-5, atol=1e-5) + # model_output.sample is a list of per-item tensors + for out, saved_out in zip(model_output.sample, saved_model_output.sample): + assert torch.allclose(out, saved_out, rtol=1e-5, atol=1e-5) def test_torch_compile(self): """Test that the quantized model works with torch.compile.""" @@ -224,8 +233,9 @@ def test_torch_compile(self): with torch.no_grad(): compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample - model_output = model_output.detach().float().cpu().numpy() - compiled_model_output = compiled_model_output.detach().float().cpu().numpy() + # model_output is a list of per-item tensors; stack for comparison + model_output = torch.stack([o.detach().float().cpu() for o in model_output]).numpy() + compiled_model_output = torch.stack([o.detach().float().cpu() for o in compiled_model_output]).numpy() max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten()) assert max_diff < 1e-3 From 677a26ea8a3d26c8bebee59b0816cdc71df3f93d Mon Sep 17 00:00:00 2001 From: Xin He Date: Fri, 24 Apr 2026 09:50:54 +0800 Subject: [PATCH 04/14] Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> --- docs/source/en/quantization/autoround.md | 36 +++++++++++------------- 1 file changed, 16 insertions(+), 20 deletions(-) diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md index 20f5d9a9be92..1a13b6ba3d26 100644 --- a/docs/source/en/quantization/autoround.md +++ b/docs/source/en/quantization/autoround.md @@ -11,31 +11,30 @@ specific language governing permissions and limitations under the License. --> # AutoRound -[AutoRound](https://github.com/intel/auto-round) is a weight-only quantization algorithm that uses **S**ign **G**radient **D**escent to jointly optimize rounding values and min-max ranges for weights. It targets the W4A16 configuration (4-bit weights, 16-bit activations), reducing model memory footprint while preserving inference accuracy. +[AutoRound](https://huggingface.co/papers/2309.05516) is a weight-only quantization algorithm. It uses sign gradient descent to jointly optimize weight rounding and min-max ranges. AutoRound targets W4A16 (4-bit weights, 16-bit activations), reducing memory usage without sacrificing inference accuracy. -> **Paper:** [AutoRound: Optimize Weight Rounding via Signed Gradient Descent for the Quantization of LLMs](https://huggingface.co/papers/2309.05516) -Before you begin, make sure you have `auto-round` installed (version ≥ 0.13.0): +Install `auto-round`(version ≥ 0.13.0): ```bash pip install "auto-round>=0.13.0" ``` -For best CUDA inference performance with the Marlin kernel, also install `gptqmodel`: +To use the Marlin kernel for faster CUDA inference, install `gptqmodel`: ```bash pip install "gptqmodel>=5.8.0" ``` -## Quickstart +## Load a quantized model -Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelMixin.from_pretrained`]. This works for any model in any modality, as long as it supports loading with [Accelerate](https://hf.co/docs/accelerate/index) and contains `torch.nn.Linear` layers. +Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelMixin.from_pretrained`]. The method works with any model that loads via [Accelerate(https://hf.co/docs/accelerate/index) and has `torch.nn.Linear` layers. ```python import torch from diffusers import AutoModel, FluxPipeline, AutoRoundConfig -model_id = "your-org/flux-autoround-w4g128" +model_id = "INCModel/Z-Image-W4A16-AutoRound" quantization_config = AutoRoundConfig(bits=4, group_size=128, sym=False) transformer = AutoModel.from_pretrained( @@ -55,11 +54,12 @@ image = pipe("A cat holding a sign that says hello world").images[0] image.save("output.png") ``` -> **Note:** AutoRound in Diffusers only supports loading **pre-quantized** models. To quantize a model from scratch, use the [AutoRound CLI or Python API](https://github.com/intel/auto-round) directly, then load the result with Diffusers. +> [!NOTE] +> AutoRound in Diffusers only supports loading *pre-quantized* models. To quantize a model from scratch, use the [AutoRound CLI or Python API](https://github.com/intel/auto-round) directly, then load the result with Diffusers. -## Inference backends +## Backends -AutoRound supports multiple inference backends. The backend determines which kernel is used for dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: +AutoRound supports multiple inference backends. The backend controls which kernel handles dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: | Backend | Value | Device | Requirements | Notes | |---------|-------|--------|--------------|-------| @@ -68,7 +68,6 @@ AutoRound supports multiple inference backends. The backend determines which ker | **Triton** | `"auto_round:tritonv2_zp"` | CUDA | `triton` | Triton-based kernel for GPU inference. | | **Marlin** | `"gptqmodel:marlin_zp"` | CUDA | `gptqmodel>=5.8.0` | Best CUDA performance via the Marlin kernel. | -### Example: specifying a backend ```python from diffusers import AutoRoundConfig @@ -96,19 +95,19 @@ The [`AutoRoundConfig`] class accepts the following parameters: - `modules_to_not_convert` (`list[str]` or `None`, defaults to `None`): List of module name patterns to exclude from quantization. Use this to keep sensitive layers in full precision. - `backend` (`str`, defaults to `"auto"`): The inference backend kernel. See the [Inference backends](#inference-backends) table above. -## Supported quantization configurations +## Quantization configurations AutoRound focuses on weight-only quantization. The primary configuration is W4A16 (4-bit weights, 16-bit activations), with flexibility in group size and symmetry: | Configuration | `bits` | `group_size` | `sym` | Description | |--------------|--------|-------------|-------|-------------| | W4G128 asymmetric | `4` | `128` | `False` | Default. Good balance of accuracy and compression. | -| W4G128 symmetric | `4` | `128` | `True` | Slightly faster dequantization, marginal accuracy loss. | -| W4G32 asymmetric | `4` | `32` | `False` | Finer granularity, better accuracy, slightly more metadata overhead. | +| W4G128 symmetric | `4` | `128` | `True` | Faster dequantization, small accuracy trade-off. | +| W4G32 asymmetric | `4` | `32` | `False` | Higher accuracy at the cost of more metadata. | -## Serializing and deserializing quantized models +## Save and load -AutoRound quantized models can be saved and reloaded using the standard [`~ModelMixin.save_pretrained`] and [`~ModelMixin.from_pretrained`] methods. +Save and reload AutoRound quantized models using the standard [`~ModelMixin.save_pretrained`] and [`~ModelMixin.from_pretrained`] methods. ### Save @@ -116,7 +115,7 @@ AutoRound quantized models can be saved and reloaded using the standard [`~Model import torch from diffusers import AutoModel, AutoRoundConfig -model_id = "your-org/flux-autoround-w4g128" +model_id = "INCModel/Z-Image-W4A16-AutoRound" quantization_config = AutoRoundConfig(bits=4, group_size=128, sym=False) model = AutoModel.from_pretrained( model_id, @@ -127,7 +126,6 @@ model = AutoModel.from_pretrained( model.save_pretrained("path/to/saved_model") ``` -### Load ```python import torch @@ -150,7 +148,5 @@ image.save("output.png") ## Resources -- [AutoRound GitHub repository](https://github.com/intel/auto-round) -- [AutoRound paper (arXiv:2309.05516)](https://arxiv.org/abs/2309.05516) - [AutoRound Hugging Face integration (Transformers)](https://huggingface.co/docs/transformers/quantization/autoround) - [Pre-quantized AutoRound models on the Hub](https://huggingface.co/models?search=autoround) From bc46f4f23acea8d38e92fc3246ff5fcca93a7b05 Mon Sep 17 00:00:00 2001 From: Xin He Date: Fri, 24 Apr 2026 05:25:52 +0000 Subject: [PATCH 05/14] update document and overwrite the default quantization_config with specified backend. Signed-off-by: Xin He --- docs/source/en/quantization/autoround.md | 98 +++++++++---------- src/diffusers/quantizers/auto.py | 13 ++- .../autoround/autoround_quantizer.py | 2 - .../quantizers/quantization_config.py | 11 +-- 4 files changed, 61 insertions(+), 63 deletions(-) diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md index 1a13b6ba3d26..0311c061fd10 100644 --- a/docs/source/en/quantization/autoround.md +++ b/docs/source/en/quantization/autoround.md @@ -11,7 +11,7 @@ specific language governing permissions and limitations under the License. --> # AutoRound -[AutoRound](https://huggingface.co/papers/2309.05516) is a weight-only quantization algorithm. It uses sign gradient descent to jointly optimize weight rounding and min-max ranges. AutoRound targets W4A16 (4-bit weights, 16-bit activations), reducing memory usage without sacrificing inference accuracy. +[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit. It achieves high accuracy at ultra-low bit widths (2-4 bits) with minimal tuning by leveraging sign-gradient descent and providing broad hardware compatibility. See our papers [SignRoundV1](https://arxiv.org/pdf/2309.05516) and [SignRoundV2](https://arxiv.org/abs/2512.04746) for more details. Install `auto-round`(version ≥ 0.13.0): @@ -32,25 +32,27 @@ Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelM ```python import torch -from diffusers import AutoModel, FluxPipeline, AutoRoundConfig +from diffusers import ZImageTransformer2DModel, ZImagePipeline, AutoRoundConfig model_id = "INCModel/Z-Image-W4A16-AutoRound" -quantization_config = AutoRoundConfig(bits=4, group_size=128, sym=False) -transformer = AutoModel.from_pretrained( +quantization_config = AutoRoundConfig(backend="marlin") +transformer = ZImageTransformer2DModel.from_pretrained( model_id, subfolder="transformer", quantization_config=quantization_config, - torch_dtype=torch.float16, + torch_dtype=torch.bfloat16, + device_map="cuda", ) -pipe = FluxPipeline.from_pretrained( + +pipe = ZImagePipeline.from_pretrained( model_id, transformer=transformer, - torch_dtype=torch.float16, + torch_dtype=torch.bfloat16, + device_map="cuda", ) -pipe.to("cuda") -image = pipe("A cat holding a sign that says hello world").images[0] +image = pipe("a cat holding a sign that says hello").images[0] image.save("output.png") ``` @@ -64,36 +66,31 @@ AutoRound supports multiple inference backends. The backend controls which kerne | Backend | Value | Device | Requirements | Notes | |---------|-------|--------|--------------|-------| | **Auto** | `"auto"` | Any | — | Default. Automatically selects the best available backend. | -| **PyTorch** | `"auto_round:torch_zp"` | CPU / CUDA | — | Pure PyTorch implementation. Broadest compatibility. | -| **Triton** | `"auto_round:tritonv2_zp"` | CUDA | `triton` | Triton-based kernel for GPU inference. | -| **Marlin** | `"gptqmodel:marlin_zp"` | CUDA | `gptqmodel>=5.8.0` | Best CUDA performance via the Marlin kernel. | +| **PyTorch** | `"torch"` | CPU / CUDA | — | Pure PyTorch implementation. Broadest compatibility. | +| **Triton** | `"tritonv2"` | CUDA | `triton` | Triton-based kernel for GPU inference. | +| **ExllamaV2** | `"exllamav2"` | CUDA | `gptqmodel>=5.8.0` | Good CUDA performance via the ExllamaV2 kernel. | +| **Marlin** | `"marlin"` | CUDA | `gptqmodel>=5.8.0` | Best CUDA performance via the Marlin kernel. | ```python from diffusers import AutoRoundConfig # Auto-select (default) -config = AutoRoundConfig(bits=4, group_size=128) +config = AutoRoundConfig() # Explicit Triton backend for CUDA -config = AutoRoundConfig(bits=4, group_size=128, backend="auto_round:tritonv2_zp") +config = AutoRoundConfig(backend="triton") # Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) -config = AutoRoundConfig(bits=4, group_size=128, backend="gptqmodel:marlin_zp") - -# PyTorch backend for CPU inference -config = AutoRoundConfig(bits=4, group_size=128, backend="auto_round:torch_zp") -``` +config = AutoRoundConfig(backend="marlin") -## AutoRoundConfig +# Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) +config = AutoRoundConfig(backend="exllamav2") -The [`AutoRoundConfig`] class accepts the following parameters: +# PyTorch backend for CPU/CUDA inference +config = AutoRoundConfig(backend="torch") +``` -- `bits` (`int`, defaults to `4`): Number of bits for weight quantization. Use `4` for W4A16. -- `group_size` (`int`, defaults to `128`): Group size for quantization. Weights in each group share the same scale and zero-point. Common values: `32`, `64`, `128`, or `-1` (per-channel). -- `sym` (`bool`, defaults to `False`): Whether to use symmetric quantization (`True`) or asymmetric quantization (`False`). Asymmetric is generally more accurate. -- `modules_to_not_convert` (`list[str]` or `None`, defaults to `None`): List of module name patterns to exclude from quantization. Use this to keep sensitive layers in full precision. -- `backend` (`str`, defaults to `"auto"`): The inference backend kernel. See the [Inference backends](#inference-backends) table above. ## Quantization configurations @@ -107,46 +104,43 @@ AutoRound focuses on weight-only quantization. The primary configuration is W4A1 ## Save and load -Save and reload AutoRound quantized models using the standard [`~ModelMixin.save_pretrained`] and [`~ModelMixin.from_pretrained`] methods. - -### Save + + ```python -import torch -from diffusers import AutoModel, AutoRoundConfig - -model_id = "INCModel/Z-Image-W4A16-AutoRound" -quantization_config = AutoRoundConfig(bits=4, group_size=128, sym=False) -model = AutoModel.from_pretrained( - model_id, - subfolder="transformer", - quantization_config=quantization_config, - torch_dtype=torch.float16, +from auto_round import AutoRound +autoround = AutoRound( + tiny_z_image_model_path, + num_inference_steps=3, + guidance_scale=7.5, + dataset="coco2014, ) -model.save_pretrained("path/to/saved_model") +autoround.quantize_and_save("Z-Image-W4A16-AutoRound") ``` + + ```python import torch -from diffusers import AutoModel, FluxPipeline +from diffusers import ZImageTransformer2DModel, ZImagePipeline -transformer = AutoModel.from_pretrained( - "path/to/saved_model", - torch_dtype=torch.float16, -) -pipe = FluxPipeline.from_pretrained( - "your-org/flux-autoround-w4g128", - transformer=transformer, - torch_dtype=torch.float16, +model_id = "INCModel/Z-Image-W4A16-AutoRound" + +# The inference backend will be automatically selected. +pipe = ZImagePipeline.from_pretrained( + model_id, + torch_dtype=torch.bfloat16, + device_map="cuda", ) -pipe.to("cuda") -image = pipe("A beautiful sunset over the ocean").images[0] +image = pipe("a cat holding a sign that says hello").images[0] image.save("output.png") ``` + + + ## Resources -- [AutoRound Hugging Face integration (Transformers)](https://huggingface.co/docs/transformers/quantization/autoround) - [Pre-quantized AutoRound models on the Hub](https://huggingface.co/models?search=autoround) diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index f61e4d8225eb..81714112a630 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -140,13 +140,24 @@ def merge_quantization_configs( ) else: warning_msg = "" - + existing_fields = set(quantization_config.keys()) if isinstance(quantization_config, dict): quantization_config = cls.from_dict(quantization_config) if isinstance(quantization_config, NVIDIAModelOptConfig): quantization_config.check_model_patching() + if quantization_config_from_args is not None: + # Only override fields that the user explicitly set. + for key, value in quantization_config_from_args.__dict__.items(): + if key not in existing_fields: + # Field does not exist in the model's quantization_config, add it. + setattr(quantization_config, key, value) + warning_msg += ( + f" Field `{key}` from `quantization_config_from_args` is not present in the model's " + f"`quantization_config`. Adding it with value: {value!r}." + ) + if warning_msg != "": warnings.warn(warning_msg) diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py index f6f46048bc66..f64c328f0261 100644 --- a/src/diffusers/quantizers/autoround/autoround_quantizer.py +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -90,8 +90,6 @@ def _process_model_before_weight_loading( model, used_backends = convert_hf_model(model, target_device) self.used_backends = used_backends - model.config.quantization_config = self.quantization_config - def _process_model_after_weight_loading(self, model, **kwargs): """ Finalizes the model after all quantized weights (qweight, scales, qzeros, etc.) have diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 53797ff4dd08..6f0bbd7bf5c7 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -768,12 +768,9 @@ class AutoRoundConfig(QuantizationConfigMixin): group_size (`int`, *optional*, defaults to `128`): The group size for weight quantization. Weights in each group share the same scale and zero-point. Common choices: 32, 64, 128, -1 (per-channel). - sym (`bool`, *optional*, defaults to `False`): + sym (`bool`, *optional*, defaults to `True`): Whether to use symmetric quantization (zero-point fixed at 0) or asymmetric - quantization (zero-point is learned). Asymmetric is generally more accurate. - modules_to_not_convert (`list[str]` or `None`, *optional*, defaults to `None`): - List of module name patterns that should NOT be quantized. Useful for keeping - certain sensitive layers (e.g. the final output projection) in full precision. + quantization (zero-point is learned). backend (`str`, *optional*, defaults to `"auto"`): The backend kernel to use for quantized inference. Available backends: - `"auto"`: Automatically select the best available backend for the current device. @@ -790,8 +787,7 @@ def __init__( self, bits: int = 4, group_size: int = 128, - sym: bool = False, - modules_to_not_convert: list[str] | None = None, + sym: bool = True, backend: str = "auto", **kwargs, ) -> None: @@ -799,7 +795,6 @@ def __init__( self.bits = bits self.group_size = group_size self.sym = sym - self.modules_to_not_convert = modules_to_not_convert self.backend = backend for k, v in kwargs.items(): setattr(self, k, v) From fb2e4c2177ea7eb304845b3ba302bf6e8f70431a Mon Sep 17 00:00:00 2001 From: Xin He Date: Fri, 24 Apr 2026 05:41:15 +0000 Subject: [PATCH 06/14] add UT and fix bug Signed-off-by: Xin He --- docs/source/en/quantization/autoround.md | 2 +- src/diffusers/quantizers/auto.py | 4 +- .../quantization/auto_round/test_autoround.py | 131 +++++++++++++----- 3 files changed, 104 insertions(+), 33 deletions(-) diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md index 0311c061fd10..cc3b3693e381 100644 --- a/docs/source/en/quantization/autoround.md +++ b/docs/source/en/quantization/autoround.md @@ -79,7 +79,7 @@ from diffusers import AutoRoundConfig config = AutoRoundConfig() # Explicit Triton backend for CUDA -config = AutoRoundConfig(backend="triton") +config = AutoRoundConfig(backend="tritonv2") # Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) config = AutoRoundConfig(backend="marlin") diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 81714112a630..fae2dfab7327 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -140,9 +140,11 @@ def merge_quantization_configs( ) else: warning_msg = "" - existing_fields = set(quantization_config.keys()) if isinstance(quantization_config, dict): + existing_fields = set(quantization_config.keys()) quantization_config = cls.from_dict(quantization_config) + else: + existing_fields = set(quantization_config.__dict__.keys()) if isinstance(quantization_config, NVIDIAModelOptConfig): quantization_config.check_model_patching() diff --git a/tests/quantization/auto_round/test_autoround.py b/tests/quantization/auto_round/test_autoround.py index 18b4b8b14ffa..7b53cd0fdaef 100644 --- a/tests/quantization/auto_round/test_autoround.py +++ b/tests/quantization/auto_round/test_autoround.py @@ -1,8 +1,11 @@ import gc import tempfile import unittest +import warnings from diffusers import AutoRoundConfig, ZImageTransformer2DModel, ZImagePipeline +from diffusers.quantizers.auto import DiffusersAutoQuantizer +from diffusers.quantizers.quantization_config import QuantizationMethod from diffusers.utils import is_auto_round_available, is_torch_available from diffusers.utils.testing_utils import ( backend_empty_cache, @@ -70,7 +73,6 @@ class AutoRoundBaseTesterMixin: pipeline_cls = ZImagePipeline torch_dtype = torch.bfloat16 expected_memory_reduction = 0.0 - modules_to_not_convert = "" _test_torch_compile = False def setUp(self): @@ -162,36 +164,6 @@ def test_autoround_memory_usage(self): assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction - def test_modules_to_not_convert(self): - """Verify that modules listed in `modules_to_not_convert` remain as standard nn.Linear.""" - from auto_round.inference.convert_model import dynamic_import_inference_linear - - init_kwargs = self.get_dummy_model_init_kwargs() - quantization_config_kwargs = self.get_dummy_init_kwargs() - quantization_config_kwargs.update({"modules_to_not_convert": self.modules_to_not_convert}) - quantization_config = AutoRoundConfig(**quantization_config_kwargs) - init_kwargs.update({"quantization_config": quantization_config}) - - model = self.model_cls.from_pretrained(**init_kwargs) - model.to(torch_device) - - # Resolve the actual backend used after model loading. - # When backend='auto', AutoRound selects the best available backend - # per-layer during convert_hf_model(); 'auto' itself is not a valid - # argument to dynamic_import_inference_linear. - used_backends = getattr(model.hf_quantizer, "used_backends", []) - resolved_backend = ( - used_backends[0] - if used_backends - else quantization_config_kwargs.get("backend", "auto_round:torch_zp") - ) - quant_linear_cls = dynamic_import_inference_linear(resolved_backend, quantization_config_kwargs) - - for name, module in model.named_modules(): - if name in self.modules_to_not_convert: - assert not isinstance( - module, quant_linear_cls - ), f"Module '{name}' should NOT have been quantized but is a {quant_linear_cls}." def test_serialization(self): """Test round-trip save and load of an AutoRound quantized model.""" @@ -392,3 +364,100 @@ def get_dummy_init_kwargs(self): "sym": False, "backend": "auto_round:torch_zp", } + + +# ============================================================================ +# Unit tests: AutoRoundConfig (no hardware required) +# ============================================================================ + + +class AutoRoundConfigTest(unittest.TestCase): + """Unit tests for AutoRoundConfig — no GPU / nightly decorator needed.""" + + def test_defaults(self): + cfg = AutoRoundConfig() + self.assertEqual(cfg.bits, 4) + self.assertEqual(cfg.group_size, 128) + self.assertTrue(cfg.sym) + self.assertEqual(cfg.backend, "auto") + self.assertEqual(cfg.quant_method, QuantizationMethod.AUTOROUND) + + def test_backend_values(self): + """All documented backend strings are stored correctly.""" + for backend in ("auto", "torch", "tritonv2", "marlin", "exllamav2"): + self.assertEqual(AutoRoundConfig(backend=backend).backend, backend) + + def test_to_dict_round_trip(self): + """to_dict → from_dict preserves all fields including backend and extra kwargs.""" + cfg = AutoRoundConfig(bits=4, group_size=32, sym=False, backend="gptqmodel:marlin_zp", + packing_format="auto_round:auto_gptq") + restored = AutoRoundConfig.from_dict(cfg.to_dict()) + self.assertEqual(restored.bits, cfg.bits) + self.assertEqual(restored.group_size, cfg.group_size) + self.assertEqual(restored.sym, cfg.sym) + self.assertEqual(restored.backend, cfg.backend) + self.assertEqual(restored.packing_format, cfg.packing_format) + self.assertEqual(restored.to_dict()["quant_method"], "auto-round") + + +# ============================================================================ +# Unit tests: DiffusersAutoQuantizer.merge_quantization_configs (no hardware) +# ============================================================================ + + +class MergeQuantizationConfigsTest(unittest.TestCase): + """Tests for the merge logic in DiffusersAutoQuantizer.merge_quantization_configs. + + Key behaviours under test: + 1. New fields in quantization_config_from_args (e.g. `backend`) are forwarded to + the merged config when they are absent from the model's saved config. + 2. Fields already present in the model's saved config are NOT overridden. + 3. A warning is emitted when quantization_config_from_args is provided. + 4. No warning when quantization_config_from_args is None. + """ + + def _model_config_dict(self, **overrides): + """Simulate a minimal saved AutoRound quantization_config dict (no 'backend' key).""" + base = { + "quant_method": "auto-round", + "bits": 4, + "group_size": 128, + "sym": True, + "autoround_version": "0.13.0", + "packing_format": "auto_round:auto_gptq", + } + base.update(overrides) + return base + + def test_new_fields_from_args_are_forwarded(self): + """Fields absent from the model config (backend, or arbitrary kwargs) are added from args.""" + for backend in ("marlin", "triton", "torch", "exllamav2"): # tritonv2 equals triton + with self.subTest(backend=backend): + merged = DiffusersAutoQuantizer.merge_quantization_configs( + self._model_config_dict(), AutoRoundConfig(backend=backend) + ) + self.assertEqual(merged.backend, backend) + + def test_existing_fields_not_overridden(self): + """Fields already in model config are NOT overridden; absent fields (backend) ARE added.""" + args_cfg = AutoRoundConfig(bits=2, group_size=32, sym=False, backend="torch") + merged = DiffusersAutoQuantizer.merge_quantization_configs(self._model_config_dict(), args_cfg) + + self.assertEqual(merged.bits, 4) # model value kept + self.assertEqual(merged.group_size, 128) # model value kept + self.assertTrue(merged.sym) # model value kept + self.assertEqual(merged.backend, "torch") # new field added + + def test_warning_behaviour(self): + """Warning emitted with args; no warning without args.""" + model_cfg = self._model_config_dict() + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + DiffusersAutoQuantizer.merge_quantization_configs(model_cfg, AutoRoundConfig(backend="auto")) + self.assertTrue(any("quantization_config" in str(x.message) for x in w)) + + with warnings.catch_warnings(record=True) as w: + warnings.simplefilter("always") + DiffusersAutoQuantizer.merge_quantization_configs(model_cfg, None) + self.assertFalse(any("quantization_config" in str(x.message).lower() for x in w)) From 418cf2e01b60a4a6701c017c04e67d3535bc4ce5 Mon Sep 17 00:00:00 2001 From: Xin He Date: Thu, 7 May 2026 12:03:40 +0000 Subject: [PATCH 07/14] update per comments Signed-off-by: Xin He --- docs/source/en/quantization/autoround.md | 36 +++-- .../transformers/transformer_z_image.py | 1 + src/diffusers/quantizers/auto.py | 1 + .../autoround/autoround_quantizer.py | 2 +- .../quantizers/quantization_config.py | 16 ++- tests/models/testing_utils/__init__.py | 5 + tests/models/testing_utils/quantization.py | 135 ++++++++++++++++-- .../test_models_transformer_z_image.py | 43 ++++++ tests/testing_utils.py | 23 +++ 9 files changed, 237 insertions(+), 25 deletions(-) diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md index cc3b3693e381..fbd0718aad87 100644 --- a/docs/source/en/quantization/autoround.md +++ b/docs/source/en/quantization/autoround.md @@ -36,7 +36,7 @@ from diffusers import ZImageTransformer2DModel, ZImagePipeline, AutoRoundConfig model_id = "INCModel/Z-Image-W4A16-AutoRound" -quantization_config = AutoRoundConfig(backend="marlin") +quantization_config = AutoRoundConfig(backend="auto") transformer = ZImageTransformer2DModel.from_pretrained( model_id, subfolder="transformer", @@ -61,7 +61,7 @@ image.save("output.png") ## Backends -AutoRound supports multiple inference backends. The backend controls which kernel handles dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: +AutoRound supports multiple inference backends for Weight-only quantized model. The backend controls which kernel handles dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: | Backend | Value | Device | Requirements | Notes | |---------|-------|--------|--------------|-------| @@ -92,16 +92,6 @@ config = AutoRoundConfig(backend="torch") ``` -## Quantization configurations - -AutoRound focuses on weight-only quantization. The primary configuration is W4A16 (4-bit weights, 16-bit activations), with flexibility in group size and symmetry: - -| Configuration | `bits` | `group_size` | `sym` | Description | -|--------------|--------|-------------|-------|-------------| -| W4G128 asymmetric | `4` | `128` | `False` | Default. Good balance of accuracy and compression. | -| W4G128 symmetric | `4` | `128` | `True` | Faster dequantization, small accuracy trade-off. | -| W4G32 asymmetric | `4` | `32` | `False` | Higher accuracy at the cost of more metadata. | - ## Save and load @@ -111,6 +101,8 @@ AutoRound focuses on weight-only quantization. The primary configuration is W4A1 from auto_round import AutoRound autoround = AutoRound( tiny_z_image_model_path, + scheme="W4A16", # W4G128 symmetric + enable_torch_compile=True, num_inference_steps=3, guidance_scale=7.5, dataset="coco2014, @@ -141,6 +133,26 @@ image.save("output.png") + +### Supported Quantization Schemes + +AutoRound supports several Schemes: + +- **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**(all Q*_K,Q*_0,Q*_1 provided by llamacpp are supported) +- **NVFP4**(Experimental feature, recommend exporting to `llm_compressor` format.data_type nvfp4,act_data_type nvfp4,static_global_scale,group_size 16) +- **MXFP4**(**Research feature, no real kernel**, Standard MXFP4, data_type mxfp,act_data_type mxfp,bits 4, act_bits 4, group_size 32) +- **MXINT4**(**Research feature, no real kernel**, Standard MXINT4, data_type mxint,act_data_type mxint,bits 4, act_bits 4, group_size 32) +- **MXFP4_RCEIL**(**Research feature,no real kernel**, NVIDIA's variant, data_type mxfp,act_data_type mxfp_rceil,bits 4, act_bits 4, group_size 32) +- **MXFP8**(**Research feature, no real kernel**, data_type mxfp,act_data_type mxfp_rceil,group_size 32) +- **FPW8A16**(**Research feature, no real kernel**, data_type fp8,group_size 0->per tensor ) +- **FP8_STATIC**(**Research feature, no real kernel**, data_type:fp8,act_data_type:fp8,group_size -1 ->per channel, act_group_size=0->per tensor) + +Besides, you could modify the `group_size`, `bits`, `sym` and many other configs you want, though there are maybe no real kernels. + ## Resources - [Pre-quantized AutoRound models on the Hub](https://huggingface.co/models?search=autoround) diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index ba401e7fdef1..99b21d1f29bf 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -15,6 +15,7 @@ import math import torch +import torch._dynamo as dynamo import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index fae2dfab7327..0aa885a55734 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -140,6 +140,7 @@ def merge_quantization_configs( ) else: warning_msg = "" + if isinstance(quantization_config, dict): existing_fields = set(quantization_config.keys()) quantization_config = cls.from_dict(quantization_config) diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py index f64c328f0261..7f02686bb04b 100644 --- a/src/diffusers/quantizers/autoround/autoround_quantizer.py +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -60,7 +60,7 @@ def validate_environment(self, *args, **kwargs): if not is_auto_round_available(): raise ImportError( "Loading an AutoRound quantized model requires the auto-round library " - "(`pip install 'auto-round>=0.5'`)" + "(`pip install 'auto-round>=0.13.0'`)" ) def _process_model_before_weight_loading( diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 6f0bbd7bf5c7..75c25540befb 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -774,9 +774,11 @@ class AutoRoundConfig(QuantizationConfigMixin): backend (`str`, *optional*, defaults to `"auto"`): The backend kernel to use for quantized inference. Available backends: - `"auto"`: Automatically select the best available backend for the current device. - - `"auto_round:torch_zp"`: Pure PyTorch kernel — works on CPU and CUDA. - - `"auto_round:tritonv2_zp"`: Triton-based kernel — requires CUDA. - - `"gptqmodel:marlin_zp"`: Marlin kernel via GPTQModel — requires CUDA and + - `"torch"`: Pure PyTorch kernel — works on CPU and CUDA. + - `"tritonv2"`: Triton-based kernel — requires CUDA. + - `"exllamav2"`: Exllamav2 kernel via GPTQModel — requires CUDA and + `gptqmodel>=5.8.0`. Offers good CUDA inference performance. + - `"marlin"`: Marlin kernel via GPTQModel — requires CUDA and `gptqmodel>=5.8.0`. Offers the best CUDA inference performance. kwargs (`dict[str, Any]`, *optional*): Additional keyword arguments forwarded to AutoRound (e.g. `iters`, `seqlen`, @@ -795,10 +797,16 @@ def __init__( self.bits = bits self.group_size = group_size self.sym = sym - self.backend = backend + self.backend = self._validate_backend(backend) for k, v in kwargs.items(): setattr(self, k, v) + def _validate_backend(self, backend): + valid_backends = ["auto","torch","tritonv2","exllamav2","marlin"] + if backend not in valid_backends: + raise ValueError(f"Invalid backend '{backend}'. Valid options are: {valid_backends}") + return backend + def to_dict(self) -> dict: """Serialize the config to a JSON-compatible dict. diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index ea076b3ec774..15a4b27c6d05 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -15,6 +15,9 @@ from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ContextParallelTesterMixin from .quantization import ( + AutoRoundCompileTesterMixin, + AutoRoundConfigMixin, + AutoRoundTesterMixin, BitsAndBytesCompileTesterMixin, BitsAndBytesConfigMixin, BitsAndBytesTesterMixin, @@ -39,6 +42,8 @@ __all__ = [ "AttentionTesterMixin", + "AutoRoundConfigMixin", + "AutoRoundTesterMixin", "BaseModelTesterConfig", "BitsAndBytesCompileTesterMixin", "BitsAndBytesConfigMixin", diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 1aab0b240148..44b768f2bd75 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -18,8 +18,9 @@ import pytest import torch -from diffusers import BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, QuantoConfig, TorchAoConfig +from diffusers import AutoRoundConfig, BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, QuantoConfig, TorchAoConfig from diffusers.utils.import_utils import ( + is_auto_round_available, is_bitsandbytes_available, is_gguf_available, is_nvidia_modelopt_available, @@ -31,6 +32,7 @@ backend_empty_cache, backend_max_memory_allocated, backend_reset_peak_memory_stats, + is_autoround, is_bitsandbytes, is_gguf, is_modelopt, @@ -40,6 +42,7 @@ is_torchao, require_accelerate, require_accelerator, + require_auto_round_version_greater_or_equal, require_bitsandbytes_version_greater, require_gguf_version_greater_or_equal, require_modelopt_version_greater_or_equal, @@ -109,11 +112,11 @@ class QuantizationTesterMixin: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ - def setup_method(self): + def setup_method(self, method=None): gc.collect() backend_empty_cache(torch_device) - def teardown_method(self): + def teardown_method(self, method=None): gc.collect() backend_empty_cache(torch_device) @@ -176,6 +179,7 @@ def _test_quantization_inference(self, config_kwargs): inputs = self.get_dummy_inputs() output = model_quantized(**inputs, return_dict=False)[0] + output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -336,6 +340,7 @@ def _test_quantization_device_map(self, config_kwargs): inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] + output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -1159,18 +1164,18 @@ class QuantizationCompileTesterMixin: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ - def setup_method(self): + def setup_method(self, method=None): gc.collect() backend_empty_cache(torch_device) torch.compiler.reset() - def teardown_method(self): + def teardown_method(self, method=None): gc.collect() backend_empty_cache(torch_device) torch.compiler.reset() @torch.no_grad() - def _test_torch_compile(self, config_kwargs): + def _test_torch_compile(self, config_kwargs, fullgraph=True, error_on_recompile=True): """ Test that torch.compile works correctly with a quantized model. @@ -1181,11 +1186,12 @@ def _test_torch_compile(self, config_kwargs): model.to(torch_device) model.eval() - model = torch.compile(model, fullgraph=True) + model = torch.compile(model, fullgraph=fullgraph) - with torch._dynamo.config.patch(error_on_recompile=True): + with torch._dynamo.config.patch(error_on_recompile=error_on_recompile): inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] + output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -1217,6 +1223,7 @@ def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] + output = output[0] if isinstance(output, (tuple, list)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -1362,3 +1369,115 @@ def test_modelopt_torch_compile(self, config_name): @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_torch_compile_with_group_offload(self, config_name): self._test_torch_compile_with_group_offload(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + +@is_quantization +@is_autoround +@require_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundConfigMixin: + """ + Base mixin providing AutoRound quantization config and model creation. + + AutoRound is a weight-only quantization method (W4A16). It supports multiple inference + + When `backend="auto"`, AutoRound selects the best available backend automatically. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - quantized_model_name_or_path: Hub repository ID for the quantized model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + config_dict = {"backend": "auto"} + + def _load_unquantized_model(self): + kwargs = getattr(self, "pretrained_model_kwargs", {}) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = AutoRoundConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs["torch_dtype"] = torch.bfloat16 + if "device_map" not in kwargs: + kwargs["device_map"] = torch_device + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.quantized_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + # AutoRound replaces linear layers with quantized linear layers + assert isinstance(module, torch.nn.Linear), f"Layer {name} is not Linear, got {type(module)}" + + +@is_autoround +@require_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundTesterMixin(AutoRoundConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing AutoRound quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - quantized_model_name_or_path: Hub repository ID for the quantized model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - AUTOROUND_CONFIGS: Dict of config name -> AutoRoundConfig kwargs to test + + Pytest mark: autoround + Use `pytest -m "not autoround"` to skip these tests + """ + config_dict = {"backend": "auto"} + + def test_autoround_quantization_memory_footprint(self): + expected = 1.5 # AutoRound is a W4A16 method, so we expect around 1.5x memory reduction + self._test_quantization_memory_footprint( + self.config_dict, expected_memory_reduction=expected + ) + + def test_autoround_quantization_inference(self): + self._test_quantization_inference(self.config_dict) + + def test_autoround_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(self.config_dict) + + +@is_autoround +@require_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundCompileTesterMixin(AutoRoundConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing `torch.compile` with AutoRound-quantized models. + + This mixin provides tests that verify `torch.compile` works correctly with models + quantized using AutoRound. Subclasses are expected to inherit from + `AutoRoundConfigMixin` (which defines `config_dict`) and to provide the + following class attributes: `model_class`, `pretrained_model_name_or_path`, and + `quantized_model_name_or_path`. + + The mixin uses `config_dict` (defaults to {"backend": "auto"}) as the + quantization configuration passed into `_create_quantized_model` when + invoking the compile-related tests. + + Provided tests: + - `test_autoround_torch_compile`: Ensures `torch.compile` runs and produces + valid, non-NaN outputs for an AutoRound-quantized model. + - `test_autoround_torch_compile_with_group_offload`: Ensures `torch.compile` + works together with group offloading when supported by the quantized + model implementation. + """ + + def test_autoround_torch_compile(self): + self._test_torch_compile(self.config_dict, fullgraph=False, error_on_recompile=False) + + def test_autoround_torch_compile_with_group_offload(self): + self._test_torch_compile_with_group_offload(self.config_dict) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 79054019f2d2..3d1a78b1b8d3 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -21,6 +21,7 @@ from diffusers import ZImageTransformer2DModel +from ..testing_utils import AutoRoundTesterMixin, AutoRoundCompileTesterMixin from ...testing_utils import IS_GITHUB_ACTIONS, torch_device from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin @@ -169,3 +170,45 @@ def test_compile_works_with_aot(self): @unittest.skip("Fullgraph is broken") def test_compile_on_different_shapes(self): super().test_compile_on_different_shapes() + + +class ZImageTransformerTesterConfig: + """Configuration class for Z-Image Transformer tests.""" + @property + def model_class(self): + return ZImageTransformer2DModel + + @property + def pretrained_model_name_or_path(self): + return "INCModel/Z-Image-tiny-for-testing" + + @property + def quantized_model_name_or_path(self): + return "INCModel/Z-Image-tiny-for-testing-W4A16-AutoRound" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + def get_dummy_inputs(self): + batch_size = 1 + in_channels = 16 + cap_feat_dim = 512 + height = width = 8 + frames = 1 + seq_len = 16 + + torch.manual_seed(0) + x = [torch.randn((in_channels, frames, height, width)).to(torch_device, dtype=torch.bfloat16) for _ in range(batch_size)] + cap_feats = [torch.randn((seq_len, cap_feat_dim)).to(torch_device, dtype=torch.bfloat16) for _ in range(batch_size)] + t = torch.tensor([0.5]).to(torch_device, dtype=torch.bfloat16) + + return {"x": x, "cap_feats": cap_feats, "t": t} + + +class TestZImageTransformerAutoRound(ZImageTransformerTesterConfig, AutoRoundTesterMixin): + """AutoRound quantization tests for Z-Image Transformer.""" + + +class TestZImageTransformerAutoRoundCompile(ZImageTransformerTesterConfig, AutoRoundCompileTesterMixin): + """AutoRound quantization + torch.compile tests for Z-Image Transformer.""" diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 060f9ee0f882..23d22702a70b 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -32,6 +32,7 @@ from diffusers.utils.import_utils import ( BACKENDS_MAPPING, is_accelerate_available, + is_auto_round_available, is_bitsandbytes_available, is_compel_available, is_flashpack_available, @@ -460,6 +461,15 @@ def is_gguf(test_case): return pytest.mark.gguf(test_case) +def is_autoround(test_case): + """ + Decorator marking a test as an AutoRound quantization test. These tests can be filtered using: + pytest -m "not autoround" to skip + pytest -m autoround to run only these tests + """ + return pytest.mark.autoround(test_case) + + def is_modelopt(test_case): """ Decorator marking a test as a NVIDIA ModelOpt quantization test. These tests can be filtered using: @@ -842,6 +852,19 @@ def decorator(test_case): return decorator +def require_auto_round_version_greater_or_equal(auto_round_version): + def decorator(test_case): + correct_auto_round_version = is_auto_round_available() and version.parse( + version.parse(importlib.metadata.version("auto_round")).base_version + ) >= version.parse(auto_round_version) + return pytest.mark.skipif( + not correct_auto_round_version, + reason=f"Test requires auto-round with version greater than {auto_round_version}.", + )(test_case) + + return decorator + + def require_kernels_version_greater_or_equal(kernels_version): def decorator(test_case): correct_kernels_version = is_kernels_available() and version.parse( From 3a7c040564ea847de114b58c75a14c8ed94ed989 Mon Sep 17 00:00:00 2001 From: Xin He Date: Tue, 19 May 2026 16:06:03 +0800 Subject: [PATCH 08/14] update per comments Signed-off-by: Xin He --- docs/source/en/_toctree.yml | 2 + docs/source/en/quantization/autoround.md | 58 ++- .../transformers/transformer_z_image.py | 1 - src/diffusers/quantizers/auto.py | 21 +- .../autoround/autoround_quantizer.py | 10 + .../quantizers/quantization_config.py | 11 +- src/diffusers/utils/testing_utils.py | 14 - tests/models/testing_utils/quantization.py | 12 +- .../test_models_transformer_z_image.py | 71 +++ tests/quantization/auto_round/__init__.py | 0 .../quantization/auto_round/test_autoround.py | 463 ------------------ 11 files changed, 156 insertions(+), 507 deletions(-) delete mode 100644 tests/quantization/auto_round/__init__.py delete mode 100644 tests/quantization/auto_round/test_autoround.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 1db7a7cc3e9f..f6befb08340f 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -180,6 +180,8 @@ title: quanto - local: quantization/modelopt title: NVIDIA ModelOpt + - local: quantization/autoround + title: AutoRound title: Quantization - isExpanded: false sections: diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md index fbd0718aad87..3e7cd9fcc541 100644 --- a/docs/source/en/quantization/autoround.md +++ b/docs/source/en/quantization/autoround.md @@ -28,7 +28,29 @@ pip install "gptqmodel>=5.8.0" ## Load a quantized model -Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelMixin.from_pretrained`]. The method works with any model that loads via [Accelerate(https://hf.co/docs/accelerate/index) and has `torch.nn.Linear` layers. +Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelMixin.from_pretrained`]. The method works with any model that loads via [Accelerate](https://hf.co/docs/accelerate/index) and has `torch.nn.Linear` layers. + +You can use [`PipelineQuantizationConfig`] to quantize specific components of a pipeline: + +```python +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, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +image = pipe("a cat holding a sign that says hello").images[0] +image.save("output.png") +``` + +Or load a quantized model component directly: ```python import torch @@ -59,6 +81,27 @@ image.save("output.png") > [!NOTE] > AutoRound in Diffusers only supports loading *pre-quantized* models. To quantize a model from scratch, use the [AutoRound CLI or Python API](https://github.com/intel/auto-round) directly, then load the result with Diffusers. +## torch.compile + +AutoRound is compatible with [`torch.compile`](../optimization/fp16#torchcompile) for faster inference. You can compile the quantized transformer (DiT) for better performance: + +```python +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, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True) +``` + ## Backends AutoRound supports multiple inference backends for Weight-only quantized model. The backend controls which kernel handles dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: @@ -84,7 +127,7 @@ config = AutoRoundConfig(backend="tritonv2") # Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) config = AutoRoundConfig(backend="marlin") -# Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) +# ExllamaV2 backend for good CUDA performance (requires gptqmodel>=5.8.0) config = AutoRoundConfig(backend="exllamav2") # PyTorch backend for CPU/CUDA inference @@ -97,22 +140,28 @@ config = AutoRoundConfig(backend="torch") +AutoRound requires data calibration to quantize a model. This is done outside of Diffusers using the [AutoRound library](https://github.com/intel/auto-round) directly: + ```python from auto_round import AutoRound + autoround = AutoRound( - tiny_z_image_model_path, + "Tongyi-MAI/Z-Image", scheme="W4A16", # W4G128 symmetric enable_torch_compile=True, num_inference_steps=3, guidance_scale=7.5, - dataset="coco2014, + dataset="coco2014", ) autoround.quantize_and_save("Z-Image-W4A16-AutoRound") ``` +For more details on calibration options, see the [AutoRound documentation](https://github.com/intel/auto-round). + + ```python import torch from diffusers import ZImageTransformer2DModel, ZImagePipeline @@ -129,7 +178,6 @@ pipe = ZImagePipeline.from_pretrained( image = pipe("a cat holding a sign that says hello").images[0] image.save("output.png") ``` - diff --git a/src/diffusers/models/transformers/transformer_z_image.py b/src/diffusers/models/transformers/transformer_z_image.py index 99b21d1f29bf..ba401e7fdef1 100644 --- a/src/diffusers/models/transformers/transformer_z_image.py +++ b/src/diffusers/models/transformers/transformer_z_image.py @@ -15,7 +15,6 @@ import math import torch -import torch._dynamo as dynamo import torch.nn as nn import torch.nn.functional as F from torch.nn.utils.rnn import pad_sequence diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 0aa885a55734..a10bf0cdcb3f 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -142,24 +142,23 @@ def merge_quantization_configs( warning_msg = "" if isinstance(quantization_config, dict): - existing_fields = set(quantization_config.keys()) quantization_config = cls.from_dict(quantization_config) - else: - existing_fields = set(quantization_config.__dict__.keys()) if isinstance(quantization_config, NVIDIAModelOptConfig): quantization_config.check_model_patching() - if quantization_config_from_args is not None: - # Only override fields that the user explicitly set. + if quantization_config_from_args is not None and isinstance(quantization_config, AutoRoundConfig): + # For AutoRound, allow overriding fields like `backend` from user args, + # since the model config may store a default value (e.g. backend="auto"). for key, value in quantization_config_from_args.__dict__.items(): - if key not in existing_fields: - # Field does not exist in the model's quantization_config, add it. - setattr(quantization_config, key, value) - warning_msg += ( - f" Field `{key}` from `quantization_config_from_args` is not present in the model's " - f"`quantization_config`. Adding it with value: {value!r}." + if key in ("quant_method",): + continue + if hasattr(quantization_config, key) and getattr(quantization_config, key) != value: + warnings.warn( + f"Overriding `{key}` in the model's quantization_config with value {value!r} " + f"from the user-provided `quantization_config`." ) + setattr(quantization_config, key, value) if warning_msg != "": warnings.warn(warning_msg) diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py index 7f02686bb04b..aa5675e83e77 100644 --- a/src/diffusers/quantizers/autoround/autoround_quantizer.py +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -62,6 +62,12 @@ def validate_environment(self, *args, **kwargs): "Loading an AutoRound quantized model requires the auto-round library " "(`pip install 'auto-round>=0.13.0'`)" ) + if not self.pre_quantized: + raise ValueError( + "AutoRound quantizer in diffusers only supports loading pre-quantized models. " + "To quantize a model from scratch, use the AutoRound CLI or Python API " + "(https://github.com/intel/auto-round) directly, then load the result with Diffusers." + ) def _process_model_before_weight_loading( self, @@ -126,3 +132,7 @@ def is_serializable(self): updated by the backend, e.g. for GPTQ/AWQ-compatible formats).""" return True + @property + def is_compileable(self) -> bool: + return True + diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 75c25540befb..dc213a6eebca 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -785,6 +785,8 @@ class AutoRoundConfig(QuantizationConfigMixin): `batch_size`, `lr`, `minmax_lr` for calibration when quantizing from scratch). """ + VALID_BACKENDS = ["auto", "torch", "tritonv2", "exllamav2", "marlin"] + def __init__( self, bits: int = 4, @@ -794,18 +796,17 @@ def __init__( **kwargs, ) -> None: self.quant_method = QuantizationMethod.AUTOROUND + self._validate_backend(backend) self.bits = bits self.group_size = group_size self.sym = sym - self.backend = self._validate_backend(backend) + self.backend = backend for k, v in kwargs.items(): setattr(self, k, v) def _validate_backend(self, backend): - valid_backends = ["auto","torch","tritonv2","exllamav2","marlin"] - if backend not in valid_backends: - raise ValueError(f"Invalid backend '{backend}'. Valid options are: {valid_backends}") - return backend + if backend not in self.VALID_BACKENDS: + raise ValueError(f"Invalid backend '{backend}'. Valid options are: {self.VALID_BACKENDS}") def to_dict(self) -> dict: """Serialize the config to a JSON-compatible dict. diff --git a/src/diffusers/utils/testing_utils.py b/src/diffusers/utils/testing_utils.py index 2ea9d039e4f0..619a37034949 100644 --- a/src/diffusers/utils/testing_utils.py +++ b/src/diffusers/utils/testing_utils.py @@ -33,7 +33,6 @@ from .import_utils import ( BACKENDS_MAPPING, is_accelerate_available, - is_auto_round_available, is_bitsandbytes_available, is_compel_available, is_flax_available, @@ -655,19 +654,6 @@ def decorator(test_case): return decorator -def require_auto_round_version_greater_or_equal(auto_round_version): - def decorator(test_case): - correct_auto_round_version = is_auto_round_available() and version.parse( - version.parse(importlib.metadata.version("auto_round")).base_version - ) >= version.parse(auto_round_version) - return unittest.skipUnless( - correct_auto_round_version, - f"Test requires auto-round with version greater than {auto_round_version}.", - )(test_case) - - return decorator - - def require_kernels_version_greater_or_equal(kernels_version): def decorator(test_case): correct_kernels_version = is_kernels_available() and version.parse( diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 44b768f2bd75..c495db6daebd 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -112,11 +112,11 @@ class QuantizationTesterMixin: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ - def setup_method(self, method=None): + def setup_method(self): gc.collect() backend_empty_cache(torch_device) - def teardown_method(self, method=None): + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) @@ -179,7 +179,6 @@ def _test_quantization_inference(self, config_kwargs): inputs = self.get_dummy_inputs() output = model_quantized(**inputs, return_dict=False)[0] - output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -340,7 +339,6 @@ def _test_quantization_device_map(self, config_kwargs): inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] - output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -1164,12 +1162,12 @@ class QuantizationCompileTesterMixin: - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass """ - def setup_method(self, method=None): + def setup_method(self): gc.collect() backend_empty_cache(torch_device) torch.compiler.reset() - def teardown_method(self, method=None): + def teardown_method(self): gc.collect() backend_empty_cache(torch_device) torch.compiler.reset() @@ -1191,7 +1189,6 @@ def _test_torch_compile(self, config_kwargs, fullgraph=True, error_on_recompile= with torch._dynamo.config.patch(error_on_recompile=error_on_recompile): inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] - output = output[0] if isinstance(output, (list, tuple)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" @@ -1223,7 +1220,6 @@ def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] - output = output[0] if isinstance(output, (tuple, list)) else output assert output is not None, "Model output is None" assert not torch.isnan(output).any(), "Model output contains NaN" diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 3d1a78b1b8d3..f8042947bb50 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -209,6 +209,77 @@ def get_dummy_inputs(self): class TestZImageTransformerAutoRound(ZImageTransformerTesterConfig, AutoRoundTesterMixin): """AutoRound quantization tests for Z-Image Transformer.""" + @torch.no_grad() + def _test_quantization_inference(self, config_kwargs): + model_quantized = self._create_quantized_model(config_kwargs) + model_quantized.to(torch_device) + + inputs = self.get_dummy_inputs() + output = model_quantized(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + @torch.no_grad() + def _test_quantization_device_map(self, config_kwargs): + model = self._create_quantized_model(config_kwargs, device_map="auto") + + assert hasattr(model, "hf_device_map"), "Model should have hf_device_map attribute" + assert model.hf_device_map is not None, "hf_device_map should not be None" + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + class TestZImageTransformerAutoRoundCompile(ZImageTransformerTesterConfig, AutoRoundCompileTesterMixin): """AutoRound quantization + torch.compile tests for Z-Image Transformer.""" + + @torch.no_grad() + def _test_torch_compile(self, config_kwargs, fullgraph=True, error_on_recompile=True): + model = self._create_quantized_model(config_kwargs) + model.to(torch_device) + model.eval() + + model = torch.compile(model, fullgraph=fullgraph) + + with torch._dynamo.config.patch(error_on_recompile=error_on_recompile): + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + @torch.no_grad() + def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False): + import pytest + + torch._dynamo.config.cache_size_limit = 1000 + + model = self._create_quantized_model(config_kwargs) + model.eval() + + if not hasattr(model, "enable_group_offload"): + pytest.skip("Model does not support group offloading") + + group_offload_kwargs = { + "onload_device": torch.device(torch_device), + "offload_device": torch.device("cpu"), + "offload_type": "leaf_level", + "use_stream": use_stream, + } + model.enable_group_offload(**group_offload_kwargs) + model = torch.compile(model) + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" diff --git a/tests/quantization/auto_round/__init__.py b/tests/quantization/auto_round/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/tests/quantization/auto_round/test_autoround.py b/tests/quantization/auto_round/test_autoround.py deleted file mode 100644 index 7b53cd0fdaef..000000000000 --- a/tests/quantization/auto_round/test_autoround.py +++ /dev/null @@ -1,463 +0,0 @@ -import gc -import tempfile -import unittest -import warnings - -from diffusers import AutoRoundConfig, ZImageTransformer2DModel, ZImagePipeline -from diffusers.quantizers.auto import DiffusersAutoQuantizer -from diffusers.quantizers.quantization_config import QuantizationMethod -from diffusers.utils import is_auto_round_available, is_torch_available -from diffusers.utils.testing_utils import ( - backend_empty_cache, - backend_reset_peak_memory_stats, - enable_full_determinism, - nightly, - numpy_cosine_similarity_distance, - require_accelerate, - require_big_accelerator, - require_auto_round_version_greater_or_equal, - require_torch_cuda_compatibility, - torch_device, -) - - -if is_torch_available(): - import torch - - from ..utils import get_memory_consumption_stat - - -def _is_gptqmodel_available(min_version="5.8.0"): - """Check if gptqmodel is installed with a minimum version.""" - try: - import importlib.metadata - - from packaging import version - - gptqmodel_version = importlib.metadata.version("gptqmodel") - return version.parse(gptqmodel_version) >= version.parse(min_version) - except importlib.metadata.PackageNotFoundError: - return False - - -enable_full_determinism() - - -@nightly -@require_big_accelerator -@require_accelerate -@require_auto_round_version_greater_or_equal("0.13.0") -class AutoRoundBaseTesterMixin: - """Base test mixin for AutoRound quantized models. - - AutoRound is a weight-only quantization method (W4A16). It supports multiple inference - backends depending on the hardware: - - CPU: `auto_round:torch_zp` backend - - CUDA: `auto_round:tritonv2_zp` backend - - CUDA + GPTQModel>=5.8.0: `gptqmodel:marlin_zp` backend (best performance) - - When `backend="auto"`, AutoRound selects the best available backend automatically. - - Key differences from ModelOpt tests: - - Only pre-quantized model loading is supported (no on-the-fly quantization). - - `is_trainable` returns False, so no LoRA training test. - - No `test_dtype_assignment` (AutoRound doesn't restrict dtype changes). - - `requires_calibration = True` means we always load pre-quantized checkpoints. - """ - - # TODO: Replace with a real tiny AutoRound-quantized checkpoint on the Hub. - # This should be a small model that has been quantized with AutoRound and uploaded - # in the standard format (qweight, scales, qzeros, g_idx). - model_id = "INCModel/Z-Image-tiny-for-testing-W4A16-AutoRound" - model_cls = ZImageTransformer2DModel - pipeline_cls = ZImagePipeline - torch_dtype = torch.bfloat16 - expected_memory_reduction = 0.0 - _test_torch_compile = False - - def setUp(self): - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - gc.collect() - - def tearDown(self): - backend_reset_peak_memory_stats(torch_device) - backend_empty_cache(torch_device) - gc.collect() - - def get_dummy_init_kwargs(self): - """Returns the default AutoRoundConfig kwargs for W4A16 quantization. - - Subclasses override this to specify backend, group_size, sym, etc. - """ - return { - "bits": 4, - "group_size": 128, - "sym": False, - } - - def get_dummy_model_init_kwargs(self): - """Returns kwargs for model_cls.from_pretrained() with AutoRound quantization.""" - return { - "pretrained_model_name_or_path": self.model_id, - "torch_dtype": self.torch_dtype, - "quantization_config": AutoRoundConfig(**self.get_dummy_init_kwargs()), - "subfolder": "transformer", - } - - def get_dummy_inputs(self): - """Creates dummy inputs matching ZImageTransformer2DModel.forward() signature. - - ZImageTransformer2DModel expects: - - x: list of (C, F, H, W) tensors, one per batch item - - t: 1-D timestep tensor of shape (batch_size,) - - cap_feats: list of (seq_len, cap_feat_dim) tensors, one per batch item - - Dimensions are chosen to match the tiny test checkpoint - (in_channels=16, cap_feat_dim=512, patch_size=2, f_patch_size=1). - """ - batch_size = 1 - in_channels = 16 # matches tiny model config - cap_feat_dim = 512 # matches tiny model config - height = width = 8 # must be divisible by patch_size=2 - frames = 1 # must be divisible by f_patch_size=1 - seq_len = 16 # caption token count (will be padded to multiple of 32) - - torch.manual_seed(0) - x = [ - torch.randn((in_channels, frames, height, width)).to(torch_device, dtype=self.torch_dtype) - for _ in range(batch_size) - ] - cap_feats = [ - torch.randn((seq_len, cap_feat_dim)).to(torch_device, dtype=self.torch_dtype) - for _ in range(batch_size) - ] - t = torch.tensor([0.5] * batch_size).to(torch_device, dtype=self.torch_dtype) - - return {"x": x, "cap_feats": cap_feats, "t": t} - - def test_autoround_memory_usage(self): - """Compare peak memory between unquantized and AutoRound-quantized model. - - The quantized model should use significantly less memory due to 4-bit weight packing. - `expected_memory_reduction` defines the minimum ratio (unquantized / quantized). - """ - inputs = self.get_dummy_inputs() - # x and cap_feats are lists of tensors; move each element individually. - inputs = { - k: [t.to(device=torch_device, dtype=self.torch_dtype) for t in v] - if isinstance(v, list) - else v.to(device=torch_device, dtype=self.torch_dtype) - for k, v in inputs.items() - if not isinstance(v, bool) - } - - unquantized_model = self.model_cls.from_pretrained( - self.model_id, torch_dtype=self.torch_dtype, subfolder="transformer" - ) - unquantized_model.to(torch_device) - unquantized_model_memory = get_memory_consumption_stat(unquantized_model, inputs) - - quantized_model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - quantized_model.to(torch_device) - quantized_model_memory = get_memory_consumption_stat(quantized_model, inputs) - - assert unquantized_model_memory / quantized_model_memory >= self.expected_memory_reduction - - - def test_serialization(self): - """Test round-trip save and load of an AutoRound quantized model.""" - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - inputs = self.get_dummy_inputs() - - model.to(torch_device) - with torch.no_grad(): - model_output = model(**inputs) - - with tempfile.TemporaryDirectory() as tmp_dir: - model.save_pretrained(tmp_dir) - saved_model = self.model_cls.from_pretrained( - tmp_dir, - torch_dtype=self.torch_dtype, - ) - - saved_model.to(torch_device) - with torch.no_grad(): - saved_model_output = saved_model(**inputs) - - # model_output.sample is a list of per-item tensors - for out, saved_out in zip(model_output.sample, saved_model_output.sample): - assert torch.allclose(out, saved_out, rtol=1e-5, atol=1e-5) - - def test_torch_compile(self): - """Test that the quantized model works with torch.compile.""" - if not self._test_torch_compile: - return - - model = self.model_cls.from_pretrained(**self.get_dummy_model_init_kwargs()) - compiled_model = torch.compile(model, mode="max-autotune", fullgraph=True, dynamic=False) - - model.to(torch_device) - with torch.no_grad(): - model_output = model(**self.get_dummy_inputs()).sample - - compiled_model.to(torch_device) - with torch.no_grad(): - compiled_model_output = compiled_model(**self.get_dummy_inputs()).sample - - # model_output is a list of per-item tensors; stack for comparison - model_output = torch.stack([o.detach().float().cpu() for o in model_output]).numpy() - compiled_model_output = torch.stack([o.detach().float().cpu() for o in compiled_model_output]).numpy() - - max_diff = numpy_cosine_similarity_distance(model_output.flatten(), compiled_model_output.flatten()) - assert max_diff < 1e-3 - - def test_model_cpu_offload(self): - """Test that the quantized model works with pipeline CPU offload.""" - init_kwargs = self.get_dummy_init_kwargs() - transformer = self.model_cls.from_pretrained( - self.model_id, - quantization_config=AutoRoundConfig(**init_kwargs), - subfolder="transformer", - torch_dtype=self.torch_dtype, - ) - pipe = self.pipeline_cls.from_pretrained(self.model_id, transformer=transformer, torch_dtype=self.torch_dtype) - pipe.enable_model_cpu_offload(device=torch_device) - _ = pipe("a cat holding a sign that says hello", num_inference_steps=2) - - -# ============================================================================ -# Backend: auto (auto-select best available backend) -# ============================================================================ - - -class AutoRoundW4G128AsymAutoBackendTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, asymmetric, backend='auto' (default — auto-selects best backend).""" - - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": False, - "backend": "auto", - } - - -class AutoRoundW4G128SymAutoBackendTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, symmetric, backend='auto'.""" - - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": True, - "backend": "auto", - } - - -class AutoRoundW4G32AsymAutoBackendTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=32, asymmetric, backend='auto' (finer granularity).""" - - expected_memory_reduction = 0.50 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 32, - "sym": False, - "backend": "auto", - } - - -# ============================================================================ -# Backend: auto_round:tritonv2_zp (CUDA, Triton-based kernel) -# ============================================================================ - - -@require_torch_cuda_compatibility(7.0) -class AutoRoundW4G128AsymTritonTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, asymmetric, backend='auto_round:tritonv2_zp' (CUDA Triton kernel).""" - - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": False, - "backend": "auto_round:tritonv2_zp", - } - - -@require_torch_cuda_compatibility(7.0) -class AutoRoundW4G128SymTritonTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, symmetric, backend='auto_round:tritonv2_zp'.""" - - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": True, - "backend": "auto_round:tritonv2_zp", - } - - -# ============================================================================ -# Backend: gptqmodel:marlin_zp (CUDA, requires GPTQModel>=5.8.0, best perf) -# ============================================================================ - - -@unittest.skipUnless(_is_gptqmodel_available("5.8.0"), "Test requires gptqmodel>=5.8.0") -@require_torch_cuda_compatibility(8.0) -class AutoRoundW4G128AsymMarlinTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, asymmetric, backend='gptqmodel:marlin_zp' (best CUDA performance).""" - - _test_torch_compile = True - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": False, - "backend": "gptqmodel:marlin_zp", - } - - -@unittest.skipUnless(_is_gptqmodel_available("5.8.0"), "Test requires gptqmodel>=5.8.0") -@require_torch_cuda_compatibility(8.0) -class AutoRoundW4G128SymMarlinTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, symmetric, backend='gptqmodel:marlin_zp'.""" - - _test_torch_compile = True - expected_memory_reduction = 0.55 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": True, - "backend": "gptqmodel:marlin_zp", - } - - -# ============================================================================ -# Backend: auto_round:torch_zp (CPU, pure PyTorch kernel) -# ============================================================================ - - -class AutoRoundW4G128AsymTorchCPUTest(AutoRoundBaseTesterMixin, unittest.TestCase): - """W4A16, group_size=128, asymmetric, backend='auto_round:torch_zp' (CPU).""" - - expected_memory_reduction = 0.50 - - def get_dummy_init_kwargs(self): - return { - "bits": 4, - "group_size": 128, - "sym": False, - "backend": "auto_round:torch_zp", - } - - -# ============================================================================ -# Unit tests: AutoRoundConfig (no hardware required) -# ============================================================================ - - -class AutoRoundConfigTest(unittest.TestCase): - """Unit tests for AutoRoundConfig — no GPU / nightly decorator needed.""" - - def test_defaults(self): - cfg = AutoRoundConfig() - self.assertEqual(cfg.bits, 4) - self.assertEqual(cfg.group_size, 128) - self.assertTrue(cfg.sym) - self.assertEqual(cfg.backend, "auto") - self.assertEqual(cfg.quant_method, QuantizationMethod.AUTOROUND) - - def test_backend_values(self): - """All documented backend strings are stored correctly.""" - for backend in ("auto", "torch", "tritonv2", "marlin", "exllamav2"): - self.assertEqual(AutoRoundConfig(backend=backend).backend, backend) - - def test_to_dict_round_trip(self): - """to_dict → from_dict preserves all fields including backend and extra kwargs.""" - cfg = AutoRoundConfig(bits=4, group_size=32, sym=False, backend="gptqmodel:marlin_zp", - packing_format="auto_round:auto_gptq") - restored = AutoRoundConfig.from_dict(cfg.to_dict()) - self.assertEqual(restored.bits, cfg.bits) - self.assertEqual(restored.group_size, cfg.group_size) - self.assertEqual(restored.sym, cfg.sym) - self.assertEqual(restored.backend, cfg.backend) - self.assertEqual(restored.packing_format, cfg.packing_format) - self.assertEqual(restored.to_dict()["quant_method"], "auto-round") - - -# ============================================================================ -# Unit tests: DiffusersAutoQuantizer.merge_quantization_configs (no hardware) -# ============================================================================ - - -class MergeQuantizationConfigsTest(unittest.TestCase): - """Tests for the merge logic in DiffusersAutoQuantizer.merge_quantization_configs. - - Key behaviours under test: - 1. New fields in quantization_config_from_args (e.g. `backend`) are forwarded to - the merged config when they are absent from the model's saved config. - 2. Fields already present in the model's saved config are NOT overridden. - 3. A warning is emitted when quantization_config_from_args is provided. - 4. No warning when quantization_config_from_args is None. - """ - - def _model_config_dict(self, **overrides): - """Simulate a minimal saved AutoRound quantization_config dict (no 'backend' key).""" - base = { - "quant_method": "auto-round", - "bits": 4, - "group_size": 128, - "sym": True, - "autoround_version": "0.13.0", - "packing_format": "auto_round:auto_gptq", - } - base.update(overrides) - return base - - def test_new_fields_from_args_are_forwarded(self): - """Fields absent from the model config (backend, or arbitrary kwargs) are added from args.""" - for backend in ("marlin", "triton", "torch", "exllamav2"): # tritonv2 equals triton - with self.subTest(backend=backend): - merged = DiffusersAutoQuantizer.merge_quantization_configs( - self._model_config_dict(), AutoRoundConfig(backend=backend) - ) - self.assertEqual(merged.backend, backend) - - def test_existing_fields_not_overridden(self): - """Fields already in model config are NOT overridden; absent fields (backend) ARE added.""" - args_cfg = AutoRoundConfig(bits=2, group_size=32, sym=False, backend="torch") - merged = DiffusersAutoQuantizer.merge_quantization_configs(self._model_config_dict(), args_cfg) - - self.assertEqual(merged.bits, 4) # model value kept - self.assertEqual(merged.group_size, 128) # model value kept - self.assertTrue(merged.sym) # model value kept - self.assertEqual(merged.backend, "torch") # new field added - - def test_warning_behaviour(self): - """Warning emitted with args; no warning without args.""" - model_cfg = self._model_config_dict() - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - DiffusersAutoQuantizer.merge_quantization_configs(model_cfg, AutoRoundConfig(backend="auto")) - self.assertTrue(any("quantization_config" in str(x.message) for x in w)) - - with warnings.catch_warnings(record=True) as w: - warnings.simplefilter("always") - DiffusersAutoQuantizer.merge_quantization_configs(model_cfg, None) - self.assertFalse(any("quantization_config" in str(x.message).lower() for x in w)) From 16a2abb3caa0d54693bfe8612c1f99747c1a9b86 Mon Sep 17 00:00:00 2001 From: Xin He Date: Tue, 19 May 2026 17:05:58 +0800 Subject: [PATCH 09/14] fix compile error in doc Signed-off-by: Xin He --- docs/source/en/quantization/autoround.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md index 3e7cd9fcc541..f4fcf1a780c3 100644 --- a/docs/source/en/quantization/autoround.md +++ b/docs/source/en/quantization/autoround.md @@ -99,7 +99,7 @@ pipe = DiffusionPipeline.from_pretrained( device_map="cuda", ) -pipe.transformer = torch.compile(pipe.transformer, mode="max-autotune", fullgraph=True) +pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=False) ``` ## Backends From 941f41dae7b88072e03fb9365a3d0af8b34e3e5d Mon Sep 17 00:00:00 2001 From: "github-actions[bot]" Date: Tue, 19 May 2026 09:48:46 +0000 Subject: [PATCH 10/14] Apply style fixes --- .../autoround/autoround_quantizer.py | 41 +++++++++---------- .../quantizers/quantization_config.py | 33 ++++++++------- src/diffusers/utils/__init__.py | 2 +- tests/models/testing_utils/quantization.py | 16 +++++--- .../test_models_transformer_z_image.py | 12 ++++-- 5 files changed, 57 insertions(+), 47 deletions(-) diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py index aa5675e83e77..4ae58b91d1ff 100644 --- a/src/diffusers/quantizers/autoround/autoround_quantizer.py +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -27,7 +27,7 @@ if is_torch_available(): - import torch + pass logger = logging.get_logger(__name__) @@ -36,12 +36,12 @@ class AutoRoundQuantizer(DiffusersQuantizer): r""" Diffusers Quantizer for AutoRound (https://github.com/intel/auto-round). - AutoRound is a weight-only quantization method that uses sign gradient descent to jointly optimize - rounding values and min-max ranges for weights. It supports W4A16 (4-bit weight, 16-bit activation) - quantization for efficient inference. + AutoRound is a weight-only quantization method that uses sign gradient descent to jointly optimize rounding values + and min-max ranges for weights. It supports W4A16 (4-bit weight, 16-bit activation) quantization for efficient + inference. - This quantizer only supports loading pre-quantized AutoRound models. On-the-fly quantization - (calibration) is not supported through this interface. + This quantizer only supports loading pre-quantized AutoRound models. On-the-fly quantization (calibration) is not + supported through this interface. """ # AutoRound requires data calibration — we only support loading pre-quantized checkpoints. @@ -53,8 +53,8 @@ def __init__(self, quantization_config, **kwargs): def validate_environment(self, *args, **kwargs): """ - Validates that the auto-round library (>= 0.5) is installed and captures the device_map - for later use during model conversion. + Validates that the auto-round library (>= 0.5) is installed and captures the device_map for later use during + model conversion. """ self.device_map = kwargs.get("device_map", None) if not is_auto_round_available(): @@ -77,17 +77,17 @@ def _process_model_before_weight_loading( **kwargs, ): """ - Replaces target nn.Linear layers with AutoRound's quantized QuantLinear layers before - weights are loaded from the checkpoint. + Replaces target nn.Linear layers with AutoRound's quantized QuantLinear layers before weights are loaded from + the checkpoint. Uses `auto_round.inference.convert_model.convert_hf_model` which: - Inspects the model architecture and the quantization config (bits, group_size, sym, backend). - - Replaces eligible nn.Linear modules with the appropriate QuantLinear variant - (the packed-weight layer that stores qweight, scales, qzeros). + - Replaces eligible nn.Linear modules with the appropriate QuantLinear variant (the packed-weight layer that + stores qweight, scales, qzeros). - Returns the converted model and a set of used backend names. - `infer_target_device` resolves the device_map into a single target device string - that AutoRound uses to select the correct kernel backend (e.g. "cuda", "cpu"). + `infer_target_device` resolves the device_map into a single target device string that AutoRound uses to select + the correct kernel backend (e.g. "cuda", "cpu"). """ from auto_round.inference.convert_model import convert_hf_model, infer_target_device @@ -98,17 +98,17 @@ def _process_model_before_weight_loading( def _process_model_after_weight_loading(self, model, **kwargs): """ - Finalizes the model after all quantized weights (qweight, scales, qzeros, etc.) have - been loaded into the QuantLinear layers. + Finalizes the model after all quantized weights (qweight, scales, qzeros, etc.) have been loaded into the + QuantLinear layers. Uses `auto_round.inference.convert_model.post_init` which: - - Performs backend-specific finalization (e.g. repacking weights into the kernel's - expected memory layout, moving buffers to the correct device). + - Performs backend-specific finalization (e.g. repacking weights into the kernel's expected memory layout, + moving buffers to the correct device). - Freezes quantized parameters (requires_grad=False). - Prepares the model for inference. - Raises ValueError if the model is not pre-quantized, since AutoRound does not support - on-the-fly quantization through this loading path. + Raises ValueError if the model is not pre-quantized, since AutoRound does not support on-the-fly quantization + through this loading path. """ if self.pre_quantized: from auto_round.inference.convert_model import post_init @@ -135,4 +135,3 @@ def is_serializable(self): @property def is_compileable(self) -> bool: return True - diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index dc213a6eebca..0c98e40ba962 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -756,9 +756,8 @@ def get_config_from_quant_type(self) -> dict[str, Any]: class AutoRoundConfig(QuantizationConfigMixin): """Configuration class for AutoRound quantization. - AutoRound is a weight-only quantization algorithm that uses sign gradient descent to - jointly optimize weight rounding and min-max values. This config targets the W4A16 - (4-bit weights, 16-bit activations) setting. + AutoRound is a weight-only quantization algorithm that uses sign gradient descent to jointly optimize weight + rounding and min-max values. This config targets the W4A16 (4-bit weights, 16-bit activations) setting. Reference: https://github.com/intel/auto-round @@ -766,23 +765,23 @@ class AutoRoundConfig(QuantizationConfigMixin): bits (`int`, *optional*, defaults to `4`): The number of bits to quantize weights to. For W4A16 this should be 4. group_size (`int`, *optional*, defaults to `128`): - The group size for weight quantization. Weights in each group share the same - scale and zero-point. Common choices: 32, 64, 128, -1 (per-channel). + The group size for weight quantization. Weights in each group share the same scale and zero-point. Common + choices: 32, 64, 128, -1 (per-channel). sym (`bool`, *optional*, defaults to `True`): - Whether to use symmetric quantization (zero-point fixed at 0) or asymmetric - quantization (zero-point is learned). + Whether to use symmetric quantization (zero-point fixed at 0) or asymmetric quantization (zero-point is + learned). backend (`str`, *optional*, defaults to `"auto"`): The backend kernel to use for quantized inference. Available backends: - `"auto"`: Automatically select the best available backend for the current device. - `"torch"`: Pure PyTorch kernel — works on CPU and CUDA. - `"tritonv2"`: Triton-based kernel — requires CUDA. - - `"exllamav2"`: Exllamav2 kernel via GPTQModel — requires CUDA and - `gptqmodel>=5.8.0`. Offers good CUDA inference performance. - - `"marlin"`: Marlin kernel via GPTQModel — requires CUDA and - `gptqmodel>=5.8.0`. Offers the best CUDA inference performance. + - `"exllamav2"`: Exllamav2 kernel via GPTQModel — requires CUDA and `gptqmodel>=5.8.0`. Offers good CUDA + inference performance. + - `"marlin"`: Marlin kernel via GPTQModel — requires CUDA and `gptqmodel>=5.8.0`. Offers the best CUDA + inference performance. kwargs (`dict[str, Any]`, *optional*): - Additional keyword arguments forwarded to AutoRound (e.g. `iters`, `seqlen`, - `batch_size`, `lr`, `minmax_lr` for calibration when quantizing from scratch). + Additional keyword arguments forwarded to AutoRound (e.g. `iters`, `seqlen`, `batch_size`, `lr`, + `minmax_lr` for calibration when quantizing from scratch). """ VALID_BACKENDS = ["auto", "torch", "tritonv2", "exllamav2", "marlin"] @@ -811,8 +810,8 @@ def _validate_backend(self, backend): def to_dict(self) -> dict: """Serialize the config to a JSON-compatible dict. - Output: A dict containing all config fields. The `quant_method` is stored as - its string value so it can be round-tripped through JSON. + Output: A dict containing all config fields. The `quant_method` is stored as its string value so it can be + round-tripped through JSON. """ output = super().to_dict() output["quant_method"] = output["quant_method"].value @@ -822,8 +821,8 @@ def to_dict(self) -> dict: def from_dict(cls, config_dict: dict, return_unused_kwargs: bool = False, **kwargs): """Instantiate an AutoRoundConfig from a dictionary. - Input: config_dict with keys like bits, group_size, sym, etc. - Output: An AutoRoundConfig instance (and optionally unused kwargs). + Input: config_dict with keys like bits, group_size, sym, etc. Output: An AutoRoundConfig instance (and + optionally unused kwargs). """ # Filter out keys that are not constructor parameters # (e.g. quant_method is set automatically) diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 392b1d7e351e..a65a43ab1a28 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -69,6 +69,7 @@ is_accelerate_version, is_aiter_available, is_aiter_version, + is_auto_round_available, is_av_available, is_better_profanity_available, is_bitsandbytes_available, @@ -88,7 +89,6 @@ is_hpu_available, is_inflect_available, is_invisible_watermark_available, - is_auto_round_available, is_kernels_available, is_kernels_version, is_kornia_available, diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index f51c77c84437..cae2397dfd9b 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -18,9 +18,15 @@ import pytest import torch -from diffusers import AutoRoundConfig, BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, QuantoConfig, TorchAoConfig +from diffusers import ( + AutoRoundConfig, + BitsAndBytesConfig, + GGUFQuantizationConfig, + NVIDIAModelOptConfig, + QuantoConfig, + TorchAoConfig, +) from diffusers.utils.import_utils import ( - is_auto_round_available, is_bitsandbytes_available, is_gguf_available, is_nvidia_modelopt_available, @@ -1399,6 +1405,7 @@ class AutoRoundConfigMixin: - quantized_model_name_or_path: Hub repository ID for the quantized model - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained """ + config_dict = {"backend": "auto"} def _load_unquantized_model(self): @@ -1443,13 +1450,12 @@ class AutoRoundTesterMixin(AutoRoundConfigMixin, QuantizationTesterMixin): Pytest mark: autoround Use `pytest -m "not autoround"` to skip these tests """ + config_dict = {"backend": "auto"} def test_autoround_quantization_memory_footprint(self): expected = 1.5 # AutoRound is a W4A16 method, so we expect around 1.5x memory reduction - self._test_quantization_memory_footprint( - self.config_dict, expected_memory_reduction=expected - ) + self._test_quantization_memory_footprint(self.config_dict, expected_memory_reduction=expected) def test_autoround_quantization_inference(self): self._test_quantization_inference(self.config_dict) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index f8042947bb50..e8d31245220e 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -21,9 +21,9 @@ from diffusers import ZImageTransformer2DModel -from ..testing_utils import AutoRoundTesterMixin, AutoRoundCompileTesterMixin from ...testing_utils import IS_GITHUB_ACTIONS, torch_device from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin +from ..testing_utils import AutoRoundCompileTesterMixin, AutoRoundTesterMixin # Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations @@ -174,6 +174,7 @@ def test_compile_on_different_shapes(self): class ZImageTransformerTesterConfig: """Configuration class for Z-Image Transformer tests.""" + @property def model_class(self): return ZImageTransformer2DModel @@ -199,8 +200,13 @@ def get_dummy_inputs(self): seq_len = 16 torch.manual_seed(0) - x = [torch.randn((in_channels, frames, height, width)).to(torch_device, dtype=torch.bfloat16) for _ in range(batch_size)] - cap_feats = [torch.randn((seq_len, cap_feat_dim)).to(torch_device, dtype=torch.bfloat16) for _ in range(batch_size)] + x = [ + torch.randn((in_channels, frames, height, width)).to(torch_device, dtype=torch.bfloat16) + for _ in range(batch_size) + ] + cap_feats = [ + torch.randn((seq_len, cap_feat_dim)).to(torch_device, dtype=torch.bfloat16) for _ in range(batch_size) + ] t = torch.tensor([0.5]).to(torch_device, dtype=torch.bfloat16) return {"x": x, "cap_feats": cap_feats, "t": t} From eaa496d09849fbc5501262a69b2a798984ca365a Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Mon, 8 Jun 2026 11:51:31 +0530 Subject: [PATCH 11/14] small nits --- .../autoround/autoround_quantizer.py | 33 +++++-------------- 1 file changed, 9 insertions(+), 24 deletions(-) diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py index 4ae58b91d1ff..f80563fed406 100644 --- a/src/diffusers/quantizers/autoround/autoround_quantizer.py +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -1,4 +1,4 @@ -# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# Copyright 2025 The Intel and The HuggingFace Inc. teams. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,11 +14,7 @@ from typing import TYPE_CHECKING -from ...utils import ( - is_auto_round_available, - is_torch_available, - logging, -) +from ...utils import is_auto_round_available, logging from ..base import DiffusersQuantizer @@ -26,9 +22,6 @@ from ...models.modeling_utils import ModelMixin -if is_torch_available(): - pass - logger = logging.get_logger(__name__) @@ -91,10 +84,9 @@ def _process_model_before_weight_loading( """ from auto_round.inference.convert_model import convert_hf_model, infer_target_device - if self.pre_quantized: - target_device = infer_target_device(self.device_map) - model, used_backends = convert_hf_model(model, target_device) - self.used_backends = used_backends + target_device = infer_target_device(self.device_map) + model, used_backends = convert_hf_model(model, target_device) + self.used_backends = used_backends def _process_model_after_weight_loading(self, model, **kwargs): """ @@ -103,22 +95,15 @@ def _process_model_after_weight_loading(self, model, **kwargs): Uses `auto_round.inference.convert_model.post_init` which: - Performs backend-specific finalization (e.g. repacking weights into the kernel's expected memory layout, - moving buffers to the correct device). + moving buffers to the correct device). - Freezes quantized parameters (requires_grad=False). - Prepares the model for inference. - Raises ValueError if the model is not pre-quantized, since AutoRound does not support on-the-fly quantization - through this loading path. """ - if self.pre_quantized: - from auto_round.inference.convert_model import post_init + from auto_round.inference.convert_model import post_init + + post_init(model, self.used_backends) - post_init(model, self.used_backends) - else: - raise ValueError( - "AutoRound quantizer in diffusers only supports pre-quantized models. " - "Please provide a model that has already been quantized with AutoRound." - ) return model @property From 018a989542b5ee08c6399e3e1c33f3dab222c818 Mon Sep 17 00:00:00 2001 From: Xin He Date: Mon, 8 Jun 2026 20:42:36 +0800 Subject: [PATCH 12/14] Add auto_round dependency to the versions table Signed-off-by: Xin He --- src/diffusers/dependency_versions_table.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 747d1011aa40..7db44efc8de9 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -37,6 +37,7 @@ "onnx": "onnx", "optimum_quanto": "optimum_quanto>=0.2.6", "gguf": "gguf>=0.10.0", + "auto_round": "auto-round>=0.13.0", "torchao": "torchao>=0.7.0", "bitsandbytes": "bitsandbytes>=0.43.3", "nvidia_modelopt[hf]": "nvidia_modelopt[hf]>=0.33.1", From 6699d3febf3b16490c214050a81e8daf9a7ad75f Mon Sep 17 00:00:00 2001 From: Xin He Date: Tue, 9 Jun 2026 13:44:50 +0800 Subject: [PATCH 13/14] fix make deps_table_check_updated Signed-off-by: Xin He --- setup.py | 1 + src/diffusers/dependency_versions_table.py | 2 +- 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/setup.py b/setup.py index bc8110bbc594..a9bafaff5399 100644 --- a/setup.py +++ b/setup.py @@ -130,6 +130,7 @@ "onnx", "optimum_quanto>=0.2.6", "gguf>=0.10.0", + "auto-round>=0.13.0", "torchao>=0.7.0", "bitsandbytes>=0.43.3", "nvidia_modelopt[hf]>=0.33.1", diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 7db44efc8de9..3aac2f280af6 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -37,7 +37,7 @@ "onnx": "onnx", "optimum_quanto": "optimum_quanto>=0.2.6", "gguf": "gguf>=0.10.0", - "auto_round": "auto-round>=0.13.0", + "auto-round": "auto-round>=0.13.0", "torchao": "torchao>=0.7.0", "bitsandbytes": "bitsandbytes>=0.43.3", "nvidia_modelopt[hf]": "nvidia_modelopt[hf]>=0.33.1", From cddfc0fbdce1679f97e043dc7de02567a34b9ad5 Mon Sep 17 00:00:00 2001 From: Xin He Date: Tue, 9 Jun 2026 13:53:14 +0800 Subject: [PATCH 14/14] fix CI Signed-off-by: Xin He --- tests/others/test_dependencies.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/others/test_dependencies.py b/tests/others/test_dependencies.py index b2e28077b131..4a33bd529c15 100644 --- a/tests/others/test_dependencies.py +++ b/tests/others/test_dependencies.py @@ -39,6 +39,8 @@ def test_backend_registration(self): backend = "opencv-python" elif backend == "nvidia_modelopt": backend = "nvidia_modelopt[hf]" + elif backend == "auto_round": + backend = "auto-round" assert backend in deps, f"{backend} is not in the deps table!" def test_pipeline_imports(self):