enable smoothquant for int8 static tensor

[jcaip]

This PR hooks up the static quant workflow added in #3442 to the prototype smoothquant API.

You can use the new flow like follows:

from torchao.quantization.quant_api import (
    Int8StaticActivationInt8WeightConfig,
)
from torchao.prototype.smoothquant import (
    SmoothQuantConfig
)

config = SmoothQuantConfig(
            base_config=Int8StaticActivationInt8Weight(granularity=PerRow()),
            step=SmoothQuantStep.PREPARE,
            alpha=0.5,
        )

quantize_(model, config)

# Perform calibration with test data
model(*x)

config.step = SmoothQuantStep.CONVERT
quantize_(model, config)

# model will now be statically quantized with the inputs used in smoothquant observer. 
model(*x)

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 2 New Failures

As of commit 2586ab6 with merge base f99105a ():

NEW FAILURES - The following jobs have failed:

• Run Regression Tests / test-nightly (CUDA Nightly, linux.g5.12xlarge.nvidia.gpu, --pre torch --index-url https://downloa... / linux-job (gh)
  test/test_low_bit_optim.py::TestFSDP2::test_fsdp2
• Run TorchAO Experimental MPS Tests / test-mps-ops / macos-job (gh)
  RuntimeError: Command bash /Users/ec2-user/runner/_work/_temp/exec_script failed with exit code 2

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

[jcaip]

cc @Xia-Weiwen and @cyxlily fyi

[jerryzh168]

we should not be using this, please check awq on how this should be implemented in the new stack:

ao/torchao/prototype/awq/api.py

Lines 108 to 113 in 08e5e20

	assert isinstance(qw, SupportsActivationPreScaling), (
	"weight must support activation scaling through implementing `SupportsActivationPreScaling`"
	)
	# since we want to do `act` * `act_pre_scale` during runtime for speed, we'll save the
	# reciprocal of the `equalization_scale`
	qw.act_pre_scale = 1.0 / equalization_scale

[jerryzh168]

nit: Optional[torch.Tensor]

[jerryzh168]

also maybe static_scale might be more descriptive I feel

[meta-codesync]

@jcaip has imported this pull request. If you are a Meta employee, you can view this in D88784212.

[cyxlily]

@jcaip Our customer needs activation quantization PerTensor and weight quantization PerRow. Will you implement it, or may I create a new PR to do it?

[jcaip]

@cyxlily feel free to open a new PR for activation per tensor x weight per row, it's not something im planning to do currently.

Thank you for your smoothquant pr btw, I used it to implement this.

[jerryzh168]

are we trying to say static_out_compile == static_out_eager == dynamic_out_compile == dynamic_out_eager here? if so I think it might be clearer just to assert all these are equal to each other

[jerryzh168]

I think we shouldn't have specific config here, maybe change this to a similar protocol like SupportsActivationPreScaling for config?

[jcaip]

I think figuring out how to do this generally will need a bit more design, we'd need to figure out how to map to the appropriate QuantizeTensorToInt/FloatXKwargs object. Agree we should be able to do this though, but can I address in a later PR?

[jerryzh168]

I guess slice doesn't work for static quant int8 before, can you add a test for that?

[jerryzh168]

same for this one, seems like select op breaks before with static quant