From 5e5b575fb3413fbaf04a949c3d0fa5796b79e4f4 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 25 Mar 2026 09:38:49 +0530 Subject: [PATCH 1/5] fix torchao tests --- tests/models/testing_utils/quantization.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 0f1fbde72485..ec74422741c3 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -177,6 +177,11 @@ def _test_quantization_inference(self, config_kwargs): model_quantized.to(torch_device) inputs = self.get_dummy_inputs() + model_dtype = next(model_quantized.parameters()).dtype + inputs = { + k: v.to(dtype=model_dtype) if torch.is_tensor(v) and torch.is_floating_point(v) else v + for k, v in inputs.items() + } output = model_quantized(**inputs, return_dict=False)[0] assert output is not None, "Model output is None" @@ -930,6 +935,7 @@ 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"]) + @pytest.mark.xfail(reason="dequantize is not implemented in torchao") def test_torchao_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"]) From 4e01e02395145cd79e258cd40ad4ec0d62f4c42c Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Wed, 25 Mar 2026 09:41:04 +0530 Subject: [PATCH 2/5] add mslk for additional dependencies. --- .github/workflows/nightly_tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 416d2af3fc2e..e242b4b57cb0 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -341,7 +341,7 @@ jobs: additional_deps: ["peft", "kernels"] - backend: "torchao" test_location: "torchao" - additional_deps: [] + additional_deps: [mslk-cuda] - backend: "optimum_quanto" test_location: "quanto" additional_deps: [] From d742b19f8b571205582a539342498277724d7887 Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 1 May 2026 12:16:22 +0530 Subject: [PATCH 3/5] add dtype --- tests/models/testing_utils/quantization.py | 5 ----- tests/models/transformers/test_models_transformer_flux.py | 4 ++++ 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index 64ab15d272d7..b08bd8d2e37d 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -175,11 +175,6 @@ def _test_quantization_inference(self, config_kwargs): model_quantized.to(torch_device) inputs = self.get_dummy_inputs() - model_dtype = next(model_quantized.parameters()).dtype - inputs = { - k: v.to(dtype=model_dtype) if torch.is_tensor(v) and torch.is_floating_point(v) else v - for k, v in inputs.items() - } output = model_quantized(**inputs, return_dict=False)[0] assert output is not None, "Model output is None" diff --git a/tests/models/transformers/test_models_transformer_flux.py b/tests/models/transformers/test_models_transformer_flux.py index e4e91e52fb80..5eaadf7a8ad2 100644 --- a/tests/models/transformers/test_models_transformer_flux.py +++ b/tests/models/transformers/test_models_transformer_flux.py @@ -362,6 +362,10 @@ def pretrained_model_kwargs(self): class TestFluxTransformerTorchAo(FluxTransformerTesterConfig, TorchAoTesterMixin): """TorchAO quantization tests for Flux Transformer.""" + @property + def torch_dtype(self): + return torch.bfloat16 + class TestFluxTransformerGGUF(FluxTransformerTesterConfig, GGUFTesterMixin): @property From e5daed105ce3517b0b3ccb88792eea38abc60dae Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Tue, 12 May 2026 17:42:56 +0900 Subject: [PATCH 4/5] more fixes --- tests/quantization/torchao/test_torchao.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/quantization/torchao/test_torchao.py b/tests/quantization/torchao/test_torchao.py index 8a811cfc1c73..53e1ed1dc2b7 100644 --- a/tests/quantization/torchao/test_torchao.py +++ b/tests/quantization/torchao/test_torchao.py @@ -29,6 +29,7 @@ FluxTransformer2DModel, TorchAoConfig, ) +from diffusers.models.attention import AttentionModuleMixin from diffusers.models.attention_processor import Attention from diffusers.quantizers import PipelineQuantizationConfig @@ -231,8 +232,8 @@ def test_quantization(self): (Int8DynamicActivationIntxWeightConfig(version=2), np.array([0.4688, 0.5195, 0.5547, 0.418, 0.4414, 0.6406, 0.4336, 0.4531, 0.5625])), (Int8WeightOnlyConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), (Int8DynamicActivationInt8WeightConfig(version=2), np.array([0.4648, 0.5195, 0.5547, 0.4199, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), - (IntxWeightOnlyConfig(dtype=torch.uint4, group_size=16, version=2), np.array([0.4609, 0.5234, 0.5508, 0.4199, 0.4336, 0.6406, 0.4316, 0.4531, 0.5625])), - (IntxWeightOnlyConfig(dtype=torch.uint7, group_size=16, version=2), np.array([0.4648, 0.5195, 0.5547, 0.4219, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), + (IntxWeightOnlyConfig(weight_dtype=torch.int4, version=2), np.array([0.4609, 0.5234, 0.5508, 0.4199, 0.4336, 0.6406, 0.4316, 0.4531, 0.5625])), + (IntxWeightOnlyConfig(weight_dtype=torch.int7, version=2), np.array([0.4648, 0.5195, 0.5547, 0.4219, 0.4414, 0.6445, 0.4316, 0.4531, 0.5625])), ] if _is_xpu_or_cuda_capability_atleast_8_9(): @@ -393,7 +394,7 @@ def test_training(self): param.data = param.data.to(torch.float32) for _, module in quantized_model.named_modules(): - if isinstance(module, Attention): + if isinstance(module, (Attention, AttentionModuleMixin)): module.to_q = LoRALayer(module.to_q, rank=4) module.to_k = LoRALayer(module.to_k, rank=4) module.to_v = LoRALayer(module.to_v, rank=4) @@ -457,7 +458,7 @@ def test_memory_footprint(self): # Will quantize all the linear layers for module in transformer_int8wo.modules(): if isinstance(module, nn.Linear): - self.assertTrue(isinstance(module.weight, Int8Tensor)) + self.assertTrue(isinstance(module.weight, TorchAOBaseTensor)) total_int4wo = get_model_size_in_bytes(transformer_int4wo) total_int4wo_gs32 = get_model_size_in_bytes(transformer_int4wo_gs32) From 7ce56c855588bf3decae9d24e832e986d4b50e79 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Fri, 5 Jun 2026 11:10:44 +0000 Subject: [PATCH 5/5] dequantization actually works in torchao --- tests/models/testing_utils/quantization.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index efd3ca5d0b30..ded5cab52268 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -933,7 +933,6 @@ 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"]) - @pytest.mark.xfail(reason="dequantize is not implemented in torchao") def test_torchao_dequantize(self): """Test that dequantize() works correctly.""" self._test_dequantize(TorchAoConfigMixin.TORCHAO_QUANT_TYPES["int8wo"])