Relax dtype requirements for int4 and float8 quants in autoquant

[jerryzh168]

Summary:
Some of the int4 quant and fp8 only works with bfloat16/float16, previously we require the model to be in correct dtype to apply these in autoquant, this PR relaxes the constraints by converting weight, bias and activation to compatible dtypes

Test Plan:
python test/integration/test_integration.py -k test_autoquant_int4wo

Reviewers:

Subscribers:

Tasks:

Tags:

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

⏳ No Failures, 3 Pending

As of commit c03760d with merge base e1cb44a ():
💚 Looks good so far! There are no failures yet. 💚

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

[jerryzh168]

@jcaip I can't run sparse marlin kernel locally: Could not run 'torchao::marlin _24_gemm' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). how do we compile to get this op? right now I'm using python setup.py develop

[jerryzh168]

also when there is one dimension that is 1, we'll do a squeeze, so the scale will become 1-d, this is a corner case
cc @jcaip @Diogo-V let me know if you have other ideas for the fix, alternatively we can pass around block_size here as well, or we can store group_size in _layout maybe

[Diogo-V]

Hey Jerry, I haven’t been actively working on this repo lately as I have been more focused on work, so I am not entirely up to date. That said, please keep in mind that my input might not fully reflect the current state of things.

My two cents here is that this should work well, as the condition seems to allow for generalization to other scenarios that might produce that corner case in the future i.e: if weight_shape[-1] == group_size == 64 yields a scale with dimension 1.

[jcaip]

@jcaip I can't run sparse marlin kernel locally: Could not run 'torchao::marlin _24_gemm' with arguments from the 'CUDA' backend. This could be because the operator doesn't exist for this backend, or was omitted during the selective/custom build process (if using custom build). how do we compile to get this op? right now I'm using python setup.py develop

You need to build with USE_CPP=1 and it should show up