Implement _dequantize for TorchAO quantizer#13538
Merged
Merged
Conversation
Contributor
Author
|
Hi @sayakpaul . Would you please review this PR? Thanks! |
sayakpaul
reviewed
May 13, 2026
Comment on lines
+826
to
+828
| assert isinstance(module.weight, TorchAOBaseTensor), ( | ||
| f"Layer {name} weight is {type(module.weight)}, expected TorchAOBaseTensor" | ||
| ) |
Member
There was a problem hiding this comment.
Can we also enable dequantization tests for TorchAO tester mixin?
- Add _dequantize() method in TorchAoHfQuantizer that dequantizes TorchAOBaseTensor weights back to standard nn.Parameter - Fix _verify_if_layer_quantized to check isinstance(weight, TorchAOBaseTensor) so dequantized layers are correctly detected as non-quantized
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
Signed-off-by: jiqing-feng <[email protected]>
a819214 to
9102fb8
Compare
vkuzo
reviewed
May 13, 2026
| from torchao.utils import TorchAOBaseTensor | ||
|
|
||
| for name, module in model.named_modules(): | ||
| if isinstance(module, nn.Linear) and isinstance(module.weight, TorchAOBaseTensor): |
Contributor
There was a problem hiding this comment.
TorchAOBaseTensor does not expose dequantize as a public API, it is defined on child classes. I agree that it would make sense to do so in the future. If you want to be safe here it might be better to check for individual tensor subclasses that do expose it.
Contributor
Author
There was a problem hiding this comment.
Thanks for the review @vkuzo! You're right that dequantize() is defined on child classes rather than on TorchAOBaseTensor itself. I've added a hasattr guard so we safely skip any subclass that doesn't expose it. In practice all quantized tensor subclasses we encounter do implement dequantize(), but this makes it future-proof.
Signed-off-by: jiqing-feng <[email protected]>
95d0118 to
df36f1a
Compare
sayakpaul
reviewed
May 14, 2026
| ], | ||
| ids=["int4wo", "int8wo", "int8dq"], | ||
| ) | ||
| def test_torchao_dequantize(self, quant_type): |
Member
There was a problem hiding this comment.
I ran the tests with the following command: ``
And there are test failures:
FAILED tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerTorchAo::test_torchao_dequantize[int4wo] - NotImplementedError: Int4Tensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload...
FAILED tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerTorchAo::test_torchao_dequantize[int8wo] - RuntimeError: mat1 and mat2 must have the same dtype, but got Float and BFloat16
FAILED tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerTorchAo::test_torchao_dequantize[int8dq] - NotImplementedError: LinearActivationQuantizedTensor dispatch: attempting to run unimplemented operator/funct...With the following diff I managed to get it down to two:
diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py
index 840eaa338..e73c31561 100644
--- a/tests/models/transformers/test_models_transformer_flux.py
+++ b/tests/models/transformers/test_models_transformer_flux.py
@@ -367,6 +367,10 @@ class TestFluxTransformerQuanto(FluxTransformerTesterConfig, QuantoTesterMixin):
class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin):
"""TorchAO quantization tests for Flux Transformer."""
+ @property
+ def torch_dtype(self):
+ return torch.bfloat16
+
class TestFluxTransformerGGUF(FluxTransformerTesterConfig, GGUFTesterMixin):
@propertyFAILED tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerTorchAo::test_torchao_dequantize[int4wo] - NotImplementedError: Int4Tensor dispatch: attempting to run unimplemented operator/function: func=<OpOverload...
FAILED tests/models/transformers/test_models_transformer_flux.py::TestFluxTransformerTorchAo::test_torchao_dequantize[int8dq] - NotImplementedError: LinearActivationQuantizedTensor dispatch: attempting to run unimplemented operator/funct...I am on an H100.
2cbe719 to
83431bf
Compare
83431bf to
450d0e4
Compare
Signed-off-by: jiqing-feng <[email protected]>
Contributor
Author
|
Hi @sayakpaul . I have fixed the dtype issue and skip |
vkuzo
reviewed
May 15, 2026
| for name, module in model.named_modules(): | ||
| if isinstance(module, nn.Linear) and isinstance(module.weight, TorchAOBaseTensor): | ||
| if not hasattr(module.weight, "dequantize"): | ||
| continue |
Contributor
There was a problem hiding this comment.
should this throw an exception then? looks like right now it's silently failing to dequantize
Signed-off-by: jiqing-feng <[email protected]>
sayakpaul
approved these changes
May 18, 2026
sayakpaul
left a comment
Member
There was a problem hiding this comment.
Thanks for working on this!
|
|
||
| # Delete quantizer and quantization config | ||
| del model.hf_quantizer | ||
| model.is_quantized = False |
Enderfga
pushed a commit
to Enderfga/diffusers
that referenced
this pull request
May 19, 2026
* Implement _dequantize for TorchAO quantizer - Add _dequantize() method in TorchAoHfQuantizer that dequantizes TorchAOBaseTensor weights back to standard nn.Parameter - Fix _verify_if_layer_quantized to check isinstance(weight, TorchAOBaseTensor) so dequantized layers are correctly detected as non-quantized * enable dequantize for TorchAO tester mixin Signed-off-by: jiqing-feng <[email protected]> * check dequantize Signed-off-by: jiqing-feng <[email protected]> * fix dequantize: clear is_quantized flag and cast dtype after dequantize * fix Signed-off-by: jiqing-feng <[email protected]> * fix error report Signed-off-by: jiqing-feng <[email protected]> --------- Signed-off-by: jiqing-feng <[email protected]> Co-authored-by: Sayak Paul <[email protected]>
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
4 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.
What does this PR do?
Implements the
_dequantize()method forTorchAoHfQuantizer, enablingmodel.dequantize()to convert TorchAO-quantized models back to standard float weights.Changes
Add
_dequantize()method: Iterates allnn.Linearmodules, callsweight.dequantize()onTorchAOBaseTensorweights, replaces them with standardnn.Parameter, and resets any overriddenextra_repr.Fix
_verify_if_layer_quantized: Addedisinstance(module.weight, TorchAOBaseTensor)check so that dequantized layers (which are stillnn.Linearbut with plain tensor weights) are correctly detected as non-quantized.