From 5990608152dd325f898829a509a791719b4414c0 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 12:45:24 -0700 Subject: [PATCH 01/11] [dev] refactor to support emerging optimizers beyond muon (#3618) --- megatron/core/optimizer/__init__.py | 192 ++++++++- .../core/optimizer/emerging_optimizers.py | 282 +++++++++++++ .../core/optimizer/layer_wise_optimizer.py | 14 +- megatron/core/optimizer/muon.py | 394 +----------------- megatron/core/optimizer/optimizer_config.py | 49 +-- megatron/core/optimizer_param_scheduler.py | 5 +- megatron/training/arguments.py | 25 +- megatron/training/checkpointing.py | 4 +- megatron/training/training.py | 55 +-- tests/unit_tests/dist_checkpointing/utils.py | 70 ++-- tests/unit_tests/test_layer_wise_optimizer.py | 71 ++-- tests/unit_tests/test_muon_optimizer.py | 71 ++-- tests/unit_tests/test_optimizer.py | 6 +- 13 files changed, 636 insertions(+), 602 deletions(-) create mode 100644 megatron/core/optimizer/emerging_optimizers.py diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index ce9bb2d30ec..9677132c685 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -2,6 +2,7 @@ import copy import logging import warnings +from collections import defaultdict from dataclasses import astuple from typing import Any, Callable, Dict, List, Optional, Tuple, Union @@ -60,7 +61,13 @@ from ..transformer.module import MegatronModule from ..utils import get_model_config, get_pg_rank, get_pg_size, is_te_min_version, log_single_rank from .distrib_optimizer import DistributedOptimizer +from .emerging_optimizers import ( + _EMERGING_OPTIMIZERS, + HAVE_EMERGING_OPTIMIZERS, + _create_emerging_optimizer, +) from .grad_scaler import ConstantGradScaler, DynamicGradScaler +from .layer_wise_optimizer import LayerWiseDistributedOptimizer from .optimizer import ( ChainedOptimizer, Float16OptimizerWithFloat16Params, @@ -68,6 +75,8 @@ MegatronOptimizer, param_group_identifier_keys, ) + +# Subclass aliases kept for backward compatibility; all are OptimizerConfig. from .optimizer_config import ( AdamOptimizerConfig, OptimizerConfig, @@ -316,14 +325,6 @@ def _get_param_groups( # Map (pg_overrides, is_expert_parallel) to params. params_map = {} - if config_overrides is None: - # TODO remove this default behavior eventually. - # This is only needed for backwards compatibility with the old config overrides API where - # the config_overrides argument by default lead to bias parameters and length 1 parameters. - # We assume that users of decoupled LR already provide config overrides so will adapt - # to the new API. - config_overrides = get_standard_config_overrides(config=config) - for model_chunk in model_chunks: for name, param in model_chunk.named_parameters(): if not param.requires_grad: @@ -458,7 +459,8 @@ def _get_megatron_optimizer_based_on_param_groups( intra_dist_opt_group: Optional[torch.distributed.ProcessGroup] = None, distributed_optimizer_instance_id: Optional[int] = 0, pg_collection: Optional[ProcessGroupCollection] = None, -) -> MegatronOptimizer: + skip_megatron_wrapping: bool = False, +) -> Union[MegatronOptimizer, Tuple[Optional[torch.optim.Optimizer], Optional[Callable]]]: """Get Megatron optimizer based on parameter groups. Args: @@ -474,12 +476,24 @@ def _get_megatron_optimizer_based_on_param_groups( optimizer. Defaults to None. distributed_optimizer_instance_id (int, optional): Distributed optimizer instance. Defaults 0. + skip_megatron_wrapping (bool): if True, return a + ``(optimizer, init_state_fn)`` tuple of the raw PyTorch optimizer + without any Megatron wrapping. Useful when the caller + (e.g. LayerWiseDistributedOptimizer) performs its own wrapping. Returns: - Instance of MegatronOptimizer. + Instance of MegatronOptimizer, or ``(optimizer, init_state_fn)`` when + *skip_megatron_wrapping=True*. """ - # TODO: Logic needs to be updated to handle different optimizer types (i.e., param_groups - # passed into this function need to correspond to the same optimizer). + # All param_groups passed here must belong to the same optimizer type (adam / sgd). + # Callers are responsible for splitting by optimizer type before calling this function. + + if skip_megatron_wrapping and config.use_precision_aware_optimizer: + raise ValueError( + "skip_megatron_wrapping=True is incompatible with use_precision_aware_optimizer." + ) + if skip_megatron_wrapping and config.optimizer_cpu_offload: + raise ValueError("skip_megatron_wrapping=True is incompatible with optimizer_cpu_offload.") # When freezing sub-models we may have no trainable parameters on a rank and # hence an empty param_groups. However, we still need to create an optimizer @@ -613,6 +627,9 @@ def init_state_fn(opt, config=None): optimizer = None init_state_fn = None + if skip_megatron_wrapping: + return optimizer, init_state_fn + # Mixed precision optimizer. # - Note: both the Float16Optimizer and the DistributedOptimizer inherit # from the MixedPrecisionOptimizer, which manages any optimizer where @@ -703,6 +720,137 @@ def check_config_overrides_consistency( return True +def _get_megatron_emerging_optimizer( + config: OptimizerConfig, + model_chunks: List[MegatronModule], + config_overrides: Optional[Dict[ParamKey, Any]] = None, + pg_collection: Optional[ProcessGroupCollection] = None, +) -> MegatronOptimizer: + """Build an emerging optimizer (e.g. Muon) for the given model chunks. + + Parameter separation (e.g., linear weights -> Muon, rest -> Adam) is expressed as a + config_override, the same mechanism used for weight-decay and learning-rate overrides. + Adam/SGD groups are delegated to _get_megatron_optimizer_based_on_param_groups so they + go through the exact same code path as the standard optimizer factory. + + When ``config.use_layer_wise_distributed_optimizer`` is True, the underlying optimizers + are wrapped with :class:`LayerWiseDistributedOptimizer`. + """ + eopt_name = config.optimizer + use_layer_wise = config.use_layer_wise_distributed_optimizer + + # Handle legacy "dist_*" optimizer names (e.g. "dist_muon" → "muon" + layer-wise). + if eopt_name.startswith('dist_'): + bare_name = eopt_name[len('dist_') :] + warnings.warn( + f"optimizer='{eopt_name}' is deprecated. " + f"Use optimizer='{bare_name}' with use_layer_wise_distributed_optimizer=True.", + DeprecationWarning, + stacklevel=3, + ) + eopt_name = bare_name + use_layer_wise = True + + if not HAVE_EMERGING_OPTIMIZERS: + raise ImportError( + f"emerging-optimizers package is required for optimizer='{eopt_name}'. " + "Install it with: pip install emerging-optimizers" + ) + if eopt_name not in _EMERGING_OPTIMIZERS: + raise ValueError(f"Unsupported emerging optimizer: {eopt_name}") + if config.fp16: + raise ValueError('emerging optimizer with fp16 is not supported.') + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + + log_single_rank(logger, logging.INFO, f'Setting up emerging optimizer with config {config}') + + # Tag parameters with optimizer-specific attributes (expert_tp, is_qkv). + for model_chunk in model_chunks: + for name, param in model_chunk.named_parameters(): + if not param.requires_grad: + continue + if 'experts' in name and 'shared' not in name: + param.expert_tp = True + # TODO(deyuf): support MLA + if 'linear_qkv.weight' in name and len(param.shape) == 2: + param.is_qkv = True + + # Apply optimizer-specific default param overrides (e.g. muon: non-linear -> adam). + config_overrides.update(_EMERGING_OPTIMIZERS[eopt_name].default_param_overrides) + + # Build param groups and bucket by (optimizer_name, is_expert_parallel). + # Layer-wise distributed optimizer handles expert params internally so we skip that split. + all_param_groups = _get_param_groups(model_chunks, config, config_overrides) + grouped_param_groups = defaultdict(list) + for group in all_param_groups: + opt_name = group.get('optimizer', eopt_name) + is_expert = group['is_expert_parallel'] and not use_layer_wise + grouped_param_groups[(opt_name, is_expert)].append(group) + + # Build an optimizer for each (optimizer_name, is_expert) bucket and combine. + results = [] + for (opt_name, is_expert), groups in grouped_param_groups.items(): + if not groups: + continue + + model_parallel_group = pg_collection.tp_ep_pp if is_expert else pg_collection.mp + + if opt_name in _EMERGING_OPTIMIZERS: + optimizer, init_state_fn = _create_emerging_optimizer( + config, groups, eopt_name, model_chunks, pg_collection + ) + if use_layer_wise: + result = (optimizer, init_state_fn) + else: + if config.bf16: + optimizer = Float16OptimizerWithFloat16Params( + optimizer, config, None, init_state_fn + ) + else: + optimizer = FP32Optimizer(optimizer, config, init_state_fn) + setattr(optimizer, 'grad_stats_parallel_group', model_parallel_group) + if pg_collection is None or not hasattr(pg_collection, 'tp'): + tp_group = parallel_state.get_tensor_model_parallel_group() + else: + tp_group = pg_collection.tp + setattr(optimizer, 'tp_group', tp_group) + result = optimizer + else: + fallback_config = copy.copy(config) + fallback_config.optimizer = opt_name + fallback_config.use_distributed_optimizer = False + result = _get_megatron_optimizer_based_on_param_groups( + config=fallback_config, + model_chunks=model_chunks, + param_groups=groups, + model_parallel_group=model_parallel_group, + pg_collection=pg_collection, + skip_megatron_wrapping=use_layer_wise, + ) + # TODO(deyuf): ChainedOptimizer currently asserts all sub-optimizers + # share the same config. Revisit this design now that emerging + # optimizers mix different optimizer types (e.g. Muon + Adam). + # For now, reset to the top-level config so the assertion holds. + if not use_layer_wise and hasattr(result, 'config'): + result.config = config + results.append(result) + + if use_layer_wise: + base_optimizers, init_fns = (), () + if results: + base_optimizers, init_fns = zip(*results) + log_single_rank( + logger, logging.INFO, f'Using LayerWiseDistributedOptimizer for {eopt_name}' + ) + return LayerWiseDistributedOptimizer( + list(base_optimizers), config, pg_collection, init_state_fn_list=list(init_fns) + ) + + return ChainedOptimizer(results) + + def get_megatron_optimizer( config: OptimizerConfig, model_chunks: List[MegatronModule], @@ -713,7 +861,10 @@ def get_megatron_optimizer( ) -> MegatronOptimizer: """Retrieve the Megatron optimizer for model chunks. + Handles both standard optimizers (Adam, SGD) and emerging optimizers (e.g. Muon). We use separate optimizers for expert parameters and non-expert parameters. + For emerging optimizers with ``config.use_layer_wise_distributed_optimizer=True``, + the optimizer is automatically wrapped with :class:`LayerWiseDistributedOptimizer`. Args: config (OptimizerConfig): optimizer configuration object. @@ -730,10 +881,25 @@ def get_megatron_optimizer( Instance of MegatronOptimizer. """ - log_single_rank(logger, logging.INFO, f'Setting up optimizer with config {config}') + # None → apply standard defaults. To extend defaults with custom overrides, + # start from get_standard_config_overrides(config) and merge yours in. + if config_overrides is None: + config_overrides = get_standard_config_overrides(config) check_config_overrides_consistency(config, config_overrides) + # TODO: the standard and emerging optimizer paths handle pg_collection differently; + # unify them so both use a single pg_collection-based flow. + if config.optimizer not in ('adam', 'sgd'): + return _get_megatron_emerging_optimizer( + config=config, + model_chunks=model_chunks, + config_overrides=config_overrides, + pg_collection=pg_collection, + ) + + log_single_rank(logger, logging.INFO, f'Setting up optimizer with config {config}') + # Separate out first model chunk if overlapping param AG with optimizer step. if config.overlap_param_gather_with_optimizer_step: all_dense_model_chunks = [[model_chunks[0]], model_chunks[1:]] diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py new file mode 100644 index 00000000000..391a247a912 --- /dev/null +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -0,0 +1,282 @@ +# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Emerging optimizer registry. + +To add a new emerging optimizer: + 1. Define its optimizer class (or import it). + 2. Write its ``__init_state_fn`` and ``__config_to_kwargs``. + 3. Add an ``EmergingOptimizerEntry`` to ``_EMERGING_OPTIMIZERS`` at the bottom. +""" + +import logging +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Literal, Optional, get_args + +import torch +from torch.optim.optimizer import ParamsT + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.utils import get_pg_size, log_single_rank + +from .optimizer_config import ParamKey, ParamPredicate + +try: + from emerging_optimizers.orthogonalized_optimizers import ( + OrthogonalizedOptimizer, + get_muon_scale_factor, + ) + from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT, newton_schulz_tp + + HAVE_EMERGING_OPTIMIZERS = True +except ImportError: + HAVE_EMERGING_OPTIMIZERS = False + + +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_EMERGING_OPTIMIZERS + ), "emerging_optimizers >= 0.2 is required for NSCoeffT. Please install or upgrade it." + return get_args(NSCoeffT) + + +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}" + ) + + +# =========================================================================== +# Registry dataclass and public API +# =========================================================================== + + +@dataclass +class EmergingOptimizerEntry: + """Everything needed to create and configure an emerging optimizer. + + Attributes: + optimizer_cls: The torch optimizer class. + init_state_fn: Lazily initialises optimizer state (needed for checkpoint formats). + config_to_kwargs: ``(config, model_chunks, pg_collection) -> dict`` of constructor kwargs. + default_param_overrides: Per-parameter config overrides applied automatically + (e.g. route non-linear params to Adam). + """ + + optimizer_cls: type + init_state_fn: Callable + config_to_kwargs: Callable + default_param_overrides: Dict[ParamKey, Dict[str, Any]] = field(default_factory=dict) + + +def _create_emerging_optimizer(config, param_groups, eopt_name, model_chunks, pg_collection): + """Instantiate an emerging optimizer and return it with its init_state_fn.""" + entry = _EMERGING_OPTIMIZERS[eopt_name] + eopt_kwargs = entry.config_to_kwargs(config, model_chunks, pg_collection) + optimizer = entry.optimizer_cls(param_groups, **eopt_kwargs) + return optimizer, entry.init_state_fn + + +# =========================================================================== +# Shared helpers +# =========================================================================== + + +def _is_nonlinear_or_embedding(param): + """True for parameters that should NOT use the emerging optimizer.""" + return getattr(param, 'is_embedding_or_output_parameter', False) or len(param.shape) != 2 + + +def _get_qkv_split_shapes(model_cfg) -> List[int]: + """Compute QKV split shapes from model config.""" + return [ + model_cfg.num_attention_heads // model_cfg.num_query_groups * model_cfg.kv_channels, + model_cfg.kv_channels, + model_cfg.kv_channels, + ] + + +# =========================================================================== +# Registry – populated below only when emerging_optimizers is installed. +# =========================================================================== + +_EMERGING_OPTIMIZERS: Dict[str, EmergingOptimizerEntry] = {} + + +# =========================================================================== +# Muon +# =========================================================================== + +if HAVE_EMERGING_OPTIMIZERS: + + class TensorParallelMuon(OrthogonalizedOptimizer): + """Tensor Parallel Muon optimizer.""" + + def __init__( + self, + params: ParamsT, + lr: float = 3e-4, + momentum_beta: float = 0.95, + use_nesterov: bool = True, + weight_decay: float = 0.01, + use_decoupled_weight_decay: bool = True, + split_qkv: bool = False, + is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, + qkv_split_shapes: tuple[int, int, int] | None = None, + fp32_matmul_prec: str = "medium", + coefficient_type: str = "quintic", + num_ns_steps: int = 5, + scale_mode: str = "spectral", + extra_scale_factor: float = 1.0, + pg_collection: Optional[ProcessGroupCollection] = None, + mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", + ) -> None: + if num_ns_steps < 1: + raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") + + def scaled_orthogonalize_fn( + grad: torch.Tensor, + tp_group: torch.distributed.ProcessGroup, + partition_dim: int | None = None, + ) -> torch.Tensor: + log_single_rank( + logger, + logging.DEBUG, + f'Orthogonalizing grad with {num_ns_steps} steps, ' + f'{coefficient_type} coefficient, ' + f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}', + ) + 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, + ) + scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode) + return orth_grad * scale_factor * extra_scale_factor + + self.pg_collection = pg_collection + self.mode = mode + self.split_qkv = split_qkv + self.is_qkv_fn = is_qkv_fn + self.qkv_split_shapes = qkv_split_shapes + + weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2" + super().__init__( + params, + lr, + momentum_beta, + use_nesterov=use_nesterov, + weight_decay=weight_decay, + weight_decay_method=weight_decay_method, + fp32_matmul_prec=fp32_matmul_prec, + scaled_orthogonalize_fn=scaled_orthogonalize_fn, + ) + + def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: + """Orthogonalize the momentum. + + Args: + p: The parameter tensor. i is necessary to pass param tensor in addition to + momentum because a lot of information is only available in the param tensor, + attributes for example. + grad: The momentum tensor. + + Returns: + The orthogonalized gradient tensor. + """ + # TODO(deyuf): switch to group + if self.pg_collection: + tp_group = ( + self.pg_collection.expt_tp + if getattr(p, 'expert_tp', False) + else self.pg_collection.tp + ) + else: + tp_group = None + partition_dim = None if self.mode == "blockwise" else getattr(p, "partition_dim", None) + if partition_dim == -1: + partition_dim = None + + if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] + grad_shape = grad.shape + log_single_rank( + logger, + logging.DEBUG, + f'qkv split grad shape {grad_shape}, ' f'split shapes {self.qkv_split_shapes}', + ) + num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes) + qkv_grads = torch.split( + grad.view(num_query_groups, sum(self.qkv_split_shapes), -1), + self.qkv_split_shapes, + dim=1, + ) + qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads] + + qkv_grads = [ + self.scaled_orthogonalize_fn(g, tp_group, partition_dim).view( + num_query_groups, -1, grad_shape[-1] + ) + for g in qkv_grads + ] + grad = torch.cat(qkv_grads, dim=1).view(grad_shape) + else: + grad = self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) + return grad + + def _muon_init_state_fn(opt, config=None): + """Initialize Muon optimizer state for torch_dist checkpoint format.""" + for group in opt.param_groups: + for p in group['params']: + if len(opt.state[p]) == 0: + opt.state[p]['momentum_buffer'] = torch.zeros_like(p.data) + + def _muon_config_to_kwargs(config, model_chunks, pg_collection) -> Dict[str, Any]: + """Convert OptimizerConfig to TensorParallelMuon constructor kwargs.""" + return { + "lr": config.lr, + "weight_decay": config.weight_decay, + "momentum_beta": config.muon_momentum, + "use_nesterov": config.muon_use_nesterov, + "fp32_matmul_prec": config.muon_fp32_matmul_prec, + "num_ns_steps": config.muon_num_ns_steps, + "scale_mode": config.muon_scale_mode, + "extra_scale_factor": config.muon_extra_scale_factor, + "mode": config.muon_tp_mode, + "split_qkv": config.muon_split_qkv, + "is_qkv_fn": lambda p: getattr(p, "is_qkv", False), + "qkv_split_shapes": _get_qkv_split_shapes(model_chunks[0].config), + "pg_collection": pg_collection, + } + + # ----------------------------------------------------------------------- + # Register Muon + # ----------------------------------------------------------------------- + _EMERGING_OPTIMIZERS['muon'] = EmergingOptimizerEntry( + optimizer_cls=TensorParallelMuon, + init_state_fn=_muon_init_state_fn, + config_to_kwargs=_muon_config_to_kwargs, + default_param_overrides={ + ParamKey( + predicate=ParamPredicate( + name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding + ) + ): {'optimizer': 'adam'} + }, + ) diff --git a/megatron/core/optimizer/layer_wise_optimizer.py b/megatron/core/optimizer/layer_wise_optimizer.py index a9fdc7ba72f..6e59e03ae42 100644 --- a/megatron/core/optimizer/layer_wise_optimizer.py +++ b/megatron/core/optimizer/layer_wise_optimizer.py @@ -76,19 +76,17 @@ def __init__( optimizers ), "init_state_fn_list must be the same length as optimizers if provided" - # wrap optimizer after sharding to avoid unnecessary master weight creation - # for higher precision, optimizers are wrapped with megatron already + # Wrap base torch optimizers with Float16 for bf16 training. + # Callers pass base optimizers; wrapping happens here *after* + # shard_params so master weights are only created for the local shard. if config.bf16: - # unwrap FP32 optimizer, possibly from reusing get_megatron_optimizer for adam for i in range(len(optimizers)): opt = optimizers[i] - if isinstance(opt, Float16OptimizerWithFloat16Params): + if isinstance(opt, (Float16OptimizerWithFloat16Params, FP32Optimizer)): raise TypeError( - 'LayerWiseDistributedOptimizer received Float16 optimizer already.' + 'LayerWiseDistributedOptimizer expects base torch optimizers, ' + f'got {type(opt).__name__}. Do not pre-wrap with Megatron optimizers.' ) - # unwrap FP32 optimizer from reusing get_megatron_optimizer for adam - if isinstance(opt, FP32Optimizer): - opt = opt.optimizer optimizers[i] = Float16OptimizerWithFloat16Params( opt, config, None, init_state_fn_list[i] if init_state_fn_list else None ) diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py index 5dfddd786f5..a3f7506f941 100644 --- a/megatron/core/optimizer/muon.py +++ b/megatron/core/optimizer/muon.py @@ -1,394 +1,16 @@ # Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -"""Megatron muon optimizer wrapper to handle tensor-parallel.""" +"""Backward-compatible shim — all code now lives in ``emerging_optimizers``.""" -import logging -from typing import Any, Callable, Dict, List, Literal, Optional, get_args +from typing import Any -import torch -from torch.optim.optimizer import ParamsT -from megatron.core.optimizer_param_scheduler import ParamGroupOverride -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.transformer.module import MegatronModule -from megatron.core.utils import get_pg_size, log_single_rank +def get_megatron_muon_optimizer(*args: Any, **kwargs: Any) -> Any: + """Backward compatible muon optimizer getter. -from . import HAVE_EMERGING_OPTIMIZERS, _get_param_groups, get_megatron_optimizer -from .layer_wise_optimizer import LayerWiseDistributedOptimizer -from .optimizer import ( - ChainedOptimizer, - Float16OptimizerWithFloat16Params, - FP32Optimizer, - MegatronOptimizer, -) -from .optimizer_config import OptimizerConfig, ParamKey - -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 -else: - OrthogonalizedOptimizer = object - -if HAVE_EMERGING_OPTIMIZERS: - 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_EMERGING_OPTIMIZERS - ), "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 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.""" - - def __init__( - self, - params: ParamsT, - lr: float = 3e-4, - momentum_beta: float = 0.95, - use_nesterov: bool = True, - weight_decay: float = 0.01, - use_decoupled_weight_decay: bool = True, - split_qkv: bool = False, - is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, - qkv_split_shapes: tuple[int, int, int] | None = None, - fp32_matmul_prec: str = "medium", - coefficient_type: str = "quintic", - num_ns_steps: int = 5, - scale_mode: str = "spectral", - extra_scale_factor: float = 1.0, - pg_collection: Optional[ProcessGroupCollection] = None, - mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", - ) -> 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, - tp_group: torch.distributed.ProcessGroup, - partition_dim: int | None = None, - ) -> torch.Tensor: - log_single_rank( - logger, - logging.DEBUG, - f'Orthogonalizing grad with {num_ns_steps} steps, {coefficient_type} coefficient, ' - f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}', - ) - size = [grad.size(-2), grad.size(-1)] - 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} - 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 - - self.pg_collection = pg_collection - self.mode = mode - self.split_qkv = split_qkv - self.is_qkv_fn = is_qkv_fn - self.qkv_split_shapes = qkv_split_shapes - - weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2" - nesterov_kwarg = {"nesterov": use_nesterov} - super().__init__( - params, - lr, - momentum_beta, - **nesterov_kwarg, - weight_decay=weight_decay, - weight_decay_method=weight_decay_method, - fp32_matmul_prec=fp32_matmul_prec, - scaled_orthogonalize_fn=scaled_orthogonalize_fn, - ) - - def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: - """Orthogonalize the momentum. - - Args: - p: The parameter tensor. i is necessary to pass param tensor in addition to momentum - because a lot of information is only available in the param tensor, - attributes for example. - grad: The momentum tensor. - - Returns: - The orthogonalized gradient tensor. - """ - # TODO(deyuf): switch to group - if self.pg_collection: - tp_group = ( - self.pg_collection.expt_tp - if getattr(p, 'expert_tp', False) - else self.pg_collection.tp - ) - else: - tp_group = None - partition_dim = None if self.mode == "blockwise" else getattr(p, "partition_dim", None) - if partition_dim == -1: - # emerging-optimizers use None instead of -1 to indicate no tensor parallel - partition_dim = None - - if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] - # split grouped attention parameters (e.g., QKV, GQA, etc.) - grad_shape = grad.shape - log_single_rank( - logger, - logging.DEBUG, - f'qkv split grad shape {grad_shape}, split shapes {self.qkv_split_shapes}', - ) - num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes) - qkv_grads = torch.split( - grad.view(num_query_groups, sum(self.qkv_split_shapes), -1), - self.qkv_split_shapes, - dim=1, - ) - qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads] - - # Apply Newton-Schulz and scales to each component, concat back - qkv_grads = [ - self.scaled_orthogonalize_fn(g, tp_group, partition_dim).view( - num_query_groups, -1, grad_shape[-1] - ) - for g in qkv_grads - ] - grad = torch.cat(qkv_grads, dim=1).view(grad_shape) - else: - grad = self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) - return grad - - -def get_megatron_muon_optimizer( - config: OptimizerConfig, - model_chunks: List[MegatronModule], - config_overrides: Optional[Dict[ParamKey, ParamGroupOverride]] = None, - use_gloo_process_groups: bool = True, - layer_wise_distributed_optimizer: bool = False, - pg_collection: Optional[ProcessGroupCollection] = None, -) -> MegatronOptimizer: - """This function is used to get the muon optimizer for the model chunks. - It is used to get the muon optimizer for the model chunks. - - Args: - config (OptimizerConfig): optimizer configuration object. - model_chunks (List[MegatronModule]): model chunks to get optimizer for. - use_gloo_process_groups (bool): if false, disable use of Gloo process groups - in underlying Megatron optimizers. - layer_wise_distributed_optimizer (bool): if true, use layer-wise distributed optimizer. - Defaults to False. + .. deprecated:: + Use :func:`megatron.core.optimizer.get_megatron_optimizer` instead. """ - # TODO: Mutating config.optimizer is a side effect; clean up after - # https://github.com/NVIDIA/Megatron-LM/pull/3638 lands. - # Set the nonlinear optimizer for muon (used for embeddings, biases, norms). - config.optimizer = config.muon_scalar_optimizer - - assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers >= 0.2 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 - if config.use_distributed_optimizer: - raise Exception('muon with dist optimizer is not supported.') - # only support bf16 w/o loss scale now - if config.fp16: - raise Exception('muon with fp16 is not supported.') - - # before this function receive properly created collection - if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups() - - log_single_rank(logger, logging.INFO, f'Setting up emerging optimizer with config {config}') - - # Needed for torch_dist ckpt_format, unlike torch ckpt_format - # For other emerging optimizers, need to implement init_state_fn as well - # TODO(boxiangw): Improve usability after optimizer refactor - # TODO(boxiangw): support precision aware optimizer - def muon_init_state_fn(opt, config=None): - for group in opt.param_groups: - for p in group['params']: - if len(opt.state[p]) == 0: - opt.state[p]['momentum_buffer'] = torch.zeros_like(p.data) - - def adam_init_state_fn(opt, config=None): - for group in opt.param_groups: - for p in group['params']: - if len(opt.state[p]) == 0: - if config is None or not config.use_precision_aware_optimizer: - opt.state[p]['exp_avg'] = torch.zeros_like(p.data) - opt.state[p]['exp_avg_sq'] = torch.zeros_like(p.data) - else: - opt.initialize_state(p) - - def lion_init_state_fn(opt, config=None): - for group in opt.param_groups: - for p in group['params']: - if len(opt.state[p]) == 0: - opt.state[p]['exp_avg'] = torch.zeros_like(p.data) - - nonlinear_init_state_fn = ( - lion_init_state_fn if config.muon_scalar_optimizer == 'lion' else adam_init_state_fn - ) - - optimizers = [] - # record list of non/linear params - linear_params = [] - nonlinear_params = [] - for model_chunk in model_chunks: - # use config to determine qkv split shapes. - # no need to check tp since tp splits by head and this is per head(group) dimension - num_attention_heads = model_chunk.config.num_attention_heads - num_query_groups = model_chunk.config.num_query_groups - kv_channels = model_chunk.config.kv_channels - qkv_split_shapes = [ - num_attention_heads // num_query_groups * kv_channels, - kv_channels, - kv_channels, - ] - for name, param in model_chunk.named_parameters(): - if not param.requires_grad: - continue - # add flag for expert weight so optimizer can figure which tp group it uses - # alternatively, create new param group and save tp_group. this require more - # change in optimizer - if 'experts' in name and 'shared' not in name: - param.expert_tp = True - # add flag for qkv parameter - # TODO(deyuf): support MLA - if 'linear_qkv.weight' in name and len(param.shape) == 2: - param.is_qkv = True - # TODO(deyuf): currently only allow 2D non-embedding weight to avoid breaking - if ( - not getattr(param, 'is_embedding_or_output_parameter', False) - and len(param.shape) == 2 - ): - linear_params.append(param) - else: - nonlinear_params.append(param) - - muon_kwargs = { - "lr": config.lr, - "momentum_beta": config.muon_momentum, - "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, - "is_qkv_fn": lambda p: getattr(p, "is_qkv", False), - "qkv_split_shapes": qkv_split_shapes, - "extra_scale_factor": config.muon_extra_scale_factor, - "pg_collection": pg_collection, - "mode": config.muon_tp_mode, - } - - # freezing nonlinear params and get param groups for muon - for param in nonlinear_params: - param.requires_grad = False - - linear_param_groups = _get_param_groups(model_chunks, config, config_overrides) - # if layerwise distributed optimizer is not used, need to handle ep params separately - expert_param_groups = [] - if not layer_wise_distributed_optimizer: - for group in linear_param_groups: - if group['is_expert_parallel']: - expert_param_groups.append(group) - linear_param_groups.remove(group) - - optimizer = TensorParallelMuon(linear_param_groups, **muon_kwargs) - - reset_config_bf16 = False - if config.bf16: - if layer_wise_distributed_optimizer: - # creating master weight before layerwise sharding will lead to unnecessary master - # weight so here we delay master weight creation into layer_wise unset config.bf16 - # will also result in all optimizers below(adam) to also not be wrapped - config.bf16 = False - reset_config_bf16 = True - else: - # if not using layer_wise wrapper, just create master weight here is fine - optimizer = Float16OptimizerWithFloat16Params( - optimizer, config, None, muon_init_state_fn - ) - else: - optimizer = FP32Optimizer(optimizer, config, muon_init_state_fn) - - optimizers.append(optimizer) - - # expert optimizer exists meaning layerwise distributed optimizer is not used - if len(expert_param_groups) > 0: - expert_optimizer = TensorParallelMuon(expert_param_groups, **muon_kwargs) - if config.bf16: - expert_optimizer = Float16OptimizerWithFloat16Params( - expert_optimizer, config, None, muon_init_state_fn - ) - else: - expert_optimizer = FP32Optimizer(expert_optimizer, config, muon_init_state_fn) - setattr(expert_optimizer, 'grad_stats_parallel_group', pg_collection.tp_ep_pp) - optimizers.append(expert_optimizer) - - # done with muon, unfreeze nonlinear and freeze linear - for param in nonlinear_params: - param.requires_grad = True - for param in linear_params: - param.requires_grad = False - - # call original get. linear params will be skipped since they're freezed - chained_adam = get_megatron_optimizer( - config, - model_chunks, - config_overrides=config_overrides, - use_gloo_process_groups=use_gloo_process_groups, - ) - - # unfreeze everything - for param in linear_params: - param.requires_grad = True - - # chain everything together - init_fns = [muon_init_state_fn] + len(chained_adam.chained_optimizers) * [ - nonlinear_init_state_fn - ] - optimizers += chained_adam.chained_optimizers + from . import get_megatron_optimizer - if layer_wise_distributed_optimizer: - log_single_rank(logger, logging.INFO, 'Using LayerWiseDistributedOptimizer for Muon') - if reset_config_bf16: - config.bf16 = True - return LayerWiseDistributedOptimizer( - optimizers, - config, - pg_collection, - init_state_fn_list=init_fns, - model_chunks=model_chunks, - async_allgather=config.overlap_param_gather, - ) - return ChainedOptimizer(optimizers) + return get_megatron_optimizer(*args, **kwargs) diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 9e6375b978c..7534227f5a0 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -207,7 +207,8 @@ class OptimizerConfig: """dtype of exp_avg_sq when enabling precision-aware-optimizer""" optimizer: str = 'adam' - """Optimizer name. NOTE: Deprecated, use individual optimizer classes instead.""" + """Optimizer name (e.g., 'adam', 'sgd', 'muon'). Can be overridden per-parameter group + via config_overrides to use different optimizers for different parameters.""" ############### # Loss scaling @@ -230,7 +231,7 @@ class OptimizerConfig: """Hysteresis for dynamic loss scaling.""" ################################################################################### - # Optimizer (NOTE: Deprecated, use individual optimizer classes instead.). + # Optimizer-specific parameters. ################################################################################### # Adam. adam_beta1: float = 0.9 @@ -255,10 +256,9 @@ class OptimizerConfig: sgd_momentum: float = 0.9 """Momentum factor for SGD optimizer.""" - # Muon. - # TODO: move muon configs to it's own `MuonConfig`. + # Muon / emerging optimizers. muon_momentum: float = 0.95 - """The momentum used by the internal SGD.""" + """The momentum used by the internal SGD in Muon optimizer.""" muon_split_qkv: bool = True """Whether to split QKV parameters for Muon optimizer.""" @@ -303,6 +303,12 @@ class OptimizerConfig: use_distributed_optimizer: bool = False """Distribute optimizer state over data-parallel replicas.""" + use_layer_wise_distributed_optimizer: bool = False + """Use :class:`LayerWiseDistributedOptimizer` for emerging optimizers (e.g. Muon). + When set via ``--use-distributed-optimizer`` with an emerging optimizer, the training + arguments layer sets this flag and resets ``use_distributed_optimizer`` to False so + that the standard distributed-optimizer path is not triggered.""" + overlap_param_gather: bool = False """If true, overlap param all-gather with forward compute. This argument is intended to have the same value as the "overlap_param_gather" argument @@ -442,33 +448,6 @@ def __post_init__(self): ), "exp_avg_sq_dtype can only be fp32 when not using precision-aware optimizer" -@dataclass -class AdamOptimizerConfig(OptimizerConfig): - """Adam optimizer configuration object.""" - - optimizer: str = 'adam' - """Optimizer name.""" - - adam_beta1: float = 0.9 - """First coefficient for computing running averages of gradient and its square in Adam - optimizer. - """ - - adam_beta2: float = 0.999 - """Second coefficient for computing running averages of gradient and its square in Adam - optimizer. - """ - - adam_eps: float = 1e-08 - """Term added to the denominator to improve numerical stability in Adam optimizer.""" - - -@dataclass -class SGDOptimizerConfig(OptimizerConfig): - """SGD optimizer configuration object.""" - - optimizer: str = 'sgd' - """Optimizer name.""" - - sgd_momentum: float = 0.9 - """Momentum factor for SGD optimizer.""" +# Backward-compatible aliases (deprecated; use OptimizerConfig directly). +AdamOptimizerConfig = OptimizerConfig +SGDOptimizerConfig = OptimizerConfig diff --git a/megatron/core/optimizer_param_scheduler.py b/megatron/core/optimizer_param_scheduler.py index e01a708ce79..91ed362b1b2 100644 --- a/megatron/core/optimizer_param_scheduler.py +++ b/megatron/core/optimizer_param_scheduler.py @@ -14,7 +14,7 @@ logger = logging.getLogger(__name__) -class ParamGroupOverride(TypedDict): +class ParamGroupOverride(TypedDict, total=False): """Override values for a parameter group. These values may be optimizer-state/scheduler related. These are the values you see later in param_group.get(...) calls in the @@ -23,7 +23,7 @@ class ParamGroupOverride(TypedDict): Example: >>> param_group_override = ParamGroupOverride(min_lr=1e-4, wd_mult=0.1) - >>> param_group_override == ParamGroupOverride(newvar=3) # this is ok too + >>> param_group_override == ParamGroupOverride(optimizer='muon') # per-param optimizer """ @@ -32,6 +32,7 @@ class ParamGroupOverride(TypedDict): start_wd: float end_wd: float wd_mult: float + optimizer: str def get_canonical_lr_for_logging(param_groups: list[dict]) -> float | None: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 286e14c477d..25e8eb02e80 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1472,14 +1472,23 @@ def validate_args(args, defaults={}): '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' \ '--rl-offload-optimizer-during-inference is incompatible (no optimizer to offload).' - # Muon optimizer check - if 'muon' in args.optimizer: + # Muon / emerging optimizer check + if args.optimizer in ('muon', 'dist_muon'): + if args.optimizer == 'dist_muon': + warn_rank_0( + "optimizer='dist_muon' is deprecated. " + "Use --optimizer muon --use-distributed-optimizer instead." + ) + args.optimizer = 'muon' + args.use_layer_wise_distributed_optimizer = True - if args.optimizer == 'muon': - assert not args.overlap_grad_reduce, "Muon optimizer does not support overlap grad reduce. Use dist_muon instead." - assert not args.overlap_param_gather, "Muon optimizer does not support overlap param gather. Use dist_muon instead." + if args.use_distributed_optimizer: + args.use_layer_wise_distributed_optimizer = True + args.use_distributed_optimizer = False - assert not args.use_distributed_optimizer, "Muon optimizer does not support distributed optimizer for now." + # TODO: remove these checks once we support them + assert not args.overlap_grad_reduce, "Muon optimizer does not support overlap grad reduce for now." + assert not args.overlap_param_gather, "Muon optimizer does not support overlap param gather for now." assert not args.use_torch_fsdp2, "Muon optimizer does not support Torch-FSDP2 for now." assert not args.use_megatron_fsdp, "Muon optimizer does not support Megatron-FSDP for now." assert args.ckpt_format in ["torch", "torch_dist"], "Muon optimizer supports torch and torch_dist checkpoint format." @@ -2477,7 +2486,9 @@ def _add_training_args(parser): 'https://arxiv.org/abs/2205.14135') group.add_argument('--optimizer', type=str, default='adam', choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion'], - help='Optimizer function') + help='Optimizer function. ' + 'Note: dist_muon is deprecated; use --optimizer muon ' + 'with --use-distributed-optimizer instead.') group.add_argument('--optimizer-cpu-offload', action='store_true', help='Offload optimizer state to CPU') group.add_argument('--optimizer-offload-fraction', type=float, default=1.0, diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 6401b6cc7bc..99328f4a010 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -571,7 +571,7 @@ def save_checkpoint(iteration, model, optimizer, opt_param_scheduler, num_floati optimizer.save_parameter_state(optim_checkpoint_name) # LayerWiseDistributedOptimizer save optimizer state to file on different ranks - if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == 'torch': + if getattr(args, "use_layer_wise_distributed_optimizer", False) and args.ckpt_format == 'torch': dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join(os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt") ensure_directory_exists(optim_checkpoint_name) @@ -1870,7 +1870,7 @@ def load_model_state_dict(module, state_dict, strict: bool): if not release and not args.finetune and not args.no_load_optim: try: # Load state dict. - if getattr(args, "optimizer", "adam").startswith("dist_") and args.ckpt_format == 'torch': + if getattr(args, "use_layer_wise_distributed_optimizer", False) and args.ckpt_format == 'torch': # LayerWiseDistributedOptimizer load optimizer state from file on different ranks dp_rank = mpu.get_data_parallel_rank() optim_checkpoint_name = os.path.join(os.path.dirname(checkpoint_name), f"layer_wise_optimizer_{dp_rank}.pt") diff --git a/megatron/training/training.py b/megatron/training/training.py index d7f7a2ff4f0..84e38cb88a6 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -175,8 +175,11 @@ def set_startup_timestamps(program_start=None, main_entry=None): from megatron.core.distributed import finalize_model_grads from megatron.core.enums import ModelType -from megatron.core.optimizer import get_megatron_optimizer, AdamOptimizerConfig, SGDOptimizerConfig, OptimizerConfig, ParamKey -from megatron.core.optimizer.muon import get_megatron_muon_optimizer +from megatron.core.optimizer import ( + get_megatron_optimizer, + OptimizerConfig, + ParamKey, +) from megatron.core.rerun_state_machine import ( get_rerun_state_machine, destroy_rerun_state_machine, @@ -1571,23 +1574,11 @@ def get_optimizer_param_scheduler(optimizer): def get_megatron_optimizer_config(args: Any) -> OptimizerConfig: """Return a Megatron optimizer config object from Megatron's arguments.""" - config = None - if args.optimizer == 'adam' or 'muon' in args.optimizer: - # TODO(deyuf): Muon needs both adam + muon but get() only receive one config - # So for now we keep using adam config that's back compat with old way - kwargs = {} - for f in dataclasses.fields(AdamOptimizerConfig): - if hasattr(args, f.name): - kwargs[f.name] = getattr(args, f.name) - config = AdamOptimizerConfig(**kwargs) - elif args.optimizer == 'sgd': - kwargs = {} - for f in dataclasses.fields(SGDOptimizerConfig): - if hasattr(args, f.name): - kwargs[f.name] = getattr(args, f.name) - config = SGDOptimizerConfig(**kwargs) - else: - raise ValueError("Invalid optimizer type!") + kwargs = {} + for f in dataclasses.fields(OptimizerConfig): + if hasattr(args, f.name): + kwargs[f.name] = getattr(args, f.name) + config = OptimizerConfig(**kwargs) # Construct the appropriate config_overrides object. This default handles many cases, but # can be added to as needed by the user, or replaced entirely with a custom override. @@ -1636,25 +1627,13 @@ def setup_model_and_optimizer( if mup_overrides: config_overrides = {**(config_overrides or {}), **mup_overrides} - if 'muon' not in config.optimizer: - # If the user is asking for a non-zero embedding init std, skip weight decay for embeddings - # to avoid embeddings from shrinking to zero as recommended in https://arxiv.org/abs/2312.16903 - # default_skip_embedding_weight_decay=args.embedding_init_method_std is not None, - optimizer = get_megatron_optimizer( - config, - model, - config_overrides=config_overrides, - use_gloo_process_groups=args.use_gloo_process_groups, - dump_param_to_param_group_map=args.dump_param_to_param_group_map, - ) - else: - optimizer = get_megatron_muon_optimizer( - config, - model, - config_overrides=config_overrides, - use_gloo_process_groups=args.use_gloo_process_groups, - layer_wise_distributed_optimizer='dist' in config.optimizer, - ) + optimizer = get_megatron_optimizer( + config, + model, + config_overrides=config_overrides, + use_gloo_process_groups=args.use_gloo_process_groups, + dump_param_to_param_group_map=args.dump_param_to_param_group_map, + ) opt_param_scheduler = get_optimizer_param_scheduler(optimizer) one_logger and one_logger.log_metrics({"app_build_optimzer_finish_time": one_logger_utils.get_timestamp_in_ms()}) diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py index ec95602b020..8fe5e28e12c 100644 --- a/tests/unit_tests/dist_checkpointing/utils.py +++ b/tests/unit_tests/dist_checkpointing/utils.py @@ -15,7 +15,7 @@ get_gpt_layer_with_transformer_engine_spec, ) from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer -from megatron.core.optimizer.muon import get_megatron_muon_optimizer +from megatron.core.optimizer.optimizer import ChainedOptimizer from megatron.core.tensor_parallel import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig from megatron.training.arguments import parse_args @@ -178,11 +178,6 @@ def init_checkpointing_mock_args(args, ckpt_dir, fully_parallel=False): def setup_model_and_optimizer( seed, tp, pp, initialize_fn=initialize_gpt_model, bf16=True, dist_opt=True, optimizer='adam' ): - if 'muon' in optimizer and dist_opt: - raise ValueError( - "Layer-wise distributed optimizer with Muon is not supported with distributed optimizer." - ) - mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) @@ -197,37 +192,39 @@ def setup_model_and_optimizer( ) ) + optimizer_type = optimizer + use_layer_wise = False + if optimizer_type == 'dist_muon': + optimizer = 'muon' + use_layer_wise = True + if optimizer_type in ('muon', 'dist_muon') and dist_opt: + use_layer_wise = True + dist_opt = False + config = OptimizerConfig( bf16=bf16, params_dtype=torch.bfloat16 if bf16 else torch.float, use_distributed_optimizer=dist_opt, + use_layer_wise_distributed_optimizer=use_layer_wise, optimizer=optimizer, ) - if 'muon' in optimizer: - # Use layer-wise distributed optimizer with Muon - optimizer_type = optimizer - # default lr None feels wrong. only change muon lr to avoid breaking old tests + if optimizer_type in ('muon', 'dist_muon'): config.lr = 0.0 - optimizer = get_megatron_muon_optimizer( - config, model, layer_wise_distributed_optimizer='dist' in optimizer_type - ) - else: - optimizer_type = optimizer - optimizer = get_megatron_optimizer(config, model) + optimizer = get_megatron_optimizer(config, model) torch.manual_seed(seed + 1) model_parallel_cuda_manual_seed(seed + 1) - if not 'muon' in optimizer_type: + if isinstance(optimizer, ChainedOptimizer): + for opt in optimizer.chained_optimizers: + opt.init_state_fn(opt) + else: for group in optimizer.optimizer.param_groups: for p in group['params']: if len(optimizer.optimizer.state[p]) == 0: optimizer.optimizer.state[p]['exp_avg'] = torch.rand_like(p.data) optimizer.optimizer.state[p]['exp_avg_sq'] = torch.rand_like(p.data) - else: - for opt in optimizer.chained_optimizers: - opt.init_state_fn(opt) optimizer.reload_model_params() CachedMetadataFileSystemReader.clear_metadata_cache() @@ -272,10 +269,6 @@ def setup_moe_model_and_optimizer( use_glu=False, optimizer='adam', ): - if 'muon' in optimizer and dist_opt: - raise ValueError( - "Layer-wise distributed optimizer with Muon is not supported with distributed optimizer." - ) mock_args = parse_args(ignore_unknown_args=True) with mock.patch('megatron.training.training.get_args', new=lambda: mock_args): init_basic_mock_args(mock_args, tp, pp, bf16=bf16) @@ -295,37 +288,40 @@ def setup_moe_model_and_optimizer( ) ) + optimizer_type = optimizer + use_layer_wise = False + if optimizer_type == 'dist_muon': + optimizer = 'muon' + use_layer_wise = True + if optimizer_type in ('muon', 'dist_muon') and dist_opt: + use_layer_wise = True + dist_opt = False + config = OptimizerConfig( bf16=bf16, params_dtype=torch.bfloat16 if bf16 else torch.float, use_distributed_optimizer=dist_opt, + use_layer_wise_distributed_optimizer=use_layer_wise, optimizer=optimizer, ) - if 'muon' in optimizer: - optimizer_type = optimizer - # default lr None feels wrong. only change muon lr to avoid breaking old tests + if optimizer_type in ('muon', 'dist_muon'): config.lr = 0.0 - optimizer = get_megatron_muon_optimizer( - config, model, layer_wise_distributed_optimizer='dist' in optimizer_type - ) - else: - optimizer_type = optimizer - optimizer = get_megatron_optimizer(config, model) + optimizer = get_megatron_optimizer(config, model) torch.manual_seed(seed + 1) model_parallel_cuda_manual_seed(seed + 1) - if not 'muon' in optimizer_type: + if optimizer_type in ('muon', 'dist_muon'): + for opt in optimizer.chained_optimizers: + opt.init_state_fn(opt) + else: for opt in optimizer.chained_optimizers: for group in opt.param_groups: for p in group['params']: if len(opt.state[p]) == 0: opt.state[p]['exp_avg'] = torch.rand_like(p.data) opt.state[p]['exp_avg_sq'] = torch.rand_like(p.data) - else: - for opt in optimizer.chained_optimizers: - opt.init_state_fn(opt) optimizer.reload_model_params() CachedMetadataFileSystemReader.clear_metadata_cache() diff --git a/tests/unit_tests/test_layer_wise_optimizer.py b/tests/unit_tests/test_layer_wise_optimizer.py index c484ca104ee..8cc13d19b13 100644 --- a/tests/unit_tests/test_layer_wise_optimizer.py +++ b/tests/unit_tests/test_layer_wise_optimizer.py @@ -10,9 +10,8 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.optimizer import OptimizerConfig +from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer from megatron.core.optimizer.layer_wise_optimizer import LayerWiseDistributedOptimizer -from megatron.core.optimizer.muon import get_megatron_muon_optimizer from megatron.core.optimizer.optimizer import Float16OptimizerWithFloat16Params from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer import TransformerConfig @@ -124,13 +123,14 @@ def create_model_and_optimizer( pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True) pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() - optimizer = get_megatron_muon_optimizer( - config=optimizer_config, - model_chunks=[model], - use_gloo_process_groups=True, - layer_wise_distributed_optimizer=use_layer_wise, - pg_collection=pg_collection, - ) + optimizer = get_megatron_optimizer(optimizer_config, [model]) + if use_layer_wise: + # Extract base torch optimizers from the FP32Optimizer wrappers. + base_optimizers = [opt.optimizer for opt in optimizer.chained_optimizers] + optimizer_config.bf16 = True + optimizer = LayerWiseDistributedOptimizer( + base_optimizers, optimizer_config, pg_collection + ) return model, optimizer, pg_collection def create_model_and_optimizer_with_overlap_param_gather( @@ -146,7 +146,7 @@ def create_model_and_optimizer_with_overlap_param_gather( """Create model, DDP wrapper, and optimizer with overlap-param-gather enabled. This variant sets overlap_param_gather=True in DDP config and uses - get_megatron_muon_optimizer with layer_wise_distributed_optimizer=True, + get_megatron_optimizer with layer_wise_distributed_optimizer=True, enabling the bucket-based async param gather path. Args: @@ -197,12 +197,8 @@ def create_model_and_optimizer_with_overlap_param_gather( pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True) pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() - optimizer = get_megatron_muon_optimizer( - config=optimizer_config, - model_chunks=[model], - use_gloo_process_groups=True, - layer_wise_distributed_optimizer=True, - pg_collection=pg_collection, + optimizer = get_megatron_optimizer( + config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True ) return model, optimizer, pg_collection @@ -350,8 +346,34 @@ def test_multiple_optimizers(self): """ model, optimizer, pg_collection = self.create_model_and_optimizer() - # get_megatron_muon_optimizer produces muon + adam chained optimizers - assert len(optimizer.chained_optimizers) >= 2, "Should have multiple chained optimizers" + ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=False) + model = DistributedDataParallel( + TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model + ) + + optimizer_config = OptimizerConfig( + optimizer='adam', lr=0.01, bf16=True, use_distributed_optimizer=False + ) + + # Split parameters into two groups for testing multiple optimizers + params = list(model.parameters()) + mid_point = len(params) // 2 + param_groups_1 = [{'params': params[:mid_point]}] + param_groups_2 = [{'params': params[mid_point:]}] + + # Create two separate plain base optimizers (LayerWise wraps them itself) + base_optimizer_1 = torch.optim.Adam(param_groups_1, lr=optimizer_config.lr) + base_optimizer_2 = torch.optim.Adam(param_groups_2, lr=optimizer_config.lr) + + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True) + pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() + + optimizer = LayerWiseDistributedOptimizer( + [base_optimizer_1, base_optimizer_2], optimizer_config, pg_collection + ) + + assert len(optimizer.chained_optimizers) == 2, "Should have two chained optimizers" # Set gradients and test optimizer step - this will trigger allgather for param in model.parameters(): @@ -391,7 +413,7 @@ def test_bf16_error(self): pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True) pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() - # Create muon optimizer (non-layer-wise) — produces Float16-wrapped chained optimizers + # Create optimizer (non-layer-wise) — produces Float16-wrapped chained optimizers optimizer_config = OptimizerConfig( optimizer='muon', lr=0.01, @@ -399,12 +421,8 @@ def test_bf16_error(self): use_distributed_optimizer=False, muon_tp_mode="duplicated", ) - muon_optimizer = get_megatron_muon_optimizer( - config=optimizer_config, - model_chunks=[model], - use_gloo_process_groups=True, - layer_wise_distributed_optimizer=False, - pg_collection=pg_collection, + muon_optimizer = get_megatron_optimizer( + config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True ) # Extract a Float16-wrapped chained optimizer @@ -412,12 +430,11 @@ def test_bf16_error(self): assert isinstance(wrapped_optimizer, Float16OptimizerWithFloat16Params) # Should raise TypeError when receiving already-wrapped Float16 optimizer - # Use a fresh config since get_megatron_muon_optimizer mutates config.optimizer lw_config = OptimizerConfig( optimizer='muon', lr=0.01, bf16=True, use_distributed_optimizer=False ) with pytest.raises( - TypeError, match='LayerWiseDistributedOptimizer received Float16 optimizer already' + TypeError, match='LayerWiseDistributedOptimizer expects base torch optimizers' ): LayerWiseDistributedOptimizer([wrapped_optimizer], lw_config, pg_collection) diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_muon_optimizer.py index 7a15674d7b2..f853ddb2a5c 100644 --- a/tests/unit_tests/test_muon_optimizer.py +++ b/tests/unit_tests/test_muon_optimizer.py @@ -10,10 +10,13 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.optimizer import HAVE_EMERGING_OPTIMIZERS, OptimizerConfig -from megatron.core.optimizer.muon import ( +from megatron.core.optimizer import ( + HAVE_EMERGING_OPTIMIZERS, + OptimizerConfig, + get_megatron_optimizer, +) +from megatron.core.optimizer.emerging_optimizers import ( TensorParallelMuon, - get_megatron_muon_optimizer, get_supported_coefficient_types, validate_coefficient_type, ) @@ -139,8 +142,8 @@ def create_ddp_model(self, model): TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model ) - def test_get_megatron_muon_optimizer_smoke(self): - """Smoke test for get_megatron_muon_optimizer function.""" + def test_get_megatron_optimizer_smoke(self): + """Smoke test for get_megatron_optimizer function.""" model = Net().bfloat16().cuda() model.requires_grad_(True) model = self.create_ddp_model(model) @@ -165,11 +168,8 @@ def test_get_megatron_muon_optimizer_smoke(self): ) # Test creating the optimizer - optimizer = get_megatron_muon_optimizer( - config=optimizer_config, - model_chunks=[model], - use_gloo_process_groups=True, - layer_wise_distributed_optimizer=False, + optimizer = get_megatron_optimizer( + config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True ) # Test basic properties @@ -214,24 +214,13 @@ def test_get_megatron_muon_optimizer_smoke(self): # Load state dict should not raise error optimizer.load_state_dict(state_dict) - def test_get_megatron_muon_optimizer_validation(self): - """Test validation logic for get_megatron_muon_optimizer.""" + def test_get_megatron_optimizer_validation(self): + """Test validation logic for get_megatron_optimizer.""" model = torch.nn.Linear(100, 50, bias=False, dtype=torch.bfloat16, device='cuda') model.requires_grad_(True) model = self.create_ddp_model(model) - # Test 1: Distributed optimizer should raise exception - optimizer_config_dist = OptimizerConfig( - optimizer='muon', - lr=0.01, - bf16=True, - use_distributed_optimizer=True, # This should cause an exception - ) - - with pytest.raises(Exception, match='muon with dist optimizer is not supported'): - get_megatron_muon_optimizer(config=optimizer_config_dist, model_chunks=[model]) - - # Test 2: FP16 should raise exception + # Test 1: FP16 should raise exception optimizer_config_fp16 = OptimizerConfig( optimizer='muon', lr=0.01, @@ -239,8 +228,8 @@ def test_get_megatron_muon_optimizer_validation(self): use_distributed_optimizer=False, ) - with pytest.raises(Exception, match='muon with fp16 is not supported'): - get_megatron_muon_optimizer(config=optimizer_config_fp16, model_chunks=[model]) + with pytest.raises(Exception, match='emerging optimizer with fp16 is not supported'): + get_megatron_optimizer(config=optimizer_config_fp16, model_chunks=[model]) # Test 3: Invalid num_ns_steps should raise exception optimizer_config_invalid_ns = OptimizerConfig( @@ -252,10 +241,10 @@ def test_get_megatron_muon_optimizer_validation(self): ) with pytest.raises(ValueError, match='num_ns_steps must be at least 1'): - get_megatron_muon_optimizer(config=optimizer_config_invalid_ns, model_chunks=[model]) + get_megatron_optimizer(config=optimizer_config_invalid_ns, model_chunks=[model]) - def test_get_megatron_muon_optimizer_layer_wise(self): - """Test get_megatron_muon_optimizer with layer-wise distributed optimizer.""" + def test_get_megatron_optimizer_layer_wise(self): + """Test get_megatron_optimizer with layer-wise distributed optimizer.""" model = Net().bfloat16().cuda() model.requires_grad_(True) model = self.create_ddp_model(model) @@ -265,7 +254,7 @@ def test_get_megatron_muon_optimizer_layer_wise(self): lr=0.01, weight_decay=0.01, bf16=True, - use_distributed_optimizer=False, + use_layer_wise_distributed_optimizer=True, muon_momentum=0.95, muon_use_nesterov=True, muon_fp32_matmul_prec="medium", @@ -274,12 +263,9 @@ def test_get_megatron_muon_optimizer_layer_wise(self): muon_tp_mode="duplicated", ) - # Test with layer_wise_distributed_optimizer=True - optimizer = get_megatron_muon_optimizer( - config=optimizer_config, - model_chunks=[model], - use_gloo_process_groups=True, - layer_wise_distributed_optimizer=True, + # use_layer_wise_distributed_optimizer=True triggers LayerWiseDistributedOptimizer + optimizer = get_megatron_optimizer( + config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True ) # Verify it's a LayerWiseDistributedOptimizer @@ -709,7 +695,7 @@ def test_muon_optimizer_invalid_coefficient_type(): 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.""" + """Test coefficient_type integration through get_megatron_optimizer.""" @pytest.fixture(autouse=True) def setup_and_teardown(self): @@ -724,8 +710,8 @@ def create_ddp_model(self, 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.""" + def test_get_megatron_optimizer_coefficient_type(self, coefficient_type): + """Test that coefficient_type flows through get_megatron_optimizer.""" model = Net().bfloat16().cuda() model.requires_grad_(True) model = self.create_ddp_model(model) @@ -741,11 +727,8 @@ def test_get_megatron_muon_optimizer_coefficient_type(self, coefficient_type): 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, + optimizer = get_megatron_optimizer( + config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True ) assert optimizer is not None diff --git a/tests/unit_tests/test_optimizer.py b/tests/unit_tests/test_optimizer.py index 2488900ba72..56af8545042 100644 --- a/tests/unit_tests/test_optimizer.py +++ b/tests/unit_tests/test_optimizer.py @@ -106,10 +106,10 @@ def test_get_param_groups_no_overrides(mock_get_world_size): def test_get_param_groups_default_overrides(mock_get_world_size): """Test that the default overrides are applied to the parameter groups.""" net = Net() - # NOTE: to get legacy default overrides, supply None. opt_config = OptimizerConfig(optimizer='adam', lr=0.01) - check_config_overrides_consistency(opt_config, None) - param_groups = _get_param_groups([net], opt_config, None) + config_overrides = get_standard_config_overrides(opt_config) + check_config_overrides_consistency(opt_config, config_overrides) + param_groups = _get_param_groups([net], opt_config, config_overrides) assert len(param_groups) == 2 pg0, pg1 = param_groups wd_mults = {pg0['wd_mult'], pg1['wd_mult']} From 9fefb3253eb8d5ba1bdee53b77c6e32766105b09 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 12:47:06 -0700 Subject: [PATCH 02/11] Continue emerging optimizer refactoring (#3737) --- .../core/optimizer/emerging_optimizers.py | 350 +++++++++-------- megatron/core/optimizer/optimizer_config.py | 13 +- megatron/training/arguments.py | 8 +- tests/unit_tests/dist_checkpointing/utils.py | 10 +- ...timizer.py => test_emerging_optimizers.py} | 357 ++++++++++++++++-- 5 files changed, 547 insertions(+), 191 deletions(-) rename tests/unit_tests/{test_muon_optimizer.py => test_emerging_optimizers.py} (69%) diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 391a247a912..107fe67fb01 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -8,6 +8,7 @@ 3. Add an ``EmergingOptimizerEntry`` to ``_EMERGING_OPTIMIZERS`` at the bottom. """ +import inspect import logging from dataclasses import dataclass, field from typing import Any, Callable, Dict, List, Literal, Optional, get_args @@ -21,15 +22,20 @@ from .optimizer_config import ParamKey, ParamPredicate try: + from emerging_optimizers import registry from emerging_optimizers.orthogonalized_optimizers import ( OrthogonalizedOptimizer, get_muon_scale_factor, ) from emerging_optimizers.orthogonalized_optimizers.muon_utils import NSCoeffT, newton_schulz_tp + # It is necessary to import SOAP for the registry to work. + from emerging_optimizers.soap import SOAP # pylint: disable=unused-import + HAVE_EMERGING_OPTIMIZERS = True except ImportError: HAVE_EMERGING_OPTIMIZERS = False + OrthogonalizedOptimizer = object logger = logging.getLogger(__name__) @@ -76,14 +82,19 @@ class EmergingOptimizerEntry: optimizer_cls: type init_state_fn: Callable - config_to_kwargs: Callable + config_to_kwargs: Callable | None default_param_overrides: Dict[ParamKey, Dict[str, Any]] = field(default_factory=dict) def _create_emerging_optimizer(config, param_groups, eopt_name, model_chunks, pg_collection): """Instantiate an emerging optimizer and return it with its init_state_fn.""" entry = _EMERGING_OPTIMIZERS[eopt_name] - eopt_kwargs = entry.config_to_kwargs(config, model_chunks, pg_collection) + if entry.config_to_kwargs is not None: + eopt_kwargs = entry.config_to_kwargs(config, model_chunks, pg_collection) + else: + eopt_kwargs = _default_adam_based_eopt_config_to_kwargs( + eopt_name, config, model_chunks, pg_collection + ) optimizer = entry.optimizer_cls(param_groups, **eopt_kwargs) return optimizer, entry.init_state_fn @@ -118,159 +129,180 @@ def _get_qkv_split_shapes(model_cfg) -> List[int]: # Muon # =========================================================================== -if HAVE_EMERGING_OPTIMIZERS: - class TensorParallelMuon(OrthogonalizedOptimizer): - """Tensor Parallel Muon optimizer.""" - - def __init__( - self, - params: ParamsT, - lr: float = 3e-4, - momentum_beta: float = 0.95, - use_nesterov: bool = True, - weight_decay: float = 0.01, - use_decoupled_weight_decay: bool = True, - split_qkv: bool = False, - is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, - qkv_split_shapes: tuple[int, int, int] | None = None, - fp32_matmul_prec: str = "medium", - coefficient_type: str = "quintic", - num_ns_steps: int = 5, - scale_mode: str = "spectral", - extra_scale_factor: float = 1.0, - pg_collection: Optional[ProcessGroupCollection] = None, - mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", - ) -> None: - if num_ns_steps < 1: - raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") - - def scaled_orthogonalize_fn( - grad: torch.Tensor, - tp_group: torch.distributed.ProcessGroup, - partition_dim: int | None = None, - ) -> torch.Tensor: - log_single_rank( - logger, - logging.DEBUG, - f'Orthogonalizing grad with {num_ns_steps} steps, ' - f'{coefficient_type} coefficient, ' - f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}', - ) - 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, - ) - scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode) - return orth_grad * scale_factor * extra_scale_factor - - self.pg_collection = pg_collection - self.mode = mode - self.split_qkv = split_qkv - self.is_qkv_fn = is_qkv_fn - self.qkv_split_shapes = qkv_split_shapes - - weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2" - super().__init__( - params, - lr, - momentum_beta, - use_nesterov=use_nesterov, - weight_decay=weight_decay, - weight_decay_method=weight_decay_method, - fp32_matmul_prec=fp32_matmul_prec, - scaled_orthogonalize_fn=scaled_orthogonalize_fn, +class TensorParallelMuon(OrthogonalizedOptimizer): + """Tensor Parallel Muon optimizer.""" + + def __init__( + self, + params: ParamsT, + lr: float = 3e-4, + momentum: float = 0.95, + nesterov: bool = True, + weight_decay: float = 0.01, + use_decoupled_weight_decay: bool = True, + split_qkv: bool = False, + is_qkv_fn: Callable[[torch.Tensor], bool] | None = None, + qkv_split_shapes: tuple[int, int, int] | None = None, + fp32_matmul_prec: str = "medium", + coefficient_type: str = "quintic", + num_ns_steps: int = 5, + scale_mode: str = "spectral", + extra_scale_factor: float = 1.0, + pg_collection: Optional[ProcessGroupCollection] = None, + tp_mode: Literal["blockwise", "duplicated", "distributed"] = "duplicated", + ) -> None: + if num_ns_steps < 1: + raise ValueError(f"num_ns_steps must be at least 1, got {num_ns_steps}") + + def scaled_orthogonalize_fn( + grad: torch.Tensor, + tp_group: torch.distributed.ProcessGroup, + partition_dim: int | None = None, + ) -> torch.Tensor: + log_single_rank( + logger, + logging.DEBUG, + f'Orthogonalizing grad with {num_ns_steps} steps, ' + f'{coefficient_type} coefficient, ' + f'{scale_mode} scale mode, extra_scale_factor={extra_scale_factor}', + ) + 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, + tp_mode="duplicated" if tp_mode == "blockwise" else tp_mode, ) + scale_factor = get_muon_scale_factor(size[0], size[1], mode=scale_mode) + return orth_grad * scale_factor * extra_scale_factor + + self.pg_collection = pg_collection + self.tp_mode = tp_mode + self.split_qkv = split_qkv + self.is_qkv_fn = is_qkv_fn + self.qkv_split_shapes = qkv_split_shapes + + weight_decay_method = "decoupled" if use_decoupled_weight_decay else "l2" + super().__init__( + params, + lr, + momentum, + nesterov=nesterov, + weight_decay=weight_decay, + weight_decay_method=weight_decay_method, + fp32_matmul_prec=fp32_matmul_prec, + scaled_orthogonalize_fn=scaled_orthogonalize_fn, + ) - def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: - """Orthogonalize the momentum. - - Args: - p: The parameter tensor. i is necessary to pass param tensor in addition to - momentum because a lot of information is only available in the param tensor, - attributes for example. - grad: The momentum tensor. - - Returns: - The orthogonalized gradient tensor. - """ - # TODO(deyuf): switch to group - if self.pg_collection: - tp_group = ( - self.pg_collection.expt_tp - if getattr(p, 'expert_tp', False) - else self.pg_collection.tp - ) - else: - tp_group = None - partition_dim = None if self.mode == "blockwise" else getattr(p, "partition_dim", None) - if partition_dim == -1: - partition_dim = None - - if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] - grad_shape = grad.shape - log_single_rank( - logger, - logging.DEBUG, - f'qkv split grad shape {grad_shape}, ' f'split shapes {self.qkv_split_shapes}', - ) - num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes) - qkv_grads = torch.split( - grad.view(num_query_groups, sum(self.qkv_split_shapes), -1), - self.qkv_split_shapes, - dim=1, + def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> torch.Tensor: + """Orthogonalize the momentum. + + Args: + p: The parameter tensor. i is necessary to pass param tensor in addition to + momentum because a lot of information is only available in the param tensor, + attributes for example. + grad: The momentum tensor. + + Returns: + The orthogonalized gradient tensor. + """ + # TODO(deyuf): switch to group + if self.pg_collection: + tp_group = ( + self.pg_collection.expt_tp + if getattr(p, 'expert_tp', False) + else self.pg_collection.tp + ) + else: + tp_group = None + partition_dim = None if self.tp_mode == "blockwise" else getattr(p, "partition_dim", None) + if partition_dim == -1: + partition_dim = None + + if self.split_qkv and self.is_qkv_fn(p): # type: ignore[misc] + grad_shape = grad.shape + log_single_rank( + logger, + logging.DEBUG, + f'qkv split grad shape {grad_shape}, ' f'split shapes {self.qkv_split_shapes}', + ) + num_query_groups = grad_shape[0] // sum(self.qkv_split_shapes) + qkv_grads = torch.split( + grad.view(num_query_groups, sum(self.qkv_split_shapes), -1), + self.qkv_split_shapes, + dim=1, + ) + qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads] + + qkv_grads = [ + self.scaled_orthogonalize_fn(g, tp_group, partition_dim).view( + num_query_groups, -1, grad_shape[-1] ) - qkv_grads = [g.reshape(-1, grad_shape[-1]) for g in qkv_grads] + for g in qkv_grads + ] + grad = torch.cat(qkv_grads, dim=1).view(grad_shape) + else: + grad = self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) + return grad - qkv_grads = [ - self.scaled_orthogonalize_fn(g, tp_group, partition_dim).view( - num_query_groups, -1, grad_shape[-1] - ) - for g in qkv_grads - ] - grad = torch.cat(qkv_grads, dim=1).view(grad_shape) - else: - grad = self.scaled_orthogonalize_fn(grad, tp_group, partition_dim) - return grad - - def _muon_init_state_fn(opt, config=None): - """Initialize Muon optimizer state for torch_dist checkpoint format.""" - for group in opt.param_groups: - for p in group['params']: - if len(opt.state[p]) == 0: - opt.state[p]['momentum_buffer'] = torch.zeros_like(p.data) - - def _muon_config_to_kwargs(config, model_chunks, pg_collection) -> Dict[str, Any]: - """Convert OptimizerConfig to TensorParallelMuon constructor kwargs.""" - return { - "lr": config.lr, - "weight_decay": config.weight_decay, - "momentum_beta": config.muon_momentum, - "use_nesterov": config.muon_use_nesterov, - "fp32_matmul_prec": config.muon_fp32_matmul_prec, - "num_ns_steps": config.muon_num_ns_steps, - "scale_mode": config.muon_scale_mode, - "extra_scale_factor": config.muon_extra_scale_factor, - "mode": config.muon_tp_mode, - "split_qkv": config.muon_split_qkv, - "is_qkv_fn": lambda p: getattr(p, "is_qkv", False), - "qkv_split_shapes": _get_qkv_split_shapes(model_chunks[0].config), - "pg_collection": pg_collection, - } - - # ----------------------------------------------------------------------- - # Register Muon - # ----------------------------------------------------------------------- - _EMERGING_OPTIMIZERS['muon'] = EmergingOptimizerEntry( + +def _eopt_init_state_fn(opt, config=None): + """Initialize emerging optimizer state for torch_dist checkpoint format.""" + for group in opt.param_groups: + opt._init_group(group) + + +def _kwargs_from_config(optimizer_cls: type, prefix: str, config) -> Dict[str, Any]: + """Match ``optimizer_cls.__init__`` parameters to config attributes. + + For each init parameter, looks for ``{prefix}_{name}`` on *config* first, + then falls back to ``{name}`` (unprefixed). ``self`` and ``params`` are + always skipped. + """ + skip_params = {"self", "params"} + sig = inspect.signature(optimizer_cls.__init__) + kwargs: Dict[str, Any] = {} + for name in sig.parameters: + if name in skip_params: + continue + prefixed = f"{prefix}_{name}" + if hasattr(config, prefixed): + kwargs[name] = getattr(config, prefixed) + elif hasattr(config, name): + kwargs[name] = getattr(config, name) + return kwargs + + +def _muon_config_to_kwargs(config, model_chunks, pg_collection) -> Dict[str, Any]: + """Convert OptimizerConfig to TensorParallelMuon constructor kwargs.""" + kwargs = _kwargs_from_config(TensorParallelMuon, "muon", config) + kwargs["is_qkv_fn"] = lambda p: getattr(p, "is_qkv", False) + kwargs["qkv_split_shapes"] = _get_qkv_split_shapes(model_chunks[0].config) + kwargs["pg_collection"] = pg_collection + return kwargs + + +def _default_adam_based_eopt_config_to_kwargs( + eopt_name, config, model_chunks, pg_collection +) -> Dict[str, Any]: + """Convert OptimizerConfig to default emerging optimizer constructor kwargs.""" + kwargs = _kwargs_from_config(registry.get_optimizer_cls(eopt_name), eopt_name, config) + kwargs["betas"] = (config.adam_beta1, config.adam_beta2) + return kwargs + + +# ----------------------------------------------------------------------- +# Register emerging optimizers +# ----------------------------------------------------------------------- +_EMERGING_OPTIMIZERS = { + 'muon': EmergingOptimizerEntry( optimizer_cls=TensorParallelMuon, - init_state_fn=_muon_init_state_fn, + init_state_fn=_eopt_init_state_fn, config_to_kwargs=_muon_config_to_kwargs, default_param_overrides={ ParamKey( @@ -280,3 +312,23 @@ def _muon_config_to_kwargs(config, model_chunks, pg_collection) -> Dict[str, Any ): {'optimizer': 'adam'} }, ) +} + +# Register soap with default config +# TODO(skyw): register all emerging optimizers. +if HAVE_EMERGING_OPTIMIZERS: + for eopt_name in ["soap"]: + if eopt_name in _EMERGING_OPTIMIZERS: + continue + _EMERGING_OPTIMIZERS[eopt_name] = EmergingOptimizerEntry( + optimizer_cls=registry.get_optimizer_cls(eopt_name), + init_state_fn=_eopt_init_state_fn, + config_to_kwargs=None, + default_param_overrides={ + ParamKey( + predicate=ParamPredicate( + name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding + ) + ): {'optimizer': 'adam'} + }, + ) diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 7534227f5a0..61dd8a38d12 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -256,14 +256,14 @@ class OptimizerConfig: sgd_momentum: float = 0.9 """Momentum factor for SGD optimizer.""" - # Muon / emerging optimizers. + # emerging optimizers. muon_momentum: float = 0.95 """The momentum used by the internal SGD in Muon optimizer.""" muon_split_qkv: bool = True """Whether to split QKV parameters for Muon optimizer.""" - muon_use_nesterov: bool = False + muon_nesterov: bool = False """Whether to use Nesterov-style momentum in the internal SGD.""" muon_scale_mode: str = "spectral" @@ -297,6 +297,15 @@ class OptimizerConfig: """Second beta coefficient for Lion optimizer (used in momentum EMA update). Defaults to 0.98.""" + soap_shampoo_beta: float = 0.95 + """The beta parameter for the Shampoo preconditioner.""" + + soap_precondition_frequency: int = 1 + """The frequency of the Shampoo preconditioner.""" + + soap_use_kl_shampoo: bool = True + """Whether to use the KL-Shampoo preconditioner.""" + ####################### # Distributed optimizer ####################### diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 25e8eb02e80..b6577541a6f 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1472,8 +1472,8 @@ def validate_args(args, defaults={}): '--no-load-optim with --skip-train --perform-rl-step skips the optimizer; ' \ '--rl-offload-optimizer-during-inference is incompatible (no optimizer to offload).' - # Muon / emerging optimizer check - if args.optimizer in ('muon', 'dist_muon'): + # emerging optimizer check + if args.optimizer not in ('sgd', 'adam'): if args.optimizer == 'dist_muon': warn_rank_0( "optimizer='dist_muon' is deprecated. " @@ -2237,7 +2237,7 @@ def _add_regularization_args(parser): group.add_argument('--muon-no-split-qkv', action='store_false', default=True, dest='muon_split_qkv', help='Whether to split QKV parameters for Muon optimizer') - group.add_argument('--muon-use-nesterov', action='store_true', + group.add_argument('--muon-nesterov', action='store_true', help='Whether to use Nesterov-style momentum in the internal SGD') group.add_argument('--muon-scale-mode', type=str, default='spectral', choices=['spectral', 'unit_rms_norm', 'shape_scaling'], @@ -2485,7 +2485,7 @@ def _add_training_args(parser): help='use FlashAttention implementation of attention. ' 'https://arxiv.org/abs/2205.14135') group.add_argument('--optimizer', type=str, default='adam', - choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion'], + choices=['adam', 'sgd', 'muon', 'dist_muon', 'lion', 'soap'], help='Optimizer function. ' 'Note: dist_muon is deprecated; use --optimizer muon ' 'with --use-distributed-optimizer instead.') diff --git a/tests/unit_tests/dist_checkpointing/utils.py b/tests/unit_tests/dist_checkpointing/utils.py index 8fe5e28e12c..0aadaee3b29 100644 --- a/tests/unit_tests/dist_checkpointing/utils.py +++ b/tests/unit_tests/dist_checkpointing/utils.py @@ -218,7 +218,10 @@ def setup_model_and_optimizer( if isinstance(optimizer, ChainedOptimizer): for opt in optimizer.chained_optimizers: - opt.init_state_fn(opt) + if not hasattr(opt, 'optimizer'): + opt.init_state_fn(opt) + else: + opt.init_state_fn(opt.optimizer) else: for group in optimizer.optimizer.param_groups: for p in group['params']: @@ -314,7 +317,10 @@ def setup_moe_model_and_optimizer( if optimizer_type in ('muon', 'dist_muon'): for opt in optimizer.chained_optimizers: - opt.init_state_fn(opt) + if not hasattr(opt, 'optimizer'): + opt.init_state_fn(opt) + else: + opt.init_state_fn(opt.optimizer) else: for opt in optimizer.chained_optimizers: for group in opt.param_groups: diff --git a/tests/unit_tests/test_muon_optimizer.py b/tests/unit_tests/test_emerging_optimizers.py similarity index 69% rename from tests/unit_tests/test_muon_optimizer.py rename to tests/unit_tests/test_emerging_optimizers.py index f853ddb2a5c..3834c5b24a3 100644 --- a/tests/unit_tests/test_muon_optimizer.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -10,25 +10,22 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig -from megatron.core.optimizer import ( - HAVE_EMERGING_OPTIMIZERS, - OptimizerConfig, - get_megatron_optimizer, -) -from megatron.core.optimizer.emerging_optimizers import ( - TensorParallelMuon, - get_supported_coefficient_types, - validate_coefficient_type, -) +from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer +from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer import TransformerConfig from tests.unit_tests.test_utilities import Utils +if HAVE_EMERGING_OPTIMIZERS: + from emerging_optimizers.soap import SOAP +else: + SOAP = None + # 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", + reason="Skip emerging optimizer tests for LTS test", ), pytest.mark.skipif( not HAVE_EMERGING_OPTIMIZERS, reason="emerging_optimizers package is not installed" @@ -54,6 +51,11 @@ def forward(self, x): return x +# =========================================================================== +# Muon optimizer tests +# =========================================================================== + + def test_muon_optimizer_smoke(): """Smoke test for TensorParallelMuon optimizer.""" # Create a simple linear model for testing @@ -65,8 +67,8 @@ def test_muon_optimizer_smoke(): optimizer = TensorParallelMuon( params=[model.weight], lr=0.01, - momentum_beta=0.95, - use_nesterov=True, + momentum=0.95, + nesterov=True, weight_decay=0.01, use_decoupled_weight_decay=True, split_qkv=False, @@ -75,7 +77,7 @@ def test_muon_optimizer_smoke(): scale_mode="spectral", extra_scale_factor=1.0, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) # Test basic properties @@ -160,7 +162,7 @@ def test_get_megatron_optimizer_smoke(self): bf16=True, use_distributed_optimizer=False, # Muon doesn't support distributed optimizer muon_momentum=0.95, - muon_use_nesterov=True, + muon_nesterov=True, muon_fp32_matmul_prec="medium", muon_num_ns_steps=5, muon_scale_mode="spectral", @@ -256,7 +258,7 @@ def test_get_megatron_optimizer_layer_wise(self): bf16=True, use_layer_wise_distributed_optimizer=True, muon_momentum=0.95, - muon_use_nesterov=True, + muon_nesterov=True, muon_fp32_matmul_prec="medium", muon_num_ns_steps=5, muon_scale_mode="spectral", @@ -305,11 +307,11 @@ def test_muon_optimizer_different_modes_single_rank(mode): optimizer = TensorParallelMuon( params=[model.weight], lr=0.01, - momentum_beta=0.95, + momentum=0.95, weight_decay=0.0, # Disable weight decay for deterministic comparison num_ns_steps=5, pg_collection=None, - mode=mode, + tp_mode=mode, ) # Use fixed input for deterministic results @@ -365,11 +367,11 @@ def create_tp_model_and_optimizer(self, mode): optimizer = TensorParallelMuon( params=[model.weight], lr=0.01, - momentum_beta=0.95, + momentum=0.95, weight_decay=0.0, num_ns_steps=5, pg_collection=pg_collection, - mode=mode, + tp_mode=mode, ) return model, optimizer @@ -441,7 +443,7 @@ def test_muon_optimizer_coefficient_types(coefficient_type): coefficient_type=coefficient_type, num_ns_steps=_DEFAULT_NS_STEPS, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) input_tensor = torch.randn(16, 80, dtype=torch.float32, device='cuda') @@ -470,7 +472,7 @@ def test_muon_optimizer_scale_modes(scale_mode): scale_mode=scale_mode, num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) input_tensor = torch.randn(16, 60, dtype=torch.float32, device='cuda') @@ -486,8 +488,8 @@ def test_muon_optimizer_scale_modes(scale_mode): ), f"Weight should be updated with scale_mode={scale_mode}" -@pytest.mark.parametrize("use_nesterov", [True, False]) -def test_muon_optimizer_nesterov(use_nesterov): +@pytest.mark.parametrize("nesterov", [True, False]) +def test_muon_optimizer_nesterov(nesterov): """Test TensorParallelMuon optimizer with and without Nesterov momentum.""" model = torch.nn.Linear(50, 25, bias=False, dtype=torch.float32, device='cuda') model.requires_grad_(True) @@ -496,11 +498,11 @@ def test_muon_optimizer_nesterov(use_nesterov): optimizer = TensorParallelMuon( params=[model.weight], lr=0.01, - momentum_beta=0.9, - use_nesterov=use_nesterov, + momentum=0.9, + nesterov=nesterov, num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) input_tensor = torch.randn(16, 50, dtype=torch.float32, device='cuda') @@ -513,7 +515,7 @@ def test_muon_optimizer_nesterov(use_nesterov): assert not torch.equal( model.weight.data, original_weight - ), f"Weight should be updated with use_nesterov={use_nesterov}" + ), f"Weight should be updated with nesterov={nesterov}" def test_muon_optimizer_multiple_steps(): @@ -525,11 +527,11 @@ def test_muon_optimizer_multiple_steps(): optimizer = TensorParallelMuon( params=[model.weight], lr=0.01, - momentum_beta=0.95, + momentum=0.95, weight_decay=0.01, num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) weights_history = [model.weight.data.clone()] @@ -575,7 +577,7 @@ def test_muon_optimizer_qkv_split(): qkv_split_shapes=qkv_split_shapes, num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) input_tensor = torch.randn(16, hidden_size, dtype=torch.float32, device='cuda') @@ -599,7 +601,7 @@ def test_muon_optimizer_qkv_split(): split_qkv=False, num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) output = model(input_tensor) @@ -631,7 +633,7 @@ def test_muon_optimizer_extra_scale_factor(): extra_scale_factor=2.0, num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) input_tensor = torch.randn(16, 80, dtype=torch.float32, device='cuda') @@ -754,7 +756,7 @@ def test_muon_optimizer_num_ns_steps(num_ns_steps): coefficient_type="quintic", num_ns_steps=num_ns_steps, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) input_tensor = torch.randn(16, 60, dtype=torch.float32, device='cuda') @@ -768,3 +770,290 @@ def test_muon_optimizer_num_ns_steps(num_ns_steps): assert not torch.equal( model.weight.data, original_weight ), f"Weight should be updated with num_ns_steps={num_ns_steps}" + + +# =========================================================================== +# SOAP optimizer tests +# =========================================================================== + +skip_no_soap = pytest.mark.skipif( + not HAVE_EMERGING_OPTIMIZERS, reason="emerging_optimizers package not installed" +) + + +@skip_no_soap +def test_soap_optimizer_smoke(): + """Smoke test for SOAP optimizer.""" + + model = torch.nn.Linear(100, 50, bias=False, dtype=torch.float32, device='cuda') + model.requires_grad_(True) + model.weight.data.fill_(1.0) + + optimizer = SOAP( + params=[model.weight], + lr=0.01, + betas=(0.9, 0.999), + shampoo_beta=0.95, + weight_decay=0.01, + precondition_frequency=1, + ) + + # Test basic properties + assert optimizer is not None, "Optimizer should not be None" + assert hasattr(optimizer, 'param_groups'), "Optimizer should have param_groups" + assert len(optimizer.param_groups) > 0, "Optimizer should have at least one parameter group" + + # Test forward and backward pass + input_tensor = torch.randn(32, 100, dtype=torch.float32, device='cuda') + output = model(input_tensor) + loss = output.sum() + loss.backward() + + # Store original weight + original_weight = model.weight.data.clone() + + # Test optimizer step + optimizer.step() + + # Verify weight was updated + assert not torch.equal( + model.weight.data, original_weight + ), "Weight should be updated after optimizer step" + + # Test zero_grad + optimizer.zero_grad() + assert model.weight.grad is None or torch.all( + model.weight.grad == 0 + ), "Gradients should be zeroed" + + # Test state_dict and load_state_dict + state_dict = optimizer.state_dict() + assert 'state' in state_dict, "State dict should contain state" + assert 'param_groups' in state_dict, "State dict should contain param_groups" + + # Load state dict should not raise error + optimizer.load_state_dict(state_dict) + + +@skip_no_soap +def test_soap_optimizer_multiple_steps(): + """Test SOAP optimizer across multiple optimization steps.""" + model = torch.nn.Linear(100, 50, bias=False, dtype=torch.float32, device='cuda') + model.requires_grad_(True) + model.weight.data.fill_(1.0) + + optimizer = SOAP( + params=[model.weight], + lr=0.01, + betas=(0.9, 0.999), + shampoo_beta=0.95, + weight_decay=0.01, + precondition_frequency=1, + ) + + weights_history = [model.weight.data.clone()] + + for i in range(3): + input_tensor = torch.randn(32, 100, dtype=torch.float32, device='cuda') + output = model(input_tensor) + loss = output.sum() + loss.backward() + + optimizer.step() + optimizer.zero_grad() + weights_history.append(model.weight.data.clone()) + + # Verify weights changed at each step + for i in range(len(weights_history) - 1): + assert not torch.equal( + weights_history[i], weights_history[i + 1] + ), f"Weight should change at step {i}" + + +@skip_no_soap +@pytest.mark.parametrize("precondition_frequency", [1, 5, 10]) +def test_soap_optimizer_precondition_frequency(precondition_frequency): + """Test SOAP optimizer with different precondition frequencies.""" + + model = torch.nn.Linear(60, 30, bias=False, dtype=torch.float32, device='cuda') + model.requires_grad_(True) + model.weight.data.fill_(1.0) + + optimizer = SOAP( + params=[model.weight], + lr=0.01, + betas=(0.9, 0.999), + shampoo_beta=0.95, + precondition_frequency=precondition_frequency, + ) + + input_tensor = torch.randn(16, 60, 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 precondition_frequency={precondition_frequency}" + + +@skip_no_soap +@pytest.mark.parametrize("use_kl_shampoo", [True, False]) +def test_soap_optimizer_kl_shampoo(use_kl_shampoo): + """Test SOAP optimizer with and without KL-Shampoo preconditioner.""" + + model = torch.nn.Linear(60, 30, bias=False, dtype=torch.float32, device='cuda') + model.requires_grad_(True) + model.weight.data.fill_(1.0) + + optimizer = SOAP( + params=[model.weight], + lr=0.01, + betas=(0.9, 0.999), + shampoo_beta=0.95, + use_kl_shampoo=use_kl_shampoo, + precondition_frequency=1, + ) + + input_tensor = torch.randn(16, 60, 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 use_kl_shampoo={use_kl_shampoo}" + + +@skip_no_soap +@pytest.mark.parametrize("shampoo_beta", [0.5, 0.9, 0.99]) +def test_soap_optimizer_shampoo_beta(shampoo_beta): + """Test SOAP optimizer with different shampoo_beta values.""" + + model = torch.nn.Linear(60, 30, bias=False, dtype=torch.float32, device='cuda') + model.requires_grad_(True) + model.weight.data.fill_(1.0) + + optimizer = SOAP( + params=[model.weight], + lr=0.01, + betas=(0.9, 0.999), + shampoo_beta=shampoo_beta, + precondition_frequency=1, + ) + + input_tensor = torch.randn(16, 60, 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 shampoo_beta={shampoo_beta}" + + +@pytest.mark.skipif( + int(os.getenv('WORLD_SIZE', '1')) == 1, reason="Multi-rank test requires WORLD_SIZE > 1" +) +class TestSoapOptimizerMultiRank: + """Test class for SOAP optimizer with multi-rank setup.""" + + @pytest.fixture(autouse=True) + def setup_and_teardown(self): + """Setup and teardown for each test.""" + Utils.initialize_model_parallel() + yield + Utils.destroy_model_parallel() + + def create_ddp_model(self, model): + """Wrap model in DDP.""" + ddp_config = DistributedDataParallelConfig(use_distributed_optimizer=False) + return DistributedDataParallel( + TransformerConfig(num_attention_heads=1, num_layers=1), ddp_config, model + ) + + def test_get_megatron_optimizer_soap_smoke(self): + """Smoke test for get_megatron_optimizer with SOAP.""" + model = Net().bfloat16().cuda() + model.requires_grad_(True) + model = self.create_ddp_model(model) + + for param in model.parameters(): + assert param.requires_grad, "All parameters should require gradients" + + optimizer_config = OptimizerConfig( + optimizer='soap', + lr=0.01, + weight_decay=0.01, + bf16=True, + use_distributed_optimizer=False, + soap_shampoo_beta=0.95, + soap_precondition_frequency=1, + soap_use_kl_shampoo=True, + ) + + optimizer = get_megatron_optimizer( + config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True + ) + + assert optimizer is not None, "Optimizer should not be None" + assert hasattr(optimizer, 'param_groups'), "Optimizer should have param_groups" + assert hasattr(optimizer, 'chained_optimizers'), "Should be a ChainedOptimizer" + assert len(optimizer.chained_optimizers) >= 1, "Should have at least one chained optimizer" + + # Test forward and backward pass + input_tensor = torch.randn(16, 80, dtype=torch.bfloat16, device='cuda') + output = model(input_tensor) + loss = output.sum() + loss.backward() + + # Store original parameters + original_params = {} + for name, param in model.named_parameters(): + original_params[name] = param.data.clone() + + # Test optimizer step + optimizer.step() + + # Verify at least some parameters were updated + params_updated = 0 + for name, param in model.named_parameters(): + if not torch.equal(param.data, original_params[name]): + params_updated += 1 + + assert params_updated > 0, "At least some parameters should be updated after optimizer step" + + # Test zero_grad + optimizer.zero_grad() + for param in model.parameters(): + assert param.grad is None or torch.all( + param.grad == 0 + ), "Gradients should be zeroed for all parameters" + + # Test state_dict and load_state_dict + state_dict = optimizer.state_dict() + assert isinstance(state_dict, list), "State dict should be a list" + optimizer.load_state_dict(state_dict) + + def test_get_megatron_optimizer_soap_validation(self): + """Test validation logic for get_megatron_optimizer with SOAP.""" + model = torch.nn.Linear(100, 50, bias=False, dtype=torch.bfloat16, device='cuda') + model.requires_grad_(True) + model = self.create_ddp_model(model) + + # FP16 should raise exception + optimizer_config_fp16 = OptimizerConfig( + optimizer='soap', lr=0.01, fp16=True, use_distributed_optimizer=False + ) + + with pytest.raises(Exception, match='emerging optimizer with fp16 is not supported'): + get_megatron_optimizer(config=optimizer_config_fp16, model_chunks=[model]) From 710faf3b0bfdb1ee5578aa4c5559fba23d155b2d Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 12:47:26 -0700 Subject: [PATCH 03/11] Fix emerging optimizer init_group for ckpt loading (#3897) --- megatron/core/optimizer/emerging_optimizers.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 107fe67fb01..2b95b598028 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -254,7 +254,8 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t def _eopt_init_state_fn(opt, config=None): """Initialize emerging optimizer state for torch_dist checkpoint format.""" for group in opt.param_groups: - opt._init_group(group) + # Checkpoint init needs state for all parameters, including those without grads yet. + opt._init_group(group, skip_non_grad_params=False) def _kwargs_from_config(optimizer_cls: type, prefix: str, config) -> Dict[str, Any]: From 1f8ab3d4230d380bb26db5968f6ed6156183f080 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 12:50:17 -0700 Subject: [PATCH 04/11] Fix missing imports in test_emerging_optimizers.py --- tests/unit_tests/test_emerging_optimizers.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index 3834c5b24a3..8458e9af1fa 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -11,7 +11,12 @@ from megatron.core import parallel_state from megatron.core.distributed import DistributedDataParallel, DistributedDataParallelConfig from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer -from megatron.core.optimizer.emerging_optimizers import HAVE_EMERGING_OPTIMIZERS, TensorParallelMuon +from megatron.core.optimizer.emerging_optimizers import ( + HAVE_EMERGING_OPTIMIZERS, + TensorParallelMuon, + 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 From c65de580eabf64eb30d9b6bf202ea62ff1b5b099 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 12:51:12 -0700 Subject: [PATCH 05/11] Fix TensorParallelMuon kwarg: mode -> tp_mode in test --- tests/unit_tests/test_emerging_optimizers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index 8458e9af1fa..f2fa732da59 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -694,7 +694,7 @@ def test_muon_optimizer_invalid_coefficient_type(): coefficient_type="nonexistent_type_xyz", num_ns_steps=5, pg_collection=None, - mode="duplicated", + tp_mode="duplicated", ) From 527569468a556624f92aaca35fc3f21e42efd8f9 Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 12:55:58 -0700 Subject: [PATCH 06/11] remove test shouldn't exit yet Signed-off-by: Hao Wu --- tests/unit_tests/test_emerging_optimizers.py | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/tests/unit_tests/test_emerging_optimizers.py b/tests/unit_tests/test_emerging_optimizers.py index f2fa732da59..d8da25c94cc 100644 --- a/tests/unit_tests/test_emerging_optimizers.py +++ b/tests/unit_tests/test_emerging_optimizers.py @@ -682,22 +682,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') - 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, - tp_mode="duplicated", - ) - - @pytest.mark.skipif( int(os.getenv('WORLD_SIZE', '1')) == 1, reason="Multi-rank test requires WORLD_SIZE > 1" ) From a4c8749dc346c9da78b9ef1f089f6224684f4f4c Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 13:56:06 -0700 Subject: [PATCH 07/11] Fix duplicate _EMERGING_OPTIMIZERS definition for autodoc2 --- .../core/optimizer/emerging_optimizers.py | 30 ++++++++++--------- 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py index 2b95b598028..5500ef2966f 100644 --- a/megatron/core/optimizer/emerging_optimizers.py +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -300,20 +300,22 @@ def _default_adam_based_eopt_config_to_kwargs( # ----------------------------------------------------------------------- # Register emerging optimizers # ----------------------------------------------------------------------- -_EMERGING_OPTIMIZERS = { - 'muon': EmergingOptimizerEntry( - optimizer_cls=TensorParallelMuon, - init_state_fn=_eopt_init_state_fn, - config_to_kwargs=_muon_config_to_kwargs, - default_param_overrides={ - ParamKey( - predicate=ParamPredicate( - name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding - ) - ): {'optimizer': 'adam'} - }, - ) -} +_EMERGING_OPTIMIZERS.update( + { + 'muon': EmergingOptimizerEntry( + optimizer_cls=TensorParallelMuon, + init_state_fn=_eopt_init_state_fn, + config_to_kwargs=_muon_config_to_kwargs, + default_param_overrides={ + ParamKey( + predicate=ParamPredicate( + name="nonlinear_or_embedding", fn=_is_nonlinear_or_embedding + ) + ): {'optimizer': 'adam'} + }, + ) + } +) # Register soap with default config # TODO(skyw): register all emerging optimizers. From c74dd1a585070cdefd1952e74da242967adee22d Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 14:13:26 -0700 Subject: [PATCH 08/11] Fix test_layer_wise_optimizer: use get_megatron_optimizer's built-in LayerWise wrapping The test was manually extracting base optimizers from Float16OptimizerWithFloat16Params wrappers and re-wrapping them in LayerWiseDistributedOptimizer, but the fp32 master param updates never propagated back to the bf16 model params. Use use_layer_wise_distributed_optimizer=True in the config instead, which lets _get_megatron_emerging_optimizer handle wrapping correctly. --- tests/unit_tests/test_layer_wise_optimizer.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/tests/unit_tests/test_layer_wise_optimizer.py b/tests/unit_tests/test_layer_wise_optimizer.py index 8cc13d19b13..3edf35462f3 100644 --- a/tests/unit_tests/test_layer_wise_optimizer.py +++ b/tests/unit_tests/test_layer_wise_optimizer.py @@ -117,20 +117,14 @@ def create_model_and_optimizer( use_distributed_optimizer=False, clip_grad=clip_grad, muon_tp_mode="duplicated", + use_layer_wise_distributed_optimizer=use_layer_wise, ) pg_collection = ProcessGroupCollection.use_mpu_process_groups() pg_collection.dp_cp = parallel_state.get_data_parallel_group(with_context_parallel=True) pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() - optimizer = get_megatron_optimizer(optimizer_config, [model]) - if use_layer_wise: - # Extract base torch optimizers from the FP32Optimizer wrappers. - base_optimizers = [opt.optimizer for opt in optimizer.chained_optimizers] - optimizer_config.bf16 = True - optimizer = LayerWiseDistributedOptimizer( - base_optimizers, optimizer_config, pg_collection - ) + optimizer = get_megatron_optimizer(optimizer_config, [model], pg_collection=pg_collection) return model, optimizer, pg_collection def create_model_and_optimizer_with_overlap_param_gather( From 762ca99df1d91f80189174c94aa501a1c0e02c2e Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 15:17:58 -0700 Subject: [PATCH 09/11] Fix test overlap_param_gather: add use_layer_wise_distributed_optimizer=True --- tests/unit_tests/test_layer_wise_optimizer.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/unit_tests/test_layer_wise_optimizer.py b/tests/unit_tests/test_layer_wise_optimizer.py index 3edf35462f3..74837a30a44 100644 --- a/tests/unit_tests/test_layer_wise_optimizer.py +++ b/tests/unit_tests/test_layer_wise_optimizer.py @@ -185,6 +185,7 @@ def create_model_and_optimizer_with_overlap_param_gather( clip_grad=clip_grad, overlap_param_gather=async_allgather, muon_tp_mode="duplicated", + use_layer_wise_distributed_optimizer=True, ) pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -192,7 +193,10 @@ def create_model_and_optimizer_with_overlap_param_gather( pg_collection.expt_dp = parallel_state.get_expert_data_parallel_group() optimizer = get_megatron_optimizer( - config=optimizer_config, model_chunks=[model], use_gloo_process_groups=True + config=optimizer_config, + model_chunks=[model], + use_gloo_process_groups=True, + pg_collection=pg_collection, ) return model, optimizer, pg_collection From 90c5747f3020416ad14f223249fe1fb61628431c Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 15:20:32 -0700 Subject: [PATCH 10/11] remove out of date assert Signed-off-by: Hao Wu --- megatron/training/arguments.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index b6577541a6f..602dd6bcd0b 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -1486,9 +1486,6 @@ def validate_args(args, defaults={}): args.use_layer_wise_distributed_optimizer = True args.use_distributed_optimizer = False - # TODO: remove these checks once we support them - assert not args.overlap_grad_reduce, "Muon optimizer does not support overlap grad reduce for now." - assert not args.overlap_param_gather, "Muon optimizer does not support overlap param gather for now." assert not args.use_torch_fsdp2, "Muon optimizer does not support Torch-FSDP2 for now." assert not args.use_megatron_fsdp, "Muon optimizer does not support Megatron-FSDP for now." assert args.ckpt_format in ["torch", "torch_dist"], "Muon optimizer supports torch and torch_dist checkpoint format." From 16f2e532c0a05050d4bb478c6f3b47c34b9ec2cc Mon Sep 17 00:00:00 2001 From: Hao Wu Date: Thu, 2 Apr 2026 15:53:52 -0700 Subject: [PATCH 11/11] Pass overlap_param_gather config to LayerWiseDistributedOptimizer as async_allgather --- megatron/core/optimizer/__init__.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 9677132c685..d99d3235aa0 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -845,7 +845,12 @@ def _get_megatron_emerging_optimizer( logger, logging.INFO, f'Using LayerWiseDistributedOptimizer for {eopt_name}' ) return LayerWiseDistributedOptimizer( - list(base_optimizers), config, pg_collection, init_state_fn_list=list(init_fns) + list(base_optimizers), + config, + pg_collection, + init_state_fn_list=list(init_fns), + model_chunks=model_chunks if config.overlap_param_gather else None, + async_allgather=config.overlap_param_gather, ) return ChainedOptimizer(results)