enable smoothquant for int8 static tensor#3468
Merged
Merged
Conversation
Summary: This PR creates a new Int8Tensor and updates the configs to use the new Int8Tensor flow Test Plan: To ensure BC: ``` pytest test/quantization/test_quant_api.py ``` To test new Int8Tensor: ``` pytest test/quantization/quantize_/workflows/int8/test_int8_tensor.py ``` Reviewers: Subscribers: Tasks: Tags:
jerryzh168
reviewed
Dec 18, 2025
jerryzh168
reviewed
Dec 18, 2025
| else: | ||
| raise ValueError(f"Unexpected step: {step}") | ||
|
|
||
| if isinstance(base_config, Int8StaticActivationInt8WeightConfig): |
Contributor
There was a problem hiding this comment.
I think we shouldn't have specific config here, maybe change this to a similar protocol like SupportsActivationPreScaling for config?
Contributor
Author
There was a problem hiding this comment.
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?
Contributor
There was a problem hiding this comment.
we can get this info from a callback I think, like base_config.get_activation_quant_kwargs(), I feel you can iterating from that, but not from the current state
f389a94 to
2586ab6
Compare
jerryzh168
reviewed
Dec 19, 2025
| block_size, | ||
| self.dtype, | ||
| act_quant_kwargs=self.act_quant_kwargs, | ||
| act_scale=self.act_scale, |
Contributor
There was a problem hiding this comment.
I guess slice doesn't work for static quant int8 before, can you add a test for that?
jerryzh168
reviewed
Dec 19, 2025
| old_int8_tensor.scale[index], | ||
| old_int8_tensor.block_size[1:], | ||
| old_int8_tensor.dtype, | ||
| old_int8_tensor.act_scale, |
Contributor
There was a problem hiding this comment.
same for this one, seems like select op breaks before with static quant
Contributor
Author
There was a problem hiding this comment.
For slice and select, they work when granularity is PerDim(=-1), but not otherwise. we now throw an exception when passing PerDim != -1 and I added a test for both.
cyxlily
reviewed
Dec 23, 2025
ea9b8e2 to
c3abc46
Compare
This PR hooks up the static quant workflow added in #3442 to the prototype smoothquant API. You can use the new flow like follows: ```python 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) ```
c3abc46 to
7b45e3e
Compare
jerryzh168
reviewed
Jan 20, 2026
jerryzh168
reviewed
Jan 20, 2026
Comment on lines
+126
to
+127
| if isinstance(base_config, Int8StaticActivationInt8WeightConfig): | ||
| base_config.static_scale = activation_scale |
Contributor
There was a problem hiding this comment.
can do the same as SupportsActivationPreScaling, like IsStaticQuantizationConfig
jerryzh168
reviewed
Jan 20, 2026
jerryzh168
reviewed
Jan 20, 2026
8764cbe to
2f05d2c
Compare
jcaip
added a commit
that referenced
this pull request
Jan 22, 2026
* Int8Tensor migration Summary: This PR creates a new Int8Tensor and updates the configs to use the new Int8Tensor flow Test Plan: To ensure BC: ``` pytest test/quantization/test_quant_api.py ``` To test new Int8Tensor: ``` pytest test/quantization/quantize_/workflows/int8/test_int8_tensor.py ``` Reviewers: Subscribers: Tasks: Tags: * ruff fixes * add init * fix ruff again * update * wip * undo update tests * fix ruff * fix varname * fix typing * add tests * fix dtype * fix ci * address granularity cr * update _choose_quant_func_and_quantize_tensor * make block size required attribute * made dtype required as well * address nits * skip per tensor weight only test for now * add static quant * add static quant * update * static quant working eager + compile * remove file * added asserts * undo smoothquant change * fix return * got smoothquant + int8 static working * generalized smoothquat code * free tests * fix static scale check * update * address cr feedback * Hook up static quant workflow to prototype smoothquant API This PR hooks up the static quant workflow added in #3442 to the prototype smoothquant API. You can use the new flow like follows: ```python 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) ``` * fix ruff * fix test to use threshold for sqnr
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR hooks up the static quant workflow added in #3442 to the prototype smoothquant API.
You can use the new flow like follows:
Test Plan: