From 1aea4c0717a6075cf015aa62cb6ea44e4f36fdc0 Mon Sep 17 00:00:00 2001
From: sayakpaul <spsayakpaul@gmail.com>
Date: Tue, 21 Jul 2026 13:20:34 +0530
Subject: [PATCH] fix torchao tests

---
 tests/models/testing_utils/quantization.py    | 47 ++++++++++++++-----
 .../test_models_transformer_flux.py           |  2 +
 2 files changed, 37 insertions(+), 12 deletions(-)

diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py
index a458fc7b2bef..a218ef929473 100644
--- a/tests/models/testing_utils/quantization.py
+++ b/tests/models/testing_utils/quantization.py
@@ -817,17 +817,23 @@ class TorchAoConfigMixin:
     }
 
     @staticmethod
-    def _get_quant_config(config_name):
+    def _get_quant_config(config_name, modules_to_not_convert=None):
         config_cls = getattr(_torchao_quantization, config_name)
         config_kwargs = {"version": 2}
-        # TorchAO int4 quantization requires plain_int32 packing format on Intel XPU
-        if config_name == "Int4WeightOnlyConfig" and torch_device == "xpu":
-            config_kwargs.setdefault("int4_packing_format", "plain_int32")
-
-        return TorchAoConfig(config_cls(**config_kwargs))
-
-    def _create_quantized_model(self, config_name, **extra_kwargs):
-        config = self._get_quant_config(config_name)
+        # version=2 int4 defaults to the "plain" packing format, which routes through the
+        # fbgemm/mslk Int4Tensor kernels. Pin the packing format to the tinygemm
+        # (_convert_weight_to_int4pack) path on CUDA and plain_int32 on Intel XPU so the tests
+        # don't require those extra kernels to be installed.
+        if config_name == "Int4WeightOnlyConfig":
+            if torch_device == "xpu":
+                config_kwargs["int4_packing_format"] = "plain_int32"
+            elif torch_device == "cuda":
+                config_kwargs["int4_packing_format"] = "tile_packed_to_4d"
+
+        return TorchAoConfig(config_cls(**config_kwargs), modules_to_not_convert=modules_to_not_convert)
+
+    def _create_quantized_model(self, config_name, modules_to_not_convert=None, **extra_kwargs):
+        config = self._get_quant_config(config_name, modules_to_not_convert=modules_to_not_convert)
         kwargs = getattr(self, "pretrained_model_kwargs", {}).copy()
         kwargs["quantization_config"] = config
         kwargs["device_map"] = str(torch_device)
@@ -848,6 +854,11 @@ def _verify_if_layer_quantized(self, name, module, config_kwargs):
     torch_device not in ["cuda", "xpu"], reason="int4wo quantization requires CUDA or XPU"
 )
 
+# The CUDA int4 tinygemm path (Int4TilePackedTo4dTensor) does not implement aten.dequantize.
+_int4wo_dequantize_skip = pytest.mark.skip(
+    reason="int4wo tinygemm packing (Int4TilePackedTo4dTensor) does not support dequantize"
+)
+
 
 @is_torchao
 @require_accelerator
@@ -901,7 +912,7 @@ def test_torchao_quantization_memory_footprint(self, quant_type):
     @pytest.mark.parametrize(
         "quant_type",
         [
-            pytest.param("int4wo", marks=_int4wo_skip),
+            pytest.param("int4wo", marks=[_int4wo_skip, _int4wo_dequantize_skip]),
             "int8wo",
             "int8dq",
         ],
@@ -978,10 +989,22 @@ def test_torchao_modules_to_not_convert(self):
         if modules_to_exclude is None:
             pytest.skip("modules_to_not_convert_for_test not defined for this model")
 
-        self._test_quantization_modules_to_not_convert(
-            TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"], modules_to_exclude
+        # TorchAoConfig takes modules_to_not_convert directly (not inside the quant_type config),
+        # so this can't reuse the dict-based QuantizationTesterMixin helper.
+        model = self._create_quantized_model(
+            TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"], modules_to_not_convert=modules_to_exclude
         )
 
+        found_excluded = False
+        for name, module in model.named_modules():
+            if isinstance(module, torch.nn.Linear) and any(excluded in name for excluded in modules_to_exclude):
+                found_excluded = True
+                assert not self._is_module_quantized(module), (
+                    f"Module {name} should not be quantized but was found to be quantized"
+                )
+
+        assert found_excluded, f"No linear layers found in excluded modules: {modules_to_exclude}"
+
     def test_torchao_device_map(self):
         """Test that device_map='auto' works correctly with quantization."""
         self._test_quantization_device_map(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"])
diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py
index 0214c1f65cab..91e7cea270e5 100644
--- a/tests/models/transformers/test_models_transformer_flux.py
+++ b/tests/models/transformers/test_models_transformer_flux.py
@@ -374,6 +374,8 @@ def pretrained_model_kwargs(self):
 class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin):
     """TorchAO quantization tests for Flux Transformer."""
 
+    modules_to_not_convert_for_test = ["proj_out"]
+
     @property
     def torch_dtype(self):
         return torch.bfloat16