[WIP] SmoothQuant using tensor subclassing

[Xia-Weiwen]

Still WIP

The implementation of SmoothQuant with tensor subclassing (AffineQuantizedTensor) is similar to that of AWQ with the following differences:

• SmoothQuant supports both static and dynamic quantization of activation while AWQ only uses dynamic
• Matmul is computed in int8 instead of floating point (at least at op level)
• The smoothing factor is calculated differently from the equalization scales of AWQ

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 1 New Failure

As of commit fa1144c with merge base d4b2f33 ():

NEW FAILURE - The following job has failed:

• Run Regression Tests / test (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://download.pytorc... / linux-job (gh)
  RuntimeError: Command docker exec -t 95559b07c7091cb50f330d25bb5604fd4430cf7f95edb1ff7928625590e1e78f /exec failed with exit code 1

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

[jerryzh168]

do we need this? or just saving the state_dict for observed model is enough?

[Xia-Weiwen]

We want to have an API to modify (tune) quantization parameters, i.e. the recipe here. Do you have any concern about adding this API?

[jerryzh168]

so the state_dict is supposed to be used by other APIs to tune quantization parameters? I think that's fine if you have this use case in mind, is the model with SmoothQuantObservedLinear not serializable by itself?

[Xia-Weiwen]

SmoothQuantObservedLinear is serializable. However, a recipe is more flexible to tune parameters. Thanks.

[Xia-Weiwen]

Hi @jerryzh168 I added a new tensor subclass LinearActivationScaleQuantizedTensor to support x -> x/scale -> quantize x for torch.compile.

If I use LinearActivationQuantizedTensor, the x/sacle is done outside the class (by input_quant_func) and there is a dynamo error about scale during torch.compile. I guess it's because the scale tensors are not on the graph in this case. Putting the scale in the weight tensor solves the problem. And WeightTensorWithLinearActivationScaleMetadata does not quantize activation.

Do you have any concern adding this new class? Thanks.

[jerryzh168]

Hi @jerryzh168 I added a new tensor subclass LinearActivationScaleQuantizedTensor to support x -> x/scale -> quantize x for torch.compile.

If I use LinearActivationQuantizedTensor, the x/sacle is done outside the class (by input_quant_func) and there is a dynamo error about scale during torch.compile. I guess it's because the scale tensors are not on the graph in this case. Putting the scale in the weight tensor solves the problem. And WeightTensorWithLinearActivationScaleMetadata does not quantize activation.

Do you have any concern adding this new class? Thanks.

I think in this case we should be composing WeightTensorWithLinearActivationScaleMetadata LinearActivationQuantizedTensor together, i.e.

weight = to_affine_quantized(float_weight, ...)
# this will quantize input
# use https://github.com/pytorch/ao/blob/c87cc9b7286a46e9dfc076fa2417eb9b64ccc807/torchao/quantization/weight_tensor_linear_activation_quantization.py#L13 for static quantization
weight = to_linear_activation_quantized_tensor(weight)  # dynamic quant
# this will do x / scale
weight = to_weight_tensor_with_linear_activation_scale_metadata(weight)

in dispatch time, we first unwrap the outer most tensor subclass, which will be WeightTensorWithLinearActivationScaleMetadata, so we'll apply scale to activation, then LinearActivationQuantizedTensor, which will quantize activation, and then AffineQuantizedTensor

would this work?

the naming for different tensor subclasses is a bit confusing right now I think, we should cleanup a bit later

[Xia-Weiwen]

I think in this case we should be composing WeightTensorWithLinearActivationScaleMetadata LinearActivationQuantizedTensor together, i.e.

weight = to_affine_quantized(float_weight, ...)
# this will quantize input
# use https://github.com/pytorch/ao/blob/c87cc9b7286a46e9dfc076fa2417eb9b64ccc807/torchao/quantization/weight_tensor_linear_activation_quantization.py#L13 for static quantization
weight = to_linear_activation_quantized_tensor(weight)  # dynamic quant
# this will do x / scale
weight = to_weight_tensor_with_linear_activation_scale_metadata(weight)

in dispatch time, we first unwrap the outer most tensor subclass, which will be WeightTensorWithLinearActivationScaleMetadata, so we'll apply scale to activation, then LinearActivationQuantizedTensor, which will quantize activation, and then AffineQuantizedTensor

would this work?

the naming for different tensor subclasses is a bit confusing right now I think, we should cleanup a bit later

It works. Thanks

[Xia-Weiwen]

Hi @jerryzh168 It's weird that if I add these lines https://github.com/pytorch/ao/blob/f595ed41b99685cc16fc480ca2218965bb812bed/torchao/kernel/intmm.py#L142C1-L146C1 to avoid overflow of float16, there will be a failure in test_spinquant.py, but the test does not use float16 at all. And there is another failure with CUDA nightly, but its log cannot be loaded. And these failures cannot be reproduced in my local environment or on an AWS instance.
So, I have to remove these line and also remove tests for float16. If users try to run in fp16, they will get overflow as well. Do you have any suggestions? Thanks.

[jerryzh168]

Hi @jerryzh168 It's weird that if I add these lines f595ed4/torchao/kernel/intmm.py#L142C1-L146C1 to avoid overflow of float16, there will be a failure in test_spinquant.py, but the test does not use float16 at all.

what is the test failure? is it possible to do the dtype conversion before calling int_scaled_matmul, e.g. before
y_dot_scaled = int_scaled_matmul(tmp, w_vals_int8_t, x_scales.reshape(-1, 1)) in affine_quantized_tensor.py?

And there is another failure with CUDA nightly, but its log cannot be loaded. And these failures cannot be reproduced in my local environment or on an AWS instance. So, I have to remove these line and also remove tests for float16. If users try to run in fp16, they will get overflow as well. Do you have any suggestions? Thanks.

I just saw the error, it is talking about some triton error:

  E1018 18:45:40.023529 436 site-packages/torch/_inductor/runtime/triton_heuristics.py:475]     assert lhs.shape[1].value >= 32, "small blocks not supported!"
  E1018 18:45:40.023529 436 site-packages/torch/_inductor/runtime/triton_heuristics.py:475] AssertionError: small blocks not supported!

[Xia-Weiwen]

what is the test failure? is it possible to do the dtype conversion before calling int_scaled_matmul, e.g. before y_dot_scaled = int_scaled_matmul(tmp, w_vals_int8_t, x_scales.reshape(-1, 1)) in affine_quantized_tensor.py?

The error is results not all close. One element exceed the tolerance by a small amount. As for dtype conversion, I didn't make such changes in affine_quantized_tensor.py 🤔

  E1018 18:45:40.023529 436 site-packages/torch/_inductor/runtime/triton_heuristics.py:475]     assert lhs.shape[1].value >= 32, "small blocks not supported!"
  E1018 18:45:40.023529 436 site-packages/torch/_inductor/runtime/triton_heuristics.py:475] AssertionError: small blocks not supported!

Thanks for the info. Did you see which test case failed?