# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import inspect import itertools import os from collections import defaultdict from contextlib import contextmanager from typing import TYPE_CHECKING, Any, Callable, Dict, Generator, Mapping, Optional, Protocol, TypeVar import torch from torch import nn from nemo.lightning.megatron_init import initialize_model_parallel_for_nemo from nemo.utils import logging NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE = "NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE" if TYPE_CHECKING: from lightning.fabric.utilities.types import Optimizable from megatron.core.model_parallel_config import ModelParallelConfig class SharedStateDictProtocol(Protocol): """ """ def sharded_state_dict(self, prefix="", metadata: Optional[dict] = None): """ """ ... def init_parallel_ranks( world_size: int, global_rank: int, local_rank: int, parallel_config: "ModelParallelConfig", seed=1234, fp8=False, ) -> None: """ Initializes the parallel ranks for distributed training. This function sets up the parallel ranks based on the provided world size, global rank, local rank, and parallel configuration. It also sets the seed for random number generation and determines whether to use fp8 precision. Args: world_size (int): The total number of processes participating in the distributed training. global_rank (int): The rank of the current process in the distributed training setup. local_rank (int): The rank of the current process within its machine. parallel_config (ModelParallelConfig): The configuration object containing settings for model parallelism. seed (int, optional): The seed for random number generation. Defaults to 1234. fp8 (bool, optional): Whether to use fp8 precision for model parameters. Defaults to False. """ from nemo.utils import AppState app_state = AppState() if os.environ.get(NEMO_MEGATRON_MODEL_PARALLEL_APPSTATE_OVERRIDE, "false").lower() == "true": init_world_size = app_state.tensor_model_parallel_size * app_state.pipeline_model_parallel_size init_global_rank = app_state.global_rank init_local_rank = app_state.local_rank else: init_world_size = world_size pp = parallel_config.pipeline_model_parallel_size or 1 if world_size < pp: raise ValueError(f"Expected world_size ({world_size}) to be greater than/equal to pipeline size ({pp})") init_global_rank = global_rank init_local_rank = local_rank initialize_model_parallel_for_nemo( world_size=init_world_size, global_rank=init_global_rank, local_rank=init_local_rank, tensor_model_parallel_size=parallel_config.tensor_model_parallel_size, expert_model_parallel_size=parallel_config.expert_model_parallel_size, expert_tensor_parallel_size=parallel_config.expert_tensor_parallel_size, pipeline_model_parallel_size=parallel_config.pipeline_model_parallel_size, pipeline_model_parallel_comm_backend=parallel_config.pipeline_model_parallel_comm_backend, virtual_pipeline_model_parallel_size=parallel_config.virtual_pipeline_model_parallel_size, context_parallel_size=parallel_config.context_parallel_size, seed=seed, use_fp8=fp8, init_mpi_proc_group=getattr(parallel_config, "tp_comm_overlap", False) and getattr(parallel_config, "tp_comm_bootstrap_backend", None) == 'mpi', use_te_rng_tracker=getattr(parallel_config, "use_te_rng_tracker", False), use_sharp=getattr(parallel_config, "use_sharp", False), use_tp_pp_dp_mapping=getattr(parallel_config, "use_tp_pp_dp_mapping", False), num_distributed_optimizer_instances=getattr(parallel_config, "num_distributed_optimizer_instances", 1), nccl_communicator_config_path=getattr(parallel_config, "nccl_communicator_config_path", None), use_gloo_process_groups=getattr(parallel_config, "use_gloo_process_groups", True), # apex_transformer_log_level=self.cfg.get('apex_transformer_log_level', 30), ) def init_model_parallel(model: Optional[nn.Module] = None) -> None: """Initializes Megatron-LM model parallel if using model parallelism.""" import torch.distributed from megatron.core import parallel_state from nemo.utils import AppState app_state = AppState() # we initialize megatron-lm model parallel and data parallel groups # after initializing DDP with PTL. if app_state.model_parallel_size is not None: # destroy groups in case they have already been created # this happens with multiple calls to trainer.test for example parallel_state.destroy_model_parallel() if torch.distributed.is_initialized(): parallel_state.initialize_model_parallel( tensor_model_parallel_size=app_state.tensor_model_parallel_size, pipeline_model_parallel_size=app_state.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=app_state.virtual_pipeline_model_parallel_size, pipeline_model_parallel_comm_backend=app_state.pipeline_model_parallel_comm_backend, context_parallel_size=app_state.context_parallel_size, expert_model_parallel_size=app_state.expert_model_parallel_size, expert_tensor_parallel_size=app_state.expert_tensor_parallel_size, use_sharp=app_state.use_sharp, order="tp-cp-ep-pp-dp" if app_state.use_tp_pp_dp_mapping else "tp-cp-ep-dp-pp", num_distributed_optimizer_instances=app_state.num_distributed_optimizer_instances, nccl_communicator_config_path=app_state.nccl_communicator_config_path, create_gloo_process_groups=app_state.use_gloo_process_groups, ) # assert that fake tp and pp rank match after model parallel init assert app_state.tensor_model_parallel_rank == parallel_state.get_tensor_model_parallel_rank() assert app_state.pipeline_model_parallel_rank == parallel_state.get_pipeline_model_parallel_rank() assert app_state.expert_tensor_parallel_rank == parallel_state.get_expert_tensor_parallel_rank() app_state.tensor_model_parallel_group = parallel_state.get_tensor_model_parallel_group() app_state.data_parallel_group = parallel_state.get_data_parallel_group() app_state.data_parallel_rank = parallel_state.get_data_parallel_rank() app_state.data_parallel_size = parallel_state.get_data_parallel_world_size() app_state.pipeline_model_parallel_group = parallel_state.get_pipeline_model_parallel_group() # create MPI process group for UCX-based communication APIs if app_state.init_mpi_proc_group: torch.distributed.new_group(backend="mpi") def set_model_parallel_attributes(model, parallelism): """ """ # Right now mcore sub-classes ModelParellelConfig, we should remove that # Given Lightning's structure it would be better if parallelism is a different object # Since then it can be passed to the Strategy # Note: Importing nemo.lightning.pytorch.strategies creates an import cycle. from megatron.core.transformer.transformer_config import TransformerConfig has_mcore_config = isinstance(getattr(model, "config", None), TransformerConfig) if has_mcore_config and hasattr(model, "configure_model"): config: TransformerConfig = model.config for attr_name in filter(lambda x: not x.startswith('__'), dir(parallelism)): if not hasattr(config, attr_name): continue setattr(config, attr_name, getattr(parallelism, attr_name)) if hasattr(config, "__io__"): setattr(config.__io__, attr_name, getattr(parallelism, attr_name)) if hasattr(config, '__post_init__'): # MCore does not use args in __post_init__ # @akoumparouli: is there a better way (e.g. reinit config)? config.__post_init__() return config return None @contextmanager def megatron_lazy_init_context(config) -> Generator[None, None, None]: """ """ try: from megatron.core.extensions import transformer_engine as _te original = _te._get_extra_te_kwargs # noqa: SLF001 def _get_extra_te_kwargs_meta(c): """Forces device to meta""" kwargs = original(c) kwargs['device'] = 'meta' return kwargs _te._get_extra_te_kwargs = _get_extra_te_kwargs_meta # noqa: SLF001 except ImportError: pass _orig_perform_initialization = config.perform_initialization _orig_use_cpu_initialization = config.use_cpu_initialization config.perform_initialization = False config.use_cpu_initialization = True yield try: from megatron.core.extensions import transformer_engine as _te _te._get_extra_te_kwargs = original # noqa: SLF001 except ImportError: pass config.perform_initialization = _orig_perform_initialization config.use_cpu_initialization = _orig_use_cpu_initialization @contextmanager def megatron_cpu_init_context(config) -> Generator[None, None, None]: """ """ _orig_use_cpu_initialization = config.use_cpu_initialization config.use_cpu_initialization = True yield config.use_cpu_initialization = _orig_use_cpu_initialization ModelT = TypeVar("ModelT", bound=nn.Module) class GradScaler(torch.cuda.amp.GradScaler): """ Gradient sclaer for model-parallel inf check. The inf in gradients are checked across tensor-parallel ranks in (1) executing optimizer step and (2) gradient scaler update. """ def __init__( self, init_scale=2.0**16, growth_factor=2.0, backoff_factor=0.5, growth_interval=2000, enabled=True, hysteresis=1, ): super().__init__( init_scale=init_scale, growth_factor=growth_factor, backoff_factor=backoff_factor, growth_interval=growth_interval, enabled=enabled, ) self.optimizer_update_skipped: Optional[bool] = None self.hysteresis = hysteresis self._hysteresis_tracker = self.hysteresis def _unscale_grads_(self, optimizer, *args): if getattr(optimizer, "_custom_amp_unscale_grads", False): return optimizer.unscale_grads(*args) else: return super()._unscale_grads_(optimizer, *args) def _maybe_opt_step(self, optimizer, optimizer_state, *args, **kwargs): from megatron.core import parallel_state retval = None found_inf = torch.cuda.FloatTensor([sum(v.item() for v in optimizer_state["found_inf_per_device"].values())]) # Update across all model parallel instances. torch.distributed.all_reduce( found_inf, op=torch.distributed.ReduceOp.MAX, group=parallel_state.get_model_parallel_group(), ) if found_inf.item() == 0: retval = optimizer.step(*args, **kwargs) self.optimizer_update_skipped = False else: self.optimizer_update_skipped = True return retval def update(self, new_scale=None): """ Updates to native grad scaler update function. 1. Check inf across model-parallel ranks. 2. Update hysteresis tracker. 3. Apply hysteresis to grad scale update. """ from megatron.core import parallel_state if not self._enabled: return _scale, _growth_tracker = self._check_scale_growth_tracker("update") if new_scale is not None: # Accept a new user-defined scale. if isinstance(new_scale, float): self._scale.fill_(new_scale) # type: ignore[union-attr] else: reason = ( "new_scale should be a float or a 1-element torch.cuda.FloatTensor with" " requires_grad=False." ) assert isinstance(new_scale, torch.cuda.FloatTensor), reason # type: ignore[attr-defined] assert new_scale.numel() == 1, reason assert new_scale.requires_grad is False, reason self._scale.copy_(new_scale) # type: ignore[union-attr] else: # Consume shared inf/nan data collected from optimizers to update the scale. # If all found_inf tensors are on the same device as self._scale, this operation is asynchronous. found_infs = [ found_inf.to(device=_scale.device, non_blocking=True) for state in self._per_optimizer_states.values() for found_inf in state["found_inf_per_device"].values() ] assert len(found_infs) > 0, "No inf checks were recorded prior to update." found_inf_combined = found_infs[0] # Update across all model parallel instances. torch.distributed.all_reduce( found_inf_combined, op=torch.distributed.ReduceOp.MAX, group=parallel_state.get_model_parallel_group(), ) if len(found_infs) > 1: for i in range(1, len(found_infs)): found_inf = found_infs[i] # Update across all model parallel instances. torch.distributed.all_reduce( found_inf, op=torch.distributed.ReduceOp.MAX, group=parallel_state.get_model_parallel_group(), ) found_inf_combined += found_inf if found_inf_combined > 0: self._hysteresis_tracker -= 1 if self._hysteresis_tracker <= 0: # When hysteresis becomes zero, follow the native grad scale update rule. # Increase scale and reset growth tracker torch._amp_update_scale_( # noqa: SLF001 _scale, _growth_tracker, found_inf_combined, self._growth_factor, self._backoff_factor, self._growth_interval, ) else: # Only reset the growth tracker when hysteresis is larger than zero _growth_tracker.fill_(0.0) else: # When no inf found, follow the native grad scale update rule. # Increment growth_tracker, update scale when growth tracker reaches the interval, and # reset the hysteresis tracker. torch._amp_update_scale_( # noqa: SLF001 _scale, _growth_tracker, found_inf_combined, self._growth_factor, self._backoff_factor, self._growth_interval, ) self._hysteresis_tracker = self.hysteresis # To prepare for next iteration, clear the data collected from optimizers this iteration. self._per_optimizer_states = defaultdict( torch.cuda.amp.grad_scaler._refresh_per_optimizer_state # noqa: SLF001 ) def state_dict(self): """ Add hysteresis_tracker to the native functions' state_dict. """ return ( { "scale": self.get_scale(), "growth_factor": self._growth_factor, "backoff_factor": self._backoff_factor, "growth_interval": self._growth_interval, "_growth_tracker": self._get_growth_tracker(), "_hysteresis_tracker": self._hysteresis_tracker, } if self._enabled else {} ) def load_state_dict(self, state_dict): """ Load hysteresis_tracker in addition to the state dict of the native function. """ if not self._enabled: return if len(state_dict) == 0: raise RuntimeError( "The source state dict is empty, possibly because it was saved " "from a disabled instance of GradScaler." ) self._init_scale = state_dict["scale"] if self._scale is not None: self._scale.fill_(state_dict["scale"]) self._growth_factor = state_dict["growth_factor"] self._backoff_factor = state_dict["backoff_factor"] self._growth_interval = state_dict["growth_interval"] self._init_growth_tracker = state_dict["_growth_tracker"] if self._growth_tracker is not None: self._growth_tracker.fill_(state_dict["_growth_tracker"]) if "_hysterisis_tracker" in state_dict: self._hysteresis_tracker = state_dict["_hysterisis_tracker"] else: self._hysteresis_tracker = 1 def enable_nvidia_optimizations() -> None: """These optimizations are present in NVIDIA NGC PyTorch Containers.""" # NVIDIA container version check nvidia_torch_version = os.getenv("NVIDIA_PYTORCH_VERSION", None) if nvidia_torch_version is not None: try: NVIDIA_TORCH_MAJOR = int(nvidia_torch_version.split(".")[0]) except Exception: NVIDIA_TORCH_MAJOR = 0 try: NVIDIA_TORCH_MINOR = int(nvidia_torch_version.split(".")[1]) except Exception: NVIDIA_TORCH_MINOR = 0 # NVFUSER available starting with 21.11 if NVIDIA_TORCH_MAJOR >= 21 or (NVIDIA_TORCH_MAJOR == 21 and NVIDIA_TORCH_MINOR >= 11): # NVFUSER torch._C._jit_set_profiling_executor(True) # noqa: SLF001 torch._C._jit_set_profiling_mode(True) # noqa: SLF001 torch._C._jit_override_can_fuse_on_cpu(False) # noqa: SLF001 torch._C._jit_override_can_fuse_on_gpu(False) # noqa: SLF001 torch._C._jit_set_texpr_fuser_enabled(False) # noqa: SLF001 # torch._C._jit_set_nvfuser_enabled(True) torch._C._debug_set_autodiff_subgraph_inlining(False) # noqa: SLF001 else: # Not a Nvidia container. NVFUSER Dependency check is on users pass def optimizer_sharded_state_dict( model: SharedStateDictProtocol, optimizer: "Optimizable", is_loading: bool = False, sharding_type: Optional[str] = None, metadata: Optional[dict] = None, ) -> Dict[str, torch.Tensor]: """ Sharded state dictionary for an MainParamsOptimizerWrapper. Used to save and load the optimizer state when training with distributed_checkpoint. Args: model (SharedStateDictProtocol): model with a `sharded_state_dict` method optimizer (Optimizable): optimizer to get the state dict of is_loading (bool, optional): set to True if the sharded state dict is intended for checkpoint loading (as opposed to saving). Defaults to False. sharding_type (str, optional): deprecated, use metadata flags instead. metadata (dict, optional): sharded state dict metadata passed from the framework. Used to control the details of sharded state dict creation, in particular the state dict format of the DistributedOptimizer with the flag `distrib_optim_sharding_type`. Defaults to None (empty metadata). Returns ------- dict: The sharded state dictionary for the optimizer Raises: ValueError: If a parameter ID does not match any model sharded parameter. """ from megatron.core.dist_checkpointing.optimizer import ( get_param_id_to_sharded_param_map, make_sharded_optimizer_tensor, optim_state_to_sharding_state, ) from nemo.core.optim import MainParamsOptimizerWrapper from nemo.core.optim.optimizers import init_optimizer_states model_sharded_state_dict = model.sharded_state_dict(metadata=metadata) # remove _extra_state model_sharded_state_dict = { key: value for key, value in model_sharded_state_dict.items() if not key.endswith("_extra_state") } if sharding_type is not None: logging.warning("sharding_type is deprecated, please use `metadata['distrib_optim_sharding_type']` instead") if metadata is None: metadata = {} if 'distrib_optim_sharding_type' not in metadata: metadata["distrib_optim_sharding_type"] = sharding_type if hasattr(optimizer, "sharded_state_dict"): return optimizer.sharded_state_dict( model_sharded_state_dict, is_loading=is_loading, metadata=metadata, ) if not isinstance(optimizer, MainParamsOptimizerWrapper): # Regular optimizer, e.g. Adam or FusedAdam init_optimizer_states(optimizer) optimizer_state_dict = optimizer.state_dict() id_to_sharded_param_map = get_param_id_to_sharded_param_map( model_sharded_state_dict=model_sharded_state_dict, optim_params_iter=itertools.chain.from_iterable(g['params'] for g in optimizer.param_groups), ) optim_state_to_sharding_state(optimizer_state_dict, id_to_sharded_param_map) return optimizer_state_dict optimizer_state_dict: Dict[str, Any] = optimizer.state_dict() id_to_sharded_param_map = get_param_id_to_sharded_param_map( model_sharded_state_dict=model_sharded_state_dict, optim_params_iter=itertools.chain.from_iterable(g for g in optimizer.float16_groups), ) # Convert fp32_from_fp16_params assert len(optimizer_state_dict["fp32_from_fp16_params"]) == len(optimizer_state_dict["optimizer"]["param_groups"]) def get_safe(param_id): try: return id_to_sharded_param_map[param_id] except KeyError as e: raise ValueError(f"Param id {param_id} does not match any model sharded param") from e optimizer_state_dict["fp32_from_fp16_params"] = [ [ make_sharded_optimizer_tensor(get_safe(param_id), fp32_param, prefix="optimizer.state.fp32_param") for param_id, fp32_param in zip(state_group["params"], fp32_group) ] for fp32_group, state_group in zip( optimizer_state_dict["fp32_from_fp16_params"], optimizer_state_dict["optimizer"]["param_groups"], ) ] # Convert state optim_state_to_sharding_state(optimizer_state_dict["optimizer"], id_to_sharded_param_map) return optimizer_state_dict def load_model_state_dict(megatron_parallel, checkpoint: Mapping[str, Any], strict: bool = True) -> None: """ """ from megatron.core import parallel_state from megatron.core.dist_checkpointing.validation import StrictHandling, parse_strict_flag # convert from StrictHandling to bool for PTL if strict is not None and not isinstance(strict, bool): strict = parse_strict_flag(strict) strict_options = [ StrictHandling.ASSUME_OK_UNEXPECTED, StrictHandling.RAISE_UNEXPECTED, StrictHandling.RAISE_ALL, ] strict = strict in strict_options try: from megatron.core.distributed.custom_fsdp import FullyShardedDataParallel have_custom_fsdp = True except ImportError or ModuleNotFoundError: have_custom_fsdp = False try: from megatron.core.distributed import FullyShardedDataParallel have_megatron_fsdp = True except ImportError or ModuleNotFoundError: have_megatron_fsdp = False for index, module in enumerate(megatron_parallel): if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None: if "state_dict" in checkpoint: checkpoint_state_dict = checkpoint["state_dict"][f"model_{index}"] else: checkpoint_state_dict = checkpoint[f"model_{index}"] else: if "state_dict" in checkpoint: checkpoint_state_dict = checkpoint["state_dict"] else: checkpoint_state_dict = checkpoint n_nesting = 0 mcore_model = megatron_parallel.module while hasattr(mcore_model, "module"): mcore_model = mcore_model.module n_nesting += 1 _state_dict = {} for key, value in checkpoint_state_dict.items(): # Count the number of "module." at the start of the key count, _key = 0, key while _key.startswith("module."): _key = _key[len("module.") :] count += 1 # Adjust the number of "module." prefixes if count < n_nesting: to_add = "module." * (n_nesting - count) _state_dict[f"{to_add}{key}"] = value elif count > n_nesting: to_remove = "module." * (count - n_nesting) _state_dict[key[len(to_remove) :]] = value else: _state_dict[key] = value if have_custom_fsdp and hasattr(module, "module") and isinstance(module.module, FullyShardedDataParallel): module.module.load_state_dict(_state_dict, strict=strict) elif have_megatron_fsdp and hasattr(module, "module") and isinstance(module.module, FullyShardedDataParallel): module.module.load_state_dict(_state_dict, strict=strict) continue try: module.load_state_dict(_state_dict, strict=strict) except RuntimeError as e: missing_keys, expected_keys = module.load_state_dict(checkpoint_state_dict, strict=False) if all(s.endswith('_extra_state') for s in missing_keys): logging.warning( f'Loding checkpoint created with Transformer Engine version lower than 1.13. ' f'Missing layers {missing_keys} will be ignored.' ) else: raise e def _sync_from_last_pipeline_stage(value: torch.Tensor, broadcast: bool = False): """ When pipeline parallelism is enabled, casts a tensor defined on the last pipeline stage to other ranks. Args: value (torch.Tensor): A tensor to be casted from the final pipeline stage of a pipeline parallelism group (e.g. loss). Note that this tensor should already be defined on the target rank(s) to fill with received data. broadcast (bool): When True, broadcasts value from the final pipeline stage rank to all ranks in its group. When False, only rank zero receives value from the final pipeline stage rank in its group. This mode exists to avoid slow one-to-many communication when not necessary. Defaults to False. """ from megatron.core import parallel_state if parallel_state.get_pipeline_model_parallel_world_size() > 1: src_rank = parallel_state.get_pipeline_model_parallel_last_rank() if not broadcast: pp_ranks = torch.distributed.get_process_group_ranks(parallel_state.get_pipeline_model_parallel_group()) if torch.distributed.get_rank() == src_rank and 0 in pp_ranks: torch.distributed.send(value, 0) elif torch.distributed.get_rank() == 0: torch.distributed.recv(value, src_rank) else: torch.distributed.broadcast( value, src_rank, group=parallel_state.get_pipeline_model_parallel_group(), ) def setup_megatron_optimizer( model, config, no_weight_decay_cond: Optional[Callable] = None, scale_lr_cond: Optional[Callable] = None, lr_mult: float = 1.0, ): """ """ from megatron.core.optimizer import OptimizerConfig, get_megatron_optimizer from nemo.core.optim import McoreDistributedOptimizer from nemo.utils import AppState app_state = AppState() assert isinstance(config, OptimizerConfig), f"Expected OptimizerConfig, got {type(config)}" class McoreOpt(McoreDistributedOptimizer): """ """ def sharded_state_dict( self, model_sharded_state_dict, optimizer_state_dict=None, is_loading=False, sharding_type='fully_sharded_model_space', metadata=None, ): mcore_optimizer_sig = inspect.signature(self.mcore_optimizer.sharded_state_dict).parameters distrib_optim_kwargs = {} if "metadata" in mcore_optimizer_sig or "kwargs" in mcore_optimizer_sig: distrib_optim_kwargs["metadata"] = metadata elif "sharding_type" in mcore_optimizer_sig: distrib_optim_kwargs["sharding_type"] = sharding_type state_dict = self.mcore_optimizer.sharded_state_dict( model_sharded_state_dict, is_loading=is_loading, **distrib_optim_kwargs ) return state_dict # megatron optimizer expects McoreDDP ddp_modules = [m.module for m in model] mcore_opt = get_megatron_optimizer( config, ddp_modules, no_weight_decay_cond=no_weight_decay_cond, scale_lr_cond=scale_lr_cond, lr_mult=lr_mult, use_gloo_process_groups=app_state.use_gloo_process_groups, ) # Pytorch does not have the concept of an `lr_mult` or a `wd_mult` but these are added to param # groups in megatron to control which sub-modules have different learning rates or weight # decays. Apply the multipliers here to each param_group's lr and wd, and to reduce confusion # change the name of these variables. We need this because nemo does not use the custom # megatron scheduler, and the megatron scheduler is what makes use of these mult parameters: # https://github.com/NVIDIA/Megatron-LM/blob/044e2ad5/megatron/core/optimizer_param_scheduler.py#L192-L193 for pg in mcore_opt.param_groups: if 'pre_lr_mult' in pg or 'pre_mult_wd' in pg: # User has already applied custom lr and wd multipliers, don't apply `lr_mult` and # `wd_mult` again. This case may be encountered when resuming training. continue pg['pre_mult_lr'] = pg["lr"] pg['pre_mult_wd'] = pg['weight_decay'] new_lr = pg["lr"] * pg.get('lr_mult', 1.0) new_wd = pg["weight_decay"] * pg.get("wd_mult", 1.0) pg['lr'] = new_lr pg['weight_decay'] = new_wd # In case a future implementation makes use of `lr_mult` and `wd_mult` directly in the # scheduler, but accidentally also uses this function, remove `lr_mult` and `wd_mult` from # the param groups so that the default value of 1.0 gets applied. if 'lr_mult' in pg: pg['pre_lr_mult'] = pg['lr_mult'] del pg['lr_mult'] # remove so downstream methods do not apply again. if 'wd_mult' in pg: pg['pre_wd_mult'] = pg['wd_mult'] del pg['wd_mult'] # remove so downstream methods do not apply again if getattr(model.ddp_config, "overlap_param_gather", False) and getattr( model.ddp_config, "align_param_gather", False ): param_sync_func = [model_chunk.start_param_sync for model_chunk in model] param_sync_func = param_sync_func[0] if len(model) == 1 else param_sync_func for module in model: module.config.param_sync_func = param_sync_func return McoreOpt(mcore_opt)