diff --git a/src/diffusers/quantizers/torchao/torchao_quantizer.py b/src/diffusers/quantizers/torchao/torchao_quantizer.py index 91aab04975aa..8e781ebb143a 100644 --- a/src/diffusers/quantizers/torchao/torchao_quantizer.py +++ b/src/diffusers/quantizers/torchao/torchao_quantizer.py @@ -51,31 +51,8 @@ import torch import torch.nn as nn - if is_torch_version(">=", "2.5"): - SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION = ( - # At the moment, only int8 is supported for integer quantization dtypes. - # In Torch 2.6, int1-int7 will be introduced, so this can be visited in the future - # to support more quantization methods, such as intx_weight_only. - torch.int8, - torch.float8_e4m3fn, - torch.float8_e5m2, - torch.uint1, - torch.uint2, - torch.uint3, - torch.uint4, - torch.uint5, - torch.uint6, - torch.uint7, - ) - else: - SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION = ( - torch.int8, - torch.float8_e4m3fn, - torch.float8_e5m2, - ) - if is_torchao_available(): - from torchao.quantization import quantize_ + from torchao.quantization import FqnToConfig, quantize_ if is_torchao_version(">=", "0.16.0"): from torchao.prototype.safetensors.safetensors_support import ( @@ -145,6 +122,41 @@ def fuzzy_match_size(config_name: str) -> str | None: return None +def _fqn_to_config_weight_sizes(config: "FqnToConfig") -> tuple[set[str | None], bool]: + """ + Summarize the configs an `FqnToConfig` holds, for the memory estimates that assume one weight size model-wide. + + Returns the size digits (as `fuzzy_match_size` reports them) of every config it maps to, along with whether the + config leaves modules unquantized -- either mapped to `None`, or unmatched with no `_default` to fall back on. + """ + fqn_to_config = config.fqn_to_config + size_digits = {fuzzy_match_size(type(c).__name__) for c in fqn_to_config.values() if c is not None} + leaves_modules_unquantized = "_default" not in fqn_to_config or any(c is None for c in fqn_to_config.values()) + return size_digits, leaves_modules_unquantized + + +def _resolve_fqn_to_config(config: "FqnToConfig", module_fqn: str, param_fqn: str): + """ + Pick the config an `FqnToConfig` assigns to a single parameter. + + `create_quantized_param` quantizes one module at a time, so `quantize_` only ever sees a lone `nn.Linear` whose fqn + is `""` and whose parameters are named `weight`/`bias`. Full-path patterns therefore never match on their own, and + the config silently falls through to `_default`. + """ + fqn_to_config = config.fqn_to_config + + for fqn in (param_fqn, module_fqn): + if fqn in fqn_to_config: + return fqn_to_config[fqn] + + for fqn in (param_fqn, module_fqn): + for pattern, pattern_config in fqn_to_config.items(): + if pattern.startswith("re:") and re.fullmatch(pattern[3:], fqn): + return pattern_config + + return fqn_to_config.get("_default", None) + + def _linear_extra_repr(self): from torchao.utils import TorchAOBaseTensor @@ -228,6 +240,15 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": from accelerate.utils import CustomDtype quant_type = self.quantization_config.quant_type + if isinstance(quant_type, FqnToConfig): + size_digits, leaves_modules_unquantized = _fqn_to_config_weight_sizes(quant_type) + if leaves_modules_unquantized: + # Modules the config skips keep `target_dtype`, so shrinking the estimate for every parameter would + # under-estimate the model and let `infer_auto_device_map` overfill a device. + return target_dtype + # Only claim int4 when nothing the config maps to is wider than that. + return CustomDtype.INT4 if size_digits == {"4"} else torch.int8 + config_name = quant_type.__class__.__name__ size_digit = fuzzy_match_size(config_name) @@ -236,18 +257,6 @@ def adjust_target_dtype(self, target_dtype: "torch.dtype") -> "torch.dtype": else: return torch.int8 - if isinstance(target_dtype, SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION): - return target_dtype - - # We need one of the supported dtypes to be selected in order for accelerate to determine - # the total size of modules/parameters for auto device placement. - possible_device_maps = ["auto", "balanced", "balanced_low_0", "sequential"] - raise ValueError( - f"You have set `device_map` as one of {possible_device_maps} on a TorchAO quantized model but a suitable target dtype " - f"could not be inferred. The supported target_dtypes are: {SUPPORTED_TORCH_DTYPES_FOR_QUANTIZATION}. If you think the " - f"dtype you are using should be supported, please open an issue at https://github.com/huggingface/diffusers/issues." - ) - def adjust_max_memory(self, max_memory: dict[str, int | str]) -> dict[str, int | str]: max_memory = {key: val * 0.9 for key, val in max_memory.items()} return max_memory @@ -372,8 +381,25 @@ def create_quantized_param( module.extra_repr = types.MethodType(_linear_extra_repr, module) else: # As we perform quantization here, the repr of linear layers is set by TorchAO, so we don't have to do it ourselves - module._parameters[tensor_name] = torch.nn.Parameter(param_value).to(device=target_device) - quantize_(module, self.quantization_config.get_apply_tensor_subclass()) + module._parameters[tensor_name] = torch.nn.Parameter(param_value.to(device=target_device)) + + retrieved_config = self.quantization_config.get_apply_tensor_subclass() + if isinstance(retrieved_config, FqnToConfig): + module_fqn = param_name.rsplit(".", 1)[0] if "." in param_name else "" + retrieved_config = _resolve_fqn_to_config(retrieved_config, module_fqn, param_name) + if retrieved_config is None: + # This module is either explicitly excluded or unmatched with no `_default`, so it stays unquantized. + return + if isinstance(retrieved_config, FqnToConfig): + # `quantize_` matches an `FqnToConfig` against the fqns of the module it is handed, which here is a + # lone `nn.Linear`. The resolution above is what keeps fqn targeting working, so a config that is + # still an `FqnToConfig` at this point would quantize nothing at all rather than erroring. + raise ValueError( + f"Nested `FqnToConfig` entries are not supported (resolved from `{param_name}`). Map each fqn " + f"to a quantization config, or to `None` to leave the matching modules unquantized." + ) + # `retrieved_config` is a plain config by this point, so `quantize_` can use its default `filter_fn`. + quantize_(module, retrieved_config) def get_cuda_warm_up_factor(self): """ @@ -391,6 +417,12 @@ def get_cuda_warm_up_factor(self): - Use a division factor of 4 for int8 weights """ quant_type = self.quantization_config.quant_type + if isinstance(quant_type, FqnToConfig): + # Pre-allocating more than the model ends up using can OOM the warmup itself, so assume the narrowest + # weights this config can produce. Modules it leaves unquantized only make the estimate safer. + size_digits, _ = _fqn_to_config_weight_sizes(quant_type) + return 8 if "4" in size_digits else 4 + config_name = quant_type.__class__.__name__ size_digit = fuzzy_match_size(config_name) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 8285704ec2a5..65f904bd4948 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -14,6 +14,7 @@ # limitations under the License. import gc +import re import pytest import safetensors.torch @@ -796,6 +797,10 @@ class TorchAoConfigMixin: @staticmethod def _get_quant_config(config_name, modules_to_not_convert=None): + # Quant types that need constructor arguments (e.g. `FqnToConfig`) are passed in already built. + if not isinstance(config_name, str): + return TorchAoConfig(config_name, modules_to_not_convert=modules_to_not_convert) + config_cls = getattr(_torchao_quantization, config_name) config_kwargs = {"version": 2} # version=2 int4 defaults to the "plain" packing format, which routes through the @@ -1011,6 +1016,59 @@ def test_torchao_training(self): def test_torchao_keep_modules_in_fp32(self): self._test_keep_modules_in_fp32(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) + def test_torchao_fqn_to_config(self): + """An `FqnToConfig` quant type must reach `quantize_` instead of tripping its `filter_fn` guard. + + See https://github.com/huggingface/diffusers/issues/14667. + """ + quant_type = _torchao_quantization.FqnToConfig( + {"_default": _torchao_quantization.Int8WeightOnlyConfig(version=2)} + ) + self._test_quantized_layers(quant_type) + + def test_torchao_fqn_to_config_targets_named_layers(self): + """Selective `FqnToConfig` targeting must land on the same layers as a plain `quantize_` call. + + The quantizer resolves fqns one `nn.Linear` at a time, so `quantization_config=` is checked against + torchao's own whole-model pass. See https://github.com/huggingface/diffusers/issues/14667. + """ + keep_in_fp32 = getattr(self.model_class, "_keep_in_fp32_modules", None) or [] + reference = self._load_unquantized_model().to(torch_device) + linear_fqns = [ + name + for name, module in reference.named_modules() + if isinstance(module, torch.nn.Linear) and not any(fp32_name in name for fp32_name in keep_in_fp32) + ] + if len(linear_fqns) < 2: + pytest.skip("Model does not have enough linear layers to test selective fqn quantization") + + # One layer selected by its exact fqn, a set of them by a wildcard regex over the model's own fqns + # (the shape from the issue https://github.com/huggingface/diffusers/issues/14667), everything else + # left alone by `_default: None`. On Flux2, this gives + # exact_fqn="time_guidance_embed.timestep_embedder.linear_1" and + # pattern=".*proj_out" (matching "proj_out"). + exact_fqn = linear_fqns[0] + pattern = f".*{re.escape(linear_fqns[-1].split('.')[-1])}" + fqn_to_config = { + exact_fqn: _torchao_quantization.Int8WeightOnlyConfig(version=2), + f"re:{pattern}": _torchao_quantization.Int8WeightOnlyConfig(version=2), + "_default": None, + } + + # `quantize_` on the whole model is the reference: `quantization_config=` must place the same tensors. + _torchao_quantization.quantize_(reference, _torchao_quantization.FqnToConfig(fqn_to_config), filter_fn=None) + expected = {name: type(reference.get_submodule(name).weight) for name in linear_fqns} + del reference + assert set(expected.values()) != {torch.nn.Parameter}, "`quantize_` did not quantize any of the named layers" + + model = self._create_quantized_model(_torchao_quantization.FqnToConfig(fqn_to_config)) + quantized = {name: type(model.get_submodule(name).weight) for name in linear_fqns} + + mismatched = { + name: (quantized[name], expected[name]) for name in linear_fqns if quantized[name] != expected[name] + } + assert not mismatched, f"Layers quantized through `quantization_config=` differ from `quantize_`: {mismatched}" + @is_quantization @is_gguf