Repository: Lightning-AI/pytorch-lightning
Pull request: Return the output of the optimizer step
## What does this PR do?

Fixes #5235

### Does your PR introduce any breaking changes? If yes, please list them.

None

## Before submitting

- [x] Was this **discussed/approved** via a GitHub issue? (not for typos and docs)
- [x] Did you read the [contributor guideline](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/.github/CONTRIBUTING.md), **Pull Request** section?
- [x] Did you make sure your **PR does only one thing**, instead of bundling different changes together?
- [x] Did you make sure to **update the documentation** with your changes? (if necessary)
- [x] Did you write any **new necessary tests**? (not for typos and docs)
- [x] Did you verify new and **existing tests pass** locally with your changes?
- [x] Did you list all the **breaking changes** introduced by this pull request?
- [x] Did you **update the [CHANGELOG](https://github.com/PyTorchLightning/pytorch-lightning/blob/master/CHANGELOG.md)**? (not for typos, docs, test updates, or internal minor changes/refactorings)

## PR review

- [x] Is this pull request ready for review? (if not, please submit in draft mode)
- [x] Check that all items from **Before submitting** are resolved
- [x] Make sure the title is self-explanatory and the description concisely explains the PR
- [x] Add labels and milestones (and optionally projects) to the PR so it can be classified

cc @borda @rohitgr7

Problem statement:
Lighting wrapped optimizers do not allow return values
This is both a bug report /feature request:

Any optimizer.step function might have a return value. However, NO return values are currently handled by the optimizer.LightningOptimizer class in the "step" function. Optimizers with return values  in their step function therefore do not operate properly.

BugFix: Add return value to to the respective line, something like:
```
return_val=None
if make_optimizer_step:
      return_val=self.__optimizer_step(*args, closure=closure, profiler_name=profiler_name, **kwargs)
....

return return_val
```
Also the `__optimizer_step` function must be able to pass on the return value

cc @borda @rohitgr7

Hints from issue discussion:
Hi! thanks for your contribution!, great first issue!
Dear @thoglu,

Thanks for reporting this bug :) 
We will address it soon.

Best regards,
T.C
Hi I wanted to take a look at this but have a couple of questions.

