Update SmoothQuant to use subtensor instead of external wrapper

[namgyu-youn]

Summary

Similar to AWQ (#2753), SmoothQuant now uses direct protocol implementation instead of an external wrapper (to_weight_tensor_with_linear_activation_scale_metadata).

This PR was inspired by @jerryzh168, at #2728 (comment).

• add attribute to LinearActivationQuantizedTensor: act_pre_scale

Test Plan

python torchao/prototype/smoothquant/example.py
python test/prototype/test_smoothquant.py
python test/integration/test_integration.py

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

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

[namgyu-youn]

Result of test/integration/test_integration.py:

----------------------------------------------------------------------
Ran 358 tests in 477.933s

OK (skipped=205)

[jerryzh168]

@namgyu-youn thanks, I think you probably want to migrate the Int8 tensor first, can you help with that?

#2752 the plain_layout for int8_act_int8_weight

ao/torchao/dtypes/uintx/plain_layout.py

Lines 269 to 315 in 122b307

	def _linear_int8_act_int8_weight_check(input_tensor, weight_tensor, bias):
	return (
	isinstance(input_tensor, AffineQuantizedTensor)
	and _aqt_is_int8_reduced_range(input_tensor)
	and isinstance(weight_tensor, AffineQuantizedTensor)
	and _aqt_is_int8(weight_tensor)
	and input_tensor.dtype == weight_tensor.dtype
	and isinstance(input_tensor._layout, PlainLayout)
	and isinstance(weight_tensor._layout, PlainLayout)
	)
	def _linear_int8_act_int8_weight_impl(input_tensor, weight_tensor, bias):
	#
	# 1. do the matrix form of dot(X_i, W_j)
	#
	#
	# 2. rescale the output
	#
	# in cases with large matrices, y_dot_int32 can grow sufficiently
	# large that y_dot_int32 * a float16 scale is greater than the maximum
	# value of a float 16, (which results in a value of inf even if multiplying
	# by the other scale would bring it within the expected range)
	x_vals_int8 = input_tensor.tensor_impl.int_data
	x_scales = input_tensor.tensor_impl.scale
	w_vals_int8_t = weight_tensor.tensor_impl.int_data.contiguous().t()
	w_scales = weight_tensor.tensor_impl.scale
	tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1])
	x_scales_dtype = x_scales.dtype
	# Cast fp16 scale to float to avoid overflow in int_scaled_matmul
	intermediate_dtype = torch.float if x_scales_dtype == torch.half else x_scales_dtype
	y_dot_scaled = int_scaled_matmul(
	tmp, w_vals_int8_t, x_scales.reshape(-1, 1).to(intermediate_dtype)
	)
	y_dot_scaled = y_dot_scaled.to(x_scales_dtype)
	y = (y_dot_scaled * w_scales).reshape(
	*x_vals_int8.shape[:-1], y_dot_scaled.shape[-1]
	)
	# can downcast only at the very end
	output_dtype = input_tensor.dtype
	y = y.to(output_dtype)
	if bias is not None:
	y += bias
	return y

[namgyu-youn]

@namgyu-youn thanks, I think you probably want to migrate the Int8 tensor first, can you help with that?

#2752 the plain_layout for int8_act_int8_weight

ao/torchao/dtypes/uintx/plain_layout.py

Lines 269 to 315 in 122b307

	def _linear_int8_act_int8_weight_check(input_tensor, weight_tensor, bias):
	return (
	isinstance(input_tensor, AffineQuantizedTensor)
	and _aqt_is_int8_reduced_range(input_tensor)
	and isinstance(weight_tensor, AffineQuantizedTensor)
	and _aqt_is_int8(weight_tensor)
	and input_tensor.dtype == weight_tensor.dtype
	and isinstance(input_tensor._layout, PlainLayout)
	and isinstance(weight_tensor._layout, PlainLayout)
	)
	def _linear_int8_act_int8_weight_impl(input_tensor, weight_tensor, bias):
	#
	# 1. do the matrix form of dot(X_i, W_j)
	#
	#
	# 2. rescale the output
	#
	# in cases with large matrices, y_dot_int32 can grow sufficiently
	# large that y_dot_int32 * a float16 scale is greater than the maximum
	# value of a float 16, (which results in a value of inf even if multiplying
	# by the other scale would bring it within the expected range)
	x_vals_int8 = input_tensor.tensor_impl.int_data
	x_scales = input_tensor.tensor_impl.scale
	w_vals_int8_t = weight_tensor.tensor_impl.int_data.contiguous().t()
	w_scales = weight_tensor.tensor_impl.scale
	tmp = x_vals_int8.reshape(-1, x_vals_int8.shape[-1])
	x_scales_dtype = x_scales.dtype
	# Cast fp16 scale to float to avoid overflow in int_scaled_matmul
	intermediate_dtype = torch.float if x_scales_dtype == torch.half else x_scales_dtype
	y_dot_scaled = int_scaled_matmul(
	tmp, w_vals_int8_t, x_scales.reshape(-1, 1).to(intermediate_dtype)
	)
	y_dot_scaled = y_dot_scaled.to(x_scales_dtype)
	y = (y_dot_scaled * w_scales).reshape(
	*x_vals_int8.shape[:-1], y_dot_scaled.shape[-1]
	)
	# can downcast only at the very end
	output_dtype = input_tensor.dtype
	y = y.to(output_dtype)
	if bias is not None:
	y += bias
	return y

I am definitely interested in it because affine transformation is still hard for me. I will look into it; thanks for your suggestion.

[jerryzh168]

this should always be true?

[jerryzh168]

oh actually this is not ready for review, we need to update int8 tensor instead, converting this to draft now