diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index c32ed77d1d7..ef23ea22244 100644 --- a/megatron/core/optimizer/__init__.py +++ b/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), diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py index ae4a1a348fa..046be78ad10 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 @@ -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, 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 d1b5c34c619..d864f0ad5fc 100644 --- a/megatron/training/arguments.py +++ b/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', 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 cc99f7a16e6..0f0a90c91ed 100644 --- a/tests/unit_tests/test_muon_optimizer.py +++ b/tests/unit_tests/test_muon_optimizer.py @@ -10,16 +10,30 @@ 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 import HAVE_EMERGING_OPTIMIZERS, HAVE_EO_V02, OptimizerConfig +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 -# 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" + ), +] + +requires_eo_v02 = pytest.mark.skipif( + not HAVE_EO_V02, reason="emerging_optimizers >= 0.2 is required" ) @@ -420,10 +434,18 @@ 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)] +# All non-custom coefficient types supported by emerging_optimizers. +_TESTABLE_COEFFICIENT_TYPES = ( + [t for t in get_supported_coefficient_types() if t != "custom"] if HAVE_EO_V02 else [] ) -def test_muon_optimizer_coefficient_types(coefficient_type_and_steps): + +# 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.""" model = torch.nn.Linear(80, 40, bias=False, dtype=torch.float32, device='cuda') model.requires_grad_(True) @@ -432,8 +454,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", ) @@ -448,7 +470,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"]) @@ -641,6 +663,106 @@ 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() + assert isinstance(supported, tuple) + assert len(supported) > 0 + for t in supported: + 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() + for expected in ("simple", "quintic", "polar_express"): + 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(): + 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") + + +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", _TESTABLE_COEFFICIENT_TYPES) + def test_get_megatron_muon_optimizer_coefficient_type(self, coefficient_type): + """Test that coefficient_type flows through get_megatron_muon_optimizer.""" + 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=_DEFAULT_NS_STEPS, + 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."""