[Fix] Fix optimizer parameter override bugs.

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)

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 ──────────────────────────────────────