1. What would be the typical use case for returning a value from the step function? Basic torch optimizers return the result of computing the closure (if passed) e.g., [Adam](https://github.com/pytorch/pytorch/blob/master/torch/optim/adam.py). 

    - But in PL, for example in `NativeMixedPrecisionPlugin` closure is called before the optimizer.step and is not passed to step as arg (I think the reason is that `torch.cuda.amp` does not support passing closure to step). If one of the basic torch optimizers was used it would probably return None in this case.

2. Code-wise in PL it might be a little harder now to determine which value should be returned. `Accelerator::optimizer_step` has hooks for precision plugins for `pre` and `post` steps. In case of `NativeMixedPrecisionPlugin` the optimizer.step is actually called in the `post_optimizer_step` step while the `Accelerator::run_optimizer_step` is not executed.

The use case (in my case, but I think more broadly) is for allowing to return optimizer-specific logging values for debug purposes. In Adam, for example, one could return first and second moments, and see how they evolve over time, as these are the central quantities of the optimizer. The logging can then happen in the main training loop function of PL.

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

This issue has been automatically marked as stale because it hasn't had any recent activity. This issue will be closed in 7 days if no further activity occurs. Thank you for your contributions, Pytorch Lightning Team!

Commit message: Skip validaiton

Patch to review:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index c9d862bfe65c1..ac98ac375c4f4 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -75,6 +75,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `opt_idx` to scheduler config if not assigned by user ([#11247](https://github.com/PyTorchLightning/pytorch-lightning/pull/11247))
 
 
+- Return the output of the `optimizer.step`. This can be useful for `LightningLite` users, manual optimization users, or users overriding `LightningModule.optimizer_step` ([#11711](https://github.com/PyTorchLightning/pytorch-lightning/pull/11711))
+
+
 - Teardown the active loop and strategy on exception ([#11620](https://github.com/PyTorchLightning/pytorch-lightning/pull/11620))
 
 
diff --git a/pytorch_lightning/core/lightning.py b/pytorch_lightning/core/lightning.py
index 1b280f12d4166..29386d7475468 100644
--- a/pytorch_lightning/core/lightning.py
+++ b/pytorch_lightning/core/lightning.py
@@ -1546,18 +1546,19 @@ def optimizer_step(
         using_lbfgs: bool = False,
     ) -> None:
         r"""
-        Override this method to adjust the default way the
-        :class:`~pytorch_lightning.trainer.trainer.Trainer` calls each optimizer.
-        By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example
-        once per optimizer. This method (and ``zero_grad()``) won't be called during the
-        accumulation phase when ``Trainer(accumulate_grad_batches != 1)``.
+        Override this method to adjust the default way the :class:`~pytorch_lightning.trainer.trainer.Trainer` calls
+        each optimizer.
+
+        By default, Lightning calls ``step()`` and ``zero_grad()`` as shown in the example once per optimizer.
+        This method (and ``zero_grad()``) won't be called during the accumulation phase when
+        ``Trainer(accumulate_grad_batches != 1)``. Overriding this hook has no benefit with manual optimization.
 
         Args:
             epoch: Current epoch
             batch_idx: Index of current batch
             optimizer: A PyTorch optimizer
             optimizer_idx: If you used multiple optimizers, this indexes into that list.
-            optimizer_closure: Closure for all optimizers. This closure must be executed as it includes the
+            optimizer_closure: The optimizer closure. This closure must be executed as it includes the
                 calls to ``training_step()``, ``optimizer.zero_grad()``, and ``backward()``.
             on_tpu: ``True`` if TPU backward is required
             using_native_amp: ``True`` if using native amp
diff --git a/pytorch_lightning/core/optimizer.py b/pytorch_lightning/core/optimizer.py
index 3825382b2fa35..2279d580ebf6e 100644
--- a/pytorch_lightning/core/optimizer.py
+++ b/pytorch_lightning/core/optimizer.py
@@ -94,13 +94,16 @@ def toggle_model(self, sync_grad: bool = True) -> Generator[None, None, None]:
             yield
             lightning_module.untoggle_optimizer(self._optimizer_idx)
 
-    def step(self, closure: Optional[Callable[[], Any]] = None, **kwargs: Any) -> None:
+    def step(self, closure: Optional[Callable[[], Any]] = None, **kwargs: Any) -> Any:
         """Performs a single optimization step (parameter update).
 
         Args:
-            closure: An optional optimizer_closure.
+            closure: An optional optimizer closure.
             kwargs: Any additional arguments to the ``optimizer.step()`` call.
 
+        Returns:
+            The output from the step call, which is generally the output of the closure execution.
+
         Example::
 
             # Scenario for a GAN using manual optimization
@@ -163,7 +166,7 @@ def closure_dis():
         assert self._strategy is not None
         assert self._strategy.lightning_module is not None
         with self._strategy.lightning_module.trainer.profiler.profile(profiler_action):
-            self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
+            return self._strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
 
 
 def _init_optimizers_and_lr_schedulers(
diff --git a/pytorch_lightning/lite/wrappers.py b/pytorch_lightning/lite/wrappers.py
index b4a1bed0aa17d..340f5d35045ff 100644
--- a/pytorch_lightning/lite/wrappers.py
+++ b/pytorch_lightning/lite/wrappers.py
@@ -56,7 +56,7 @@ def state_dict(self) -> Dict[str, Tensor]:
 
     def step(self, closure: Optional[Callable] = None) -> None:
         closure = closure or _do_nothing_closure
-        self._strategy.optimizer_step(
+        return self._strategy.optimizer_step(
             self.optimizer,
             opt_idx=0,
             closure=closure,
diff --git a/pytorch_lightning/plugins/precision/apex_amp.py b/pytorch_lightning/plugins/precision/apex_amp.py
index 1e86ec2633fe9..9e8b35289bce5 100644
--- a/pytorch_lightning/plugins/precision/apex_amp.py
+++ b/pytorch_lightning/plugins/precision/apex_amp.py
@@ -80,7 +80,7 @@ def optimizer_step(
         optimizer_idx: int,
         closure: Callable[[], Any],
         **kwargs: Any,
-    ) -> None:
+    ) -> Any:
         if isinstance(optimizer, LBFGS):
             raise MisconfigurationException(
                 f"apex AMP and the LBFGS optimizer are not compatible (optimizer {optimizer_idx})."
@@ -90,7 +90,8 @@ def optimizer_step(
         skipped_backward = closure_result is None
         # in manual optimization, the closure does not return a value
         if not isinstance(model, pl.LightningModule) or not model.automatic_optimization or not skipped_backward:
-            optimizer.step(**kwargs)
+            return optimizer.step(**kwargs)
+        return closure_result
 
     def on_load_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
         if "amp_scaling_state" in checkpoint:
diff --git a/pytorch_lightning/plugins/precision/deepspeed.py b/pytorch_lightning/plugins/precision/deepspeed.py
index 3a6eb85769559..b629c931d2fe9 100644
--- a/pytorch_lightning/plugins/precision/deepspeed.py
+++ b/pytorch_lightning/plugins/precision/deepspeed.py
@@ -61,7 +61,7 @@ def optimizer_step(
         optimizer_idx: int,
         closure: Callable[[], Any],
         **kwargs: Any,
-    ) -> None:
+    ) -> Any:
         if isinstance(optimizer, LBFGS):
             raise MisconfigurationException(
                 f"DeepSpeed and the LBFGS optimizer are not compatible (optimizer {optimizer_idx})."
@@ -76,7 +76,7 @@ def optimizer_step(
             )
         # DeepSpeed handles the optimizer step internally
         deepspeed_engine = model.trainer.model if isinstance(model, pl.LightningModule) else model
-        deepspeed_engine.step(**kwargs)
+        return deepspeed_engine.step(**kwargs)
 
     def clip_gradients(
         self,
diff --git a/pytorch_lightning/plugins/precision/ipu.py b/pytorch_lightning/plugins/precision/ipu.py
index 1b321d2c5acd5..cb267000ad551 100644
--- a/pytorch_lightning/plugins/precision/ipu.py
+++ b/pytorch_lightning/plugins/precision/ipu.py
@@ -47,7 +47,7 @@ def optimizer_step(
         optimizer_idx: int,
         closure: Callable[[], Any],
         **kwargs: Any,
-    ) -> None:
+    ) -> Any:
         """IPUs handle the optimizer step internally."""
         if isinstance(optimizer, LBFGS):
             raise MisconfigurationException(
@@ -64,6 +64,7 @@ def optimizer_step(
                 " Please, open an issue in `https://github.com/PyTorchLightning/pytorch-lightning/issues`"
                 " requesting this feature."
             )
+        return closure_result
 
     def clip_gradients(
         self,
diff --git a/pytorch_lightning/plugins/precision/native_amp.py b/pytorch_lightning/plugins/precision/native_amp.py
index f6cb28c76c867..2c9174e5517cb 100644
--- a/pytorch_lightning/plugins/precision/native_amp.py
+++ b/pytorch_lightning/plugins/precision/native_amp.py
@@ -74,7 +74,7 @@ def optimizer_step(
         optimizer_idx: int,
         closure: Callable[[], Any],
         **kwargs: Any,
-    ) -> None:
+    ) -> Any:
         if self.scaler is None:
             # skip scaler logic, as bfloat16 does not require scaler
             return super().optimizer_step(model, optimizer, optimizer_idx, closure, **kwargs)
@@ -90,8 +90,10 @@ def optimizer_step(
         # in manual optimization, the closure does not return a value
         if not isinstance(model, pl.LightningModule) or not model.automatic_optimization or not skipped_backward:
             # note: the scaler will skip the `optimizer.step` if nonfinite gradients are found
-            self.scaler.step(optimizer, **kwargs)
+            step_output = self.scaler.step(optimizer, **kwargs)
             self.scaler.update()
+            return step_output
+        return closure_result
 
     def autocast_context_manager(self) -> Union["old_autocast", "new_autocast"]:
         if _TORCH_GREATER_EQUAL_1_10:
diff --git a/pytorch_lightning/plugins/precision/precision_plugin.py b/pytorch_lightning/plugins/precision/precision_plugin.py
index e875fe51f19e7..6bdfe12dc2a2d 100644
--- a/pytorch_lightning/plugins/precision/precision_plugin.py
+++ b/pytorch_lightning/plugins/precision/precision_plugin.py
@@ -146,11 +146,11 @@ def optimizer_step(
         optimizer_idx: int,
         closure: Callable[[], Any],
         **kwargs: Any,
-    ) -> None:
+    ) -> Any:
         """Hook to run the optimizer step."""
         if isinstance(model, pl.LightningModule):
             closure = partial(self._wrap_closure, model, optimizer, optimizer_idx, closure)
-        optimizer.step(closure=closure, **kwargs)
+        return optimizer.step(closure=closure, **kwargs)
 
     def _track_grad_norm(self, trainer: "pl.Trainer") -> None:
         if trainer.track_grad_norm == -1:
diff --git a/pytorch_lightning/plugins/precision/tpu.py b/pytorch_lightning/plugins/precision/tpu.py
index 8b8871b3006de..1afd34264c60c 100644
--- a/pytorch_lightning/plugins/precision/tpu.py
+++ b/pytorch_lightning/plugins/precision/tpu.py
@@ -36,7 +36,7 @@ def optimizer_step(
         optimizer_idx: int,
         closure: Callable[[], Any],
         **kwargs: Any
-    ) -> None:
+    ) -> Any:
         if isinstance(model, pl.LightningModule):
             closure = partial(self._wrap_closure, model, optimizer, optimizer_idx, closure)
         closure_result = xm.optimizer_step(optimizer, optimizer_args={"closure": closure, **kwargs})
@@ -49,3 +49,4 @@ def optimizer_step(
                 " Please, open an issue in `https://github.com/PyTorchLightning/pytorch-lightning/issues`"
                 " requesting this feature."
             )
+        return closure_result
diff --git a/pytorch_lightning/strategies/strategy.py b/pytorch_lightning/strategies/strategy.py
index 4dd4ca135bcf1..629911911b780 100644
--- a/pytorch_lightning/strategies/strategy.py
+++ b/pytorch_lightning/strategies/strategy.py
@@ -178,7 +178,7 @@ def optimizer_step(
         closure: Callable[[], Any],
         model: Optional[Union["pl.LightningModule", Module]] = None,
         **kwargs: Any,
-    ) -> None:
+    ) -> Any:
         """performs the actual optimizer step.
 
         Args:
@@ -189,7 +189,7 @@ def optimizer_step(
             **kwargs: Any extra arguments to ``optimizer.step``
         """
         model = model or self.lightning_module
-        self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
+        return self.precision_plugin.optimizer_step(model, optimizer, opt_idx, closure, **kwargs)
 
     def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]) -> Tuple[Module, List[Optimizer]]:
         """Setup a model and multiple optimizers together.
diff --git a/tests/core/test_lightning_module.py b/tests/core/test_lightning_module.py
index 63d5ecf670f91..70abaed4492d7 100644
--- a/tests/core/test_lightning_module.py
+++ b/tests/core/test_lightning_module.py
@@ -76,47 +76,6 @@ def test_property_logger(tmpdir):
     assert model.logger == logger
 
 
-def test_params_groups_and_state_are_accessible(tmpdir):
-    class TestModel(BoringModel):
-        def training_step(self, batch, batch_idx, optimizer_idx):
-            output = self.layer(batch)
-            loss = self.loss(batch, output)
-            return {"loss": loss}
-
-        def configure_optimizers(self):
-            optimizer = SGD(self.layer.parameters(), lr=0.1)
-            optimizer_2 = Adam(self.layer.parameters(), lr=0.1)
-            return [optimizer, optimizer_2]
-
-        def optimizer_step(
-            self,
-            epoch,
-            batch_idx,
-            optimizer,
-            optimizer_idx,
-            optimizer_closure,
-            on_tpu=False,
-            using_native_amp=False,
-            using_lbfgs=False,
-        ):
-            # warm up lr
-            if self.trainer.global_step < 500:
-                lr_scale = min(1.0, float(self.trainer.global_step + 1) / 500.0)
-                for pg in optimizer.param_groups:
-                    pg["lr"] = lr_scale * 0.01
-
-            optimizer.step(closure=optimizer_closure)
-
-    model = TestModel()
-    model.training_epoch_end = None
-
-    trainer = Trainer(
-        max_epochs=1, default_root_dir=tmpdir, limit_train_batches=8, limit_val_batches=1, accumulate_grad_batches=1
-    )
-
-    trainer.fit(model)
-
-
 def test_toggle_untoggle_2_optimizers_no_shared_parameters(tmpdir):
     class TestModel(BoringModel):
         def training_step(self, batch, batch_idx, optimizer_idx=None):
diff --git a/tests/core/test_lightning_optimizer.py b/tests/core/test_lightning_optimizer.py
index 970acafa37a41..3b2fe7bfde8cc 100644
--- a/tests/core/test_lightning_optimizer.py
+++ b/tests/core/test_lightning_optimizer.py
@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from unittest.mock import DEFAULT, patch
+from unittest.mock import DEFAULT, Mock, patch
 
 import pytest
 import torch
@@ -95,7 +95,10 @@ def closure(opt):
                 opt_1.step()
 
             closure(opt_2)
-            opt_2.step()
+            step_output = opt_2.step()
+            # check that the step output is returned with manual optimization
+            # since the optimizer is mocked, the step output is a Mock
+            assert isinstance(step_output, Mock)
 
         def configure_optimizers(self):
             optimizer_1 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
@@ -314,3 +317,31 @@ def test_lightning_optimizer_keeps_hooks(tmpdir):
     assert len(optimizer._fwd_handles) == 1
     del lightning_optimizer
     assert len(optimizer._fwd_handles) == 1
+
+
+def test_params_groups_and_state_are_accessible(tmpdir):
+    class TestModel(BoringModel):
+        def training_step(self, batch, batch_idx, optimizer_idx):
+            output = self.layer(batch)
+            loss = self.loss(batch, output)
+            self.__loss = loss
+            return loss
+
+        def configure_optimizers(self):
+            optimizer = SGD(self.layer.parameters(), lr=0.1)
+            optimizer_2 = Adam(self.layer.parameters(), lr=0.1)
+            return [optimizer, optimizer_2]
+
+        def optimizer_step(self, epoch, batch_idx, optimizer, optimizer_idx, optimizer_closure, **__):
+            # check attributes are accessible
+            assert all("lr" in pg for pg in optimizer.param_groups)
+            assert optimizer.state is optimizer._optimizer.state
+            assert optimizer.defaults is optimizer._optimizer.defaults
+
+            loss = optimizer.step(closure=optimizer_closure)
+            # the optimizer step still returns the loss
+            assert loss == self.__loss
+
+    model = TestModel()
+    trainer = Trainer(max_epochs=1, default_root_dir=tmpdir, limit_train_batches=1, limit_val_batches=0)
+    trainer.fit(model)
diff --git a/tests/lite/test_wrappers.py b/tests/lite/test_wrappers.py
index 65522b74936c8..f70a6a9e38230 100644
--- a/tests/lite/test_wrappers.py
+++ b/tests/lite/test_wrappers.py
@@ -155,7 +155,9 @@ def test_lite_optimizer_steps():
     """Test that the LiteOptimizer forwards the step() and zero_grad() calls to the wrapped optimizer."""
     optimizer = Mock()
     strategy = Mock()
+    strategy.optimizer_step.return_value = 123
     lite_optimizer = _LiteOptimizer(optimizer=optimizer, strategy=strategy)
-    lite_optimizer.step()
+    step_output = lite_optimizer.step()
+    assert step_output == 123
     strategy.optimizer_step.assert_called_once()
     strategy.optimizer_step.assert_called_with(optimizer, opt_idx=0, closure=ANY, model=strategy.model)


Code review comments:
diff_hunk: @@ -56,7 +56,7 @@ def state_dict(self) -> Dict[str, Tensor]:
 
     def step(self, closure: Optional[Callable] = None) -> None:
path: pytorch_lightning/lite/wrappers.py
text: @user1:
```suggestion
    def step(self, closure: Optional[Callable] = None) -> Any:
```