Implement an AWQ algorithm with dynamic activation quantization for ExecuTorch

[metascroy]

TorchAO has implemented AWQ, and we'd like to extend the implementation to cover dynamic activation quantization to support lowering to ExecuTorch.

1. We should modify the existing AWQ algorithm to support QDQLayout here to support ExecuTorch

2. The scales should have 8-bit dynamic activation quantization applied before computing the AWQ scaling.

With the above changes, we should be able to quantize a model with AWQ and lower it to ExecuTorch following instructions similar to here.

[metascroy]

cc @kimishpatel as FYI

[kimishpatel]

@metascroy this one https://github.com/pytorch/ao/blob/main/torchao/prototype/awq/core.py#L118 is really in the observer. You need to replace it like here https://github.com/pytorch/ao/blob/main/torchao/prototype/awq/api.py#L150 or here https://github.com/pytorch/ao/blob/main/torchao/prototype/awq/example.py#L239, right?

[metascroy]

@metascroy this one https://github.com/pytorch/ao/blob/main/torchao/prototype/awq/core.py#L118 is really in the observer. You need to replace it like here https://github.com/pytorch/ao/blob/main/torchao/prototype/awq/api.py#L150 or here https://github.com/pytorch/ao/blob/main/torchao/prototype/awq/example.py#L239, right?

Yes, I think those need to be modified as well.

It's actually a bit strange to me that the observer is not using the same quantization as the "convert" method. Seems like they should be aligned?

[jerryzh168]

actually as I thought through this more, it may not always work for dynamic quantization, the reasoning is the follows:

The way it works is it preserves the salient weights, in weight only quant, the only accuracy loss comes from weight quantization, so preserving the salient weights can help reduce the only source of error without affecting other things. In dynamic quant, activation quantization is also important, so if we scale the activation it could have some effect on the accuracy as well, I think this could explain why we can't always get better results with applying awq to dynamic quant.

[kimishpatel]

actually as I thought through this more, it may not always work for dynamic quantization, the reasoning is the follows:

The way it works is it preserves the salient weights, in weight only quant, the only accuracy loss comes from weight quantization, so preserving the salient weights can help reduce the only source of error without affecting other things. In dynamic quant, activation quantization is also important, so if we scale the activation it could have some effect on the accuracy as well, I think this could explain why we can't always get better results with applying awq to dynamic quant.

Jerry, what do you mean by salient weight. I dont quite follow

[jerryzh168]

Jerry, what do you mean by salient weight. I dont quite follow

oh it means that weights that has high impact on accuracy loss of the output will be preserved by applying a large scale to it,

for example, just for demo purpose:

output = input * weight (elementwise)
input = [1, 2, 3, 4]
weight = [0.1, 0.3, 0.3, 0.2]

let's say quantizing weight[3] has a large impact on the final output accuracy, then we'll apply a scale to it:

weight = [0.1, 0.3, 0.3, 0.2 * 10] # apply a scale to reduce the accuaracy loss of weight[3]
input = [1, 2, 3, 4 / 10]  # apply 1/scale to input[3]

if it's weight only quant, then this monotonically reduces the source of error, but with dynamic quant, since input is also quantized, scaling input may impact the accuracy loss from quantizing input

Why does it explains that awq doesn't always help dynamic quant? I think it's because it's related to the measurement we use at figuring out the scale (trying to reduce the accuracy loss of output given the clibration dataset) is not the same as the final accuracy eval measurement (task specific accuaracy measurement). This is OK for weight only quant because there is actually a guarantee in weight only quant that the result is not going to be worse than before because we reduce the output loss by reducing the loss of quantizing weight I think. It's different in the case of dynamic quant because there are accuracy loss of quantizing input activation as well, and reducing quantization error for weight may increase the quantization error of input, and when the final accuracy loss measurement is different at calibration time v.s. eval time, there won't be a guarantee anymore on awq + dynamic quant being better than dynamic quant only.

[gau-nernst]

For dynamic activation quants, I think SmoothQuant is more suitable. Actually, I think AWQ and SmoothQuant are 2 sides of the same coin: transfer outliers from one side to other other.

[jerryzh168]

makes sense, @namgyu-youn is helping us to update smoothquant here #2728 we can try it on int8 dynamic activation + int4 weight quant when it's ready

[namgyu-youn]

Actually, I think AWQ and SmoothQuant are 2 sides of the same coin: transfer outliers from one side to other other.

Same thought, because both start from activation distribution (outliers); SmoothQuant focuses on smoothing outliers, and AWQ focuses on finding salient weights using outliers. I hope to show SmoothQuant with W4A8 in the foreseeable future!