From 444f17e681a27c3424013e182fa776bfaaaf197a Mon Sep 17 00:00:00 2001
From: root <root@pool0-0084.cm.cluster>
Date: Wed, 18 Mar 2026 09:25:16 -0700
Subject: [PATCH 01/20] Add --muon-coefficient-type argument for Muon optimizer

Allow users to select the Newton-Schulz polynomial coefficient set
(e.g. simple, quintic, polar_express, aol) via the new
--muon-coefficient-type CLI flag. Supported types are discovered
dynamically from the installed emerging_optimizers package so that
upstream additions are picked up automatically without code changes.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py             |  33 ++++-
 megatron/core/optimizer/optimizer_config.py |   4 +
 megatron/training/arguments.py              |   5 +
 tests/unit_tests/test_muon_optimizer.py     | 147 +++++++++++++++++++-
 4 files changed, 187 insertions(+), 2 deletions(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index ae4a1a348fa..a141f3c7465 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -28,12 +28,16 @@
         OrthogonalizedOptimizer,
         get_muon_scale_factor,
     )
-    from emerging_optimizers.orthogonalized_optimizers.muon_utils import newton_schulz_tp
+    from emerging_optimizers.orthogonalized_optimizers.muon_utils import (
+        _COEFFICIENT_SETS,
+        newton_schulz_tp,
+    )
 
     HAVE_EMERGING_OPTIMIZERS = True
 except ImportError:
     HAVE_EMERGING_OPTIMIZERS = False
     OrthogonalizedOptimizer = object
+    _COEFFICIENT_SETS = {}
 
 # TODO: Remove this separate try/except once the next version of emerging_optimizers
 # (which includes Lion) is released. Then Lion can be imported in the block above.
@@ -47,6 +51,31 @@
 
 logger = logging.getLogger(__name__)
 
+# Fallback choices if emerging_optimizers is not installed or _COEFFICIENT_SETS is empty.
+_FALLBACK_COEFFICIENT_TYPES = ("simple", "quintic", "polar_express")
+
+
+def get_supported_coefficient_types() -> tuple[str, ...]:
+    """Return the coefficient types supported by the installed emerging_optimizers.
+
+    Dynamically reads the keys from ``_COEFFICIENT_SETS`` so that new types
+    added upstream are automatically available without code changes here.
+    Falls back to a hardcoded list when the package is not installed.
+    """
+    if _COEFFICIENT_SETS:
+        return tuple(_COEFFICIENT_SETS.keys())
+    return _FALLBACK_COEFFICIENT_TYPES
+
+
+def validate_coefficient_type(coefficient_type: str) -> None:
+    """Raise ``ValueError`` if *coefficient_type* is not supported."""
+    supported = get_supported_coefficient_types()
+    if coefficient_type not in supported:
+        raise ValueError(
+            f"Unsupported muon coefficient type '{coefficient_type}'. "
+            f"Supported types: {supported}"
+        )
+
 
 class TensorParallelMuon(OrthogonalizedOptimizer):
     """Tensor Parallel Muon optimizer."""
@@ -72,6 +101,7 @@ def __init__(
     ) -> None:
         if num_ns_steps < 1:
             raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}")
+        validate_coefficient_type(coefficient_type)
 
         def scaled_orthogonalize_fn(
             grad: torch.Tensor,
@@ -288,6 +318,7 @@ def lion_init_state_fn(opt, config=None):
         "use_nesterov": config.muon_use_nesterov,
         "weight_decay": config.weight_decay,
         "fp32_matmul_prec": config.muon_fp32_matmul_prec,
+        "coefficient_type": config.muon_coefficient_type,
         "num_ns_steps": config.muon_num_ns_steps,
         "scale_mode": config.muon_scale_mode,
         "split_qkv": config.muon_split_qkv,
diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py
index 16b0a54cb6b..9e6375b978c 100644
--- a/megatron/core/optimizer/optimizer_config.py
+++ b/megatron/core/optimizer/optimizer_config.py
@@ -272,6 +272,10 @@ class OptimizerConfig:
     muon_fp32_matmul_prec: str = "medium"
     """The precision to use for the fp32 matmul. Defaults to "medium"."""
 
+    muon_coefficient_type: str = "quintic"
+    """Newton-Schulz coefficient type for the Muon optimizer. Valid types are discovered
+    dynamically from the installed ``emerging_optimizers`` package. Defaults to "quintic"."""
+
     muon_num_ns_steps: int = 5
     """The number of iteration steps to use in the Newton-Schulz iteration."""
 
diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py
index 3d6451f3168..cd495ec817e 100644
--- a/megatron/training/arguments.py
+++ b/megatron/training/arguments.py
@@ -2191,6 +2191,11 @@ def _add_regularization_args(parser):
     group.add_argument('--muon-fp32-matmul-prec', type=str, default='medium',
                        choices=['low', 'medium', 'high'],
                        help='FP32 matmul precision for Newton-Schulz iteration')
+    group.add_argument('--muon-coefficient-type', type=str, default='quintic',
+                       help='Newton-Schulz coefficient type for the Muon optimizer. '
+                       'Valid types are discovered from the installed emerging_optimizers '
+                       'package (e.g. simple, quintic, polar_express, aol). '
+                       'Validated at optimizer creation time.')
     group.add_argument('--muon-num-ns-steps', type=int, default=5,
                        help='Number of Newton-Schulz steps for Muon optimizer')
     group.add_argument('--muon-tp-mode', type=str, default='blockwise',
diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index cc99f7a16e6..53c63e4f70b 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -11,7 +11,12 @@
 from megatron.core import parallel_state
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
 from megatron.core.optimizer import OptimizerConfig
-from megatron.core.optimizer.muon import TensorParallelMuon, get_megatron_muon_optimizer
+from megatron.core.optimizer.muon import (
+    TensorParallelMuon,
+    get_megatron_muon_optimizer,
+    get_supported_coefficient_types,
+    validate_coefficient_type,
+)
 from megatron.core.process_groups_config import ProcessGroupCollection
 from megatron.core.transformer import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
@@ -641,6 +646,146 @@ def test_muon_optimizer_extra_scale_factor():
     ), "Weight should be updated with extra_scale_factor"
 
 
+def test_get_supported_coefficient_types_returns_tuple():
+    """Test that get_supported_coefficient_types returns a non-empty tuple of strings."""
+    supported = get_supported_coefficient_types()
+    assert isinstance(supported, tuple)
+    assert len(supported) > 0
+    for t in supported:
+        assert isinstance(t, str)
+
+
+def test_get_supported_coefficient_types_contains_known_types():
+    """Test that the known coefficient types are present in the supported set."""
+    supported = get_supported_coefficient_types()
+    for expected in ("simple", "quintic", "polar_express"):
+        assert expected in supported, f"Expected '{expected}' in supported types {supported}"
+
+
+def test_validate_coefficient_type_accepts_valid():
+    """Test that validate_coefficient_type does not raise for valid types."""
+    for t in get_supported_coefficient_types():
+        validate_coefficient_type(t)  # should not raise
+
+
+def test_validate_coefficient_type_rejects_invalid():
+    """Test that validate_coefficient_type raises ValueError for an invalid type."""
+    with pytest.raises(ValueError, match="Unsupported muon coefficient type"):
+        validate_coefficient_type("nonexistent_type_xyz")
+
+
+@pytest.mark.parametrize("coefficient_type", list(get_supported_coefficient_types()))
+def test_muon_optimizer_all_supported_coefficient_types(coefficient_type):
+    """Test TensorParallelMuon optimizer with every dynamically-discovered coefficient type.
+
+    This ensures that when emerging_optimizers adds new coefficient types, they
+    are automatically exercised by these tests.
+    """
+    from emerging_optimizers.orthogonalized_optimizers.muon_utils import _COEFFICIENT_SETS
+
+    # Determine a compatible num_ns_steps for this coefficient type.
+    num_coeffs = len(_COEFFICIENT_SETS[coefficient_type])
+    num_ns_steps = num_coeffs  # 1x the set length is always valid
+
+    model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')
+    model.requires_grad_(True)
+    model.weight.data.fill_(1.0)
+
+    optimizer = TensorParallelMuon(
+        params=[model.weight],
+        lr=0.01,
+        coefficient_type=coefficient_type,
+        num_ns_steps=num_ns_steps,
+        pg_collection=None,
+        mode="duplicated",
+    )
+
+    input_tensor = torch.randn(16, 80, dtype=torch.float32, device='cuda')
+    output = model(input_tensor)
+    loss = output.sum()
+    loss.backward()
+
+    original_weight = model.weight.data.clone()
+    optimizer.step()
+
+    assert not torch.equal(
+        model.weight.data, original_weight
+    ), f"Weight should be updated with coefficient_type={coefficient_type}"
+
+
+def test_muon_optimizer_invalid_coefficient_type():
+    """Test that TensorParallelMuon raises ValueError for an invalid coefficient_type."""
+    model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')
+    model.requires_grad_(True)
+
+    with pytest.raises(ValueError, match="Unsupported muon coefficient type"):
+        TensorParallelMuon(
+            params=[model.weight],
+            lr=0.01,
+            coefficient_type="nonexistent_type_xyz",
+            num_ns_steps=5,
+            pg_collection=None,
+            mode="duplicated",
+        )
+
+
+@pytest.mark.skipif(
+    int(os.getenv('WORLD_SIZE', '1')) == 1, reason="Multi-rank test requires WORLD_SIZE > 1"
+)
+class TestMuonCoefficientTypeMultiRank:
+    """Test coefficient_type integration through get_megatron_muon_optimizer."""
+
+    @pytest.fixture(autouse=True)
+    def setup_and_teardown(self):
+        Utils.initialize_model_parallel()
+        yield
+        Utils.destroy_model_parallel()
+
+    def create_ddp_model(self, model):
+        ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=False)
+        return DistributedDataParallel(
+            TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model
+        )
+
+    @pytest.mark.parametrize("coefficient_type", ["simple", "quintic", "polar_express"])
+    def test_get_megatron_muon_optimizer_coefficient_type(self, coefficient_type):
+        """Test that coefficient_type flows through get_megatron_muon_optimizer."""
+        from emerging_optimizers.orthogonalized_optimizers.muon_utils import _COEFFICIENT_SETS
+
+        num_coeffs = len(_COEFFICIENT_SETS[coefficient_type])
+
+        model = Net().bfloat16().cuda()
+        model.requires_grad_(True)
+        model = self.create_ddp_model(model)
+
+        optimizer_config = OptimizerConfig(
+            optimizer='muon',
+            lr=0.01,
+            weight_decay=0.01,
+            bf16=True,
+            use_distributed_optimizer=False,
+            muon_coefficient_type=coefficient_type,
+            muon_num_ns_steps=num_coeffs,
+            muon_tp_mode="duplicated",
+        )
+
+        optimizer = get_megatron_muon_optimizer(
+            config=optimizer_config,
+            model_chunks=[model],
+            use_gloo_process_groups=True,
+            layer_wise_distributed_optimizer=False,
+        )
+
+        assert optimizer is not None
+
+        input_tensor = torch.randn(16, 80, dtype=torch.bfloat16, device='cuda')
+        output = model(input_tensor)
+        loss = output.sum()
+        loss.backward()
+
+        optimizer.step()
+
+
 @pytest.mark.parametrize("num_ns_steps", [5, 15, 25])
 def test_muon_optimizer_num_ns_steps(num_ns_steps):
     """Test TensorParallelMuon optimizer with different numbers of Newton-Schulz steps."""

