Support optimizer step progress tracking with manual optimization

Author: carmocca

pytorch_lightning/core/optimizer.py

@@ -161,10 +161,22 @@ def closure_dis():
         profiler_action += f"_{self._optimizer_idx}"
 
         assert self._strategy is not None
-        assert self._strategy.lightning_module is not None
-        with self._strategy.lightning_module.trainer.profiler.profile(profiler_action):
+        lightning_module = self._strategy.lightning_module
+        assert lightning_module is not None
+        trainer = lightning_module.trainer
+        batch_loop = trainer.fit_loop.epoch_loop.batch_loop
+        if lightning_module.automatic_optimization:
+            progress = batch_loop.optimizer_loop.optim_progress
+        else:
+            progress = batch_loop.manual_loop.optim_progress
+
+        progress.optimizer.step.increment_ready()
+
+        with trainer.profiler.profile(profiler_action):
             self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
 
+        progress.optimizer.step.increment_completed()
+
 
 def _init_optimizers_and_lr_schedulers(
     model: "pl.LightningModule",

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -89,7 +89,10 @@ def batch_idx(self) -> int:
 
     @property
     def global_step(self) -> int:
-        return self.batch_loop.optimizer_loop.optim_progress.optimizer_steps
+        lightning_module = self.trainer.lightning_module
+        if lightning_module is None or lightning_module.automatic_optimization:
+            return self.batch_loop.optimizer_loop.optim_progress.optimizer_steps
+        return self.batch_loop.manual_loop.optim_progress.optimizer_steps
 
     @property
     def _is_training_done(self) -> bool:

pytorch_lightning/loops/optimization/manual_loop.py

@@ -19,6 +19,7 @@
 from pytorch_lightning.loops import Loop
 from pytorch_lightning.loops.optimization.closure import OutputResult
 from pytorch_lightning.loops.utilities import _build_training_step_kwargs, _extract_hiddens
+from pytorch_lightning.trainer.progress import OptimizationProgress
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 
@@ -74,6 +75,9 @@ class ManualOptimization(Loop[_OUTPUTS_TYPE]):
 
     def __init__(self) -> None:
         super().__init__()
+        # FIXME: should this be a simpler progress? lr schedulers are not wrapped anyways
+        self.optim_progress = OptimizationProgress()
+
         self._done: bool = False
         self._hiddens: Optional[Any] = None
         self._output: _OUTPUTS_TYPE = {}

pytorch_lightning/loops/optimization/optimizer_loop.py

@@ -359,8 +359,6 @@ def _optimizer_step(
         else:
             optimizer = self.trainer.strategy._lightning_optimizers[opt_idx]
 
-        self.optim_progress.optimizer.step.increment_ready()
-
         # model hook
         self.trainer._call_lightning_module_hook(
             "optimizer_step",
@@ -374,8 +372,6 @@ def _optimizer_step(
             using_lbfgs=is_lbfgs,
         )
 
-        self.optim_progress.optimizer.step.increment_completed()
-
     def _on_before_zero_grad(self, optimizer: torch.optim.Optimizer) -> None:
         """Calls the ``on_before_zero_grad`` hook.
 

pytorch_lightning/trainer/trainer.py

@@ -2010,6 +2010,7 @@ def data_parallel_device_ids(self) -> Optional[List[int]]:
 
     @property
     def lightning_module(self) -> "pl.LightningModule":
+        # TODO: this is actually an optional return
         return self.strategy.lightning_module
 
     @property

tests/loops/test_loop_state_dict.py

@@ -56,6 +56,16 @@ def test_loops_state_dict_structure():
             },
             "epoch_loop.batch_loop.state_dict": {},
             "epoch_loop.batch_loop.manual_loop.state_dict": {},
+            "epoch_loop.batch_loop.manual_loop.optim_progress": {
+                "optimizer": {
+                    "step": {"total": {"ready": 0, "completed": 0}, "current": {"ready": 0, "completed": 0}},
+                    "zero_grad": {
+                        "total": {"ready": 0, "started": 0, "completed": 0},
+                        "current": {"ready": 0, "started": 0, "completed": 0},
+                    },
+                },
+                "optimizer_position": 0,
+            },
             "epoch_loop.batch_loop.optimizer_loop.state_dict": {},
             "epoch_loop.batch_loop.optimizer_loop.optim_progress": {
                 "optimizer": {

tests/loops/test_loops.py

@@ -512,6 +512,16 @@ def configure_optimizers_multiple(self):
         },
         "epoch_loop.batch_loop.state_dict": ANY,
         "epoch_loop.batch_loop.manual_loop.state_dict": ANY,
+        "epoch_loop.batch_loop.manual_loop.optim_progress": {
+            "optimizer_position": 0,
+            "optimizer": {
+                "step": {"total": {"ready": 0, "completed": 0}, "current": {"ready": 0, "completed": 0}},
+                "zero_grad": {
+                    "total": {"ready": 0, "started": 0, "completed": 0},
+                    "current": {"ready": 0, "started": 0, "completed": 0},
+                },
+            },
+        },
         "epoch_loop.batch_loop.optimizer_loop.state_dict": {},
         "epoch_loop.batch_loop.optimizer_loop.optim_progress": {
             "optimizer_position": stop_optimizer,
@@ -680,6 +690,16 @@ def train_dataloader(self):
         },
         "epoch_loop.batch_loop.state_dict": ANY,
         "epoch_loop.batch_loop.manual_loop.state_dict": ANY,
+        "epoch_loop.batch_loop.manual_loop.optim_progress": {
+            "optimizer_position": 0,
+            "optimizer": {
+                "step": {"total": {"ready": 0, "completed": 0}, "current": {"ready": 0, "completed": 0}},
+                "zero_grad": {
+                    "total": {"ready": 0, "started": 0, "completed": 0},
+                    "current": {"ready": 0, "started": 0, "completed": 0},
+                },
+            },
+        },
         "epoch_loop.batch_loop.optimizer_loop.state_dict": {},
         "epoch_loop.batch_loop.optimizer_loop.optim_progress": {
             "optimizer_position": n_optimizers,

tests/models/test_hooks.py

@@ -558,7 +558,6 @@ def training_step(self, batch, batch_idx):
         dict(name="on_validation_model_train"),
         dict(name="training_epoch_end", args=([dict(loss=ANY)] * train_batches,)),
         dict(name="Callback.on_train_epoch_end", args=(trainer, model)),
-        # FIXME: there seems to be a problem with manual here
         # `ModelCheckpoint.save_checkpoint` is called here from `Callback.on_train_epoch_end`
         dict(name="Callback.on_save_checkpoint", args=(trainer, model, saved_ckpt)),
         dict(name="on_save_checkpoint", args=(saved_ckpt,)),

tests/trainer/optimization/test_manual_optimization.py

@@ -165,6 +165,7 @@ def training_epoch_end(self, outputs) -> None:
     with mock.patch.object(Strategy, "backward", wraps=trainer.strategy.backward) as bwd_mock:
         trainer.fit(model)
     assert bwd_mock.call_count == limit_train_batches * 3
+    assert trainer.global_step == limit_train_batches * 2
 
 
 def test_multiple_optimizers_manual_log(tmpdir):
@@ -524,18 +525,14 @@ def optimizer_closure():
             weight_after = self.layer.weight.clone()
             assert not torch.equal(weight_before, weight_after)
 
-        def configure_optimizers(self):
-            return torch.optim.SGD(self.layer.parameters(), lr=0.1)
-
     model = TestModel()
-    model.val_dataloader = None
     model.training_epoch_end = None
 
     limit_train_batches = 2
     trainer = Trainer(
         default_root_dir=tmpdir,
         limit_train_batches=limit_train_batches,
-        limit_val_batches=2,
+        limit_val_batches=0,
         max_epochs=1,
         log_every_n_steps=1,
     )
@@ -547,58 +544,45 @@ def configure_optimizers(self):
     assert trainer.progress_bar_metrics["train_loss_epoch"] == torch.stack(model._losses).mean()
 
 
-def test_step_with_optimizer_closure_and_accumulated_grad(tmpdir):
-    """Tests that `step` works with optimizer_closure and accumulated_grad."""
-
+def test_step_with_optimizer_closure_2(tmpdir):
     class TestModel(BoringModel):
         def __init__(self):
             super().__init__()
             self.automatic_optimization = False
 
         def training_step(self, batch, batch_idx):
-            # manual
             opt = self.optimizers()
             x = batch[0]
-
-            loss_1 = self(x)
-            loss_1 = self.loss(loss_1, loss_1)
+            loss = self(x).sum()
 
             def optimizer_closure():
                 # emulate bayesian optimization.
                 num_backward = 1
                 for backward_idx in range(num_backward + 1):
                     retain_graph = num_backward != backward_idx
-                    self.manual_backward(loss_1, retain_graph=retain_graph)
+                    self.manual_backward(loss, retain_graph=retain_graph)
 
             weight_before = self.layer.weight.clone()
-
             opt.step(closure=optimizer_closure)
-
             weight_after = self.layer.weight.clone()
-            if not self.trainer.fit_loop._should_accumulate():
-                assert not torch.equal(weight_before, weight_after)
-            else:
-                assert self.layer.weight.grad is not None
-
-        def configure_optimizers(self):
-            return torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            assert not torch.equal(weight_before, weight_after)
 
     model = TestModel()
-    model.val_dataloader = None
     model.training_epoch_end = None
 
     limit_train_batches = 4
     trainer = Trainer(
         default_root_dir=tmpdir,
         limit_train_batches=limit_train_batches,
-        limit_val_batches=2,
+        limit_val_batches=0,
         max_epochs=1,
         log_every_n_steps=1,
     )
 
     with mock.patch.object(Strategy, "backward", wraps=trainer.strategy.backward) as bwd_mock:
         trainer.fit(model)
     assert bwd_mock.call_count == limit_train_batches * 2
+    assert trainer.global_step == limit_train_batches
 
 
 @patch("torch.optim.SGD.step")
@@ -614,41 +598,23 @@ def on_train_start(self) -> None:
             step_mock.reset_mock()
 
         def training_step(self, batch, batch_idx):
-            # manual
             opt = self.optimizers()
-            x = batch[0]
-
-            loss_1 = self(x)
-            loss_1 = self.loss(loss_1, loss_1)
-
-            def optimizer_closure():
-                # emulate bayesian optimization.
-                num_backward = 1
-                for backward_idx in range(num_backward + 1):
-                    retain_graph = num_backward != backward_idx
-                    self.manual_backward(loss_1, retain_graph=retain_graph)
-
-            opt.step(closure=optimizer_closure)
-            opt.zero_grad()
-
-        def configure_optimizers(self):
-            return torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            opt.step(closure=lambda: ..., foo=123)
 
     model = TestModel()
-    model.val_dataloader = None
     model.training_epoch_end = None
 
-    limit_train_batches = 4
+    limit_train_batches = 2
     trainer = Trainer(
         default_root_dir=tmpdir,
         limit_train_batches=limit_train_batches,
-        limit_val_batches=2,
+        limit_val_batches=0,
         max_epochs=1,
-        log_every_n_steps=1,
     )
 
     trainer.fit(model)
-    assert step_mock.mock_calls == [call(closure=ANY) for _ in range(limit_train_batches)]
+    assert step_mock.mock_calls == [call(closure=ANY, foo=123) for _ in range(limit_train_batches)]
+    assert trainer.global_step == limit_train_batches
 
 
 @patch("torch.optim.Adam.step")
@@ -724,6 +690,7 @@ def configure_optimizers(self):
     trainer.fit(model)
     assert mock_sgd_step.mock_calls == [call(closure=ANY, optim="sgd") for _ in range(4)]
     assert mock_adam_step.mock_calls == [call(closure=ANY) for _ in range(2)]
+    assert trainer.global_step == 4 + 2
 
 
 class TesManualOptimizationDDPModel(BoringModel):