Skip to content

PARQ quantizer support for torchao's weight-only configs#2091

Merged
lisjin merged 13 commits into
pytorch:mainfrom
lisjin:parq
Apr 28, 2025
Merged

PARQ quantizer support for torchao's weight-only configs#2091
lisjin merged 13 commits into
pytorch:mainfrom
lisjin:parq

Conversation

@lisjin

@lisjin lisjin commented Apr 21, 2025

Copy link
Copy Markdown
Contributor

This is the first step in supporting torchao.quantize_ for PARQ trained models. I target only Int4WeightOnlyConfig and IntxWeightOnlyConfig for now since PARQ does not have activation quantization.

Instead of converting the state (e.g., scale, zero point) from PARQ's existing quantizers to torchao format, I decided to create a new quantizer UnifTorchaoQuantizer. This quantizer calls torchao's quantization primitives choose_qparams_affine, quantize_affine, dequantize_affine to ensure parity between the two QAT methods.

@metascroy It would be great if you could check the correctness of how the quantizer in TestUnifTorchaoQuantizer.test_intx_weight_only is initialized. I'm not sure if I missed any subtleties with int8.

@pytorch-bot

pytorch-bot Bot commented Apr 21, 2025 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

✅ No Failures

As of commit 073e1fa with merge base e3db2b2 (image):
💚 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 facebook-github-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 Apr 21, 2025
@lisjin lisjin added the topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories) label Apr 21, 2025
metascroy
metascroy reviewed Apr 22, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py
metascroy
metascroy reviewed Apr 22, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py Outdated
metascroy
metascroy reviewed Apr 22, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py Outdated
@metascroy

metascroy commented Apr 22, 2025

Copy link
Copy Markdown
Contributor

@lisjin can you give a little code snippet of our QAT prepare/convert would work for this API?

I'm having trouble following. Here are some example code snippets from other APIs: https://fb.workplace.com/groups/pytorch.edge2.team/permalink/1186139489308568/

lisjin
lisjin commented Apr 23, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py
lisjin
lisjin commented Apr 23, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py Outdated
@lisjin lisjin requested a review from metascroy April 23, 2025 21:27
lisjin
lisjin commented Apr 24, 2025
View reviewed changes
Comment thread torchao/prototype/parq/optim/quantopt.py Outdated
@andrewor14

andrewor14 commented Apr 24, 2025

Copy link
Copy Markdown
Contributor

Hi @lisjin, do you mind adding a code snippet on the main README on what the end-to-end flow would look like? My understanding is you can just replace the LSBQuantizer there with your new UnifTorchaoQuantizer. Then what happens after training? Do we call quantize_(model, Int4WeightOnlyConfig) as before? Would be good to clarify

andrewor14
andrewor14 reviewed Apr 24, 2025
View reviewed changes
Comment thread torchao/prototype/parq/quant/uniform_torchao.py Outdated
@lisjin

lisjin commented Apr 25, 2025 •
edited
Loading

Copy link
Copy Markdown
Contributor Author

@andrewor14 Thanks for the feedback—I removed config from UnifTorchaoQuantizer. In the README, I've also added a side-by-side comparison of PARQ vs. torchao prepare and convert steps. After PARQ training, we call optimizer.torchao_quantize_(model, config). Let me know if there's anything missing.

@lisjin lisjin requested a review from andrewor14 April 25, 2025 17:48
@andrewor14

andrewor14 commented Apr 25, 2025

Copy link
Copy Markdown
Contributor

Looks great, thanks @lisjin! The README is very clear.

One thing I want to discuss is whether we can just use a new PARQConfig instead so the PARQ flow looks more like the existing torchao QAT flow. This is current convert flow in the PR now:

config = IntXWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32))
optimizer.torchao_quantize_(model, config)

What do you think about something like this instead?

inner_config = IntXWeightOnlyConfig(weight_dtype=torch.int4, granularity=PerGroup(32))
parq_config = PARQConfig(optimizer, inner_config)
quantize_(model, parq_config)

Also curious if @metascroy has any thoughts on this

lisjin
lisjin commented Apr 25, 2025
View reviewed changes
Comment thread torchao/prototype/parq/optim/quantopt.py Outdated
andrewor14
andrewor14 reviewed Apr 25, 2025
View reviewed changes
Comment thread torchao/prototype/parq/optim/quantopt.py Outdated
andrewor14
andrewor14 approved these changes Apr 25, 2025
View reviewed changes

@andrewor14 andrewor14 left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Looks good to me other than the recursion comment. @metascroy any other thoughts?

@metascroy

metascroy commented Apr 26, 2025

Copy link
Copy Markdown
Contributor

Looks good to me! Thanks @lisjin!

Can we add an end-to-end test_intx_weight_only_e2e for intx (with various x-values), similar to test_int4_weight_only_e2e?

metascroy
metascroy reviewed Apr 26, 2025
View reviewed changes
Comment thread torchao/prototype/parq/quant/uniform_torchao.py Outdated
metascroy
metascroy reviewed Apr 26, 2025
View reviewed changes
Comment thread torchao/prototype/parq/quant/uniform_torchao.py Outdated
lisjin
lisjin commented Apr 28, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py Outdated
lisjin
lisjin commented Apr 28, 2025
View reviewed changes
Comment thread test/prototype/test_parq.py
Comment on lines +246 to +247
@common_utils.parametrize("b", [2, 3, 4, 8])
def test_intx_weight_only_e2e(self, b: int = 2, group_size: int = 32):

@lisjin lisjin Apr 28, 2025

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@metascroy Thanks for looking it over! I've added this end-to-end test, along with mapping_type=MappingType.SYMMETRIC and target_dtype=torch.int8 defaults for UnifTorchaoQuantizer

@lisjin lisjin requested review from andrewor14 and metascroy April 28, 2025 14:14
andrewor14
andrewor14 reviewed Apr 28, 2025
View reviewed changes
Comment thread torchao/prototype/parq/README.md Outdated
@facebook-github-bot

facebook-github-bot commented Apr 28, 2025

Copy link
Copy Markdown
Contributor

@lisjin has imported this pull request. If you are a Meta employee, you can view this diff on Phabricator.

@lisjin lisjin merged commit 8334340 into pytorch:main Apr 28, 2025
@lisjin lisjin deleted the parq branch April 28, 2025 17:36
@lisjin lisjin mentioned this pull request Jul 21, 2025
Merged
liangel-02 pushed a commit that referenced this pull request Aug 25, 2025
@lisjin @liangel-02
PARQ quantizer support for torchao's weight-only configs (#2091)
740fc42 
* Add parq.quant.UnifTorchaoQuantizer for quantize_ API equivalence

* Test IntxWeightOnlyConfig

* Formatting fix

* Per-row IntxWeightOnlyConfig test

* Add end-to-end QAT prepare/convert test case

* Pass explicit layout to int4_weight_only

* Add QuantOptimizer.torchao_quantize_

* Update README, add Int4UnifTorchaoQuantizer

* Add test_intx_weight_only_e2e, set UnifTorchaoQuantizer defaults

* Update PARQ README
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@andrewor14 andrewor14 andrewor14 approved these changes

@jerryzh168 jerryzh168 Awaiting requested review from jerryzh168 jerryzh168 is a code owner

@metascroy metascroy Awaiting requested review from metascroy

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. topic: improvement Use this tag if this PR is an improvement (doesn't fit into any of the other categories)

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@lisjin @metascroy @andrewor14 @facebook-github-bot