[torchao]: fix how FqnConfig is resolved and handled

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)
 

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