feat: introduce new callback on_before_on_before_optimizer_setup

Author: deependujha

src/lightning/pytorch/callbacks/callback.py

@@ -277,6 +277,14 @@ def on_before_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModu
     def on_after_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         """Called after ``loss.backward()`` and before optimizers are stepped."""
 
+    def on_before_optimizer_setup(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
+        """Called after :meth:`~lightning.pytorch.core.hooks.ModelHooks.configure_model` but before
+        :meth:`~lightning.pytorch.core.hooks.ModelHooks.configure_optimizers`.
+
+        Useful when you need to make changes to the model before the optimizers are set up (e.g. freezing layers).
+
+        """
+
     def on_before_optimizer_step(
         self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", optimizer: Optimizer
     ) -> None:

src/lightning/pytorch/core/hooks.py

@@ -295,6 +295,29 @@ def on_after_backward(self) -> None:
 
         """
 
+    def on_before_optimizer_setup(self) -> None:
+        """Called after :meth:`~lightning.pytorch.core.hooks.ModelHooks.configure_model` but before
+        :meth:`~lightning.pytorch.core.hooks.ModelHooks.configure_optimizers`.
+
+        This hook provides a safe point to modify, freeze, or inspect model parameters before optimizers are created.
+        It’s particularly useful for callbacks such as
+        :class:`~lightning.pytorch.callbacks.finetuning.BaseFinetuning`, where parameters must be frozen
+        prior to optimizer setup.
+
+        This hook runs once in fit stage, after the model
+        has been fully instantiated by ``configure_model``, but before optimizers are created by
+        ``configure_optimizers``.
+
+        Example::
+
+            class MyFinetuneCallback(Callback):
+                def on_before_optimizer_setup(self, trainer, pl_module):
+                    # freeze the backbone before optimizers are created
+                    for param in pl_module.backbone.parameters():
+                        param.requires_grad = False
+
+        """
+
     def on_before_optimizer_step(self, optimizer: Optimizer) -> None:
         """Called before ``optimizer.step()``.
 

src/lightning/pytorch/trainer/trainer.py

@@ -989,6 +989,11 @@ def _run(
         log.debug(f"{self.__class__.__name__}: configuring model")
         call._call_configure_model(self)
 
+        # run hook `on_before_optimizer_setup` before optimizers are set up & after model is configured
+        if self.state.fn == TrainerFn.FITTING:
+            call._call_callback_hooks(self, "on_before_optimizer_setup")
+            call._call_lightning_module_hook(self, "on_before_optimizer_setup")
+
         # check if we should delay restoring checkpoint till later
         if not self.strategy.restore_checkpoint_after_setup:
             log.debug(f"{self.__class__.__name__}: restoring module and callbacks from checkpoint path: {ckpt_path}")

tests/tests_pytorch/callbacks/test_callback_hook_outputs.py renamed to tests/tests_pytorch/callbacks/test_callback_hooks.py

@@ -56,3 +56,44 @@ def on_test_batch_end(self, outputs, *_):
     assert any(isinstance(c, CB) for c in trainer.callbacks)
 
     trainer.fit(model)
+
+
+def test_callback_on_before_optimizer_setup(tmp_path):
+    """Tests that on_before_optimizer_step is called as expected."""
+
+    class CB(Callback):
+        def setup(self, trainer, pl_module, stage=None):
+            assert len(trainer.optimizers) == 0
+            assert pl_module.layer is None  # setup is called before `LightningModule.configure_model`
+
+        def on_before_optimizer_setup(self, trainer, pl_module):
+            assert len(trainer.optimizers) == 0
+            assert pl_module.layer is not None  # called after `LightningModule.configure_model`
+
+        def on_fit_start(self, trainer, pl_module):
+            assert len(trainer.optimizers) == 1
+            assert pl_module.layer is not None  # called after `LightningModule.configure_model`
+
+    class DemoModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.layer = None  # initialize layer in `configure_model`
+
+        def configure_model(self):
+            import torch.nn as nn
+
+            self.layer = nn.Linear(32, 2)
+
+    model = DemoModel()
+
+    trainer = Trainer(
+        callbacks=CB(),
+        default_root_dir=tmp_path,
+        limit_train_batches=2,
+        limit_val_batches=2,
+        max_epochs=1,
+        log_every_n_steps=1,
+        enable_model_summary=False,
+    )
+
+    trainer.fit(model)