remove str option for quantization config in torchao by howardzhang-cv · Pull Request #13291 · huggingface/diffusers

howardzhang-cv · 2026-03-19T20:31:22Z

What does this PR do?

Remove the deprecated string-based quant_type path from TorchAoConfig, requiring AOBaseConfig instances instead.

TorchAoConfig.init now only accepts AOBaseConfig subclass instances (e.g. Int8WeightOnlyConfig()) and raises TypeError for strings
Deleted ~200 lines of dead code: _get_torchao_quant_type_to_method, _is_xpu_or_cuda_capability_atleast_8_9, TorchAoJSONEncoder, and all string-parsing branches in post_init, to_dict, from_dict, get_apply_tensor_subclass
Simplified torchao_quantizer.py: removed string-based branches in update_torch_dtype, adjust_target_dtype, get_cuda_warm_up_factor; fixed is_trainable which would crash on AOBaseConfig objects
Converted all test cases from string quant types to their AOBaseConfig equivalents; removed test_floatx_quantization (no replacement for fpx_weight_only)
Updated docs to show only AOBaseConfig-based usage

Testing

python -m pytest tests/quantization/torchao/test_torchao.py -xvs

Who can review?

@sayakpaul

jerryzh168 · 2026-03-19T20:37:56Z

src/diffusers/quantizers/torchao/torchao_quantizer.py

+        if isinstance(quant_type, AOBaseConfig):
+            # Extract size digit using fuzzy match on the class name
+            config_name = quant_type.__class__.__name__
+            size_digit = fuzzy_match_size(config_name)
+
+            # Map the extracted digit to appropriate dtype
+            if size_digit == "4":
                return CustomDtype.INT4
-            elif quant_type == "uintx_weight_only":
-                return self.quantization_config.quant_type_kwargs.get("dtype", torch.uint8)
-            elif quant_type.startswith("uint"):
-                return {
-                    1: torch.uint1,
-                    2: torch.uint2,
-                    3: torch.uint3,
-                    4: torch.uint4,
-                    5: torch.uint5,
-                    6: torch.uint6,
-                    7: torch.uint7,
-                }[int(quant_type[4])]
-            elif quant_type.startswith("float") or quant_type.startswith("fp"):
-                return torch.bfloat16
-
-        elif is_torchao_version(">", "0.9.0"):
-            from torchao.core.config import AOBaseConfig
-
-            quant_type = self.quantization_config.quant_type
-            if isinstance(quant_type, AOBaseConfig):
-                # Extract size digit using fuzzy match on the class name
-                config_name = quant_type.__class__.__name__
-                size_digit = fuzzy_match_size(config_name)
-
-                # Map the extracted digit to appropriate dtype
-                if size_digit == "4":
-                    return CustomDtype.INT4
-                else:
-                    # Default to int8
-                    return torch.int8
+            else:
+                # Default to int8
+                return torch.int8


this seems a bit fragile, I think it's from transformers originally, not sure if this is still needed in transformers though, it might have been refactored after 5.0 update

Sorry might be a dumb question, what part are you referring to that's from transformers originally? Is it the entire adjust_target_dtype function?

jerryzh168 · 2026-03-19T20:38:49Z

src/diffusers/quantizers/torchao/torchao_quantizer.py

-        for pattern, target_dtype in map_to_target_dtype.items():
-            if fnmatch(quant_type, pattern):
-                return target_dtype
+        if isinstance(quant_type, AOBaseConfig):


same here, when do we use this? cc @sayakpaul

src/diffusers/quantizers/torchao/torchao_quantizer.py

jerryzh168 · 2026-03-19T20:40:20Z

src/diffusers/quantizers/quantization_config.py

-                    f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the "
-                    f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues."
-                )
+        if is_torchao_version("<=", "0.9.0"):


separate PR: I feel we should just have a single assertion for torchao to be a relatively recent version (e.g. 0.15) and remove all these version checks

Actually yeah I was going to ask you about that as well. There's a couple version checks scattered around right now. Would be cleaner to just remove all of them.

Yeah we should mandate a minimum version requirement here.

Do we want to wait to do this in a separate PR? I changed and set it to 0.9.0 because that's when AOBaseConfig was supported. Moving to 0.15.0 might be cleaner in a separate PR in case we need to revert for whatever reason?

sayakpaul

Thanks a lot for starting this work!

I think we can merge this fairly soon.

sayakpaul · 2026-03-20T04:37:45Z

src/diffusers/quantizers/torchao/torchao_quantizer.py

-                return torch.int8
-            elif quant_type.startswith("int4"):
+        quant_type = self.quantization_config.quant_type
+        if isinstance(quant_type, AOBaseConfig):


What's the fallback path? Would it be better to mandate that the users pass a valid AOBaseConfig instance?

sayakpaul · 2026-03-20T04:38:12Z

src/diffusers/quantizers/torchao/torchao_quantizer.py

-        for pattern, target_dtype in map_to_target_dtype.items():
-            if fnmatch(quant_type, pattern):
-                return target_dtype
+        if isinstance(quant_type, AOBaseConfig):


sayakpaul · 2026-03-20T04:39:12Z

src/diffusers/quantizers/quantization_config.py

-                    f"Requested quantization type: {self.quant_type} is not supported or is an incorrect `quant_type` name. If you think the "
-                    f"provided quantization type should be supported, please open an issue at https://github.com/huggingface/diffusers/issues."
-                )
+        if is_torchao_version("<=", "0.9.0"):


Yeah we should mandate a minimum version requirement here.

sayakpaul · 2026-03-20T04:39:27Z

src/diffusers/quantizers/quantization_config.py

+        if not isinstance(self.quant_type, AOBaseConfig):
+            raise TypeError(f"quant_type must be an AOBaseConfig instance, got {type(self.quant_type).__name__}")


sayakpaul · 2026-03-20T04:40:59Z

src/diffusers/quantizers/quantization_config.py

+        # For now we assume there is 1 config per Transformer, however in the future
+        # We may want to support a config per fqn.
+        d["quant_type"] = {"default": config_to_dict(self.quant_type)}


This would be really nice! I think we should also provide a reference link to the TorchAO docs to remind ourselves what that granularity means.

sayakpaul · 2026-03-20T04:48:53Z

@bot /style

github-actions · 2026-03-20T04:49:19Z

Style bot fixed some files and pushed the changes.

HuggingFaceDocBuilderDev · 2026-03-20T04:56:04Z

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

jerryzh168 reviewed Mar 19, 2026

View reviewed changes

src/diffusers/quantizers/torchao/torchao_quantizer.py Outdated Show resolved Hide resolved

jerryzh168 reviewed Mar 19, 2026

View reviewed changes

remove str option for quantization config in torchao

d49eb2e

howardzhang-cv force-pushed the update_torchao_test branch from ed295a7 to d49eb2e Compare March 19, 2026 21:42

howardzhang-cv marked this pull request as ready for review March 19, 2026 21:43

sayakpaul reviewed Mar 20, 2026

View reviewed changes

github-actions bot and others added 2 commits March 20, 2026 04:49

Apply style fixes

b8cb4b2

Merge branch 'main' into update_torchao_test

2449132

minor fixes

292e65b

howardzhang-cv requested a review from sayakpaul March 20, 2026 21:29

		if not isinstance(self.quant_type, AOBaseConfig):
		raise TypeError(f"quant_type must be an AOBaseConfig instance, got {type(self.quant_type).__name__}")

Conversation

howardzhang-cv commented Mar 19, 2026

What does this PR do?

Testing

Who can review?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

jerryzh168 Mar 19, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sayakpaul commented Mar 20, 2026

Uh oh!

github-actions bot commented Mar 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

HuggingFaceDocBuilderDev commented Mar 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

jerryzh168 Mar 19, 2026 •

edited

Loading

github-actions bot commented Mar 20, 2026 •

edited

Loading