From e357a411c9ad56e5233b67cf23928ab376975f77 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0084.cm.cluster>
Date: Wed, 18 Mar 2026 09:34:29 -0700
Subject: [PATCH 02/20] Use NSCoeffT Literal instead of private
 _COEFFICIENT_SETS

Derive supported coefficient types from the public NSCoeffT Literal
type via typing.get_args() rather than reading keys from the private
_COEFFICIENT_SETS dict. Tests likewise avoid importing _COEFFICIENT_SETS.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py         | 16 +++++++------
 tests/unit_tests/test_muon_optimizer.py | 31 +++++++++++++++++--------
 2 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index a141f3c7465..643e8de01c0 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -3,7 +3,7 @@
 """Megatron muon optimizer wrapper to handle tensor-parallel."""
 
 import logging
-from typing import Any, Callable, Dict, List, Literal, Optional
+from typing import Any, Callable, Dict, List, Literal, Optional, get_args
 
 import torch
 from torch.optim.optimizer import ParamsT
@@ -29,7 +29,7 @@
         get_muon_scale_factor,
     )
     from emerging_optimizers.orthogonalized_optimizers.muon_utils import (
-        _COEFFICIENT_SETS,
+        NSCoeffT,
         newton_schulz_tp,
     )
 
@@ -37,7 +37,7 @@
 except ImportError:
     HAVE_EMERGING_OPTIMIZERS = False
     OrthogonalizedOptimizer = object
-    _COEFFICIENT_SETS = {}
+    NSCoeffT = None  # type: ignore[assignment, misc]
 
 # TODO: Remove this separate try/except once the next version of emerging_optimizers
 # (which includes Lion) is released. Then Lion can be imported in the block above.
@@ -51,19 +51,21 @@
 
 logger = logging.getLogger(__name__)
 
-# Fallback choices if emerging_optimizers is not installed or _COEFFICIENT_SETS is empty.
+# Fallback choices if emerging_optimizers is not installed.
 _FALLBACK_COEFFICIENT_TYPES = ("simple", "quintic", "polar_express")
 
 
 def get_supported_coefficient_types() -> tuple[str, ...]:
     """Return the coefficient types supported by the installed emerging_optimizers.
 
-    Dynamically reads the keys from ``_COEFFICIENT_SETS`` so that new types
+    Reads the members of the ``NSCoeffT`` Literal type so that new types
     added upstream are automatically available without code changes here.
     Falls back to a hardcoded list when the package is not installed.
     """
-    if _COEFFICIENT_SETS:
-        return tuple(_COEFFICIENT_SETS.keys())
+    if NSCoeffT is not None:
+        args = get_args(NSCoeffT)
+        if args:
+            return args
     return _FALLBACK_COEFFICIENT_TYPES
 
 
diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index 53c63e4f70b..f8aed7e96d3 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -674,18 +674,31 @@ def test_validate_coefficient_type_rejects_invalid():
         validate_coefficient_type("nonexistent_type_xyz")
 
 
-@pytest.mark.parametrize("coefficient_type", list(get_supported_coefficient_types()))
+# Newton-Schulz step counts known to be compatible with each coefficient type.
+# "custom" is excluded because it requires user-supplied coefficient sets.
+_NS_STEPS_FOR_COEFF_TYPE = {
+    "simple": 3,
+    "quintic": 5,
+    "polar_express": 8,
+    "aol": 4,
+}
+
+# Parametrize over all supported types that we know how to configure.
+_TESTABLE_COEFFICIENT_TYPES = [
+    t for t in get_supported_coefficient_types()
+    if t in _NS_STEPS_FOR_COEFF_TYPE
+]
+
+
+@pytest.mark.parametrize("coefficient_type", _TESTABLE_COEFFICIENT_TYPES)
 def test_muon_optimizer_all_supported_coefficient_types(coefficient_type):
     """Test TensorParallelMuon optimizer with every dynamically-discovered coefficient type.
 
     This ensures that when emerging_optimizers adds new coefficient types, they
-    are automatically exercised by these tests.
+    are automatically exercised by these tests.  New types need a corresponding
+    entry in ``_NS_STEPS_FOR_COEFF_TYPE`` to be picked up.
     """
