AWQ Apply Scales Bugfix when smooth layer output length doesn't match balance layer input length

[brian-dellabetta]

Summary

We are hitting an edge case in AWQ we had not previously hit with the initial Llama/Qwen testing models. When a smooth layer's # of output_features does not match a balance layer's # of input_features, the code as it is currently will error out when trying to update the smooth layer's weights with weights.div(scales), due to a shape mismatch error. We are hitting this in #1440 for Phi3 models, which include a mapping between the fused qkv_proj smooth layer and o_proj balance layer in AutoAWQ (see here).

The resolution in AutoAWQ is to only use the last rows of the smooth layer so that the shapes line up, as shown here. This PR includes that update, and with #1440 will allow Phi3 models to be quantizable with AWQModifier. Like with v_proj -> o_proj, if shapes don't match up, they will be excluded from resolved mappings. This allows phi-3-mini to include the mapping because qkv_proj out_features == 3*o_proj in_features == 9216, but excludes it from phi-3-medium which has qkv_proj out_features == 7680 and o_proj in_features==5120. If the mapping is included for phi-3-medium, the model blows up with wikitext eval perplexities >2000. This implementation was agreed upon with @anmarques .

PS: I also moved mul & div to mul_ & div_, to avoid unnecessary memory allocation.

---

Test Plan

With these changes and with #1440 , examples/awq/llama_example.py works with "microsoft/Phi-3-mini-128k-instruct" and produces similar results as when qkv_proj to o_proj mapping is included

Without mapping:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
wikitext	2	none	5	bits_per_byte	↓	0.6474	±	N/A
		none	5	byte_perplexity	↓	1.5664	±	N/A
		none	5	word_perplexity	↓	11.0201	±	N/A

With mapping:

Tasks	Version	Filter	n-shot	Metric		Value		Stderr
wikitext	2	none	5	bits_per_byte	↓	0.6482	±	N/A
		none	5	byte_perplexity	↓	1.5672	±	N/A
		none	5	word_perplexity	↓	11.0527	±	N/A

I also confirmed re-running with meta-llama/Llama-3.2-3B-Instruct and meta-llama/Llama-2-7b-hf does not deviate in PPL scores from what is currently on main

[github-actions]

👋 Hi! Thank you for contributing to llm-compressor. Please add the ready label when the PR is ready for review.

Note: This is required to complete the testing suite, please only add the label once the PR is code complete and local testing has been performed.

[shanjiaz]

Thanks for the test results! Looks good to me.

[rahul-tuli]

🚀