Update

Author: zeshengzong

test/float8/test_dtensor.py

@@ -11,14 +11,14 @@
 """
 
 import os
+import unittest
 
-import pytest
 import torch
 
 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)
+    raise unittest.SkipTest("Unsupported PyTorch version")
 
 from torch.distributed._tensor import DTensor, Replicate, Shard, distribute_tensor
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh

test/float8/test_float8_utils.py

@@ -5,7 +5,6 @@
 # LICENSE file in the root directory of this source tree.
 import unittest
 
-import pytest
 import torch
 from torch.testing._internal.common_utils import (
     TestCase,
@@ -76,7 +75,7 @@ def test_round_scale_down_to_power_of_2_valid_inputs(
     )
     def test_non_float32_input(self, invalid_dtype: torch.dtype):
         non_float32_tensor = torch.tensor([3.0], dtype=invalid_dtype)
-        with pytest.raises(AssertionError, match="scale must be float32 tensor"):
+        with self.assertRaisesRegex(AssertionError, "scale must be float32 tensor"):
             _round_scale_down_to_power_of_2(non_float32_tensor)
 
 

test/float8/test_fsdp.py

@@ -13,15 +13,15 @@
 
 import copy
 import os
+import unittest
 import warnings
 
 import fire
-import pytest
 
 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)
+    raise unittest.SkipTest("Unsupported PyTorch version")
 
 import torch
 import torch.distributed as dist

test/float8/test_fsdp2_tp.py

@@ -12,14 +12,14 @@
 
 import copy
 import os
+import unittest
 
-import pytest
 import torch
 
 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)
+    raise unittest.SkipTest("Unsupported PyTorch version")
 
 from torch.distributed._composable.fsdp import fully_shard
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh

test/float8/test_fsdp_compile.py

@@ -9,15 +9,15 @@
 """
 
 import os
+import unittest
 import warnings
 
 import fire
-import pytest
 
 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)
+    raise unittest.SkipTest("Unsupported PyTorch version")
 
 import torch
 import torch.distributed as dist

test/float8/test_numerics_integration.py

@@ -7,9 +7,15 @@
 # Tests LLaMa FeedForward numerics with float8
 
 import copy
+import unittest
 from typing import Optional
 
-import pytest
+from torch.testing._internal.common_utils import (
+    TestCase,
+    instantiate_parametrized_tests,
+    parametrize,
+    run_tests,
+)
 
 from torchao.utils import (
     TORCH_VERSION_AT_LEAST_2_5,
@@ -18,7 +24,7 @@
 )
 
 if not TORCH_VERSION_AT_LEAST_2_5:
-    pytest.skip("Unsupported PyTorch version", allow_module_level=True)
+    raise unittest.SkipTest("Unsupported PyTorch version")
 
 import torch
 import torch.nn as nn
@@ -83,7 +89,7 @@ def init_weights(self, init_std: float):
             nn.init.trunc_normal_(linear.weight, mean=0.0, std=init_std)
 
 
-class TestFloat8NumericsIntegrationTest:
+class TestFloat8NumericsIntegrationTest(TestCase):
     def _test_impl(self, config: Float8LinearConfig) -> None:
         data_dtype = torch.bfloat16
         # LLaMa 3 70B shapes
@@ -147,22 +153,20 @@ def _test_impl(self, config: Float8LinearConfig) -> None:
             sqnr = compute_error(ref_grad, cur_grad)
             assert sqnr > grad_sqnr_threshold
 
-    @pytest.mark.parametrize(
+    @parametrize(
         "scaling_type_input",
         [ScalingType.DYNAMIC],
     )
-    @pytest.mark.parametrize(
+    @parametrize(
         "scaling_type_weight",
         [ScalingType.DYNAMIC],
     )
-    @pytest.mark.parametrize(
+    @parametrize(
         "scaling_type_grad_output",
         [ScalingType.DYNAMIC],
     )
-    @pytest.mark.skipif(
-        not is_sm_at_least_89(), reason="requires SM89 compatible machine"
-    )
-    @pytest.mark.skipif(IS_ROCM, reason="test doesn't currently work on the ROCm stack")
+    @unittest.skipIf(not is_sm_at_least_89(), reason="requires SM89 compatible machine")
+    @unittest.skipIf(IS_ROCM, reason="test doesn't currently work on the ROCm stack")
     def test_encoder_fw_bw_from_config_params(
         self,
         scaling_type_input: ScalingType,
@@ -177,17 +181,15 @@ def test_encoder_fw_bw_from_config_params(
         )
         self._test_impl(config)
 
-    @pytest.mark.parametrize(
+    @parametrize(
         "recipe_name",
         [
             Float8LinearRecipeName.ROWWISE,
             Float8LinearRecipeName.ROWWISE_WITH_GW_HP,
         ],
     )
-    @pytest.mark.skipif(
-        not is_sm_at_least_90(), reason="requires SM90 compatible machine"
-    )
-    @pytest.mark.skipif(IS_ROCM, reason="test doesn't currently work on the ROCm stack")
+    @unittest.skipIf(not is_sm_at_least_90(), reason="requires SM90 compatible machine")
+    @unittest.skipIf(IS_ROCM, reason="test doesn't currently work on the ROCm stack")
     def test_encoder_fw_bw_from_recipe(
         self,
         recipe_name: str,
@@ -196,5 +198,7 @@ def test_encoder_fw_bw_from_recipe(
         self._test_impl(config)
 
 
+instantiate_parametrized_tests(TestFloat8NumericsIntegrationTest)
+
 if __name__ == "__main__":
-    pytest.main([__file__])
+    run_tests()