Implement _dequantize for TorchAO quantizer

[jiqing-feng]

What does this PR do?

Implements the _dequantize() method for TorchAoHfQuantizer, enabling model.dequantize() to convert TorchAO-quantized models back to standard float weights.

Changes

1. Add _dequantize() method: Iterates all nn.Linear modules, calls weight.dequantize() on TorchAOBaseTensor weights, replaces them with standard nn.Parameter, and resets any overridden extra_repr.

2. Fix _verify_if_layer_quantized: Added isinstance(module.weight, TorchAOBaseTensor) check so that dequantized layers (which are still nn.Linear but with plain tensor weights) are correctly detected as non-quantized.

[jiqing-feng]

Hi @sayakpaul . Would you please review this PR? Thanks!

[sayakpaul]

@vkuzo could you review too?

[sayakpaul]

Can we also enable dequantization tests for TorchAO tester mixin?

[jiqing-feng]

done.

[HuggingFaceDocBuilderDev]

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.

[sayakpaul]

Will also let @vkuzo review this.

[vkuzo]

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.

[jiqing-feng]

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.

[sayakpaul]

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):
     @property

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[int8dq] - NotImplementedError: LinearActivationQuantizedTensor dispatch: attempting to run unimplemented operator/funct...

I am on an H100.

[jiqing-feng]

Hi @sayakpaul . I have fixed the dtype issue and skip [int4wo] and [int8dq] since torchao didn't implement them.