unify NPU and XPU test cases into a single class

Author: orangeH25

test/quantization/quantize_/workflows/int4/test_int4_plain_int32_tensor.py

@@ -5,12 +5,12 @@
 # LICENSE file in the root directory of this source tree.
 
 import tempfile
-import unittest
 
+import pytest
 import torch
+from torch.testing._internal.common_device_type import instantiate_device_type_tests
 from torch.testing._internal.common_utils import (
     TestCase,
-    instantiate_parametrized_tests,
     parametrize,
     run_tests,
 )
@@ -33,9 +33,19 @@ def get_config(group_size):
     )
 
 
-@unittest.skipIf(not torch_version_at_least("2.8.0"), "Need pytorch 2.8+")
-@unittest.skipIf(not torch.xpu.is_available(), "XPU not available")
-class Int4PlainInt32TensorXPU(TestCase):
+class Int4PlainInt32Tensor(TestCase):
+    _MIN_VER = {
+        "xpu": "2.8.0",
+        "npu": "2.7.1",
+    }
+
+    def setUp(self):
+        min_req = type(self)._MIN_VER.get(self.device_type)
+        if not torch_version_at_least(min_req):
+            self.skipTest(
+                f"{self.device_type} requires torch >= {min_req}, current {torch.__version__}"
+            )
+
     @parametrize(
         "sizes",
         [
@@ -46,90 +56,36 @@ class Int4PlainInt32TensorXPU(TestCase):
     )
     @parametrize("dtype", [torch.bfloat16, torch.half])
     @parametrize("group_size", [32, 64, 128])
-    def test_linear(self, sizes, dtype, group_size):
-        device = "xpu"
+    @parametrize("thresholds", [{"xpu": 20, "npu": 10}])
+    def test_linear(self, device, sizes, dtype, group_size, thresholds):
         M, N, K = sizes
+        if "npu" in device and group_size == K:
+            pytest.skip(
+                f"{device} does not support group_size equal to K dimension ({group_size} == {K})"
+            )
+        threshold = thresholds.get(device.split(":")[0])
+
         input = torch.randn(*M, K, dtype=dtype, device=device)
         linear = torch.nn.Linear(K, N, dtype=dtype, device=device)
         original = linear(input)
         quantize_(linear, get_config(group_size))
         quantized = linear(input)
-        self.assertTrue(compute_error(original, quantized) > 20)
+        self.assertTrue(compute_error(original, quantized) > threshold)
 
-        compiled_linear = torch.compile(linear)
-        quantized_and_compiled = compiled_linear(input)
-        self.assertTrue(compute_error(original, quantized_and_compiled) > 20)
+        if "xpu" in device:
+            compiled_linear = torch.compile(linear)
+            quantized_and_compiled = compiled_linear(input)
+            self.assertTrue(compute_error(original, quantized_and_compiled) > threshold)
 
     @parametrize("dtype", [torch.bfloat16, torch.half])
-    def test_module_path(self, dtype):
-        linear = torch.nn.Linear(128, 256, dtype=dtype, device="xpu")
-        quantize_(linear, get_config(group_size=128))
-        self.assertEqual(
-            str(type(linear.weight)),
-            "<class 'torchao.quantization.Int4PlainInt32Tensor'>",
-        )
-
-        with tempfile.NamedTemporaryFile() as f:
-            torch.save(linear.state_dict(), f)
-            f.seek(0)
-            state_dict = torch.load(f)
-            self.assertEqual(
-                str(type(state_dict["weight"])),
-                "<class 'torchao.quantization.Int4PlainInt32Tensor'>",
-            )
-
-    def test_activation_prescaling(self):
-        dtype = torch.bfloat16
-        device = "xpu"
-        input = torch.randn(1, 128, dtype=dtype, device=device)
-        linear = torch.nn.Linear(128, 256, bias=False, dtype=dtype, device=device)
-        original = linear(input)
-        quantize_(linear, get_config(128))
-        qw = linear.weight
-        assert isinstance(qw, SupportsActivationPreScaling), (
-            "Expected int4 tensor supports activation prescaling"
-        )
-        assert qw.act_pre_scale is None, "Default `act_pre_scale` is None"
-        _ACT_PRE_SCALE = 2
-        qw.act_pre_scale = _ACT_PRE_SCALE
-        quantized = linear(input)
-
-        # making sure activation pre scaling is successfully applied to the activation
-        self.assertTrue(compute_error(original * _ACT_PRE_SCALE, quantized) > 20)
-
+    def test_module_path(self, device, dtype):
+        device = self.device_type
+        K, N, group_size = 128, 256, 128
+        if "npu" in device:
+            group_size = 64
 
-@unittest.skipIf(not torch_version_at_least("2.7.1"), "Need pytorch 2.7.1+")
-@unittest.skipIf(
-    torch.accelerator.current_accelerator().type != "npu"
-    or not torch.accelerator.is_available(),
-    "NPU not available",
-)
-class Int4PlainInt32TensorNPU(TestCase):
-    @parametrize("device", ["npu"])
-    @parametrize(
-        "sizes",
-        [
-            ((128,), 256, 128),
-            ((32, 128), 512, 128),
-            ((2, 32, 128), 256, 128),
-        ],
-    )
-    @parametrize("dtype", [torch.float16, torch.bfloat16])
-    @parametrize("group_size", [32, 64])
-    def test_linear(self, device, sizes, dtype, group_size):
-        M, N, K = sizes
-        input = torch.randn(*M, K, dtype=dtype, device=device)
         linear = torch.nn.Linear(K, N, dtype=dtype, device=device)
-        orig_output = linear(input)
         quantize_(linear, get_config(group_size))
-        quantized_output = linear(input)
-        self.assertTrue(compute_error(orig_output, quantized_output) > 10)
-
-    @parametrize("device", ["npu"])
-    @parametrize("dtype", [torch.float16, torch.bfloat16])
-    def test_module_path(self, device, dtype):
-        linear = torch.nn.Linear(128, 256, dtype=dtype, device=device)
-        quantize_(linear, get_config(group_size=64))
         self.assertEqual(
             str(type(linear.weight)),
             "<class 'torchao.quantization.Int4PlainInt32Tensor'>",
@@ -144,13 +100,22 @@ def test_module_path(self, device, dtype):
                 "<class 'torchao.quantization.Int4PlainInt32Tensor'>",
             )
 
-    @parametrize("device", ["npu"])
     @parametrize("dtype", [torch.float16, torch.bfloat16])
-    def test_activation_prescaling(self, device, dtype):
-        input = torch.randn(1, 128, dtype=dtype, device=device)
-        linear = torch.nn.Linear(128, 256, bias=False, dtype=dtype, device=device)
+    @parametrize("thresholds", [{"xpu": 20, "npu": 10}])
+    def test_activation_prescaling(self, device, dtype, thresholds):
+        device = self.device_type
+        if "xpu" in device and dtype == torch.float16:
+            pytest.skip(f"{device} test_activation_prescaling don't test {dtype}")
+
+        threshold = thresholds.get(device.split(":")[0])
+        K, N, group_size = 128, 256, 128
+        if "npu" in device:
+            group_size = 64
+
+        input = torch.randn(1, K, dtype=dtype, device=device)
+        linear = torch.nn.Linear(K, N, bias=False, dtype=dtype, device=device)
         original = linear(input)
-        quantize_(linear, get_config(64))
+        quantize_(linear, get_config(group_size))
         qw = linear.weight
         assert isinstance(qw, SupportsActivationPreScaling), (
             "Expected int4 tensor supports activation prescaling"
@@ -161,11 +126,13 @@ def test_activation_prescaling(self, device, dtype):
         quantized = linear(input)
 
         # making sure activation pre scaling is successfully applied to the activation
-        self.assertTrue(compute_error(original * _ACT_PRE_SCALE, quantized) > 10)
+        self.assertTrue(compute_error(original * _ACT_PRE_SCALE, quantized) > threshold)
+
 
+instantiate_device_type_tests(
+    Int4PlainInt32Tensor, globals(), only_for=("xpu", "npu"), allow_xpu=True
+)
 
-instantiate_parametrized_tests(Int4PlainInt32TensorXPU)
-instantiate_parametrized_tests(Int4PlainInt32TensorNPU)
 
 if __name__ == "__main__":
     run_tests()