Add TorchAO wrapper config to allow filter_fn for quantize_ #13264

abhinaykukkadapu · 2025-08-10T03:45:04Z

Changes:

Support filter function in quantize_ function when using torchao quantize.
Update unittests accordingly
Use ComposableQuantizer if there are multiple quantizers and is of type torchao, for legacy quantizers use them directly with prepare_pt2e.
Source transform modifies model inplace, so deep copy first to avoid modifying user provided model.

[ghstack-poisoned]

abhinaykukkadapu · 2025-08-10T03:45:06Z

Stack from ghstack (oldest at bottom):

pytorch-bot · 2025-08-10T03:45:07Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/13264

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 5 New Failures

As of commit 7dab762 with merge base a84b3c9 ():

NEW FAILURES - The following jobs have failed:

Build documentation / build (buck2) / Build doc (gh)
At least one of the pre-conditions you specified did not hold
pull / unittest / linux / linux-job (gh)
examples/models/llama/tests/test_ring_attention.py::TestRingAttention::test_sliding_window_attention
pull / unittest / macos / macos-job (gh)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1
pull / unittest-editable / linux / linux-job (gh)
examples/models/llama/tests/test_ring_attention.py::TestRingAttention::test_sliding_window_attention
pull / unittest-editable / macos / macos-job (gh)
RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 1

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2025-08-10T03:45:43Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

[ghstack-poisoned]

digantdesai · 2025-08-12T12:12:14Z

backends/xnnpack/test/recipes/test_xnnpack_recipes.py

+        eager_quantized_model = source_transform_output.data["forward"]
+        output = session.run_method("forward", example_inputs[0])[0]
+        expected = eager_quantized_model(*example_inputs[0])
+        self.assertTrue(torch.allclose(output, expected, atol=atol))


You might want to print more stats if this fails - see https://github.com/pytorch/executorch/blob/main/backends/test/harness/tester.py#L337

digantdesai · 2025-08-12T12:12:58Z

backends/xnnpack/test/recipes/test_xnnpack_recipes.py

-                            atol=1e-1,
-                        )
+                    self._compare_eager_quantized_model_outputs(
+                        session, example_inputs, 1e-1


atol? Why is this so high for two linears?

Suggested change

session, example_inputs, 1e-1

session, example_inputs, atol=1e-1

Yeah, i think 1e-2 is working on my mac, will check if linux passes on CI. Nevertheless, i'm updating the tolerance tests similar to CoreML (let me know if there is any objection) to use sqnr to compare eager model vs lowered model output.

But use tolerance checks to compare post quantized model and lowered model.

digantdesai · 2025-08-12T12:18:06Z

backends/xnnpack/recipes/xnnpack_recipe_provider.py

-        config = Int8DynamicActivationIntxWeightConfig(
-            weight_dtype=weight_dtype,
-            weight_granularity=weight_granularity,
+        config = AOQuantizationConfig(


assert weight_dtype == torch.int4 because XNNPACK doesn't support anything but 4 for groupwise. For channelwise it can do both int4 and int8.

digantdesai · 2025-08-12T12:19:08Z

export/stages.py

+            # Multiple torchao quantizers - use ComposableQuantizer
+            from torchao.quantization.pt2e.quantizer import ComposableQuantizer
+
+            return ComposableQuantizer(torchao_pt2e_quantizers)


test for this?

metascroy · 2025-08-12T18:35:13Z

export/stages.py

+            raise ValueError("Mixed quantizer types are not supported")
+        if len(torch_ao_quantizers) > 1:
+            raise ValueError(
+                "Multiple quantizers of torch.ao.quantization.quantizer not supported"


Doesn't torchao already detect this and give an error if mixing? I thought I added that

May be the torchao version is different?

metascroy · 2025-08-12T18:36:37Z

export/stages.py

@@ -331,6 +330,38 @@ def valid_predecessor_stages(self) -> List["StageType"]:
    def can_start_pipeline(self) -> bool:
        return True

+    def _get_quantizer_for_prepare_pt2e(self, quantizers: List[Any]):
+        torch_ao_quantizers = []


I want to make sure this API does not invite people to use torch.ao quantizers.

We've migrated to torchao quantizers in ET and torch.ao will be deprecated. The only backend that uses torch.ao in ET right now is CoreML (because it has an external depedency on coremltools adopting the torchao quantizer)

I can create a follow up ticket to deprecate torch.ao support, once coreml moves to torchao. Currently we already support both torch.ao and torchao via prepare_pt2e's backward compatibility anyway.

I guess this is a private method, so maybe it's OK. I just want to make sure we aren't exposing this publicly

metascroy · 2025-08-12T20:21:17Z

Overall, looks good. Address comments before merging

Update

59affe6

[ghstack-poisoned]

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 10, 2025

abhinaykukkadapu mentioned this pull request Aug 10, 2025

Add coreml quant recipes #13265

Open

abhinaykukkadapu requested review from digantdesai and metascroy August 10, 2025 03:56

Update

37bdc0b

[ghstack-poisoned]

abhinaykukkadapu requested a review from mcr229 as a code owner August 12, 2025 00:23

abhinaykukkadapu added 2 commits August 11, 2025 21:37

Update

4eb6a03

[ghstack-poisoned]

Update

7dab762

[ghstack-poisoned]

digantdesai reviewed Aug 12, 2025

View reviewed changes

metascroy reviewed Aug 12, 2025

View reviewed changes

metascroy approved these changes Aug 12, 2025

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Add TorchAO wrapper config to allow filter_fn for quantize_ #13264

Add TorchAO wrapper config to allow filter_fn for quantize_ #13264

abhinaykukkadapu commented Aug 10, 2025 •

edited

Loading

Uh oh!

abhinaykukkadapu commented Aug 10, 2025 •

edited

Loading

Uh oh!

pytorch-bot bot commented Aug 10, 2025 •

edited

Loading

Uh oh!

github-actions bot commented Aug 10, 2025

Uh oh!

digantdesai Aug 12, 2025

Uh oh!

digantdesai Aug 12, 2025 •

edited

Loading

Uh oh!

abhinaykukkadapu Aug 12, 2025 •

edited

Loading

Uh oh!

digantdesai Aug 12, 2025

Uh oh!

digantdesai Aug 12, 2025

Uh oh!

metascroy Aug 12, 2025

Uh oh!

abhinaykukkadapu Aug 12, 2025

Uh oh!

metascroy Aug 12, 2025

Uh oh!

abhinaykukkadapu Aug 12, 2025

Uh oh!

metascroy Aug 12, 2025

Uh oh!

metascroy commented Aug 12, 2025

Uh oh!

Uh oh!

	session, example_inputs, 1e-1
	session, example_inputs, atol=1e-1

Add TorchAO wrapper config to allow filter_fn for quantize_ #13264

Are you sure you want to change the base?

Add TorchAO wrapper config to allow filter_fn for quantize_ #13264

Conversation

abhinaykukkadapu commented Aug 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

abhinaykukkadapu commented Aug 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Aug 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/13264

❌ 5 New Failures

Uh oh!

github-actions bot commented Aug 10, 2025

This PR needs a release notes: label

Uh oh!

Choose a reason for hiding this comment

Uh oh!

digantdesai Aug 12, 2025 • 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!

abhinaykukkadapu Aug 12, 2025 • 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!

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!

metascroy commented Aug 12, 2025

Uh oh!

Uh oh!

abhinaykukkadapu commented Aug 10, 2025 •

edited

Loading

abhinaykukkadapu commented Aug 10, 2025 •

edited

Loading

pytorch-bot bot commented Aug 10, 2025 •

edited

Loading

This PR needs a `release notes:` label

digantdesai Aug 12, 2025 •

edited

Loading

abhinaykukkadapu Aug 12, 2025 •

edited

Loading