diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index e434c9f36252..38dc4ded4b68 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -866,17 +866,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) @@ -897,6 +903,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 @@ -950,7 +961,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", ], @@ -985,10 +996,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 b6a3dfae9e3d..719429526945 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -376,6 +376,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