move decorators to testing/utils.py (#1761)

jcaip · web-flow · commit c72ebc65225c · 2025-02-21T19:06:45.000-08:00
* move decorators to testing/utils.py

* add import

* fix import

* fix ruff formatting error

* ruff fixes

* ruff format

* compute_capability test

* update

* update rest of tests

* fix ruff
diff --git a/test/dtypes/test_affine_quantized.py b/test/dtypes/test_affine_quantized.py
@@ -20,12 +20,12 @@
     quantize_,
 )
 from torchao.quantization.quant_primitives import MappingType, ZeroPointDomain
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
     TORCH_VERSION_AT_LEAST_2_6,
     is_fbcode,
     is_sm_at_least_89,
-    skip_if_rocm,
 )
 
 is_cusparselt_available = (
diff --git a/test/dtypes/test_floatx.py b/test/dtypes/test_floatx.py
@@ -27,7 +27,8 @@
     fpx_weight_only,
     quantize_,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode, skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode
 
 _DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
 _Floatx_DTYPES = [(3, 2), (2, 2)]
diff --git a/test/dtypes/test_nf4.py b/test/dtypes/test_nf4.py
@@ -33,7 +33,7 @@
     nf4_weight_only,
     to_nf4,
 )
-from torchao.utils import skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
 
 bnb_available = False
 
diff --git a/test/dtypes/test_uint4.py b/test/dtypes/test_uint4.py
@@ -28,7 +28,8 @@
 from torchao.quantization.quant_api import (
     _replace_with_custom_fn_if_matches_filter,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 
 def _apply_weight_only_uint4_quant(model):
diff --git a/test/float8/test_base.py b/test/float8/test_base.py
@@ -14,11 +14,11 @@
 import torch
 import torch.nn as nn
 
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
     is_sm_at_least_89,
     is_sm_at_least_90,
-    skip_if_rocm,
 )
 
 if not TORCH_VERSION_AT_LEAST_2_5:
diff --git a/test/float8/test_float8_utils.py b/test/float8/test_float8_utils.py
@@ -4,7 +4,8 @@
 import torch
 
 from torchao.float8.float8_utils import _round_scale_down_to_power_of_2
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 if not TORCH_VERSION_AT_LEAST_2_5:
     pytest.skip("Unsupported PyTorch version", allow_module_level=True)
diff --git a/test/hqq/test_hqq_affine.py b/test/hqq/test_hqq_affine.py
@@ -9,9 +9,9 @@
     quantize_,
     uintx_weight_only,
 )
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_3,
-    skip_if_rocm,
 )
 
 cuda_available = torch.cuda.is_available()
diff --git a/test/integration/test_integration.py b/test/integration/test_integration.py
@@ -76,6 +76,7 @@
 from torchao.quantization.utils import (
     compute_error as SQNR,
 )
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_3,
     TORCH_VERSION_AT_LEAST_2_4,
@@ -85,7 +86,6 @@
     benchmark_model,
     is_fbcode,
     is_sm_at_least_90,
-    skip_if_rocm,
     unwrap_tensor_subclass,
 )
 
diff --git a/test/kernel/test_fused_kernels.py b/test/kernel/test_fused_kernels.py
@@ -11,7 +11,7 @@
 import torch
 from galore_test_utils import get_kernel, make_copy, make_data
 
-from torchao.utils import skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
 
 torch.manual_seed(0)
 MAX_DIFF_no_tf32 = 1e-5
diff --git a/test/kernel/test_galore_downproj.py b/test/kernel/test_galore_downproj.py
@@ -11,7 +11,7 @@
 
 from torchao.prototype.galore.kernels.matmul import set_tuner_top_k as matmul_tuner_topk
 from torchao.prototype.galore.kernels.matmul import triton_mm_launcher
-from torchao.utils import skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
 
 torch.manual_seed(0)
 
diff --git a/test/prototype/test_awq.py b/test/prototype/test_awq.py
@@ -5,10 +5,10 @@
 import torch
 
 from torchao.quantization import quantize_
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_3,
     TORCH_VERSION_AT_LEAST_2_5,
-    skip_if_rocm,
 )
 
 if TORCH_VERSION_AT_LEAST_2_3:
diff --git a/test/prototype/test_low_bit_optim.py b/test/prototype/test_low_bit_optim.py
@@ -26,11 +26,11 @@
 from torchao.prototype.low_bit_optim.subclass_4bit import OptimState4bit
 from torchao.prototype.low_bit_optim.subclass_8bit import OptimState8bit
 from torchao.prototype.low_bit_optim.subclass_fp8 import OptimStateFp8
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_4,
     TORCH_VERSION_AT_LEAST_2_5,
     get_available_devices,
-    skip_if_rocm,
 )
 
 try:
diff --git a/test/prototype/test_splitk.py b/test/prototype/test_splitk.py
@@ -14,7 +14,7 @@
     triton_available = False
 
 