-    from emerging_optimizers.orthogonalized_optimizers.muon_utils import _COEFFICIENT_SETS
-
-    # Determine a compatible num_ns_steps for this coefficient type.
-    num_coeffs = len(_COEFFICIENT_SETS[coefficient_type])
-    num_ns_steps = num_coeffs  # 1x the set length is always valid
+    num_ns_steps = _NS_STEPS_FOR_COEFF_TYPE[coefficient_type]
 
     model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')
     model.requires_grad_(True)
@@ -750,9 +763,7 @@ def create_ddp_model(self, model):
     @pytest.mark.parametrize("coefficient_type", ["simple", "quintic", "polar_express"])
     def test_get_megatron_muon_optimizer_coefficient_type(self, coefficient_type):
         """Test that coefficient_type flows through get_megatron_muon_optimizer."""
-        from emerging_optimizers.orthogonalized_optimizers.muon_utils import _COEFFICIENT_SETS
-
-        num_coeffs = len(_COEFFICIENT_SETS[coefficient_type])
+        num_coeffs = _NS_STEPS_FOR_COEFF_TYPE[coefficient_type]
 
         model = Net().bfloat16().cuda()
         model.requires_grad_(True)

From bcac48165153df42e1c9f44f40f85f5731266802 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0084.cm.cluster>
Date: Wed, 18 Mar 2026 09:37:47 -0700
Subject: [PATCH 03/20] Remove fallback coefficient types; require
 emerging_optimizers

emerging_optimizers must be installed to use Muon, so there is no
need for a hardcoded fallback list. get_supported_coefficient_types()
now asserts the package is present and reads NSCoeffT directly.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py | 19 +++++--------------
 1 file changed, 5 insertions(+), 14 deletions(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 643e8de01c0..c2028c181aa 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -28,16 +28,12 @@
         OrthogonalizedOptimizer,
         get_muon_scale_factor,
     )
-    from emerging_optimizers.orthogonalized_optimizers.muon_utils import (
-        NSCoeffT,
-        newton_schulz_tp,
-    )
+    from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT, newton_schulz_tp
 
     HAVE_EMERGING_OPTIMIZERS = True
 except ImportError:
     HAVE_EMERGING_OPTIMIZERS = False
     OrthogonalizedOptimizer = object
-    NSCoeffT = None  # type: ignore[assignment, misc]
 
 # TODO: Remove this separate try/except once the next version of emerging_optimizers
 # (which includes Lion) is released. Then Lion can be imported in the block above.
@@ -51,22 +47,17 @@
 
 logger = logging.getLogger(__name__)
 
-# Fallback choices if emerging_optimizers is not installed.
-_FALLBACK_COEFFICIENT_TYPES = ("simple", "quintic", "polar_express")
-
 
 def get_supported_coefficient_types() -> tuple[str, ...]:
     """Return the coefficient types supported by the installed emerging_optimizers.
 
     Reads the members of the ``NSCoeffT`` Literal type so that new types
     added upstream are automatically available without code changes here.
-    Falls back to a hardcoded list when the package is not installed.
     """
-    if NSCoeffT is not None:
-        args = get_args(NSCoeffT)
-        if args:
-            return args
-    return _FALLBACK_COEFFICIENT_TYPES
+    assert (
+        HAVE_EMERGING_OPTIMIZERS
+    ), "emerging_optimizers is required for the Muon optimizer. Please install it."
+    return get_args(NSCoeffT)
 
 
 def validate_coefficient_type(coefficient_type: str) -> None:

From 17597acb74ca051d00caf0e6f2f23fa402c10892 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0084.cm.cluster>
Date: Wed, 18 Mar 2026 09:47:24 -0700
Subject: [PATCH 04/20] Run autoformatter on test file

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 tests/unit_tests/test_muon_optimizer.py | 10 ++--------
 1 file changed, 2 insertions(+), 8 deletions(-)

diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index f8aed7e96d3..0a84258a5ca 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -676,17 +676,11 @@ def test_validate_coefficient_type_rejects_invalid():
 
 # Newton-Schulz step counts known to be compatible with each coefficient type.
 # "custom" is excluded because it requires user-supplied coefficient sets.
-_NS_STEPS_FOR_COEFF_TYPE = {
-    "simple": 3,
-    "quintic": 5,
-    "polar_express": 8,
-    "aol": 4,
-}
+_NS_STEPS_FOR_COEFF_TYPE = {"simple": 3, "quintic": 5, "polar_express": 8, "aol": 4}
 
 # Parametrize over all supported types that we know how to configure.
 _TESTABLE_COEFFICIENT_TYPES = [
-    t for t in get_supported_coefficient_types()
-    if t in _NS_STEPS_FOR_COEFF_TYPE
+    t for t in get_supported_coefficient_types() if t in _NS_STEPS_FOR_COEFF_TYPE
 ]
 
 

From bf834c15c80807d1c48c8a86e900ebd06c12f466 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0393.cm.cluster>
Date: Wed, 18 Mar 2026 15:36:47 -0700
Subject: [PATCH 05/20] Derive test coefficient types dynamically from
 emerging_optimizers

Replace hardcoded _NS_STEPS_FOR_COEFF_TYPE mapping with dynamic discovery
via get_supported_coefficient_types() (backed by NSCoeffT). Since
get_coefficient_iterator cycles/repeats coefficients, a single default
step count works for all types. Remove redundant duplicate test.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 tests/unit_tests/test_muon_optimizer.py | 65 +++++--------------------
 1 file changed, 12 insertions(+), 53 deletions(-)

diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index 0a84258a5ca..1fec2af594b 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -425,10 +425,8 @@ def test_muon_optimizer_blockwise_mode_different_result(self):
         ), "Weight should be updated with mode=blockwise"
 
 
-@pytest.mark.parametrize(
-    "coefficient_type_and_steps", [("simple", 3), ("quintic", 5), ("polar_express", 8)]
-)
-def test_muon_optimizer_coefficient_types(coefficient_type_and_steps):
+@pytest.mark.parametrize("coefficient_type", _TESTABLE_COEFFICIENT_TYPES)
+def test_muon_optimizer_coefficient_types(coefficient_type):
     """Test TensorParallelMuon optimizer with different coefficient types."""
     model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')
     model.requires_grad_(True)
@@ -437,8 +435,8 @@ def test_muon_optimizer_coefficient_types(coefficient_type_and_steps):
     optimizer = TensorParallelMuon(
         params=[model.weight],
         lr=0.01,
-        coefficient_type=coefficient_type_and_steps[0],
-        num_ns_steps=coefficient_type_and_steps[1],
+        coefficient_type=coefficient_type,
+        num_ns_steps=_DEFAULT_NS_STEPS,
         pg_collection=None,
         mode="duplicated",
     )
@@ -453,7 +451,7 @@ def test_muon_optimizer_coefficient_types(coefficient_type_and_steps):
 
     assert not torch.equal(
         model.weight.data, original_weight
-    ), f"Weight should be updated with coefficient_type={coefficient_type_and_steps[0]} and num_ns_steps={coefficient_type_and_steps[1]}"
+    ), f"Weight should be updated with coefficient_type={coefficient_type}"
 
 
 @pytest.mark.parametrize("scale_mode", ["spectral", "unit_rms_norm", "shape_scaling"])
@@ -674,50 +672,13 @@ def test_validate_coefficient_type_rejects_invalid():
         validate_coefficient_type("nonexistent_type_xyz")
 
 
