Pass QAT learned qparams in convert #3022

andrewor14 · 2025-09-17T17:25:12Z

Summary: Add support to pass scales and zero points learned during QAT range learning to the PTQ base config. Currently only the following configs support this feature:

IntxWeightOnlyConfig
Int8DynamicActivationInt4WeightConfig
Int8DynamicActivationIntxWeightConfig

During the convert phase, QAT will detect if range learning was used during training, and pass the learned scales and zero points as custom qparams to the quantized tensor subclass, so PTQ will produce more consistent numerics.

Fixes part of #2271.

Test Plan:

python test/quantization/test_qat.py -k test_range_learning_convert_pass_qparams

pytorch-bot · 2025-09-17T17:25:14Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/3022

📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit 75c6c8d with merge base 9e5059e ():
💚 Looks good so far! There are no failures yet. 💚

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

facebook-github-bot · 2025-09-17T17:33:12Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D82654987.

torchao/quantization/qat/api.py

torchao/quantization/quantize_/workflows/intx/intx_unpacked_to_int8_tensor.py

metascroy · 2025-09-17T17:55:23Z

torchao/quantization/quant_api.py

+                **kwargs,
            )
            if scale_dtype is not None and scale_dtype != weight.dtype:
                _adjust_scale_dtype_in_intx_unpacked_tensor(


Presumably you don't want this to run if you passed custom scales?

I looked into this a bit and I think it's actually fine to run this? Basically we use the custom scale first and then cast it to the configured scale dtype (otherwise custom_scale will have to leak into this file, which I hope to avoid)

I think you're right

facebook-github-bot · 2025-09-17T18:46:06Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D82654987.

facebook-github-bot · 2025-09-17T20:21:49Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D82654987.

metascroy · 2025-09-17T22:35:50Z

test/quantization/test_qat.py

+            is_dynamic=False,
+            range_learning=True,
+            scale_precision=torch.float32,
+            zero_point_precision=torch.float32,


We should not use float32 zeros we we want to lower to XNNPACK.

We should convert the int32 zeros to int8 when passing in custom_zero_point

jerryzh168 · 2025-09-17T22:36:36Z

torchao/quantization/quant_api.py


-def _int8_dynamic_activation_intx_weight_quantize_tensor(weight, bias, config):
+def _int8_dynamic_activation_intx_weight_quantize_tensor(
+    weight, bias, config, **kwargs


spell out the kwargs might be better?

like tensor_subclass_kwargs?

no, I mean use custom_scale and custom_zero_point, since kwargs is more complicated than this

torchao/quantization/qat/api.py

facebook-github-bot · 2025-09-18T20:30:46Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D82654987.

jerryzh168 · 2025-09-18T21:24:17Z

torchao/quantization/quant_api.py

+        if custom_zero_point.dtype == torch.int32:
+            custom_zero_point = custom_zero_point.to(torch.int8)


should we do assert instead?

this is to make the flow work by default without user having to specify the zero point dtype (QAT's default is int32), so adding an assert will fail that

jerryzh168 · 2025-09-18T21:24:25Z

torchao/quantization/quant_api.py


    if config.version == 2:
        if config.intx_packing_format == IntxPackingFormat.UNPACKED_TO_INT8:
+            if custom_zero_point.dtype == torch.int32:


**Summary:** Add support to pass scales and zero points learned during QAT range learning to the PTQ base config. Currently only the following configs support this feature: ``` IntxWeightOnlyConfig Int8DynamicActivationInt4WeightConfig Int8DynamicActivationIntxWeightConfig ``` During the convert phase, QAT will detect if range learning was used during training, and pass the learned scales and zero points as custom qparams to the quantized tensor subclass, so PTQ will produce more consistent numerics. Fixes part of #2271. **Test Plan:** ``` python test/quantization/test_qat.py -k test_range_learning_convert_pass_qparams ```

facebook-github-bot · 2025-09-19T17:12:22Z

@andrewor14 has imported this pull request. If you are a Meta employee, you can view this in D82654987.

andrewor14 marked this pull request as draft September 17, 2025 17:25

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 Sep 17, 2025

andrewor14 requested a review from metascroy September 17, 2025 17:25

andrewor14 force-pushed the try-range-learning-pass-qparams branch from 2f0b504 to f061ae7 Compare September 17, 2025 17:26

andrewor14 added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Sep 17, 2025

andrewor14 force-pushed the try-range-learning-pass-qparams branch from f061ae7 to c6b0fe7 Compare September 17, 2025 17:34

metascroy reviewed Sep 17, 2025

View reviewed changes

torchao/quantization/qat/api.py Outdated Show resolved Hide resolved

metascroy reviewed Sep 17, 2025

View reviewed changes

torchao/quantization/quantize_/workflows/intx/intx_unpacked_to_int8_tensor.py Outdated Show resolved Hide resolved

metascroy reviewed Sep 17, 2025

View reviewed changes

andrewor14 force-pushed the try-range-learning-pass-qparams branch from c6b0fe7 to adb966c Compare September 17, 2025 18:44

andrewor14 force-pushed the try-range-learning-pass-qparams branch from adb966c to 9d0aadd Compare September 17, 2025 20:20

andrewor14 changed the title ~~[draft] Pass QAT learned qparams in convert~~ Pass QAT learned qparams in convert Sep 17, 2025

andrewor14 requested review from jerryzh168 and metascroy September 17, 2025 20:20

andrewor14 marked this pull request as ready for review September 17, 2025 20:52

metascroy approved these changes Sep 17, 2025

View reviewed changes

metascroy reviewed Sep 17, 2025

View reviewed changes

jerryzh168 reviewed Sep 17, 2025

View reviewed changes

torchao/quantization/qat/api.py Outdated Show resolved Hide resolved

andrewor14 force-pushed the try-range-learning-pass-qparams branch 3 times, most recently from aae7eee to ecd9349 Compare September 18, 2025 20:30

jerryzh168 reviewed Sep 18, 2025

View reviewed changes

jerryzh168 approved these changes Sep 18, 2025

View reviewed changes

andrewor14 force-pushed the try-range-learning-pass-qparams branch from ecd9349 to 75c6c8d Compare September 19, 2025 17:10

andrewor14 merged commit ae12e42 into main Sep 19, 2025
19 of 20 checks passed

		if custom_zero_point.dtype == torch.int32:
		custom_zero_point = custom_zero_point.to(torch.int8)

Pass QAT learned qparams in convert #3022

Pass QAT learned qparams in convert #3022

Uh oh!

Conversation

andrewor14 commented Sep 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Sep 17, 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/ao/3022

✅ No Failures

Uh oh!

facebook-github-bot commented Sep 17, 2025

Uh oh!

Uh oh!

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!

facebook-github-bot commented Sep 17, 2025

Uh oh!

facebook-github-bot commented Sep 17, 2025

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!

jerryzh168 Sep 17, 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!

Uh oh!

facebook-github-bot commented Sep 18, 2025

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!

facebook-github-bot commented Sep 19, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

andrewor14 commented Sep 17, 2025 •

edited

Loading

pytorch-bot bot commented Sep 17, 2025 •

edited

Loading

jerryzh168 Sep 17, 2025 •

edited

Loading