diff --git a/megatron/core/optimizer_param_scheduler.py b/megatron/core/optimizer_param_scheduler.py index 53a2c1a3951..13045e13a8b 100644 --- a/megatron/core/optimizer_param_scheduler.py +++ b/megatron/core/optimizer_param_scheduler.py @@ -15,6 +15,8 @@ logger = logging.getLogger(__name__) +_OPT_PARAM_SCHEDULER_OVERRIDE_KEYS = ('max_lr', 'min_lr', 'start_wd', 'end_wd') + class ParamGroupOverride(TypedDict, total=False): """Override values for a parameter group. These values may be optimizer-state/scheduler related. @@ -143,6 +145,14 @@ def __init__( # Class values. self.optimizer = optimizer + self._param_group_scheduler_overrides = [ + { + key: param_group[key] + for key in _OPT_PARAM_SCHEDULER_OVERRIDE_KEYS + if key in param_group + } + for param_group in self.optimizer.param_groups + ] self.init_lr = init_lr self.max_lr = float(max_lr) @@ -338,6 +348,32 @@ def _check_and_set(self, cls_value: float, sd_value: float, name: str) -> float: log_single_rank(logger, logging.INFO, f" > using checkpoint value {sd_value} for {name}") return sd_value + def _restore_param_group_scheduler_overrides(self) -> None: + """Restore runtime scheduler fields after optimizer param groups are loaded. + + Optimizer checkpoints may carry per-param-group scheduler fields such as + max_lr/min_lr/start_wd/end_wd. These fields take precedence over the scheduler's class + values in get_lr()/get_wd(), so override mode must restore the values created + from the current run's arguments before reapplying the schedule. + """ + if not self.override_opt_param_scheduler: + return + + if len(self.optimizer.param_groups) != len(self._param_group_scheduler_overrides): + log_single_rank( + logger, + logging.WARNING, + " > unable to restore optimizer param group scheduler overrides due to " + "param group count mismatch", + ) + return + + for param_group, scheduler_overrides in zip( + self.optimizer.param_groups, self._param_group_scheduler_overrides + ): + for key, value in scheduler_overrides.items(): + param_group[key] = value + def load_state_dict(self, state_dict: dict) -> None: """Load the state dict. @@ -383,12 +419,6 @@ def load_state_dict(self, state_dict: dict) -> None: self.lr_decay_style, lr_decay_style_, 'learning rate decay style' ) - if 'num_iters' in state_dict: - num_steps = state_dict['num_iters'] - else: - num_steps = state_dict['num_steps'] - self.step(increment=num_steps) - if 'start_wd' in state_dict: self.start_wd = self._check_and_set( self.start_wd, state_dict['start_wd'], "start weight decay" @@ -402,3 +432,10 @@ def load_state_dict(self, state_dict: dict) -> None: self.wd_incr_style = self._check_and_set( self.wd_incr_style, state_dict['wd_incr_style'], "weight decay incr style" ) + + if 'num_iters' in state_dict: + num_steps = state_dict['num_iters'] + else: + num_steps = state_dict['num_steps'] + self._restore_param_group_scheduler_overrides() + self.step(increment=num_steps) diff --git a/tests/unit_tests/test_optimizer_param_scheduler.py b/tests/unit_tests/test_optimizer_param_scheduler.py index 670ca92c992..3c030ec20ff 100644 --- a/tests/unit_tests/test_optimizer_param_scheduler.py +++ b/tests/unit_tests/test_optimizer_param_scheduler.py @@ -297,6 +297,88 @@ def test_load_state_dict(mock_optimizer): assert scheduler.end_wd == 0.2 assert scheduler.wd_incr_steps == 500 assert scheduler.wd_incr_style == 'cosine' + assert mock_optimizer.param_groups[0]['weight_decay'] == pytest.approx(0.2) + + +def test_load_state_dict_override_restores_runtime_param_group_scheduler_values(): + optimizer = MagicMock() + optimizer.param_groups = [ + { + 'lr': 0.0, + 'weight_decay': 0.0, + 'max_lr': 0.2, + 'min_lr': 0.02, + 'start_wd': 0.0, + 'end_wd': 0.1, + 'default_config': True, + }, + { + 'lr': 0.0, + 'weight_decay': 0.0, + 'max_lr': 0.05, + 'min_lr': 0.005, + 'start_wd': 0.01, + 'end_wd': 0.05, + 'default_config': False, + }, + ] + scheduler = OptimizerParamScheduler( + optimizer=optimizer, + init_lr=0.0, + max_lr=0.2, + min_lr=0.02, + lr_warmup_steps=20, + lr_decay_steps=200, + lr_decay_style='linear', + start_wd=0.0, + end_wd=0.1, + wd_incr_steps=200, + wd_incr_style='linear', + use_checkpoint_opt_param_scheduler=False, + override_opt_param_scheduler=True, + ) + + # Simulate optimizer.load_state_dict restoring scheduler-owned fields from + # an older checkpoint before opt_param_scheduler.load_state_dict is called. + optimizer.param_groups[0].update( + {'lr': 1.0, 'max_lr': 1.0, 'min_lr': 0.1, 'start_wd': 0.2, 'end_wd': 0.3} + ) + optimizer.param_groups[1].update( + {'lr': 0.4, 'max_lr': 0.4, 'min_lr': 0.04, 'start_wd': 0.2, 'end_wd': 0.3} + ) + + scheduler.load_state_dict( + { + 'max_lr': 1.0, + 'min_lr': 0.1, + 'lr_warmup_steps': 5, + 'lr_decay_steps': 50, + 'lr_decay_style': 'linear', + 'num_steps': 10, + 'start_wd': 0.01, + 'end_wd': 0.2, + 'wd_incr_steps': 50, + 'wd_incr_style': 'linear', + } + ) + + assert scheduler.max_lr == 0.2 + assert scheduler.min_lr == 0.02 + assert scheduler.lr_warmup_steps == 20 + assert scheduler.lr_decay_steps == 200 + assert scheduler.num_steps == 10 + assert optimizer.param_groups[0]['max_lr'] == 0.2 + assert optimizer.param_groups[0]['min_lr'] == 0.02 + assert optimizer.param_groups[0]['start_wd'] == 0.0 + assert optimizer.param_groups[0]['end_wd'] == 0.1 + assert optimizer.param_groups[0]['lr'] == pytest.approx(0.1) + assert optimizer.param_groups[0]['weight_decay'] == pytest.approx(0.005) + assert optimizer.param_groups[1]['max_lr'] == 0.05 + assert optimizer.param_groups[1]['min_lr'] == 0.005 + assert optimizer.param_groups[1]['start_wd'] == 0.01 + assert optimizer.param_groups[1]['end_wd'] == 0.05 + assert optimizer.param_groups[1]['lr'] == pytest.approx(0.025) + assert optimizer.param_groups[1]['weight_decay'] == pytest.approx(0.012) # ── get_canonical_lr_for_logging tests ──────────────────────────────────────