-from torchao.utils import skip_if_compute_capability_less_than, skip_if_rocm
+from torchao.testing.utils import skip_if_compute_capability_less_than, skip_if_rocm
 
 
 @unittest.skipIf(not triton_available, "Triton is required but not available")
diff --git a/test/quantization/test_galore_quant.py b/test/quantization/test_galore_quant.py
@@ -18,7 +18,7 @@
     triton_dequant_blockwise,
     triton_quantize_blockwise,
 )
-from torchao.utils import skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
 
 SEED = 0
 torch.manual_seed(SEED)
diff --git a/test/quantization/test_marlin_qqq.py b/test/quantization/test_marlin_qqq.py
@@ -18,7 +18,8 @@
     MappingType,
     choose_qparams_and_quantize_affine_qqq,
 )
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 
 @skip_if_rocm("ROCm enablement in progress")
diff --git a/test/quantization/test_quant_api.py b/test/quantization/test_quant_api.py
@@ -48,14 +48,14 @@
     Int8WeightOnlyQuantizedLinearWeight,
 )
 from torchao.quantization.utils import compute_error
+from torchao.testing.utils import skip_if_rocm
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_3,
     TORCH_VERSION_AT_LEAST_2_4,
     TORCH_VERSION_AT_LEAST_2_5,
     TORCH_VERSION_AT_LEAST_2_6,
     is_sm_at_least_89,
     is_sm_at_least_90,
-    skip_if_rocm,
     unwrap_tensor_subclass,
 )
 
diff --git a/test/sparsity/test_marlin.py b/test/sparsity/test_marlin.py
@@ -15,7 +15,8 @@
 )
 from torchao.sparsity.marlin import inject_24, pack_to_marlin_24, unpack_from_marlin_24
 from torchao.sparsity.sparse_api import apply_fake_sparsity
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm
+from torchao.testing.utils import skip_if_rocm
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5
 
 
 class SparseMarlin24(TestCase):
diff --git a/torchao/testing/utils.py b/torchao/testing/utils.py
@@ -14,7 +14,7 @@
 from torchao.dtypes import AffineQuantizedTensor, to_affine_quantized_intx
 from torchao.quantization import int8_weight_only, quantize_
 from torchao.quantization.quant_primitives import MappingType
-from torchao.utils import TORCH_VERSION_AT_LEAST_2_6
+from torchao.utils import TORCH_VERSION_AT_LEAST_2_6, get_compute_capability
 
 """
 How to use:
@@ -41,6 +41,50 @@ class MyTestCase(TorchAOBasicTestCase):
 """
 
 
+def skip_if_compute_capability_less_than(min_capability):
+    import unittest
+
+    def decorator(test_func):
+        def wrapper(*args, **kwargs):
+            if get_compute_capability() < min_capability:
+                raise unittest.SkipTest(
+                    f"Compute capability is less than {min_capability}"
+                )
+            return test_func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
+def skip_if_rocm(message=None):
+    """Decorator to skip tests on ROCm platform with custom message.
+
+    Args:
+        message (str, optional): Additional information about why the test is skipped.
+    """
+    import pytest
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):
+            if torch.version.hip is not None:
+                skip_message = "Skipping the test in ROCm"
+                if message:
+                    skip_message += f": {message}"
+                pytest.skip(skip_message)
+            return func(*args, **kwargs)
+
+        return wrapper
+
+    # Handle both @skip_if_rocm and @skip_if_rocm() syntax
+    if callable(message):
+        func = message
+        message = None
+        return decorator(func)
+    return decorator
+
+
 # copied from https://github.com/pytorch/pytorch/blob/941d094dd1b507dacf06ddc6ed3485a9537e09b7/test/inductor/test_torchinductor.py#L11389
 def copy_tests(my_cls, other_cls, suffix, test_failures=None, xfail_prop=None):  # noqa: B902
     for name, value in my_cls.__dict__.items():
diff --git a/torchao/utils.py b/torchao/utils.py
@@ -7,7 +7,6 @@
 from math import gcd
 from typing import Any, Callable, Tuple
 
-import pytest
 import torch
 import torch.nn.utils.parametrize as parametrize
 
@@ -16,7 +15,6 @@
     "profiler_runner",
     "get_available_devices",
     "get_compute_capability",
-    "skip_if_compute_capability_less_than",
     "benchmark_torch_function_in_microseconds",
     "find_multiple",
     "_register_custom_op",
@@ -146,49 +144,6 @@ def get_compute_capability():
     return 0.0
 
 
