fix for edge case lr lambda handling surfaced w/ PT 2.10

Author: speediedan

src/finetuning_scheduler/fts_supporters.py

@@ -1595,8 +1595,13 @@ def add_optimizer_groups(
                         scheduler.min_lrs.extend([scheduler.min_lrs[0]] * added_pgs)  # type: ignore[attr-defined]
                     else:
                         scheduler.base_lrs.extend([orig_lr_factor] * added_pgs)
-                        if hasattr(scheduler, "lr_lambdas"):
-                            scheduler.lr_lambdas.extend([scheduler.lr_lambdas[-1]] * added_pgs)
+                        if hasattr(scheduler, "lr_lambdas") and scheduler.lr_lambdas:
+                            # due to PyTorch lr scheduler state_dict peculiarities wrt lr_lambdas, lr_lambdas may
+                            # already be pg-aligned (since lr_lambdas are only conditionally saved/restored) see:
+                            # https://bit.ly/lr_lambda_state_dict_special_handling
+                            lambdas_to_sync = max(len(scheduler.base_lrs) - len(scheduler.lr_lambdas), 0)
+                            if lambdas_to_sync:
+                                scheduler.lr_lambdas.extend([scheduler.lr_lambdas[-1]] * lambdas_to_sync)
             else:
                 _ = ScheduleImplMixin._add_groups(no_decay, optimizer, module, thawed_pl, phase_lr)
 

src/finetuning_scheduler/strategy_adapters/base.py

@@ -224,6 +224,8 @@ def _clean_optim_lr_pgs(trainer: Trainer) -> List:
             lrs_cfg.scheduler.last_epoch = -1  # type: ignore[union-attr]
             if not isinstance(lrs_cfg.scheduler, ReduceLROnPlateau):
                 lrs_cfg.scheduler.base_lrs = []
+            # if hasattr(lrs_cfg.scheduler, "lr_lambdas"):
+            #     lrs_cfg.scheduler.lr_lambdas = []
         return orig_num_pgs
 
     def _reconfigure_optimizer_for_phase0(self, trainer: Trainer) -> None:
@@ -250,21 +252,23 @@ def _reconfigure_lrs_for_phase0(self, trainer: Trainer, orig_num_pgs: List) -> N
         Args:
             trainer (Trainer): The :external+pl:class:`~lightning.pytorch.trainer.trainer.Trainer` object.
             orig_num_pgs (List): A list of the number of parameter groups pruned for each optimizer (since only a single
-                optimizer is currently supported by FTS, this list will have only a single element in this verison.)
+                optimizer is currently supported by FTS, this list will have only a single element in this version.)
         """
         # since we may have added parameter groups (e.g. implementing ``no_decay`` for user), we need to reinitialize
         # certain lr_scheduler variables (including type-dependent ones like ``min_lrs`` and ``lr_lambdas``)
         if trainer.lr_scheduler_configs:
             for lrs_cfg in trainer.lr_scheduler_configs:
+                # if hasattr(lrs_cfg.scheduler, "lr_lambdas"):
+                #     lrs_cfg.scheduler.lr_lambdas = lrs_cfg.scheduler.lr_lambdas[orig_num_pgs[0] :]
                 if not isinstance(lrs_cfg.scheduler, ReduceLROnPlateau):
                     lrs_cfg.scheduler._initial_step()
                 lrs_cfg.scheduler._last_lr = [  # type: ignore[union-attr]
                     group["lr"] for group in lrs_cfg.scheduler.optimizer.param_groups
                 ]
                 if isinstance(lrs_cfg.scheduler, ReduceLROnPlateau):
                     lrs_cfg.scheduler.min_lrs = lrs_cfg.scheduler.min_lrs[orig_num_pgs[0] :]
-                elif hasattr(lrs_cfg.scheduler, "lr_lambdas"):
-                    lrs_cfg.scheduler.lr_lambdas = lrs_cfg.scheduler.lr_lambdas[orig_num_pgs[0] :]
+                # elif hasattr(lrs_cfg.scheduler, "lr_lambdas"):
+                #     lrs_cfg.scheduler.lr_lambdas = lrs_cfg.scheduler.lr_lambdas[orig_num_pgs[0] :]
 
     def phase0_optimizer_override(self) -> None:
         """Reconfigure the user-configured optimizer (configured via `configure_optimizers`) to optimize the