[AWQ] add option to take smooth layer quantization into accout

[HDCharles]

normally the way AWQ works is to pick a layer that is going to be quantized, try a bunch of scale factors to find the one that minimizes quantization error when that layer is quantized and then do an inverse rescale on the preceeding layer which is normally not quantized. However a problem arises for the up_proj -> down_proj mapping
because both the smooth and balance layers are targeted for quantization. Since we only take into account the quantization of the balance layers in our current AWQ implementation, we could be making the smooth layer harder to quantize with our choice of scale factor for the balance layer since the smooth layer is basically ignored during the quantization error calculation for quantizing the balance layer.

We should

1. test if this has a significant impact
2. add an option to enable this feature if its beneficial

STEPS

A) add a check here for whether smooth_name is in targeted_names and if so, change the get_lowest_common...etc search to include the smooth layer (this is how we determine what module is run to determine the quantization error, so we need smooth layer to be run if we're taking its quantization into accout)
B) add a flag to compute_best_scale for if the smooth layer is targeted
C) if necessary add the smooth layer to this dict
D) move the rescale weight code into a function which is called for each balance layer
E) if necessary, call the rescale weight code for 1/_scalesview on the smooth_layer
F) check whether this has an impact on lm_eval performance on some small set of models.
G) check how this affects the runtime of AWQ for those models.

if its beneficial then put up a PR with those changes demonstrating what was tested and how it affects things.

[Etelis]

Assign me here as well

[Ramshankar07]

Can I work on it @Etelis , I'm eager to learn and contribute

[Etelis]

@Ramshankar07 Sure go ahead.

[Etelis]

I'd wait a bit for me to finish running tests to see if it is beneficial (#2295

[Ramshankar07]

Sure thing, Let me know

[HDCharles]

this is all dealing with the same type of thing but in orthagonal ways. even if #2295 isn't important this one could be.

They're all kind of heuristic ways of turning the joint problem
$\min_{s1,s2} Loss(MLP(X,W1,W2) - MLP(X/s1,Q(s1 * W1)/s2, Q(s2 * W2)))$

into 2 disjoint problems $\min_{s1} Loss(MLP(X,W1,W2) - MLP(X/s1,Q(s1 * W1), W2))$ and $\min_{s2} Loss(MLP(X,W1,W2) - MLP(X/s1,s1 * W1/s2, Q(s2 * W2)))$

2295 deals with the ordering, since it makes more sense to look at $Q(\frac{s1 * W1}{s2})$ rather than $\frac{Q(s1 * W1)}{s2}$ whereas this one deals with taking some of problem 2 and taking it into account in problem 1. So it'd be mroe like

$\min_{s1} Loss(MLP(X,W1,W2) - MLP(X/s1,Q(s1 * W1), W2))$ and $\min_{s2} Loss(MLP(X,W1,W2) - MLP(X/s1,\textcolor{red}{Q(}s1 * W1/s2), Q(s2 * W2)))$

i'll assign you @Ramshankar07

[Ramshankar07]

Thanks, @HDCharles. The provided steps are very elaborate for me to check. I'll start a PR in couple of hours.