Fix LightningOptimizer step and toggling logic (#9958)

Author: carmocca

pytorch_lightning/core/optimizer.py

@@ -93,14 +93,6 @@ def _to_lightning_optimizer(cls, optimizer, trainer, opt_idx):
             optimizer = trainer.lightning_optimizers[opt_idx]
         return optimizer
 
-    def _toggle_model(self):
-        model_ref = self._trainer.lightning_module
-        model_ref.toggle_optimizer(self, self._optimizer_idx)
-
-    def _untoggle_model(self):
-        model_ref = self._trainer.lightning_module
-        model_ref.untoggle_optimizer(self)
-
     @contextmanager
     def toggle_model(self, sync_grad: bool = True):
         """This function is just a helper for advanced users.
@@ -116,16 +108,12 @@ def toggle_model(self, sync_grad: bool = True):
         # local import here to avoid circular import
         from pytorch_lightning.loops.utilities import _block_parallel_sync_behavior
 
+        lightning_module = self._trainer.lightning_module
+
         with _block_parallel_sync_behavior(self._trainer, block=(not sync_grad)):
-            self._toggle_model()
+            lightning_module.toggle_optimizer(self, self._optimizer_idx)
             yield
-            self._untoggle_model()
-
-    def __optimizer_step(self, closure: Callable, profiler_name: str = None, **kwargs):
-        trainer = self._trainer
-
-        with trainer.profiler.profile(profiler_name):
-            trainer.accelerator.optimizer_step(self._optimizer, self._optimizer_idx, lambda_closure=closure, **kwargs)
+            lightning_module.untoggle_optimizer(self._optimizer_idx)
 
     def step(self, closure: Optional[Callable] = None, **kwargs):
         """Call this directly from your training_step when doing optimizations manually. By using this we can
@@ -193,14 +181,17 @@ def closure_dis():
                     opt_dis.step(closure=closure_dis)
         """
         if closure is None:
-            profiler_name = f"closure_{self._optimizer_idx}"
             closure = do_nothing_closure
+            profiler_action = "optimizer_step_without_closure"
+        elif not callable(closure):
+            raise MisconfigurationException("When `optimizer.step(closure)` is called, the closure should be callable")
         else:
-            if not callable(closure):
-                raise MisconfigurationException("When closure is provided, it should be a function")
-            profiler_name = f"optimizer_step_and_closure_{self._optimizer_idx}"
+            profiler_action = "optimizer_step_with_closure"
+        profiler_action += f"_{self._optimizer_idx}"
 
-        self.__optimizer_step(closure=closure, profiler_name=profiler_name, **kwargs)
+        trainer = self._trainer
+        with trainer.profiler.profile(profiler_action):
+            trainer.accelerator.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
         self._total_optimizer_step_calls += 1
 
     def __repr__(self):

pytorch_lightning/profiler/pytorch.py

@@ -117,11 +117,11 @@ def pre_step(self, current_action: str) -> None:
 
     def reset(self):
         # handle properly `fast_dev_run`. PyTorch Profiler will fail otherwise.
-        self._num_optimizer_step_and_closure = 0
+        self._num_optimizer_step_with_closure = 0
         self._num_validation_step = 0
         self._num_test_step = 0
         self._num_predict_step = 0
-        self._optimizer_step_and_closure_reached_end = False
+        self._optimizer_step_with_closure_reached_end = False
         self._validation_step_reached_end = False
         self._test_step_reached_end = False
         self._predict_step_reached_end = False
@@ -132,13 +132,13 @@ def reset(self):
     @property
     def is_training(self) -> bool:
         return self._current_action is not None and (
-            self._current_action.startswith("optimizer_step_and_closure_") or self._current_action == "training_step"
+            self._current_action.startswith("optimizer_step_with_closure_") or self._current_action == "training_step"
         )
 
     @property
     def num_step(self) -> int:
         if self.is_training:
-            return self._num_optimizer_step_and_closure
+            return self._num_optimizer_step_with_closure
         if self._current_action == "validation_step":
             return self._num_validation_step
         if self._current_action == "test_step":
@@ -149,10 +149,10 @@ def num_step(self) -> int:
 
     def _step(self) -> None:
         if self.is_training:
-            self._num_optimizer_step_and_closure += 1
+            self._num_optimizer_step_with_closure += 1
         elif self._current_action == "validation_step":
             if self._start_action_name == "on_fit_start":
-                if self._num_optimizer_step_and_closure > 0:
+                if self._num_optimizer_step_with_closure > 0:
                     self._num_validation_step += 1
             else:
                 self._num_validation_step += 1
@@ -164,7 +164,7 @@ def _step(self) -> None:
     @property
     def has_finished(self) -> bool:
         if self.is_training:
-            return self._optimizer_step_and_closure_reached_end
+            return self._optimizer_step_with_closure_reached_end
         if self._current_action == "validation_step":
             return self._validation_step_reached_end
         if self._current_action == "test_step":
@@ -182,7 +182,7 @@ def __call__(self, num_step: int) -> "ProfilerAction":
         action = self._schedule(max(self.num_step, 0))
         if action == ProfilerAction.RECORD_AND_SAVE:
             if self.is_training:
-                self._optimizer_step_and_closure_reached_end = True
+                self._optimizer_step_with_closure_reached_end = True
             elif self._current_action == "validation_step":
                 self._validation_step_reached_end = True
             elif self._current_action == "test_step":
@@ -202,9 +202,9 @@ class PyTorchProfiler(BaseProfiler):
         "test_step",
         "predict_step",
     }
-    RECORD_FUNCTION_PREFIX = "optimizer_step_and_closure_"
+    RECORD_FUNCTION_PREFIX = "optimizer_step_with_closure_"
     STEP_FUNCTIONS = {"training_step", "validation_step", "test_step", "predict_step"}
-    STEP_FUNCTION_PREFIX = "optimizer_step_and_closure_"
+    STEP_FUNCTION_PREFIX = "optimizer_step_with_closure_"
     AVAILABLE_SORT_KEYS = {
         "cpu_time",
         "cuda_time",
@@ -383,8 +383,8 @@ def start(self, action_name: str) -> None:
                 self._register.__enter__()
 
         if self._lightning_module is not None:
-            # when the model is used in automatic optimization,
-            # we use `optimizer_step_and_closure` to step the model.
+            # when the model is used in automatic optimization, we use `optimizer_step_with_closure` to step the model.
+            # this profiler event is generated in the `LightningOptimizer.step` method
             if self._lightning_module.automatic_optimization and "training_step" in self.STEP_FUNCTIONS:
                 self.STEP_FUNCTIONS.remove("training_step")
 

tests/profiler/test_profiler.py

@@ -285,7 +285,7 @@ def test_pytorch_profiler_trainer_ddp(tmpdir, pytorch_profiler):
         files = [file for file in files if file.endswith(".json")]
         assert len(files) == 2, files
         local_rank = trainer.local_rank
-        assert any(f"{local_rank}-optimizer_step_and_closure_" in f for f in files)
+        assert any(f"{local_rank}-optimizer_step_with_closure_" in f for f in files)
         assert any(f"{local_rank}-validation_step" in f for f in files)