[WIP] Activation Aware Weight Quantization (AWQ)

[vayuda]

Adds AWQ per #530

To do:

• Verify correctness of implementation and add tests for this
• Fold activation scaling into previous layer if applicable
• Model works with compile
• Move implementation into the right files

[pytorch-bot]

🔗 Helpful Links

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

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

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

[jerryzh168]

can you just remove weight_quant_fn and add use_hqq to the function?

[jerryzh168]

also when it's uint4, I think it's fine to just use TensorCoreTiledLayout

[vayuda]

I added that as a feature so maybe people can find ways to compose on top of AWQ esp if they come up with new kernels or what not, but I would agree that its not a necessary feature for an initial release

[jerryzh168]

looks good, thanks for addressing all the comments! I think the main thing remaining is just to use use_hqq flag for awq_uintx so we don't take a random weight_quant_fn

also make sure to fix the CI issues as well

[HDCharles]

this shouldn't be in generate.py, it should be in eval so we can actually see the accuracy impact