-# Newton-Schulz step counts known to be compatible with each coefficient type.
-# "custom" is excluded because it requires user-supplied coefficient sets.
-_NS_STEPS_FOR_COEFF_TYPE = {"simple": 3, "quintic": 5, "polar_express": 8, "aol": 4}
-
-# Parametrize over all supported types that we know how to configure.
-_TESTABLE_COEFFICIENT_TYPES = [
-    t for t in get_supported_coefficient_types() if t in _NS_STEPS_FOR_COEFF_TYPE
-]
-
-
-@pytest.mark.parametrize("coefficient_type", _TESTABLE_COEFFICIENT_TYPES)
-def test_muon_optimizer_all_supported_coefficient_types(coefficient_type):
-    """Test TensorParallelMuon optimizer with every dynamically-discovered coefficient type.
-
-    This ensures that when emerging_optimizers adds new coefficient types, they
-    are automatically exercised by these tests.  New types need a corresponding
-    entry in ``_NS_STEPS_FOR_COEFF_TYPE`` to be picked up.
-    """
-    num_ns_steps = _NS_STEPS_FOR_COEFF_TYPE[coefficient_type]
-
-    model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')
-    model.requires_grad_(True)
-    model.weight.data.fill_(1.0)
-
-    optimizer = TensorParallelMuon(
-        params=[model.weight],
-        lr=0.01,
-        coefficient_type=coefficient_type,
-        num_ns_steps=num_ns_steps,
-        pg_collection=None,
-        mode="duplicated",
-    )
-
-    input_tensor = torch.randn(16, 80, dtype=torch.float32, device='cuda')
-    output = model(input_tensor)
-    loss = output.sum()
-    loss.backward()
+# All non-custom coefficient types supported by emerging_optimizers.
+_TESTABLE_COEFFICIENT_TYPES = [t for t in get_supported_coefficient_types() if t != "custom"]
 
-    original_weight = model.weight.data.clone()
-    optimizer.step()
+# A reasonable default NS step count for testing; get_coefficient_iterator
+# cycles/repeats coefficients so any step count works with any type.
+_DEFAULT_NS_STEPS = 5
 
-    assert not torch.equal(
-        model.weight.data, original_weight
-    ), f"Weight should be updated with coefficient_type={coefficient_type}"
 
 
 def test_muon_optimizer_invalid_coefficient_type():
@@ -754,11 +715,9 @@ def create_ddp_model(self, model):
             TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model
         )
 
-    @pytest.mark.parametrize("coefficient_type", ["simple", "quintic", "polar_express"])
+    @pytest.mark.parametrize("coefficient_type", _TESTABLE_COEFFICIENT_TYPES)
     def test_get_megatron_muon_optimizer_coefficient_type(self, coefficient_type):
         """Test that coefficient_type flows through get_megatron_muon_optimizer."""
-        num_coeffs = _NS_STEPS_FOR_COEFF_TYPE[coefficient_type]
-
         model = Net().bfloat16().cuda()
         model.requires_grad_(True)
         model = self.create_ddp_model(model)
@@ -770,7 +729,7 @@ def test_get_megatron_muon_optimizer_coefficient_type(self, coefficient_type):
             bf16=True,
             use_distributed_optimizer=False,
             muon_coefficient_type=coefficient_type,
-            muon_num_ns_steps=num_coeffs,
+            muon_num_ns_steps=_DEFAULT_NS_STEPS,
             muon_tp_mode="duplicated",
         )
 

From 7b8b0428111c6e1344abce186f41fc75ac8bf92d Mon Sep 17 00:00:00 2001
From: root <root@pool0-0154.cm.cluster>
Date: Thu, 19 Mar 2026 15:28:17 -0700
Subject: [PATCH 06/20] Add emerging_optimizers >= 0.2 version gate for Lion
 import

The Lion class moved in emerging_optimizers 0.2. Gate the import
behind an explicit version check so users get a clear error instead
of a silent ImportError on older versions.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/__init__.py | 10 +++++++++-
 megatron/core/optimizer/muon.py     | 12 +++++++++---
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 8cfb22620bb..bb1376950ac 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -34,10 +34,18 @@
         USING_PYTORCH_OPTIMIZER = True
 
 try:
+    from importlib.metadata import version as _pkg_version
+
+    _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
+    if _eo_ver < (0, 2):
+        raise ImportError(
+            f"Lion optimizer requires emerging_optimizers >= 0.2, "
+            f"found {_pkg_version('emerging-optimizers')}"
+        )
     from emerging_optimizers.scalar_optimizers import Lion
 
     HAVE_LION = True
-except ImportError:
+except (ImportError, Exception):
     HAVE_LION = False
 
 from megatron.core import parallel_state
diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index c2028c181aa..9dca752f54f 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -35,13 +35,19 @@
     HAVE_EMERGING_OPTIMIZERS = False
     OrthogonalizedOptimizer = object
 
-# TODO: Remove this separate try/except once the next version of emerging_optimizers
-# (which includes Lion) is released. Then Lion can be imported in the block above.
 try:
+    from importlib.metadata import version as _pkg_version
+
+    _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
+    if _eo_ver < (0, 2):
+        raise ImportError(
+            f"Lion optimizer requires emerging_optimizers >= 0.2, "
+            f"found {_pkg_version('emerging-optimizers')}"
+        )
     from emerging_optimizers.scalar_optimizers import Lion  # pylint: disable=unused-import
 
     HAVE_LION = True
-except ImportError:
+except (ImportError, Exception):
     HAVE_LION = False
 
 

From 02bad44c4caefe4761afe46e31da7e8f7dfbf52a Mon Sep 17 00:00:00 2001
From: root <root@pool0-0154.cm.cluster>
Date: Thu, 19 Mar 2026 15:40:05 -0700
Subject: [PATCH 07/20] Narrow broad Exception catch to PackageNotFoundError

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/__init__.py | 3 ++-
 megatron/core/optimizer/muon.py     | 3 ++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index bb1376950ac..6b7bf34330a 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -34,6 +34,7 @@
         USING_PYTORCH_OPTIMIZER = True
 
 try:
+    from importlib.metadata import PackageNotFoundError
     from importlib.metadata import version as _pkg_version
 
     _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
@@ -45,7 +46,7 @@
     from emerging_optimizers.scalar_optimizers import Lion
 
     HAVE_LION = True
-except (ImportError, Exception):
+except (ImportError, PackageNotFoundError):
     HAVE_LION = False
 
 from megatron.core import parallel_state
diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 9dca752f54f..89b6b4fbdea 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -36,6 +36,7 @@
     OrthogonalizedOptimizer = object
 
 try:
+    from importlib.metadata import PackageNotFoundError
     from importlib.metadata import version as _pkg_version
 
     _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
@@ -47,7 +48,7 @@
     from emerging_optimizers.scalar_optimizers import Lion  # pylint: disable=unused-import
 
     HAVE_LION = True
-except (ImportError, Exception):
+except (ImportError, PackageNotFoundError):
     HAVE_LION = False
 
 

From 07ee6fa0af8818817eae3b0578e2db3f3f927f3c Mon Sep 17 00:00:00 2001
From: root <root@pool0-0021.cm.cluster>
Date: Tue, 24 Mar 2026 11:48:53 -0700
Subject: [PATCH 08/20] Add emerging_optimizers >= 0.2 version gate for
 NSCoeffT and fix v0.2 API compat

