Add --muon-coefficient-type argument for Muon optimizer

megatron/core/optimizer/__init__.py

@@ -34,11 +34,18 @@
         USING_PYTORCH_OPTIMIZER = True
 
 try:
-    from emerging_optimizers.scalar_optimizers import Lion
+    from importlib.metadata import PackageNotFoundError
+    from importlib.metadata import version as _pkg_version
 
-    HAVE_LION = True
-except ImportError:
-    HAVE_LION = False
+    _eo_ver = tuple(int(x) for x in _pkg_version('emerging-optimizers').split('.')[:2])
+except (ImportError, PackageNotFoundError):
+    _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
@@ -575,12 +582,12 @@ 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."
+                    "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),

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
@@ -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_EMERGING_OPTIMIZERS, HAVE_EO_V02, _get_param_groups, get_megatron_optimizer
 from .layer_wise_optimizer import LayerWiseDistributedOptimizer
 from .optimizer import (
     ChainedOptimizer,
@@ -23,31 +23,44 @@
 )
 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
 
-# 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 emerging_optimizers.scalar_optimizers import Lion  # pylint: disable=unused-import
-
-    HAVE_LION = True
-except ImportError:
-    HAVE_LION = False
+if HAVE_EO_V02:
+    from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT
 
 
 logger = logging.getLogger(__name__)
 
 
+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.
+    """
+    assert (
+        HAVE_EO_V02
+    ), "emerging_optimizers >= 0.2 is required for NSCoeffT. Please install or upgrade it."
+    return get_args(NSCoeffT)  # pylint: disable=possibly-used-before-assignment
+
+
+def validate_coefficient_type(coefficient_type: str) -> None:
+    """Raise ``ValueError`` if *coefficient_type* is not supported."""
+    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}'. "
+            f"Supported types: {supported}"
+        )
+
+
 class TensorParallelMuon(OrthogonalizedOptimizer):
     """Tensor Parallel Muon optimizer."""
 
@@ -72,6 +85,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,
@@ -87,14 +101,15 @@ 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,
-                steps=num_ns_steps,
-                coefficient_type=coefficient_type,
-                tp_group=tp_group,
-                partition_dim=partition_dim,
-                mode="duplicated" if mode == "blockwise" else mode,
+            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, tp_group=tp_group, partition_dim=partition_dim, **mode_kwarg
             )
+            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
 
@@ -105,11 +120,14 @@ 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,
@@ -195,12 +213,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_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."
         )
+    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
@@ -288,6 +307,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,

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."""
 

megatron/training/arguments.py

@@ -2231,6 +2231,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',

tests/unit_tests/test_lion_optimizer.py

@@ -14,15 +14,15 @@
 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,
 )
 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,17 +97,17 @@ 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)
 
             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