Add int8 static quantization workflow

[jcaip]

This PR adds in a new static quant workflow based off of Int8Tensor (#3407).

It introduces a new config, Int8StaticActivationInt8WeightConfig which requires a scale tensor and granularity

static_config = Int8StaticActivationInt8WeightConfig(
    scale=int8_input.scale.detach(), granularity=PerRow
)
quantize_(model_static_quant, static_config)

Currently PerRow and PerTensor symmetric quant is support only.

This scale tensor is stored on the weight Int8Tensor under activation_scale, and is used to create a new activation Int8Tensor for static quantization.

It would be nice to store this scale tensor in QuantizeTensorToInt8Kwargs but unfortunately this breaks dynamo tracing, as we store the quant kwargs as an object for the weight tensor and we are unable to fakeify them properly.

As a result, we need to keep track and pass scale outside of this Kwargs object.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ No Failures

As of commit b5309eb with merge base c4273fe ():
💚 Looks good so far! There are no failures yet. 💚

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

[jcaip]

note: I changed these checks in #3468

[jerryzh168]

I think we typically also check the shape of scale tensor as well, like these

ao/torchao/float8/inference.py

Lines 181 to 198 in 08e5e20

	def _is_rowwise_scaled(x: torch.Tensor) -> bool:
	"""Checks if a quantized tensor is rowwise scaled
	Args:
	x: quantized tensor (should have `block_size` attribute)
	"""
	assert hasattr(x, "block_size"), "Expecting input to have `block_size` attribute"
	return tuple(x.block_size) == (1,) * (x.dim() - 1) + (x.shape[-1],)
	def _is_tensorwise_scaled(x: torch.Tensor) -> bool:
	"""Checks if a quantized tensor is rowwise scaled
	Args:
	x: quantized tensor (should have `block_size` attribute)
	"""
	assert hasattr(x, "block_size"), "Expecting input to have `block_size` attribute"
	return all(
	x.block_size[i] == -1 or x.block_size[i] == x.shape[i] for i in range(x.ndim)
	)

[jerryzh168]

nit: also check the shapes, and check PerRow as well

[jcaip]

I think we might want to enable scales to be None, for passing Int8StaticActivationInt8WeightConfig() as a base config, we can discuss on #3468

[jerryzh168]

OK makes sense, sounds good, this is easier for user I think, otherwise they have to do a separate flow to get the scale here

[jerryzh168]

lg, see some comments inline