Add HQQ support

[mobicham]

Adds support for HQQ quantization without using the hqq lib.

Note:

The dequantized output produced by AffineQuantizedTensor is a bit worse than that produced by the hqq-lib. You can check that by setting raw_output=True. The problem has to do with the midpoint used in the dequantization logic by AffineQuantizedTensor which produces zero-point values with very low magnitude.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 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 0b511f5 with merge base ffa88a4 ():
💚 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]

Hi @mobicham!

Thank you for your pull request and welcome to our community.

Action Required

In order to merge any pull request (code, docs, etc.), we require contributors to sign our Contributor License Agreement, and we don't seem to have one on file for you.

Process

In order for us to review and merge your suggested changes, please sign at https://code.facebook.com/cla. If you are contributing on behalf of someone else (eg your employer), the individual CLA may not be sufficient and your employer may need to sign the corporate CLA.

Once the CLA is signed, our tooling will perform checks and validations. Afterwards, the pull request will be tagged with CLA signed. The tagging process may take up to 1 hour after signing. Please give it that time before contacting us about it.

If you have received this in error or have any questions, please contact us at cla@meta.com. Thanks!

[facebook-github-bot]

Thank you for signing our Contributor License Agreement. We can now accept your code for this (and any) Meta Open Source project. Thanks!

[mobicham]

Llama3.1 8B Instruct lm-eval, the individual scores for each benchmark are a bit different (probably because of the mid-point thing), but the average score is the same

[HDCharles]

Can you modify https://github.com/pytorch/ao/blob/main/torchao/quantization/quant_api.py#L384

to have a new argument use_hqq default to False that applies the hqq quantization
so that it aligns with the existing API?

[jerryzh168]

The problem has to do with the midpoint used in the dequantization logic by AffineQuantizedTensor which produces zero-point values with very low magnitude.

this is talking about accuracy? I can see you'd need this conversion for int4 tinygemm, but it's not required for other cases, if this is significant, maybe you can expose raw_output as a flag to the top level quant_api for hqq

[mobicham]

@jerryzh168 we can keep it this way for the moment. The other solution would be to modify the alternate minimization and update the quant parameters based ont he mid-point format instead of (Wq - zero)*scale. For the moment, the average lm-eval is about the same as hqq-lib, so I guess it's mostly fine. If there are some performance issues I can get back to it and try to change the steps to improve performance with the mid-point approach which is needed for using the tinygemm kernel.

[jerryzh168]

@mobicham thanks, for tests right now we support both unittest.TestCase and pytest (e.g.

ao/test/dtypes/test_uintx.py

Line 39 in 174e630

@pytest.mark.parametrize("bit_width", bit_widths)

), what you have in the file makes sense I think, since you have hardcoded errors for each dtype I feel it's fine to just separate separate them into different tests

[mobicham]

@mobicham thanks, for tests right now we support both unittest.TestCase and pytest (e.g.

ao/test/dtypes/test_uintx.py

Line 39 in 174e630

@pytest.mark.parametrize("bit_width", bit_widths)

), what you have in the file makes sense I think, since you have hardcoded errors for each dtype I feel it's fine to just separate separate them into different tests

I separated the tests in separate classes and added a check regarding the device, otherwise the tests would fail if no gpu is available.

[mobicham]

@jerryzh168 anything missing for the merge?

[msaroufim]

IIRC there was some test failure, just retriggered CI if all g ill merge

[mobicham]

@msaroufim the tests fail because of triton. I was about to delete torchao/prototype/hqq/core.py since we moved the hqq core stuff to quantization, but just saw this pr uses it : #679

  test/hqq/test_hqq_affine.py:3: in <module>
      from torchao.prototype.hqq.core import HQQQuantizer
  /opt/conda/envs/venv/lib/python3.9/site-packages/torchao/prototype/hqq/__init__.py:1: in <module>
      from .mixed_mm import triton_mixed_mm, pack_2xint4
  /opt/conda/envs/venv/lib/python3.9/site-packages/torchao/prototype/hqq/mixed_mm.py:2: in <module>
      import triton.language as tl
  E   ModuleNotFoundError: No module named 'triton'

[msaroufim]

I do see some legit looking errors for GPU https://github.com/pytorch/ao/actions/runs/10403091060/job/28816346879?pr=605#step:12:2570

And for CPU not sure I follow the connection to the version guard fix PR, basically we need to ensure HQQ Quantizer doesnt get imported by accident on CPU machines that way we dont crash on the triton import. And that HQQ quantizer tests are skipped on CPU instnaces

Also I gave you access to trigger CI yourself that way you'll get signal per commit (its just an annoying thing for first time contributors)

Some of the changes you might see on local vs CI runs are due to us running multiple pytorch versions in CI so feel free to add skips for older pytorch versions in case your code isn't working for that

EDIT: Discussed offline to skip def test_dynamic_quant_per_channel_numerics_cuda(self):

[mobicham]

It looks like it's the tensorcore test that is failing not the rest:

  test/hqq/test_hqq_affine.py::TestHQQ4bit::test_hqq_tensorcore_4bit FAILED

It was working fine just 2 days ago I think, let me clone and recheck on an instance

[mobicham]

And for CPU not sure I follow the connection to the version guard fix PR, basically we need to ensure HQQ Quantizer doesnt get imported by accident on CPU machines that way we dont crash on the triton import

It's not supposed to be imported, I wanted to delete torchao/prototype/hqq/core.py because we moved all the stuff there to torchao/quantization/quant_primitives .py. Will take a look at it

[jerryzh168]

I just found that this flag is not used, so we don't really expose hqq to users right now, are you planning to create a new function for hqq? cc @mobicham

[mobicham]

My understand was that @HDCharles suggested putting it there and later turning it on by default.
It is exposed though via to_affine_quantized https://github.com/pytorch/ao/pull/605/files#diff-a9708dc28f15bb9cf665417e6c66601f9e8e2f1f672d1858603b74fa879a3357R62
Let me know if there's another way of exposing it

[jerryzh168]

commented in #786 (comment)