Remove warnings in favor of skiptests for Moe code (#2654)

* Remove warnings in favor of skiptests for Moe code

* update

* lint

Author: msaroufim

test/float8/test_base.py

@@ -8,7 +8,6 @@
 import random
 import re
 import unittest
-import warnings
 
 import pytest
 import torch
@@ -381,6 +380,9 @@ def test_linear_from_config_params(
         "linear_dtype", [torch.bfloat16, torch.float16, torch.float32]
     )
     @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+    @unittest.skipIf(
+        torch.cuda.is_available() and not is_sm_at_least_90(), "CUDA capability < 9.0"
+    )
     @skip_if_rocm("ROCm enablement in progress")
     def test_linear_from_recipe(
         self,
@@ -389,12 +391,6 @@ def test_linear_from_recipe(
         linear_dtype: torch.dtype,
         linear_bias: bool,
     ):
-        if torch.cuda.get_device_capability() < (9, 0):
-            warnings.warn(
-                f"CUDA capability {torch.cuda.get_device_capability()} < (9.0)"
-            )
-            pytest.skip()
-
         x = torch.randn(*x_shape, device="cuda", dtype=linear_dtype)
         m_ref = nn.Linear(16, 32, bias=linear_bias, device="cuda", dtype=linear_dtype)
         config = Float8LinearConfig.from_recipe_name(recipe_name)

test/prototype/moe_training/test_fsdp.py

@@ -38,10 +38,9 @@
     from torchtitan.experiments.llama4.model.args import TransformerModelArgs
     from torchtitan.experiments.llama4.model.moe import MoE
 except ImportError:
-    import warnings
-
-    warnings.warn("torchtitan not installed, skipping MoE tests.")
-    pytest.skip(allow_module_level=True)
+    pytest.skip(
+        "torchtitan not installed, skipping MoE tests.", allow_module_level=True
+    )
 
 
 def test_moe_float8_training_fsdp():

test/prototype/moe_training/test_fsdp_tp.py

@@ -30,12 +30,10 @@
         parallelize_module,
     )
 except ImportError:
-    import warnings
-
-    warnings.warn(
-        "torch version is too old, these tests require nightly build. Skipping MoE training tests."
+    pytest.skip(
+        "torch version is too old, these tests require nightly build. Skipping MoE training tests.",
+        allow_module_level=True,
     )
-    pytest.skip(allow_module_level=True)
 
 # this feature requires CUDA and SM89+
 if not torch.cuda.is_available() or torch.cuda.get_device_capability() < (8, 9):
@@ -60,10 +58,9 @@
     from torchtitan.experiments.llama4.model.args import TransformerModelArgs
     from torchtitan.experiments.llama4.model.moe import MoE
 except ImportError:
-    import warnings
-
-    warnings.warn("torchtitan not installed, skipping MoE tests.")
-    pytest.skip(allow_module_level=True)
+    pytest.skip(
+        "torchtitan not installed, skipping MoE tests.", allow_module_level=True
+    )
 
 
 @pytest.mark.parametrize(

test/prototype/moe_training/test_tp.py

@@ -29,12 +29,10 @@
         parallelize_module,
     )
 except ImportError:
-    import warnings
-
-    warnings.warn(
-        "torch version is too old, these tests require nightly build. Skipping MoE training tests."
+    pytest.skip(
+        "torch version is too old, these tests require nightly build. Skipping MoE training tests.",
+        allow_module_level=True,
     )
-    pytest.skip(allow_module_level=True)
 
 
 # this feature requires CUDA and SM89+
@@ -60,10 +58,9 @@
     from torchtitan.experiments.llama4.model.args import TransformerModelArgs
     from torchtitan.experiments.llama4.model.moe import MoE
 except ImportError:
-    import warnings
-
-    warnings.warn("torchtitan not installed, skipping MoE tests.")
-    pytest.skip(allow_module_level=True)
+    pytest.skip(
+        "torchtitan not installed, skipping MoE tests.", allow_module_level=True
+    )
 
 
 @pytest.mark.parametrize(

test/prototype/moe_training/test_training.py

@@ -22,10 +22,9 @@
     from torchtitan.experiments.llama4.model.args import TransformerModelArgs
     from torchtitan.experiments.llama4.model.moe import MoE
 except ImportError:
-    import warnings
-
-    warnings.warn("torchtitan not installed, skipping MoE tests.")
-    pytest.skip(allow_module_level=True)
+    pytest.skip(
+        "torchtitan not installed, skipping MoE tests.", allow_module_level=True
+    )
 
 
 @pytest.mark.parametrize(