`FqnToConfig` raises inside `TorchAoConfig`

[Beinsezii]

Describe the bug

seems like it's not applied to quantize_ properly

Reproduction

from collections import OrderedDict

import torch
from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
from diffusers.quantizers.quantization_config import TorchAoConfig
from torchao.quantization import FqnToConfig, Int8WeightOnlyConfig, PerGroup, PerTensor

with torch.no_grad():
    _t = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-klein-4B",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
        low_cpu_mem_usage=True,
        quantization_config=TorchAoConfig(
            FqnToConfig(
                OrderedDict(
                    [
                        (
                            r"re:.*transformer_blocks\.\d+\.ff.*",
                            Int8WeightOnlyConfig(version=2, granularity=PerTensor()),
                        ),
                        (
                            r"re:.*transformer_blocks\.\d+\.attn.*",
                            Int8WeightOnlyConfig(version=2, granularity=PerGroup(64)),
                        ),
                        ("_default", None),
                    ]
                )
            )
        ),
    )

Logs

Traceback (most recent call last):
  File "/tmp/reproduce.py", line 9, in <module>
    _t = Flux2Transformer2DModel.from_pretrained(
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/repro/.venv/lib/python3.12/site-packages/huggingface_hub/utils/_validators.py", line 88, in _inner_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/tmp/repro/.venv/lib/python3.12/site-packages/diffusers/models/modeling_utils.py", line 1398, in from_pretrained
    ) = cls._load_pretrained_model(
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/repro/.venv/lib/python3.12/site-packages/diffusers/models/modeling_utils.py", line 1792, in _load_pretrained_model
    offload_index, state_dict_index, _mismatched_keys, _error_msgs = load_fn(shard_file)
                                                                     ^^^^^^^^^^^^^^^^^^^
  File "/tmp/repro/.venv/lib/python3.12/site-packages/diffusers/models/model_loading_utils.py", line 371, in _load_shard_file
    offload_index, state_dict_index = load_model_dict_into_meta(
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/tmp/repro/.venv/lib/python3.12/site-packages/diffusers/models/model_loading_utils.py", line 303, in load_model_dict_into_meta
    hf_quantizer.create_quantized_param(
  File "/tmp/repro/.venv/lib/python3.12/site-packages/diffusers/quantizers/torchao/torchao_quantizer.py", line 366, in create_quantized_param
    quantize_(module, self.quantization_config.get_apply_tensor_subclass())
  File "/tmp/repro/.venv/lib/python3.12/site-packages/torchao/quantization/quant_api.py", line 422, in quantize_
    raise ValueError(
ValueError: Custom filter_fn and FqnToConfig were both specified. Only filter_fn=None is supported when FqnToConfig is specified.

System Info

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- 🤗 Diffusers version: 0.40.0.dev0
- Platform: Linux-7.2.2-arch1-1-x86_64-with-glibc2.44
- Running on Google Colab?: No
- Python version: 3.12.13
- PyTorch version (GPU?): 2.13.0+rocm7.2 (True)
- Huggingface_hub version: 1.24.0
- Transformers version: 5.5.4
- Accelerate version: 1.12.0
- PEFT version: 0.18.1
- bitsandbytes version: 0.49.2
- optimum-quanto version: 0.2.7
- torchao version: 0.17.0+rocm7.2
- Safetensors version: 0.8.0
- xFormers version: not installed
- Accelerator: NA
- Using GPU in script?: optional
- Using distributed or parallel set-up in script?: no

Who can help?

@sayakpaul @DN6

[DN6]

@Beinsezii The traceback seems to point to an issue with TorchAO? Are you able to quantize the model with TorchAO directly?

[Beinsezii]

@Beinsezii The traceback seems to point to an issue with TorchAO? Are you able to quantize the model with TorchAO directly?

@DN6, Yes, the issue is diffusers callsite of quantize_. Transformers already handles this, so passing FqnToConfig to ModularPipeline.load_components only works on text encoders.

direct_quantization.py

from collections import OrderedDict

import torch
from diffusers.models.transformers.transformer_flux2 import Flux2Transformer2DModel
from diffusers.quantizers.quantization_config import TorchAoConfig
from torchao.quantization import (
    FqnToConfig,
    Int8WeightOnlyConfig,
    PerGroup,
    PerTensor,
    quantize_,
)

with torch.no_grad():
    t = Flux2Transformer2DModel.from_pretrained(
        "black-forest-labs/FLUX.2-klein-4B",
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
        device_map="cuda",
        low_cpu_mem_usage=True,
        # # raises because filter_fn is set
        # quantization_config=TorchAoConfig(
        #     FqnToConfig(
        #         OrderedDict(
        #             [
        #                 (
        #                     r"re:.*transformer_blocks\.\d+\.ff.*.weight",
        #                     Int8WeightOnlyConfig(version=2, granularity=PerTensor()),
        #                 ),
        #                 (
        #                     r"re:.*transformer_blocks\.\d+\.attn.*.weight",
        #                     Int8WeightOnlyConfig(version=2, granularity=PerGroup(64)),
        #                 ),
        #                 ("_default", None),
        #             ]
        #         )
        #     )
        # ),
    )
    # direct quantization works fine
    quantize_(
        t,
        FqnToConfig(
            OrderedDict(
                [
                    (
                        r"re:.*transformer_blocks\.\d+\.ff.*.weight",
                        Int8WeightOnlyConfig(version=2, granularity=PerTensor()),
                    ),
                    (
                        r"re:.*transformer_blocks\.\d+\.attn.*.weight",
                        Int8WeightOnlyConfig(version=2, granularity=PerGroup(64)),
                    ),
                    ("_default", None),
                ]
            )
        ),
        filter_fn=None,
    )

i made a hacky patch on a fork that just adds a special case for FqnToConfig but the Transformers source code seems to have a more sophisticated solution so I thought it would be better for a HuggingFace maintainer to implement this.

[sayakpaul]

Checking.

[sayakpaul]

@Beinsezii can you check #14686 and see if that fixes the problem?