guard triton version and sm version better

Signed-off-by: Hao Wu <skyw@nvidia.com>

Author: skyw

emerging_optimizers/orthogonalized_optimizers/muon.py

@@ -16,8 +16,10 @@
 from typing import Callable
 
 import torch
+from absl import logging
 from torch.optim.optimizer import ParamsT
 
+from emerging_optimizers import triton_kernels
 from emerging_optimizers.orthogonalized_optimizers.muon_utils import newton_schulz
 from emerging_optimizers.orthogonalized_optimizers.orthogonalized_optimizer import OrthogonalizedOptimizer, _args_doc
 
@@ -80,6 +82,18 @@ def __init__(
         if num_ns_steps < 1:
             raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
 
+        if torch.cuda.is_available():
+            sm_version = torch.cuda.get_device_capability()
+        else:
+            sm_version = (0, 0)
+        if not triton_kernels.HAS_TRITON_340:  # type: ignore[attr-defined]
+            logging.error("Triton 3.4.0 or higher is required for use_syrk to be True.")
+            use_syrk = False
+        elif sm_version not in ((8, 0), (9, 0), (10, 0), (10, 3)):
+            logging.error(
+                f"Correctness of Triton kernel on SM {sm_version} cannot be guaranteed. Setting use_syrk to False."
+            )
+            use_syrk = False
         orthogonalize_fn = partial(
             newton_schulz, steps=num_ns_steps, coefficient_type=coefficient_type, use_syrk=use_syrk
         )

emerging_optimizers/orthogonalized_optimizers/muon_utils.py

@@ -278,6 +278,9 @@ def newton_schulz_step_tsyrk(
     Returns:
         The orthogonalization of X.
     """
+    assert triton_kernels.HAS_TRITON_340, (  # type: ignore[attr-defined]
+        "Triton version doesn't support tensor descriptor API. Minimum required version is 3.4.0."
+    )
     A = triton_kernels.tsyrk_ex(X)  # type: ignore[attr-defined]
     if tp_group is not None:
         torch.distributed.all_reduce(A, op=torch.distributed.ReduceOp.SUM, group=tp_group)

emerging_optimizers/triton_kernels/syrk.py

@@ -20,12 +20,13 @@
 
 try:
     from triton.tools.tensor_descriptor import TensorDescriptor
+
+    HAS_TRITON_340 = True
 except ImportError:
-    raise ImportError(
-        f"Triton version ({triton.__version__}) doesn't support tensor descriptor API. Minimum required version is 3.4.0."
-    )
+    HAS_TRITON_340 = False
+
 
-__all__ = ["ssyrk", "tsyrk_ex"]
+__all__ = ["ssyrk", "tsyrk_ex", "HAS_TRITON_340"]
 
 
 @triton.jit

tests/test_muon_utils.py

@@ -22,7 +22,7 @@
 from emerging_optimizers.orthogonalized_optimizers import muon, muon_utils
 
 
-_SM_VERSION = torch.cuda.get_device_capability() if torch.cuda.is_available() else None
+_SM_VERSION = torch.cuda.get_device_capability() if torch.cuda.is_available() else (0, 0)
 
 
 def newton_schulz_ref(x: torch.Tensor, coefficient_sets: list[tuple[float, float, float]]) -> torch.Tensor:
@@ -212,7 +212,7 @@ def test_qkv_split_shapes_validation(self):
 
 
 @absltest.skipIf(
-    _SM_VERSION is None or _SM_VERSION not in ((8, 0), (9, 0), (10, 0), (11, 0)),
+    _SM_VERSION not in ((8, 0), (9, 0), (10, 0), (10, 3)),
     f"Correctness of Triton kernel on SM {_SM_VERSION} cannot be guaranteed.",
 )
 class TestNewtonSchulzStepWithTsyrk(parameterized.TestCase):