PARQ quantizer support for torchao's weight-only configs

[lisjin]

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]

🔗 Helpful Links

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

• 📄 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 073e1fa with merge base e3db2b2 ():
💚 Looks good so far! There are no failures yet. 💚

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

[metascroy]

@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/

[andrewor14]

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

[lisjin]

@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.

[andrewor14]

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

[andrewor14]

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

[metascroy]

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?

[lisjin]

@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

[facebook-github-bot]

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