Initial PARQ addition and testing

[lisjin]

This PR imports PARQ code from an internal repo which will soon be released under https://github.com/facebookresearch. I exclude the image classification examples for now to minimize the code footprint.

PARQ is a QAT method based on a convex regularization framework that converges to hard quantization (i.e., STE) at its asymptotic limit. It splits quantization into: (1) estimating quantization values and (2) projecting full-precision weights onto the values. More details can be found in torchao/prototype/parq/README.md and we will add an arXiv link to the paper soon.

PARQ supports choices of:

• bit-width: 1, ternary, 2, 3, 4, ...
• quantization function: uniform, least squares binary quantization
• proximal mapping: hard, PARQ, BinaryRelax

Test plan: python test/prototype/test_parq.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[andrewor14]

Hi @lisjin, thanks for submitting this! If you have any high level results from your torchtune experiments so far, can you share them in the PR description as well? Also, are there any significant differences from the code in the PARQ repo or is it mostly just copy and paste?

[lisjin]

Definitely! I'll add some preliminary torchtune results (and fix the failed tests) when I'm back early next week. The code was mostly copy and pasted from the PARQ repo.

[lisjin]

@andrewor14 As we discussed in a different thread, we'll plan to include the torchtune experiment results in a separate PR. The code should be ready now, outside of any new suggestions you have!

[andrewor14]

Thanks. I think we can merge this since it's not touching anything outside prototype. As discussed offline, we'll merge the distributed benchmark results separately in the README. Merging!