Gate NSCoeffT import behind the >= 0.2 version check alongside Lion.
Unify HAVE_LION into a single HAVE_EO_V02 flag for all v0.2 features.
Handle v0.2 API renames (use_nesterov -> nesterov, mode -> tp_mode).
Fix forward reference bug for _TESTABLE_COEFFICIENT_TYPES in tests.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/__init__.py     |  6 ++--
 megatron/core/optimizer/muon.py         | 39 ++++++++++++++++---------
 tests/unit_tests/test_lion_optimizer.py | 10 +++----
 tests/unit_tests/test_muon_optimizer.py | 16 +++++-----
 4 files changed, 41 insertions(+), 30 deletions(-)

diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 6b7bf34330a..a026ededa5e 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -45,9 +45,9 @@
         )
     from emerging_optimizers.scalar_optimizers import Lion
 
-    HAVE_LION = True
+    HAVE_EO_V02 = True
 except (ImportError, PackageNotFoundError):
-    HAVE_LION = False
+    HAVE_EO_V02 = False
 
 from megatron.core import parallel_state
 from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer
@@ -560,7 +560,7 @@ def init_state_fn(opt, config=None):
                                 opt.initialize_state(p)
 
         elif config.optimizer == 'lion':
-            if not HAVE_LION:
+            if not HAVE_EO_V02:
                 raise ImportError(
                     "Lion optimizer requires the 'emerging_optimizers' package. "
                     "Please install it to use --optimizer lion."
diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 89b6b4fbdea..221b739d6dd 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -28,7 +28,7 @@
         OrthogonalizedOptimizer,
         get_muon_scale_factor,
     )
-    from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT, newton_schulz_tp
+    from emerging_optimizers.orthogonalized_optimizers.muon_utils import newton_schulz_tp
 
     HAVE_EMERGING_OPTIMIZERS = True
 except ImportError:
@@ -42,14 +42,15 @@
     _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
     if _eo_ver < (0, 2):
         raise ImportError(
-            f"Lion optimizer requires emerging_optimizers >= 0.2, "
+            f"NSCoeffT and Lion require emerging_optimizers >= 0.2, "
             f"found {_pkg_version('emerging-optimizers')}"
         )
+    from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
     from emerging_optimizers.scalar_optimizers import Lion  # pylint: disable=unused-import
 
-    HAVE_LION = True
+    HAVE_EO_V02 = True
 except (ImportError, PackageNotFoundError):
-    HAVE_LION = False
+    HAVE_EO_V02 = False
 
 
 logger = logging.getLogger(__name__)
