diff --git a/src/diffusers/loaders/single_file_model.py b/src/diffusers/loaders/single_file_model.py index 610e377d3ab9..7e2f934e1b63 100644 --- a/src/diffusers/loaders/single_file_model.py +++ b/src/diffusers/loaders/single_file_model.py @@ -416,7 +416,9 @@ def from_single_file(cls, pretrained_model_link_or_path_or_dict: str | None = No if quantization_config is not None: if not low_cpu_mem_usage: raise ValueError("`low_cpu_mem_usage` cannot be False when using quantization.") - hf_quantizer = DiffusersAutoQuantizer.from_config(quantization_config) + hf_quantizer = DiffusersAutoQuantizer.from_config( + quantization_config, pre_quantized=getattr(quantization_config, "pre_quantized", False) + ) hf_quantizer.validate_environment() torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype) checkpoint = hf_quantizer.maybe_update_state_dict(checkpoint) diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 377b46b7b361..8b6df8c8e099 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -160,11 +160,10 @@ def merge_quantization_configs( 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 + # Only `backend` is runtime-only; bits/group_size/sym/etc. describe how the + # checkpoint was quantized and must come from the checkpoint's own config. + for key in ("backend",): + value = getattr(quantization_config_from_args, key, None) 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} " diff --git a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py index 63d2680b3d33..883bc1236369 100644 --- a/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py +++ b/src/diffusers/quantizers/bitsandbytes/bnb_quantizer.py @@ -95,7 +95,7 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": if target_dtype != torch.int8: from accelerate.utils import CustomDtype - logger.info("target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization") + logger.info(f"target_dtype {target_dtype} is replaced by `CustomDtype.INT4` for 4-bit BnB quantization") return CustomDtype.INT4 else: raise ValueError(f"Wrong `target_dtype` ({target_dtype}) provided.") @@ -207,7 +207,7 @@ def check_quantized_param_shape(self, param_name, current_param, loaded_param): inferred_shape = (n,) if "bias" in param_name else ((n + 1) // 2, 1) if loaded_param_shape != inferred_shape: raise ValueError( - f"Expected the flattened shape of the current param ({param_name}) to be {loaded_param_shape} but is {inferred_shape}." + f"Expected the flattened shape of the current param ({param_name}) to be {inferred_shape} but is {loaded_param_shape}." ) else: return True @@ -241,8 +241,7 @@ def update_device_map(self, device_map): device_map = {"": current_device} logger.info( "The device_map was not initialized. " - "Setting device_map to {" - ": {current_device}}. " + f"Setting device_map to {{'': {current_device}}}. " "If you want to use the model for inference, please set device_map ='auto' " ) return device_map @@ -263,6 +262,8 @@ def _process_model_before_weight_loading( if not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] + else: + self.modules_to_not_convert = list(self.modules_to_not_convert) self.modules_to_not_convert.extend(keep_in_fp32_modules) @@ -441,8 +442,7 @@ def update_device_map(self, device_map): device_map = {"": current_device} logger.info( "The device_map was not initialized. " - "Setting device_map to {" - ": {current_device}}. " + f"Setting device_map to {{'': {current_device}}}. " "If you want to use the model for inference, please set device_map ='auto' " ) return device_map @@ -465,7 +465,7 @@ def check_if_quantized_param( module, tensor_name = get_module_from_name(model, param_name) if isinstance(module._parameters.get(tensor_name, None), bnb.nn.Int8Params): if self.pre_quantized: - if param_name.replace("weight", "SCB") not in state_dict.keys(): + if param_name[: -len("weight")] + "SCB" not in state_dict.keys(): raise ValueError("Missing quantization component `SCB`") if param_value.dtype != torch.int8: raise ValueError( @@ -486,8 +486,8 @@ def create_quantized_param( ): import bitsandbytes as bnb - fp16_statistics_key = param_name.replace("weight", "SCB") - fp16_weights_format_key = param_name.replace("weight", "weight_format") + fp16_statistics_key = param_name[: -len("weight")] + "SCB" + fp16_weights_format_key = param_name[: -len("weight")] + "weight_format" fp16_statistics = state_dict.get(fp16_statistics_key, None) fp16_weights_format = state_dict.get(fp16_weights_format_key, None) @@ -550,6 +550,8 @@ def _process_model_before_weight_loading( if not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] + else: + self.modules_to_not_convert = list(self.modules_to_not_convert) self.modules_to_not_convert.extend(keep_in_fp32_modules) diff --git a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py index b7cadb50e54b..d64e18014ea6 100644 --- a/src/diffusers/quantizers/modelopt/modelopt_quantizer.py +++ b/src/diffusers/quantizers/modelopt/modelopt_quantizer.py @@ -161,8 +161,10 @@ def _process_model_before_weight_loading( if modules_to_not_convert is None: modules_to_not_convert = [] - if isinstance(modules_to_not_convert, str): + elif isinstance(modules_to_not_convert, str): modules_to_not_convert = [modules_to_not_convert] + else: + modules_to_not_convert = list(modules_to_not_convert) modules_to_not_convert.extend(keep_in_fp32_modules) if self.quantization_config.disable_conv_quantization: modules_to_not_convert.extend(self.get_conv_param_names(model)) diff --git a/src/diffusers/quantizers/pipe_quant_config.py b/src/diffusers/quantizers/pipe_quant_config.py index cfb5803ee0ca..ee4751cae89e 100644 --- a/src/diffusers/quantizers/pipe_quant_config.py +++ b/src/diffusers/quantizers/pipe_quant_config.py @@ -101,7 +101,7 @@ def _validate_init_kwargs_in_backends(self): init_kwargs_diffusers = inspect.signature(quant_config_mapping_diffusers[quant_backend].__init__) init_kwargs_diffusers = {name for name in init_kwargs_diffusers.parameters if name != "self"} - if init_kwargs_transformers != init_kwargs_diffusers: + if init_kwargs_transformers is not None and init_kwargs_transformers != init_kwargs_diffusers: raise ValueError( "The signatures of the __init__ methods of the quantization config classes in `diffusers` and `transformers` don't match. " f"Please provide a `quant_mapping` instead, in the {self.__class__.__name__} class. Refer to [the docs](https://huggingface.co/docs/diffusers/main/en/quantization/overview#pipeline-level-quantization) to learn more about how " diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index 5777ab45fa7a..21bda08ffb76 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -27,7 +27,6 @@ import json import os import warnings -from dataclasses import dataclass from enum import Enum from typing import Any, Callable @@ -54,7 +53,6 @@ class QuantizationMethod(str, Enum): SDNQ = "sdnq" -@dataclass class QuantizationConfigMixin: """ Mixin class for quantization config @@ -120,6 +118,13 @@ def to_dict(self) -> dict[str, Any]: """ return copy.deepcopy(self.__dict__) + def to_diff_dict(self) -> dict[str, Any]: + """ + Serializes this instance to a Python dictionary. Returns: + `dict[str, Any]`: Dictionary of all the attributes that make up this configuration instance. + """ + return self.to_dict() + def __iter__(self): """allows `dict(obj)` for situations where obj may be a dict or QuantizationConfigMixin""" for attr, value in copy.deepcopy(self.__dict__).items(): @@ -169,7 +174,6 @@ def update(self, **kwargs): return unused_kwargs -@dataclass class BitsAndBytesConfig(QuantizationConfigMixin): """ This is a wrapper class about all possible attributes and features that you can play with a model that has been @@ -409,7 +413,6 @@ def to_diff_dict(self) -> dict[str, Any]: return serializable_config_dict -@dataclass class GGUFQuantizationConfig(QuantizationConfigMixin): """This is a config class for GGUF Quantization techniques. @@ -431,8 +434,12 @@ def __init__(self, compute_dtype: "torch.dtype" | None = None): if self.compute_dtype is None: self.compute_dtype = torch.float32 + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + output["compute_dtype"] = str(output["compute_dtype"]).split(".")[1] + return output + -@dataclass class BFLQuantizationConfig(QuantizationConfigMixin): """Config for loading checkpoints quantized with the Black Forest Labs `fp8r` and `nvfp4` schemes. @@ -455,8 +462,12 @@ def __init__(self, compute_dtype: "torch.dtype" | None = None): self.compute_dtype = compute_dtype if compute_dtype is not None else torch.bfloat16 self.pre_quantized = True + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + output["compute_dtype"] = str(output["compute_dtype"]).split(".")[1] + return output + -@dataclass class NunchakuLiteQuantizationConfig(QuantizationConfigMixin): """Configuration for loading Nunchaku Lite checkpoints. @@ -570,6 +581,8 @@ def post_init(self): ) elif precision != "int4": raise ValueError("Nunchaku AWQ target requires precision='int4'.") + elif group_size != 64: + raise ValueError(f"Nunchaku AWQ config requires group_size=64, got {group_size}.") def to_dict(self) -> dict[str, Any]: output = super().to_dict() @@ -577,7 +590,6 @@ def to_dict(self) -> dict[str, Any]: return output -@dataclass class TorchAoConfig(QuantizationConfigMixin): """This is a config class for torchao quantization/sparsity techniques. @@ -672,7 +684,6 @@ def __repr__(self): return f"{self.__class__.__name__} {json.dumps(config_dict, indent=2, sort_keys=True)}\n" -@dataclass class QuantoConfig(QuantizationConfigMixin): """ This is a wrapper class about all possible attributes and features that you can play with a model that has been @@ -716,7 +727,6 @@ def post_init(self): raise ValueError(f"Only support weights in {accepted_weights} but found {self.weights_dtype}") -@dataclass class NVIDIAModelOptConfig(QuantizationConfigMixin): """This is a config class to use nvidia modelopt for quantization. @@ -797,6 +807,11 @@ def __init__( self.modelopt_config = self.get_config_from_quant_type() if not modelopt_config else modelopt_config self.disable_conv_quantization = disable_conv_quantization + def to_dict(self) -> dict[str, Any]: + output = super().to_dict() + output.pop("forward_loop", None) + return output + def check_model_patching(self, operation: str = "loading"): # ModelOpt imports diffusers internally. This is here to prevent circular imports from modelopt.torch.opt.plugins.huggingface import _PATCHED_CLASSES @@ -907,7 +922,6 @@ def get_config_from_quant_type(self) -> dict[str, Any]: return BASE_CONFIG -@dataclass class AutoRoundConfig(QuantizationConfigMixin): """Configuration class for AutoRound quantization. diff --git a/src/diffusers/quantizers/quanto/quanto_quantizer.py b/src/diffusers/quantizers/quanto/quanto_quantizer.py index 92bbecbb1b53..858210311ac0 100644 --- a/src/diffusers/quantizers/quanto/quanto_quantizer.py +++ b/src/diffusers/quantizers/quanto/quanto_quantizer.py @@ -165,6 +165,8 @@ def _process_model_before_weight_loading( if not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] + else: + self.modules_to_not_convert = list(self.modules_to_not_convert) self.modules_to_not_convert.extend(keep_in_fp32_modules) diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index e31db9faa068..a254febffed9 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -443,6 +443,8 @@ def _process_model_before_weight_loading( if not isinstance(self.modules_to_not_convert, list): self.modules_to_not_convert = [self.modules_to_not_convert] + else: + self.modules_to_not_convert = list(self.modules_to_not_convert) self.modules_to_not_convert.extend(keep_in_fp32_modules) diff --git a/tests/quantization/bnb/test_mixed_int8.py b/tests/quantization/bnb/test_mixed_int8.py index d37099ff623c..f73726ab217b 100644 --- a/tests/quantization/bnb/test_mixed_int8.py +++ b/tests/quantization/bnb/test_mixed_int8.py @@ -69,3 +69,36 @@ def test_bnb_8bit_logs_warning_for_no_quantization(self): "You are loading your model in 8bit or 4bit but no linear modules were found in your model." in cap_logger.out ) + + def test_bnb_8bit_modules_to_not_convert_does_not_mutate_config(self): + from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB8BitDiffusersQuantizer + + skip_modules = ["proj_out"] + quantization_config = BitsAndBytesConfig(load_in_8bit=True, llm_int8_skip_modules=skip_modules) + model = torch.nn.Sequential(torch.nn.Linear(4, 4), torch.nn.Linear(4, 4)) + + for _ in range(2): + quantizer = BnB8BitDiffusersQuantizer(quantization_config) + quantizer._process_model_before_weight_loading(model, device_map=None, keep_in_fp32_modules=["norm_out"]) + + assert quantization_config.llm_int8_skip_modules == skip_modules + + def test_bnb_8bit_scb_lookup_is_suffix_safe(self): + import bitsandbytes as bnb + + from diffusers.quantizers.bitsandbytes.bnb_quantizer import BnB8BitDiffusersQuantizer + + quantization_config = BitsAndBytesConfig(load_in_8bit=True) + quantizer = BnB8BitDiffusersQuantizer(quantization_config) + quantizer.pre_quantized = True + + # "weight" appears outside the trailing ".weight" param, which a naive str.replace corrupts. + model = torch.nn.Module() + model.weight_proj = torch.nn.Linear(4, 4, bias=False) + model.weight_proj.weight = bnb.nn.Int8Params(torch.zeros(4, 4, dtype=torch.int8), requires_grad=False) + + param_name = "weight_proj.weight" + state_dict = {"weight_proj.SCB": torch.zeros(4, dtype=torch.float32)} + assert quantizer.check_if_quantized_param( + model, param_value=torch.zeros(4, 4, dtype=torch.int8), param_name=param_name, state_dict=state_dict + ) diff --git a/tests/quantization/test_quantization_config.py b/tests/quantization/test_quantization_config.py new file mode 100644 index 000000000000..5d6639f3ee3b --- /dev/null +++ b/tests/quantization/test_quantization_config.py @@ -0,0 +1,88 @@ +# coding=utf-8 +# Copyright 2026 The HuggingFace Team Inc. +# +# 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. +"""CPU-only tests for quantization config logic that doesn't need a backend installed.""" + +import json + +import pytest + +from diffusers.quantizers.auto import DiffusersAutoQuantizer +from diffusers.quantizers.quantization_config import ( + AutoRoundConfig, + GGUFQuantizationConfig, + NunchakuLiteQuantizationConfig, + NVIDIAModelOptConfig, + QuantoConfig, +) + + +class TestAutoRoundConfigMerge: + def test_only_backend_is_overridden_from_args(self): + checkpoint_config = AutoRoundConfig(bits=4, group_size=32, sym=False, backend="torch") + args_config = AutoRoundConfig(backend="auto") + + merged = DiffusersAutoQuantizer.merge_quantization_configs(checkpoint_config, args_config) + + assert merged.backend == "auto" + assert merged.bits == 4 + assert merged.group_size == 32 + assert merged.sym is False + + +class TestNVIDIAModelOptConfigToDict: + def test_to_dict_drops_forward_loop_callable(self): + config = NVIDIAModelOptConfig( + quant_type="FP8", forward_loop=lambda model: model, modelopt_config={"quant_cfg": {}} + ) + + output = config.to_dict() + + assert "forward_loop" not in output + json.dumps(output) + + +class TestQuantizationConfigEqAndRepr: + def test_gguf_config_eq_does_not_ignore_fields(self): + import torch + + a = GGUFQuantizationConfig(compute_dtype=torch.float16) + b = GGUFQuantizationConfig(compute_dtype=torch.bfloat16) + + assert a != b + + def test_gguf_config_repr_shows_fields(self): + import torch + + config = GGUFQuantizationConfig(compute_dtype=torch.float16) + assert "float16" in repr(config) + + def test_quanto_config_to_json_string_does_not_raise(self): + config = QuantoConfig(weights_dtype="int8") + json_string = config.to_json_string() + assert "int8" in json_string + + +class TestNunchakuAwqGroupSize: + def test_awq_rejects_group_size_other_than_64(self): + with pytest.raises(ValueError, match="group_size=64"): + NunchakuLiteQuantizationConfig( + awq_w4a16={"precision": "int4", "group_size": 128, "targets": ["final_linear"]} + ) + + def test_awq_accepts_group_size_64(self): + config = NunchakuLiteQuantizationConfig( + awq_w4a16={"precision": "int4", "group_size": 64, "targets": ["final_linear"]} + ) + assert config.awq_w4a16["group_size"] == 64 diff --git a/tests/quantization/torchao/test_torchao.py b/tests/quantization/torchao/test_torchao.py index 6666569187cb..4239f4b2f50e 100644 --- a/tests/quantization/torchao/test_torchao.py +++ b/tests/quantization/torchao/test_torchao.py @@ -14,9 +14,12 @@ # limitations under the License. +from types import SimpleNamespace + import pytest from diffusers import TorchAoConfig +from diffusers.quantizers.torchao.torchao_quantizer import TorchAoHfQuantizer from ...testing_utils import ( enable_full_determinism, @@ -73,3 +76,13 @@ def test_repr(self): quantization_repr = repr(quantization_config) assert "TorchAoConfig" in quantization_repr assert "torchao" in quantization_repr + + def test_process_model_before_weight_loading_does_not_mutate_config_list(self): + # The same config can be reused across multiple loads. + quantization_config = TorchAoConfig(Int4WeightOnlyConfig(), modules_to_not_convert=["a"]) + quantizer = TorchAoHfQuantizer(quantization_config) + model = SimpleNamespace(config=SimpleNamespace()) + + quantizer._process_model_before_weight_loading(model, device_map=None, keep_in_fp32_modules=["b"]) + + assert quantization_config.modules_to_not_convert == ["a"]