diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 1d6dbb4a301c..25b9f0ec2fbe 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 new file mode 100644 index 000000000000..f4fcf1a780c3 --- /dev/null +++ b/docs/source/en/quantization/autoround.md @@ -0,0 +1,206 @@ + + +# AutoRound + +[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): + +```bash +pip install "auto-round>=0.13.0" +``` + +To use the Marlin kernel for faster CUDA inference, install `gptqmodel`: + +```bash +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. + +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 +from diffusers import ZImageTransformer2DModel, ZImagePipeline, AutoRoundConfig + +model_id = "INCModel/Z-Image-W4A16-AutoRound" + +quantization_config = AutoRoundConfig(backend="auto") +transformer = ZImageTransformer2DModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=quantization_config, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +pipe = ZImagePipeline.from_pretrained( + model_id, + transformer=transformer, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +image = pipe("a cat holding a sign that says hello").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. + +## 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="default", fullgraph=False) +``` + +## 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: + +| Backend | Value | Device | Requirements | Notes | +|---------|-------|--------|--------------|-------| +| **Auto** | `"auto"` | Any | — | Default. Automatically selects the best available backend. | +| **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() + +# Explicit Triton backend for CUDA +config = AutoRoundConfig(backend="tritonv2") + +# Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) +config = AutoRoundConfig(backend="marlin") + +# ExllamaV2 backend for good CUDA performance (requires gptqmodel>=5.8.0) +config = AutoRoundConfig(backend="exllamav2") + +# PyTorch backend for CPU/CUDA inference +config = AutoRoundConfig(backend="torch") +``` + + +## Save and load + + + + +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( + "Tongyi-MAI/Z-Image", + scheme="W4A16", # W4G128 symmetric + enable_torch_compile=True, + num_inference_steps=3, + guidance_scale=7.5, + dataset="coco2014", +) +autoround.quantize_and_save("Z-Image-W4A16-AutoRound") +``` + +For more details on calibration options, see the [AutoRound documentation](https://github.com/intel/auto-round). + + + + + +```python +import torch +from diffusers import ZImageTransformer2DModel, ZImagePipeline + +model_id = "INCModel/Z-Image-W4A16-AutoRound" + +# The inference backend will be automatically selected. +pipe = ZImagePipeline.from_pretrained( + model_id, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +image = pipe("a cat holding a sign that says hello").images[0] +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/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/__init__.py b/src/diffusers/__init__.py index 4a2c3bca5bcc..0f4eb50a709a 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, @@ -123,6 +124,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() @@ -982,6 +995,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/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 747d1011aa40..3aac2f280af6 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", diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 6cd24c459c9d..a10bf0cdcb3f 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, } @@ -143,6 +147,19 @@ def merge_quantization_configs( if isinstance(quantization_config, NVIDIAModelOptConfig): quantization_config.check_model_patching() + 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 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/__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..f80563fed406 --- /dev/null +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -0,0 +1,122 @@ +# 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. +# 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, logging +from ..base import DiffusersQuantizer + + +if TYPE_CHECKING: + from ...models.modeling_utils import ModelMixin + + +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.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, + 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 + + 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): + """ + 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. + + """ + from auto_round.inference.convert_model import post_init + + post_init(model, self.used_backends) + + 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 + + @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 c3d829fde8cf..0c98e40ba962 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,81 @@ 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 `True`): + 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. + 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). + """ + + VALID_BACKENDS = ["auto", "torch", "tritonv2", "exllamav2", "marlin"] + + def __init__( + self, + bits: int = 4, + group_size: int = 128, + sym: bool = True, + backend: str = "auto", + **kwargs, + ) -> None: + self.quant_method = QuantizationMethod.AUTOROUND + self._validate_backend(backend) + self.bits = bits + self.group_size = group_size + self.sym = sym + self.backend = backend + for k, v in kwargs.items(): + setattr(self, k, v) + + def _validate_backend(self, 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. + + 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 10ad75d92f17..5cd6885e0364 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, 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 ce439bfecbf2..a0fa882d2705 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -232,6 +232,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") _flashpack_available, _flashpack_version = _is_package_available("flashpack") _av_available, _av_version = _is_package_available("av") @@ -404,6 +405,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 diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 0b31342ffd4a..728a7ac80248 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -19,6 +19,9 @@ from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin from .quantization import ( + AutoRoundCompileTesterMixin, + AutoRoundConfigMixin, + AutoRoundTesterMixin, BitsAndBytesCompileTesterMixin, BitsAndBytesConfigMixin, BitsAndBytesTesterMixin, @@ -44,6 +47,8 @@ __all__ = [ "AttentionBackendTesterMixin", "AttentionTesterMixin", + "AutoRoundConfigMixin", + "AutoRoundTesterMixin", "BaseModelTesterConfig", "BitsAndBytesCompileTesterMixin", "BitsAndBytesConfigMixin", diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index ded5cab52268..cae2397dfd9b 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -18,7 +18,14 @@ 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_bitsandbytes_available, is_gguf_available, @@ -31,6 +38,7 @@ backend_empty_cache, backend_max_memory_allocated, backend_reset_peak_memory_stats, + is_autoround, is_bitsandbytes, is_gguf, is_modelopt, @@ -40,6 +48,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, @@ -1183,7 +1192,7 @@ def teardown_method(self): 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. @@ -1196,7 +1205,7 @@ def _test_torch_compile(self, config_kwargs): model.compile(fullgraph=True) - 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] assert output is not None, "Model output is None" @@ -1375,3 +1384,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..e8d31245220e 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -23,6 +23,7 @@ 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 @@ -169,3 +170,122 @@ 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.""" + + @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/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): diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 6d6df8b24d1e..a8306b3d65f8 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, @@ -449,6 +450,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: @@ -836,6 +846,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(