@@ -61,9 +62,9 @@ def get_supported_coefficient_types() -> tuple[str, ...]:
     Reads the members of the ``NSCoeffT`` Literal type so that new types
     added upstream are automatically available without code changes here.
     """
-    assert (
-        HAVE_EMERGING_OPTIMIZERS
-    ), "emerging_optimizers is required for the Muon optimizer. Please install it."
+    assert HAVE_EO_V02, (
+        "emerging_optimizers >= 0.2 is required for NSCoeffT. Please install or upgrade it."
+    )
     return get_args(NSCoeffT)
 
 
@@ -117,14 +118,19 @@ def scaled_orthogonalize_fn(
             size = [grad.size(-2), grad.size(-1)]
             if partition_dim is not None:
                 size[partition_dim] *= get_pg_size(tp_group)
-            orth_grad = newton_schulz_tp(
-                grad,
+            mode_value = "duplicated" if mode == "blockwise" else mode
+            mode_kwarg = (
+                {"tp_mode": mode_value} if HAVE_EO_V02 else {"mode": mode_value}
+            )
+            ns_kwargs = dict(
                 steps=num_ns_steps,
-                coefficient_type=coefficient_type,
                 tp_group=tp_group,
                 partition_dim=partition_dim,
-                mode="duplicated" if mode == "blockwise" else mode,
+                **mode_kwarg,
             )
+            if HAVE_EO_V02:
+                ns_kwargs["coefficient_type"] = coefficient_type
+            orth_grad = newton_schulz_tp(grad, **ns_kwargs)
             scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode)
             return orth_grad * scale_factor * extra_scale_factor
 
@@ -135,11 +141,16 @@ def scaled_orthogonalize_fn(
         self.qkv_split_shapes = qkv_split_shapes
 
         weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2"
+        nesterov_kwarg = (
+            {"nesterov": use_nesterov}
+            if HAVE_EO_V02
+            else {"use_nesterov": use_nesterov}
+        )
         super().__init__(
             params,
             lr,
             momentum_beta,
-            use_nesterov=use_nesterov,
+            **nesterov_kwarg,
             weight_decay=weight_decay,
             weight_decay_method=weight_decay_method,
             fp32_matmul_prec=fp32_matmul_prec,
@@ -227,8 +238,8 @@ def get_megatron_muon_optimizer(
 
     assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed."
     if config.muon_scalar_optimizer == 'lion':
-        assert HAVE_LION, (
-            "Lion optimizer requires a version of 'emerging_optimizers' that includes Lion. "
+        assert HAVE_EO_V02, (
+            "Lion optimizer requires emerging_optimizers >= 0.2. "
             "Please upgrade to use --muon-scalar-optimizer lion."
         )
 
diff --git a/tests/unit_tests/test_lion_optimizer.py b/tests/unit_tests/test_lion_optimizer.py
index 5cd479e655a..589ed82764c 100644
--- a/tests/unit_tests/test_lion_optimizer.py
+++ b/tests/unit_tests/test_lion_optimizer.py
@@ -14,7 +14,7 @@
 import torch.nn as nn
 
 from megatron.core.optimizer import (
-    HAVE_LION,
+    HAVE_EO_V02,
     OptimizerConfig,
     _get_megatron_optimizer_based_on_param_groups,
     _get_param_groups,
@@ -22,7 +22,7 @@
 from megatron.core.optimizer.optimizer import FP32Optimizer
 
 requires_emerging_optimizers = pytest.mark.skipif(
-    not HAVE_LION, reason="emerging_optimizers package not installed"
+    not HAVE_EO_V02, reason="emerging_optimizers package not installed"
 )
 
 
@@ -97,9 +97,9 @@ def test_lion_import_error_without_package(self):
         """Should raise ImportError with helpful message if emerging_optimizers not installed."""
         import megatron.core.optimizer as opt_module
 
-        original_have_lion = opt_module.HAVE_LION
+        original_have_lion = opt_module.HAVE_EO_V02
         try:
-            opt_module.HAVE_LION = False
+            opt_module.HAVE_EO_V02 = False
 
             model = SimpleModel()
             config = OptimizerConfig(optimizer="lion", lr=1e-4)
@@ -107,7 +107,7 @@ def test_lion_import_error_without_package(self):
             with pytest.raises(ImportError, match="emerging_optimizers"):
                 _create_lion_optimizer(model, config)
         finally:
-            opt_module.HAVE_LION = original_have_lion
+            opt_module.HAVE_EO_V02 = original_have_lion
 
 
 @requires_emerging_optimizers
diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index 1fec2af594b..4e5d8245638 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -425,6 +425,14 @@ def test_muon_optimizer_blockwise_mode_different_result(self):
         ), "Weight should be updated with mode=blockwise"
 
 
+# All non-custom coefficient types supported by emerging_optimizers.
+_TESTABLE_COEFFICIENT_TYPES = [t for t in get_supported_coefficient_types() if t != "custom"]
+
+# A reasonable default NS step count for testing; get_coefficient_iterator
+# cycles/repeats coefficients so any step count works with any type.
+_DEFAULT_NS_STEPS = 5
+
+
 @pytest.mark.parametrize("coefficient_type", _TESTABLE_COEFFICIENT_TYPES)
 def test_muon_optimizer_coefficient_types(coefficient_type):
     """Test TensorParallelMuon optimizer with different coefficient types."""
@@ -672,14 +680,6 @@ def test_validate_coefficient_type_rejects_invalid():
         validate_coefficient_type("nonexistent_type_xyz")
 
 
-# All non-custom coefficient types supported by emerging_optimizers.
-_TESTABLE_COEFFICIENT_TYPES = [t for t in get_supported_coefficient_types() if t != "custom"]
-
-# A reasonable default NS step count for testing; get_coefficient_iterator
-# cycles/repeats coefficients so any step count works with any type.
-_DEFAULT_NS_STEPS = 5
-
-
 
 def test_muon_optimizer_invalid_coefficient_type():
     """Test that TensorParallelMuon raises ValueError for an invalid coefficient_type."""

From 4e929d28b104ba0a206ddef85198d3d49641b12e Mon Sep 17 00:00:00 2001
From: root <root@pool0-0021.cm.cluster>
Date: Tue, 24 Mar 2026 11:57:05 -0700
Subject: [PATCH 09/20] Consolidate HAVE_EO_V02 version check into __init__.py

Remove duplicate emerging_optimizers >= 0.2 version check from muon.py
and import HAVE_EO_V02 from __init__.py instead. Update Lion error
message to mention the >= 0.2 requirement.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/__init__.py |  6 +++---
 megatron/core/optimizer/muon.py     | 18 ++----------------
 2 files changed, 5 insertions(+), 19 deletions(-)

diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index a026ededa5e..1ea171d269b 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -40,7 +40,7 @@
     _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
     if _eo_ver < (0, 2):
         raise ImportError(
-            f"Lion optimizer requires emerging_optimizers >= 0.2, "
+            f"emerging_optimizers >= 0.2 is required, "
             f"found {_pkg_version('emerging-optimizers')}"
         )
     from emerging_optimizers.scalar_optimizers import Lion
@@ -562,8 +562,8 @@ def init_state_fn(opt, config=None):
         elif config.optimizer == 'lion':
             if not HAVE_EO_V02:
                 raise ImportError(
-                    "Lion optimizer requires the 'emerging_optimizers' package. "
-                    "Please install it to use --optimizer lion."
+                    "Lion optimizer requires emerging_optimizers >= 0.2. "
+                    "Please install or upgrade it to use --optimizer lion."
                 )
             optimizer = Lion(
                 param_groups,
diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 221b739d6dd..9762d1ea2bc 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -13,7 +13,7 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_pg_size, log_single_rank
 
-from . import _get_param_groups, get_megatron_optimizer
+from . import HAVE_EO_V02, _get_param_groups, get_megatron_optimizer
 from .layer_wise_optimizer import LayerWiseDistributedOptimizer
 from .optimizer import (
     ChainedOptimizer,
@@ -35,22 +35,8 @@
     HAVE_EMERGING_OPTIMIZERS = False
     OrthogonalizedOptimizer = object
 
-try:
-    from importlib.metadata import PackageNotFoundError
-    from importlib.metadata import version as _pkg_version
-
-    _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
-    if _eo_ver < (0, 2):
-        raise ImportError(
-            f"NSCoeffT and Lion require emerging_optimizers >= 0.2, "
-            f"found {_pkg_version('emerging-optimizers')}"
-        )
+if HAVE_EO_V02:
     from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
-    from emerging_optimizers.scalar_optimizers import Lion  # pylint: disable=unused-import
-
-    HAVE_EO_V02 = True
-except (ImportError, PackageNotFoundError):
-    HAVE_EO_V02 = False
 
 
 logger = logging.getLogger(__name__)

From 662880183bb8860dd819cf4538ba93a353a77118 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0021.cm.cluster>
Date: Tue, 24 Mar 2026 11:59:59 -0700
Subject: [PATCH 10/20] Skip redundant HAVE_EMERGING_OPTIMIZERS check when
 HAVE_EO_V02 is true

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 9762d1ea2bc..2aa9832015a 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -222,12 +222,13 @@ def get_megatron_muon_optimizer(
     # Set the nonlinear optimizer for muon (used for embeddings, biases, norms).
     config.optimizer = config.muon_scalar_optimizer
 
-    assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed."
     if config.muon_scalar_optimizer == 'lion':
         assert HAVE_EO_V02, (
             "Lion optimizer requires emerging_optimizers >= 0.2. "
             "Please upgrade to use --muon-scalar-optimizer lion."
         )
+    else:
+        assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed."
 
     # Dist-opt is not supported due to strong coupling with how DDP init grad buffer
     # In theory we can change DDP to enable use muon and dist-opt-adam together

From 3de021ebeaa5448d3dadc2102c1cfbcbffa0da35 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0021.cm.cluster>
Date: Tue, 24 Mar 2026 12:04:12 -0700
Subject: [PATCH 11/20] Consolidate emerging_optimizers version detection into
 __init__.py

Single version check sets HAVE_EMERGING_OPTIMIZERS (>= 0.1) and
HAVE_EO_V02 (>= 0.2). muon.py imports both flags instead of
doing its own try/except detection.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/__init__.py | 16 +++++++---------
 megatron/core/optimizer/muon.py     |  9 +++------
 2 files changed, 10 insertions(+), 15 deletions(-)

diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 1ea171d269b..686563b1fdc 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -38,16 +38,14 @@
     from importlib.metadata import version as _pkg_version
 
     _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
-    if _eo_ver < (0, 2):
-        raise ImportError(
-            f"emerging_optimizers >= 0.2 is required, "
-            f"found {_pkg_version('emerging-optimizers')}"
-        )
-    from emerging_optimizers.scalar_optimizers import Lion
-
-    HAVE_EO_V02 = True
 except (ImportError, PackageNotFoundError):
-    HAVE_EO_V02 = False
+    _eo_ver = (0, 0)
+
+HAVE_EMERGING_OPTIMIZERS = _eo_ver >= (0, 1)
+HAVE_EO_V02 = _eo_ver >= (0, 2)
+
+if HAVE_EO_V02:
+    from emerging_optimizers.scalar_optimizers import Lion
 
 from megatron.core import parallel_state
 from megatron.core.optimizer.cpu_offloading.hybrid_optimizer import HybridDeviceOptimizer
diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 2aa9832015a..5fec2d9eb37 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -13,7 +13,7 @@
 from megatron.core.transformer.module import MegatronModule
 from megatron.core.utils import get_pg_size, log_single_rank
 
-from . import HAVE_EO_V02, _get_param_groups, get_megatron_optimizer
+from . import HAVE_EMERGING_OPTIMIZERS, HAVE_EO_V02, _get_param_groups, get_megatron_optimizer
 from .layer_wise_optimizer import LayerWiseDistributedOptimizer
 from .optimizer import (
     ChainedOptimizer,
@@ -23,16 +23,13 @@
 )
 from .optimizer_config import OptimizerConfig, ParamKey
 
-try:
+if HAVE_EMERGING_OPTIMIZERS:
     from emerging_optimizers.orthogonalized_optimizers import (
         OrthogonalizedOptimizer,
         get_muon_scale_factor,
     )
     from emerging_optimizers.orthogonalized_optimizers.muon_utils import newton_schulz_tp
-
-    HAVE_EMERGING_OPTIMIZERS = True
-except ImportError:
-    HAVE_EMERGING_OPTIMIZERS = False
+else:
     OrthogonalizedOptimizer = object
 
 if HAVE_EO_V02:

From 420ebbd4b20017123d730a17c16c27e333bdeadf Mon Sep 17 00:00:00 2001
From: root <root@pool0-0331.cm.cluster>
Date: Tue, 24 Mar 2026 18:26:04 -0700
Subject: [PATCH 12/20] Run black and isort autoformatting on PR files

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py         | 19 ++++++-------------
 megatron/core/transformer/attention.py  |  8 ++++++--
 tests/unit_tests/test_muon_optimizer.py |  1 -
 3 files changed, 12 insertions(+), 16 deletions(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 5fec2d9eb37..473e905d561 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -45,9 +45,9 @@ def get_supported_coefficient_types() -> tuple[str, ...]:
     Reads the members of the ``NSCoeffT`` Literal type so that new types
     added upstream are automatically available without code changes here.
     """
