diff --git a/megatron/core/optimizer/__init__.py b/megatron/core/optimizer/__init__.py index 11aa6c49585..3bd5fbd9ca5 100644 --- a/megatron/core/optimizer/__init__.py +++ b/megatron/core/optimizer/__init__.py @@ -57,6 +57,7 @@ ) from .optimizer_config import ( AdamOptimizerConfig, + MuonOptimizerConfig, OptimizerConfig, ParamKey, ParamPredicate, diff --git a/megatron/core/optimizer/emerging_optimizers.py b/megatron/core/optimizer/emerging_optimizers.py new file mode 100644 index 00000000000..2c5a89c4fd1 --- /dev/null +++ b/megatron/core/optimizer/emerging_optimizers.py @@ -0,0 +1,243 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Megatron emerging optimizers (e.g. Muon) factory and wiring.""" + +import logging +from copy import copy +from typing import Any, Dict, List, Optional + +import torch + +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.module import MegatronModule +from megatron.core.utils import log_single_rank + +from . import _get_param_groups, get_megatron_optimizer +from .layer_wise_optimizer import LayerWiseDistributedOptimizer +from .muon import TensorParallelMuon +from .optimizer import ( + ChainedOptimizer, + Float16OptimizerWithFloat16Params, + FP32Optimizer, + MegatronOptimizer, +) +from .optimizer_config import OptimizerConfig, ParamGroupOverride, ParamKey + +logger = logging.getLogger(__name__) + + +def get_megatron_emerging_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: + """Get the emerging 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. + """ + eopt_name = copy(config.optimizer) + # FIXME(skyw): Confirm whether we still need this hack with MuonOptimizerConfig + config.optimizer = 'adam' + + # Dist-opt is not supported due to strong coupling with how DDP init grad buffer + # In theory we can change DDP to enable use emerging optimizer and dist-opt-adam together + if config.use_distributed_optimizer: + raise Exception('emerging optimizer with dist optimizer is not supported.') + # only support bf16 w/o loss scale now + if config.fp16: + raise Exception('emerging optimizer 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) + + if eopt_name == 'muon': + eopt_cls = TensorParallelMuon + eopt_init_state_fn = muon_init_state_fn + else: + raise ValueError(f"Unsupported emerging optimizer: {eopt_name}") + + 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) + + eopt_kwargs = _config_to_eopt_kwargs(config) + if eopt_kwargs.get("split_qkv", False): + eopt_kwargs["is_qkv_fn"] = lambda p: getattr(p, "is_qkv", False) + eopt_kwargs["qkv_split_shapes"] = qkv_split_shapes + eopt_kwargs['pg_collection'] = pg_collection + + # freezing nonlinear params and get param groups for emerging optimizer + 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 = eopt_cls(linear_param_groups, **eopt_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, eopt_init_state_fn + ) + else: + optimizer = FP32Optimizer(optimizer, config, eopt_init_state_fn) + + optimizers.append(optimizer) + + # expert optimizer exists meaning layerwise distributed optimizer is not used + if len(expert_param_groups) > 0: + expert_optimizer = eopt_cls(expert_param_groups, **eopt_kwargs) + if config.bf16: + expert_optimizer = Float16OptimizerWithFloat16Params( + expert_optimizer, config, None, eopt_init_state_fn + ) + else: + expert_optimizer = FP32Optimizer(expert_optimizer, config, eopt_init_state_fn) + setattr(expert_optimizer, 'grad_stats_parallel_group', pg_collection.tp_ep_pp) + optimizers.append(expert_optimizer) + + # done with emerging optimizer, 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 = [eopt_init_state_fn] + len(chained_adam.chained_optimizers) * [adam_init_state_fn] + optimizers += chained_adam.chained_optimizers + + if layer_wise_distributed_optimizer: + log_single_rank( + logger, logging.INFO, f'Using LayerWiseDistributedOptimizer for {eopt_name}' + ) + if reset_config_bf16: + config.bf16 = True + return LayerWiseDistributedOptimizer( + optimizers, config, pg_collection, init_state_fn_list=init_fns + ) + return ChainedOptimizer(optimizers) + + +def _config_to_eopt_kwargs(config: OptimizerConfig) -> Dict[str, Any]: + """Convert a Megatron optimizer config to a dictionary of kwargs for the emerging optimizer. + + Note: + Ideally this can be a method of OptimizerConfig class, which would require a naming + convention being enforced. Currently because optimizer implementation is coming from + EmergingOptimizers, we don't have any mechanism to enforce same naming convention between + MegatronCore and EmergingOptimizers. Before a reliable syncing mechanism is in place, + we'll have this dispatch function. + + Args: + config (OptimizerConfig): the Megatron optimizer config. + + Returns: + Dict[str, Any]: a dictionary of kwargs for the emerging optimizer. + """ + if config.optimizer == 'muon': + eopt_kwargs = { + "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, + } + else: + raise ValueError(f"Unsupported emerging optimizer: {config.optimizer}") + return eopt_kwargs diff --git a/megatron/core/optimizer/muon.py b/megatron/core/optimizer/muon.py index 57eb1e94478..267800f5300 100644 --- a/megatron/core/optimizer/muon.py +++ b/megatron/core/optimizer/muon.py @@ -3,26 +3,14 @@ """Megatron muon optimizer wrapper to handle tensor-parallel.""" import logging -from typing import Any, Callable, Dict, List, Literal, Optional +from typing import Any, Callable, Literal, Optional 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 -from . import _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 - try: from emerging_optimizers.orthogonalized_optimizers import ( OrthogonalizedOptimizer, @@ -162,189 +150,10 @@ def orthogonalize(self, p: torch.Tensor, grad: torch.Tensor, **kwargs: Any) -> t 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. - """ - # Muon currently use adam config. setting str here to call regular get for adam creation - # side effect is muon optimizer will have wrong name, i.e. config.optimizer == 'adam' - config.optimizer = 'adam' - - assert HAVE_EMERGING_OPTIMIZERS, "Emerging Optimizers is not installed." - - # Dist-opt is not supported due to strong coupling with how DDP init grad buffer - # In theory we can change DDP to enable use muon and dist-opt-adam together - 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 get_megatron_muon_optimizer(*args: Any, **kwargs: Any) -> Any: + """Backward compatible muon optimizer getter.""" - 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) + # Re-export from emerging_optimizer to avoid circular imports + from .emerging_optimizers import get_megatron_emerging_optimizer as _get_megatron_muon_optimizer - 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, - "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) * [adam_init_state_fn] - optimizers += chained_adam.chained_optimizers - - 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 - ) - return ChainedOptimizer(optimizers) + return _get_megatron_muon_optimizer(*args, **kwargs) diff --git a/megatron/core/optimizer/optimizer_config.py b/megatron/core/optimizer/optimizer_config.py index 94163102eb3..1542a8fc265 100644 --- a/megatron/core/optimizer/optimizer_config.py +++ b/megatron/core/optimizer/optimizer_config.py @@ -254,32 +254,6 @@ class OptimizerConfig: sgd_momentum: float = 0.9 """Momentum factor for SGD optimizer.""" - # Muon. - # TODO: move muon configs to it's own `MuonConfig`. - muon_momentum: float = 0.95 - """The momentum used by the internal SGD.""" - - muon_split_qkv: bool = True - """Whether to split QKV parameters for Muon optimizer.""" - - muon_use_nesterov: bool = False - """Whether to use Nesterov-style momentum in the internal SGD.""" - - muon_scale_mode: str = "spectral" - """The mode to use for the scale factor. Defaults to "spectral".""" - - muon_fp32_matmul_prec: str = "medium" - """The precision to use for the fp32 matmul. Defaults to "medium".""" - - muon_num_ns_steps: int = 5 - """The number of iteration steps to use in the Newton-Schulz iteration.""" - - muon_tp_mode: str = "blockwise" - """How to perform NS calculation for tensor parallel weights. Defaults to "blockwise".""" - - muon_extra_scale_factor: float = 1.0 - """Additional scale factor for the muon update.""" - ####################### # Distributed optimizer ####################### @@ -461,3 +435,35 @@ class SGDOptimizerConfig(OptimizerConfig): sgd_momentum: float = 0.9 """Momentum factor for SGD optimizer.""" + + +@dataclass +class MuonOptimizerConfig(OptimizerConfig): + """Muon optimizer configuration object.""" + + optimizer: str = 'muon' + """Optimizer name.""" + + muon_momentum: float = 0.95 + """The momentum used by the internal SGD.""" + + muon_split_qkv: bool = True + """Whether to split QKV parameters for Muon optimizer.""" + + muon_use_nesterov: bool = False + """Whether to use Nesterov-style momentum in the internal SGD.""" + + muon_scale_mode: str = "spectral" + """The mode to use for the scale factor. Defaults to "spectral".""" + + muon_fp32_matmul_prec: str = "medium" + """The precision to use for the fp32 matmul. Defaults to "medium".""" + + muon_num_ns_steps: int = 5 + """The number of iteration steps to use in the Newton-Schulz iteration.""" + + muon_tp_mode: str = "blockwise" + """How to perform NS calculation for tensor parallel weights. Defaults to "blockwise".""" + + muon_extra_scale_factor: float = 1.0 + """Additional scale factor for the muon update.""" diff --git a/megatron/training/training.py b/megatron/training/training.py index e9736ac085c..f8db9eb02dd 100644 --- a/megatron/training/training.py +++ b/megatron/training/training.py @@ -124,7 +124,14 @@ 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 import ( + get_megatron_optimizer, + AdamOptimizerConfig, + MuonOptimizerConfig, + SGDOptimizerConfig, + OptimizerConfig, + ParamKey, +) from megatron.core.optimizer.muon import get_megatron_muon_optimizer from megatron.core.rerun_state_machine import ( get_rerun_state_machine, @@ -1412,9 +1419,7 @@ 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 + if args.optimizer == 'adam': kwargs = {} for f in dataclasses.fields(AdamOptimizerConfig): if hasattr(args, f.name): @@ -1426,6 +1431,12 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig: if hasattr(args, f.name): kwargs[f.name] = getattr(args, f.name) config = SGDOptimizerConfig(**kwargs) + elif 'muon' in args.optimizer: + kwargs = {} + for f in dataclasses.fields(MuonOptimizerConfig): + if hasattr(args, f.name): + kwargs[f.name] = getattr(args, f.name) + config = MuonOptimizerConfig(**kwargs) else: raise ValueError("Invalid optimizer type!")