fx validator fix

Author: deependujha

src/lightning/pytorch/core/hooks.py

@@ -295,28 +295,28 @@ 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`.
+    # 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 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``.
+    #     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::
+    #     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
+    #         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/connectors/logger_connector/fx_validator.py

@@ -35,6 +35,7 @@ class _LogOptions(TypedDict):
         "on_after_backward": _LogOptions(
             allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
         ),
+        "on_before_optimizer_setup": None,
         "on_before_optimizer_step": _LogOptions(
             allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
         ),

src/lightning/pytorch/trainer/trainer.py

@@ -992,7 +992,7 @@ def _run(
         # 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")
+            # 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:

tests/tests_pytorch/trainer/logging_/test_logger_connector.py

@@ -43,6 +43,7 @@ def test_fx_validator():
     callbacks_func = {
         "on_before_backward",
         "on_after_backward",
+        "on_before_optimizer_setup",
         "on_before_optimizer_step",
         "on_before_zero_grad",
         "on_fit_end",
@@ -83,6 +84,7 @@ def test_fx_validator():
     }
 
     not_supported = {
+        "on_before_optimizer_setup",
         "on_fit_end",
         "on_fit_start",
         "on_exception",
@@ -198,6 +200,7 @@ def test_fx_validator_integration(tmp_path):
         "setup": "You can't",
         "configure_model": "You can't",
         "configure_optimizers": "You can't",
+        "on_before_optimizer_setup": "You can't",
         "on_fit_start": "You can't",
         "train_dataloader": "You can't",
         "val_dataloader": "You can't",