-def skip_if_compute_capability_less_than(min_capability):
-    import unittest
-
-    def decorator(test_func):
-        def wrapper(*args, **kwargs):
-            if get_compute_capability() < min_capability:
-                raise unittest.SkipTest(
-                    f"Compute capability is less than {min_capability}"
-                )
-            return test_func(*args, **kwargs)
-
-        return wrapper
-
-    return decorator
-
-
-def skip_if_rocm(message=None):
-    """Decorator to skip tests on ROCm platform with custom message.
-
-    Args:
-        message (str, optional): Additional information about why the test is skipped.
-    """
-
-    def decorator(func):
-        @functools.wraps(func)
-        def wrapper(*args, **kwargs):
-            if torch.version.hip is not None:
-                skip_message = "Skipping the test in ROCm"
-                if message:
-                    skip_message += f": {message}"
-                pytest.skip(skip_message)
-            return func(*args, **kwargs)
-
-        return wrapper
-
-    # Handle both @skip_if_rocm and @skip_if_rocm() syntax
-    if callable(message):
-        func = message
-        message = None
-        return decorator(func)
-    return decorator
-
-
 def compute_max_diff(output: torch.Tensor, output_ref: torch.Tensor) -> torch.Tensor:
     return torch.mean(torch.abs(output - output_ref)) / torch.mean(
         torch.abs(output_ref)

Original file line number	Diff line number	Diff line change
`@@ -20,12 +20,12 @@`
`20`	`20`	`quantize_,`
`21`	`21`	`)`
`22`	`22`	`from torchao.quantization.quant_primitives import MappingType, ZeroPointDomain`
	`23`	`+from torchao.testing.utils import skip_if_rocm`
`23`	`24`	`from torchao.utils import (`
`24`	`25`	`TORCH_VERSION_AT_LEAST_2_5,`
`25`	`26`	`TORCH_VERSION_AT_LEAST_2_6,`
`26`	`27`	`is_fbcode,`
`27`	`28`	`is_sm_at_least_89,`
`28`		`- skip_if_rocm,`
`29`	`29`	`)`
`30`	`30`
`31`	`31`	`is_cusparselt_available = (`
Original file line number	Diff line number	Diff line change
`@@ -27,7 +27,8 @@`
`27`	`27`	`fpx_weight_only,`
`28`	`28`	`quantize_,`
`29`	`29`	`)`
`30`		`-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode, skip_if_rocm`
	`30`	`+from torchao.testing.utils import skip_if_rocm`
	`31`	`+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, is_fbcode`
`31`	`32`
`32`	`33`	`_DEVICES = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])`
`33`	`34`	`_Floatx_DTYPES = [(3, 2), (2, 2)]`
Original file line number	Diff line number	Diff line change
`@@ -33,7 +33,7 @@`
`33`	`33`	`nf4_weight_only,`
`34`	`34`	`to_nf4,`
`35`	`35`	`)`
`36`		`-from torchao.utils import skip_if_rocm`
	`36`	`+from torchao.testing.utils import skip_if_rocm`
`37`	`37`
`38`	`38`	`bnb_available = False`
`39`	`39`
Original file line number	Diff line number	Diff line change
`@@ -28,7 +28,8 @@`
`28`	`28`	`from torchao.quantization.quant_api import (`
`29`	`29`	`_replace_with_custom_fn_if_matches_filter,`
`30`	`30`	`)`
`31`		`-from torchao.utils import TORCH_VERSION_AT_LEAST_2_5, skip_if_rocm`
	`31`	`+from torchao.testing.utils import skip_if_rocm`
	`32`	`+from torchao.utils import TORCH_VERSION_AT_LEAST_2_5`
`32`	`33`
`33`	`34`
`34`	`35`	`def _apply_weight_only_uint4_quant(model):`
Original file line number	Diff line number	Diff line change
`@@ -9,9 +9,9 @@`
`9`	`9`	`quantize_,`
`10`	`10`	`uintx_weight_only,`
`11`	`11`	`)`
	`12`	`+from torchao.testing.utils import skip_if_rocm`
`12`	`13`	`from torchao.utils import (`
`13`	`14`	`TORCH_VERSION_AT_LEAST_2_3,`
`14`		`- skip_if_rocm,`
`15`	`15`	`)`
`16`	`16`
`17`	`17`	`cuda_available = torch.cuda.is_available()`
Original file line number	Diff line number	Diff line change
`@@ -76,6 +76,7 @@`
`76`	`76`	`from torchao.quantization.utils import (`
`77`	`77`	`compute_error as SQNR,`
`78`	`78`	`)`
	`79`	`+from torchao.testing.utils import skip_if_rocm`
`79`	`80`	`from torchao.utils import (`
`80`	`81`	`TORCH_VERSION_AT_LEAST_2_3,`
`81`	`82`	`TORCH_VERSION_AT_LEAST_2_4,`
`@@ -85,7 +86,6 @@`
`85`	`86`	`benchmark_model,`
`86`	`87`	`is_fbcode,`
`87`	`88`	`is_sm_at_least_90,`
`88`		`- skip_if_rocm,`
`89`	`89`	`unwrap_tensor_subclass,`
`90`	`90`	`)`
`91`	`91`