-    assert HAVE_EO_V02, (
-        "emerging_optimizers >= 0.2 is required for NSCoeffT. Please install or upgrade it."
-    )
+    assert (
+        HAVE_EO_V02
+    ), "emerging_optimizers >= 0.2 is required for NSCoeffT. Please install or upgrade it."
     return get_args(NSCoeffT)
 
 
@@ -102,14 +102,9 @@ def scaled_orthogonalize_fn(
             if partition_dim is not None:
                 size[partition_dim] *= get_pg_size(tp_group)
             mode_value = "duplicated" if mode == "blockwise" else mode
-            mode_kwarg = (
-                {"tp_mode": mode_value} if HAVE_EO_V02 else {"mode": mode_value}
-            )
+            mode_kwarg = {"tp_mode": mode_value} if HAVE_EO_V02 else {"mode": mode_value}
             ns_kwargs = dict(
-                steps=num_ns_steps,
-                tp_group=tp_group,
-                partition_dim=partition_dim,
-                **mode_kwarg,
+                steps=num_ns_steps, tp_group=tp_group, partition_dim=partition_dim, **mode_kwarg
             )
             if HAVE_EO_V02:
                 ns_kwargs["coefficient_type"] = coefficient_type
@@ -125,9 +120,7 @@ def scaled_orthogonalize_fn(
 
         weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2"
         nesterov_kwarg = (
-            {"nesterov": use_nesterov}
-            if HAVE_EO_V02
-            else {"use_nesterov": use_nesterov}
+            {"nesterov": use_nesterov} if HAVE_EO_V02 else {"use_nesterov": use_nesterov}
         )
         super().__init__(
             params,
diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index ea028f9ac47..bfab4348520 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -61,7 +61,9 @@
     rearrange = None
 
 try:
-    from flash_attn_3.flash_attn_interface import _flash_attn_forward
+    from flash_attn_3.flash_attn_interface import (
+        _flash_attn_forward,
+    )
     from flash_attn_3.flash_attn_interface import (
         flash_attn_with_kvcache as flash_attn3_with_kvcache,
     )
@@ -72,7 +74,9 @@
 
 if not HAVE_FA3:
     try:
-        from flashattn_hopper.flash_attn_interface import _flash_attn_forward
+        from flashattn_hopper.flash_attn_interface import (
+            _flash_attn_forward,
+        )
         from flashattn_hopper.flash_attn_interface import (
             flash_attn_with_kvcache as flash_attn3_with_kvcache,
         )
diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index 4e5d8245638..44e59294149 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -680,7 +680,6 @@ def test_validate_coefficient_type_rejects_invalid():
         validate_coefficient_type("nonexistent_type_xyz")
 
 
-
 def test_muon_optimizer_invalid_coefficient_type():
     """Test that TensorParallelMuon raises ValueError for an invalid coefficient_type."""
     model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda')

From 99e8ad4cbbe6a8028703519344bccf1f665cbc68 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 10:43:14 -0700
Subject: [PATCH 13/20] Fix pylint E0606 for conditional imports of
 emerging_optimizers

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/__init__.py | 2 +-
 megatron/core/optimizer/muon.py     | 4 +++-
 2 files changed, 4 insertions(+), 2 deletions(-)

diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py
index 7c37a8cced5..ef23ea22244 100644
--- a/megatron/core/optimizer/__init__.py
+++ b/megatron/core/optimizer/__init__.py
@@ -587,7 +587,7 @@ def init_state_fn(opt, config=None):
                     "Lion optimizer requires emerging_optimizers >= 0.2. "
                     "Please install or upgrade it to use --optimizer lion."
                 )
-            optimizer = Lion(
+            optimizer = Lion(  # pylint: disable=possibly-used-before-assignment
                 param_groups,
                 lr=config.lr,
                 betas=(config.lion_beta1, config.lion_beta2),
diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 473e905d561..453232f1d6c 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -48,7 +48,7 @@ def get_supported_coefficient_types() -> tuple[str, ...]:
     assert (
         HAVE_EO_V02
     ), "emerging_optimizers >= 0.2 is required for NSCoeffT. Please install or upgrade it."
-    return get_args(NSCoeffT)
+    return get_args(NSCoeffT)  # pylint: disable=possibly-used-before-assignment
 
 
 def validate_coefficient_type(coefficient_type: str) -> None:
@@ -108,7 +108,9 @@ def scaled_orthogonalize_fn(
             )
             if HAVE_EO_V02:
                 ns_kwargs["coefficient_type"] = coefficient_type
+            # pylint: disable-next=possibly-used-before-assignment
             orth_grad = newton_schulz_tp(grad, **ns_kwargs)
+            # pylint: disable-next=possibly-used-before-assignment
             scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode)
             return orth_grad * scale_factor * extra_scale_factor
 

From 8d9f7c5fc4d9cc00a6a5a54b0953e0cf6a154acd Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 10:49:21 -0700
Subject: [PATCH 14/20] Fix black formatting for flash_attn imports in
 attention.py

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/transformer/attention.py | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py
index bfab4348520..ea028f9ac47 100644
--- a/megatron/core/transformer/attention.py
+++ b/megatron/core/transformer/attention.py
@@ -61,9 +61,7 @@
     rearrange = None
 
 try:
-    from flash_attn_3.flash_attn_interface import (
-        _flash_attn_forward,
-    )
+    from flash_attn_3.flash_attn_interface import _flash_attn_forward
     from flash_attn_3.flash_attn_interface import (
         flash_attn_with_kvcache as flash_attn3_with_kvcache,
     )
@@ -74,9 +72,7 @@
 
 if not HAVE_FA3:
     try:
-        from flashattn_hopper.flash_attn_interface import (
-            _flash_attn_forward,
-        )
+        from flashattn_hopper.flash_attn_interface import _flash_attn_forward
         from flashattn_hopper.flash_attn_interface import (
             flash_attn_with_kvcache as flash_attn3_with_kvcache,
         )

From 8d39d23f4f11e2ddd66aac7030453e897c562793 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 13:03:45 -0700
Subject: [PATCH 15/20] Skip muon optimizer tests when emerging_optimizers is
 not installed

Tests were failing with AssertionError because they unconditionally
called get_supported_coefficient_types() which requires emerging_optimizers >= 0.2.
Added pytestmark skip conditions and guarded module-level collection code.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 tests/unit_tests/test_muon_optimizer.py | 26 ++++++++++++++++++-------
 1 file changed, 19 insertions(+), 7 deletions(-)

diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index 44e59294149..de63a491472 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -10,7 +10,7 @@
 
 from megatron.core import parallel_state
 from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig
-from megatron.core.optimizer import OptimizerConfig
+from megatron.core.optimizer import HAVE_EMERGING_OPTIMIZERS, HAVE_EO_V02, OptimizerConfig
 from megatron.core.optimizer.muon import (
     TensorParallelMuon,
     get_megatron_muon_optimizer,
@@ -21,11 +21,21 @@
 from megatron.core.transformer import TransformerConfig
 from tests.unit_tests.test_utilities import Utils
 
-# Skip all tests in this file for LTS versions
-pytestmark = pytest.mark.skipif(
-    Version(os.getenv('NVIDIA_PYTORCH_VERSION', "24.01")) <= Version("25.05"),
-    reason="Skip muon optimizer for LTS test",
-)
+# Skip all tests in this file for LTS versions or when emerging_optimizers is missing
+pytestmark = [
+    pytest.mark.skipif(
+        Version(os.getenv('NVIDIA_PYTORCH_VERSION', "24.01")) <= Version("25.05"),
+        reason="Skip muon optimizer for LTS test",
+    ),
+    pytest.mark.skipif(
+        not HAVE_EMERGING_OPTIMIZERS,
+        reason="emerging_optimizers package is not installed",
+    ),
+    pytest.mark.skipif(
+        not HAVE_EO_V02,
+        reason="emerging_optimizers >= 0.2 is required for coefficient type support",
+    ),
+]
 
 
 class Net(nn.Module):
@@ -426,7 +436,9 @@ def test_muon_optimizer_blockwise_mode_different_result(self):
 
 
 # All non-custom coefficient types supported by emerging_optimizers.
-_TESTABLE_COEFFICIENT_TYPES = [t for t in get_supported_coefficient_types() if t != "custom"]
+_TESTABLE_COEFFICIENT_TYPES = (
+    [t for t in get_supported_coefficient_types() if t != "custom"] if HAVE_EO_V02 else []
+)
 
 # A reasonable default NS step count for testing; get_coefficient_iterator
 # cycles/repeats coefficients so any step count works with any type.

From a1dd63d83507a7d54e5c2b04cd9ea46379cd7d3e Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 13:43:44 -0700
Subject: [PATCH 16/20] Run black autoformatting on test_muon_optimizer.py

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 tests/unit_tests/test_muon_optimizer.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index de63a491472..d6a1f3c7ecf 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -28,8 +28,7 @@
         reason="Skip muon optimizer for LTS test",
     ),
     pytest.mark.skipif(
-        not HAVE_EMERGING_OPTIMIZERS,
-        reason="emerging_optimizers package is not installed",
+        not HAVE_EMERGING_OPTIMIZERS, reason="emerging_optimizers package is not installed"
     ),
     pytest.mark.skipif(
         not HAVE_EO_V02,

From a8e82d04d52322773c418a61806ea1dbb507455a Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 14:27:08 -0700
Subject: [PATCH 17/20] Fix validate_coefficient_type crash when
 emerging_optimizers < 0.2

validate_coefficient_type was unconditionally calling
get_supported_coefficient_types() which asserts HAVE_EO_V02. This
caused dist_checkpointing tests to fail when emerging_optimizers < 0.2
was installed, since TensorParallelMuon.__init__ always calls
validate_coefficient_type. Skip validation when HAVE_EO_V02 is False,
as the coefficient_type kwarg is not passed to newton_schulz_tp in
that case anyway.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 453232f1d6c..04a3b0f15eb 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -53,6 +53,8 @@ def get_supported_coefficient_types() -> tuple[str, ...]:
 
 def validate_coefficient_type(coefficient_type: str) -> None:
     """Raise ``ValueError`` if *coefficient_type* is not supported."""
+    if not HAVE_EO_V02:
+        return
     supported = get_supported_coefficient_types()
     if coefficient_type not in supported:
         raise ValueError(

From 0bc8f486c1c4ad1aa4a353fd107f5591c678855c Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 14:29:18 -0700
Subject: [PATCH 18/20] Fall back to ("quintic",) for coefficient type
 validation without eo >= 0.2

"quintic" is the default coefficient type supported before
emerging_optimizers 0.2, so validate against it rather than
skipping validation entirely.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index 04a3b0f15eb..a3a4649f9bb 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -53,9 +53,7 @@ def get_supported_coefficient_types() -> tuple[str, ...]:
 
 def validate_coefficient_type(coefficient_type: str) -> None:
     """Raise ``ValueError`` if *coefficient_type* is not supported."""
-    if not HAVE_EO_V02:
-        return
-    supported = get_supported_coefficient_types()
+    supported = get_supported_coefficient_types() if HAVE_EO_V02 else ("quintic",)
     if coefficient_type not in supported:
         raise ValueError(
             f"Unsupported muon coefficient type '{coefficient_type}'. "

From 0d613ec896750ce32618fc090866bca97a9db62c Mon Sep 17 00:00:00 2001
From: root <root@pool0-0079.cm.cluster>
Date: Wed, 25 Mar 2026 14:33:39 -0700
Subject: [PATCH 19/20] Narrow HAVE_EO_V02 skip to only tests that call
 get_supported_coefficient_types

Now that validate_coefficient_type falls back to ("quintic",) without
emerging_optimizers >= 0.2, most tests only need HAVE_EMERGING_OPTIMIZERS.
Only the three tests that directly call get_supported_coefficient_types()
still require HAVE_EO_V02.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 tests/unit_tests/test_muon_optimizer.py | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py
index d6a1f3c7ecf..0f0a90c91ed 100644
--- a/tests/unit_tests/test_muon_optimizer.py
+++ b/tests/unit_tests/test_muon_optimizer.py
@@ -30,12 +30,12 @@
     pytest.mark.skipif(
         not HAVE_EMERGING_OPTIMIZERS, reason="emerging_optimizers package is not installed"
     ),
-    pytest.mark.skipif(
-        not HAVE_EO_V02,
-        reason="emerging_optimizers >= 0.2 is required for coefficient type support",
-    ),
 ]
 
+requires_eo_v02 = pytest.mark.skipif(
+    not HAVE_EO_V02, reason="emerging_optimizers >= 0.2 is required"
+)
+
 
 class Net(nn.Module):
     def __init__(self):
@@ -663,6 +663,7 @@ def test_muon_optimizer_extra_scale_factor():
     ), "Weight should be updated with extra_scale_factor"
 
 
+@requires_eo_v02
 def test_get_supported_coefficient_types_returns_tuple():
     """Test that get_supported_coefficient_types returns a non-empty tuple of strings."""
     supported = get_supported_coefficient_types()
@@ -672,6 +673,7 @@ def test_get_supported_coefficient_types_returns_tuple():
         assert isinstance(t, str)
 
 
+@requires_eo_v02
 def test_get_supported_coefficient_types_contains_known_types():
     """Test that the known coefficient types are present in the supported set."""
     supported = get_supported_coefficient_types()
@@ -679,6 +681,7 @@ def test_get_supported_coefficient_types_contains_known_types():
         assert expected in supported, f"Expected '{expected}' in supported types {supported}"
 
 
+@requires_eo_v02
 def test_validate_coefficient_type_accepts_valid():
     """Test that validate_coefficient_type does not raise for valid types."""
     for t in get_supported_coefficient_types():

From 10bf5eb9285246b51e3b4a53d724c69cc7063a03 Mon Sep 17 00:00:00 2001
From: root <root@pool0-0137.cm.cluster>
Date: Wed, 25 Mar 2026 20:22:17 -0700
Subject: [PATCH 20/20] Always pass coefficient_type to newton_schulz_tp
 regardless of EO version

The coefficient_type parameter is supported in emerging_optimizers >= 0.1,
so there's no need to gate it behind HAVE_EO_V02 (>= 0.2).

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 megatron/core/optimizer/muon.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py
index a3a4649f9bb..046be78ad10 100644
--- a/megatron/core/optimizer/muon.py
+++ b/megatron/core/optimizer/muon.py
@@ -106,8 +106,7 @@ def scaled_orthogonalize_fn(
             ns_kwargs = dict(
                 steps=num_ns_steps, tp_group=tp_group, partition_dim=partition_dim, **mode_kwarg
             )
-            if HAVE_EO_V02:
-                ns_kwargs["coefficient_type"] = coefficient_type
+            ns_kwargs["coefficient_type"] = coefficient_type
             # pylint: disable-next=possibly-used-before-assignment
             orth_grad = newton_schulz_tp(grad, **ns_kwargs)
             # pylint: disable-next=possibly-used-before-assignment