[Feature] Add CuTeDSL Gram Newton-Schulz support for Muon

examples/v1/config/sft_qwen3p5.py

@@ -102,6 +102,7 @@ def _get_bool_env(name: str, default: bool = False) -> bool:
 optim_cfg = MuonConfig(
     lr=float(os.environ.get("LR", "2e-5")),
     weight_decay=float(os.environ.get("WEIGHT_DECAY", "0.05")),
+    use_gram_newton_schulz=_get_bool_env("USE_GRAM_NEWTON_SCHULZ", False),
 )
 lr_cfg = LRConfig(
     lr_type="cosine",

pyproject.toml

@@ -81,6 +81,11 @@ video = [
   "decord",
   "av",
 ]
+gram-ns = [
+  "torch>=2.7.1",
+  "quack-kernels==0.5.0",
+  "nvidia-cutlass-dsl==4.5.2",
+]
 all = [
   "jsonlines",
   "decord",

tests/optim/test_gram_newton_schulz.py

@@ -0,0 +1,107 @@
+# Copyright (c) OpenMMLab. All rights reserved.
+
+from typing import Any
+
+import pytest
+import torch
+
+from xtuner.v1.config.optim import MuonConfig
+from xtuner.v1.optim import Muon
+from xtuner.v1.optim import gram_newton_schulz as gram_ns_module
+from xtuner.v1.optim.gram_newton_schulz import (
+    POLAR_EXPRESS_COEFFICIENTS,
+    POLAR_EXPRESS_SAFETY_FACTOR,
+    GramNewtonSchulz,
+)
+
+
+def _cpu_orthogonalizer(**kwargs: Any) -> GramNewtonSchulz:
+    return GramNewtonSchulz(use_kernels=False, compile_kwargs=None, **kwargs)
+
+
+class TestGramNewtonSchulz:
+    def test_polar_express_coefficients_include_recommended_safety_factor(self) -> None:
+        unmodified_first = (8.28721201814563, -23.595886519098837, 17.300387312530933)
+        expected = (
+            unmodified_first[0] / POLAR_EXPRESS_SAFETY_FACTOR,
+            unmodified_first[1] / POLAR_EXPRESS_SAFETY_FACTOR**3,
+            unmodified_first[2] / POLAR_EXPRESS_SAFETY_FACTOR**5,
+        )
+
+        assert len(POLAR_EXPRESS_COEFFICIENTS) == 5
+        torch.testing.assert_close(torch.tensor(POLAR_EXPRESS_COEFFICIENTS[0]), torch.tensor(expected))
+
+    def test_adaptive_selection_falls_back_for_square_and_low_aspect_ratio(self) -> None:
+        orthogonalizer = _cpu_orthogonalizer(min_aspect_ratio=2.0)
+
+        assert orthogonalizer._selection(256, 256) == ("standard_fallback", "torch")
+        assert orthogonalizer._selection(256, 512) == ("standard_fallback", "torch")
+        assert orthogonalizer._selection(256, 768) == ("gram", "torch")
+        assert orthogonalizer._selection(768, 256) == ("gram", "torch")
+
+        orthogonalizer._kernel_backend = object()
+        assert orthogonalizer._selection(256, 768) == ("gram", "cutedsl")
+
+    def test_gram_matches_recommended_standard_path_on_cpu(self) -> None:
+        torch.manual_seed(0)
+        inputs = (torch.randn(12, 4), torch.randn(4, 12))
+        gram = _cpu_orthogonalizer(min_aspect_ratio=2.0)
+        standard = _cpu_orthogonalizer(min_aspect_ratio=100.0)
+
+        for input_ in inputs:
+            gram_output = gram(input_)
+            standard_output = standard(input_)
+            assert gram_output.dtype == input_.dtype
+            assert torch.isfinite(gram_output).all()
+            torch.testing.assert_close(gram_output, standard_output, atol=1e-2, rtol=1e-2)
+
+    def test_concatenated_experts_are_orthogonalized_independently(self) -> None:
+        torch.manual_seed(1)
+        orthogonalizer = _cpu_orthogonalizer(min_aspect_ratio=2.0)
+        experts = torch.randn(3, 12, 4, dtype=torch.float32)
+        concatenated = experts.view(-1, experts.size(-1))
+
+        batched = orthogonalizer(concatenated, num_experts=experts.size(0))
+        independent = torch.cat([orthogonalizer(expert) for expert in experts])
+
+        assert batched.shape == concatenated.shape
+        assert batched.dtype == concatenated.dtype
+        torch.testing.assert_close(batched, independent)
+
+    def test_rejects_invalid_expert_layout(self) -> None:
+        orthogonalizer = _cpu_orthogonalizer()
+
+        with pytest.raises(ValueError, match="Invalid num_experts"):
+            orthogonalizer(torch.randn(12, 4), num_experts=5)
+
+
+class TestMuonGramNewtonSchulz:
+    def test_config_keeps_gram_ns_opt_in(self) -> None:
+        config = MuonConfig()
+
+        assert config.use_gram_newton_schulz is False
+        assert config.gram_restart_iterations == (2,)
+        assert config.gram_ns_epsilon == 1e-7
+        assert config.gram_ns_min_aspect_ratio == 2.0
+        assert config.gram_ns_torch_compile is True
+
+    def test_switch_selects_gram_ns_callback_without_cuda(self, monkeypatch: pytest.MonkeyPatch) -> None:
+        monkeypatch.setattr(gram_ns_module, "_make_kernel_backend", lambda: None)
+        parameter = torch.nn.Parameter(torch.randn(12, 4))
+
+        optimizer = Muon(
+            [parameter],
+            use_gram_newton_schulz=True,
+            gram_restart_iterations=(2,),
+            gram_ns_min_aspect_ratio=2.0,
+            gram_ns_torch_compile=False,
+        )
+
+        assert isinstance(optimizer._newton_schulz_func, GramNewtonSchulz)
+        assert optimizer._newton_schulz_func.restart_iterations == (2,)
+
+    def test_switch_rejects_other_newton_schulz_implementations(self) -> None:
+        parameter = torch.nn.Parameter(torch.randn(12, 4))
+
+        with pytest.raises(ValueError, match="cannot be combined"):
+            Muon([parameter], use_gram_newton_schulz=True, use_triton=True)

xtuner/v1/config/optim.py

@@ -78,6 +78,21 @@ class MuonConfig(OptimConfig):
     enable_all2all: Annotated[
         bool, Parameter(help="Allow all-to-all comm strategy; set False on topologies where all-to-all is unreliable")
     ] = True
+    use_gram_newton_schulz: Annotated[
+        bool, Parameter(help="Enable CuTeDSL Gram Newton-Schulz with automatic standard-NS fallback")
+    ] = False
+    gram_ns_epsilon: Annotated[
+        float, Parameter(help="Normalization epsilon used by Gram NS and its standard fallback")
+    ] = 1e-7
+    gram_restart_iterations: Annotated[
+        Tuple[int, ...], Parameter(help="Zero-based Gram NS iterations before which to restart")
+    ] = (2,)
+    gram_ns_min_aspect_ratio: Annotated[
+        float, Parameter(help="Use Gram NS only when a matrix aspect ratio is greater than this value")
+    ] = 2.0
+    gram_ns_torch_compile: Annotated[
+        bool, Parameter(help="Compile Gram NS with torch.compile fullgraph/reduce-overhead mode")
+    ] = True
 
     def build(self, model):
         trainable_parameters_names = model.trainable_parameters()
@@ -205,6 +220,11 @@ def build(self, model):
             flatten=True,
             adjust_lr=self.adjust_lr,
             use_triton=False,
+            use_gram_newton_schulz=self.use_gram_newton_schulz,
+            gram_ns_epsilon=self.gram_ns_epsilon,
+            gram_restart_iterations=self.gram_restart_iterations,
+            gram_ns_min_aspect_ratio=self.gram_ns_min_aspect_ratio,
+            gram_ns_torch_compile=self.gram_ns_torch_compile,
             epsilon=self.eps,
             enable_all2all=self.enable_all2all,
             muon_split_sizes=muon_split_sizes,