Implement double-closure

Author: carmocca

CHANGELOG.md

@@ -665,6 +665,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Fixed gradients not being unscaled when clipping or logging the gradient norm ([#9287](https://github.com/PyTorchLightning/pytorch-lightning/pull/9287))
 
 
+- Fixed `on_before_optimizer_step` getting called before the optimizer closure (including backward) has run ([#10167](https://github.com/PyTorchLightning/pytorch-lightning/pull/10167))
+
+
 - Fixed monitor value in `ModelCheckpoint` getting moved to the wrong device in a special case where it becomes NaN ([#10118](https://github.com/PyTorchLightning/pytorch-lightning/pull/10118))
 
 

pytorch_lightning/accelerators/accelerator.py

@@ -333,6 +333,11 @@ def optimizer_step(
         """
         model = model or self.lightning_module
         self.precision_plugin.optimizer_step(model, optimizer, opt_idx, lambda_closure, **kwargs)
+
+        if not isinstance(model, pl.LightningModule):
+            # gradient clipping and norm tracking only available with a LightingModule/Trainer
+            return
+
         trainer = model.trainer
         assert isinstance(trainer, pl.Trainer)
         # TODO: this is done for the entire model but should be changed to per-optimizer

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -110,6 +110,28 @@ def _run_backward(self, tensor: Tensor, model: Optional[Module], *args: Any, **k
         """
         tensor.backward(*args, **kwargs)
 
+    def _wrap_closure(
+        self,
+        model: Union["pl.LightningModule", Module],
+        optimizer: Optimizer,
+        optimizer_idx: int,
+        lambda_closure: Callable[[], Any],
+    ) -> Callable[[], Any]:
+        """This double-closure allows makes sure the ``lambda_closure`` is executed before the
+        ``on_before_optimizer_step`` hook is called.
+
+        The closure (generally) runs ``backward`` so this allows inspecting gradients in this hook. This structure is
+        consistent with the ``PrecisionPlugin`` subclasses that cannot pass ``optimizer.step(lambda_closure)`` directly.
+        """
+
+        def inner() -> Any:
+            closure_result = lambda_closure()
+            if isinstance(model, pl.LightningModule):
+                model.trainer.call_hook("on_before_optimizer_step", optimizer, optimizer_idx)
+            return closure_result
+
+        return inner
+
     def optimizer_step(
         self,
         model: Union["pl.LightningModule", Module],
@@ -119,12 +141,11 @@ def optimizer_step(
         **kwargs: Any,
     ) -> None:
         """Hook to run the optimizer step."""
-        if isinstance(model, pl.LightningModule):
-            model.trainer.call_hook("on_before_optimizer_step", optimizer, optimizer_idx)
-        optimizer.step(closure=lambda_closure, **kwargs)
+        closure = self._wrap_closure(model, optimizer, optimizer_idx, lambda_closure)
+        optimizer.step(closure=closure, **kwargs)
 
     def _track_grad_norm(self, trainer: "pl.Trainer") -> None:
-        if float(trainer.track_grad_norm) == -1:
+        if trainer.track_grad_norm == -1:
             return
         grad_norm_dict = grad_norm(trainer.lightning_module, trainer.track_grad_norm, trainer.logger.group_separator)
         if grad_norm_dict:

pytorch_lightning/plugins/precision/tpu.py

@@ -34,9 +34,8 @@ def optimizer_step(
         lambda_closure: Callable[[], Any],
         **kwargs: Any
     ) -> None:
-        if isinstance(model, pl.LightningModule):
-            model.trainer.call_hook("on_before_optimizer_step", optimizer, optimizer_idx)
-        closure_result = xm.optimizer_step(optimizer, optimizer_args={"closure": lambda_closure, **kwargs})
+        closure = self._wrap_closure(model, optimizer, optimizer_idx, lambda_closure)
+        closure_result = xm.optimizer_step(optimizer, optimizer_args={"closure": closure, **kwargs})
         skipped_backward = closure_result is None
         # in manual optimization, the closure does not return a value
         if isinstance(model, pl.LightningModule) and model.automatic_optimization and skipped_backward:

tests/models/test_hooks.py

@@ -275,12 +275,7 @@ def _train_batch(self, *args, **kwargs):
     def _auto_train_batch(trainer, model, batches, device=torch.device("cpu"), current_epoch=0, **kwargs):
         using_native_amp = kwargs.get("amp_backend") == "native"
         using_deepspeed = kwargs.get("strategy") == "deepspeed"
-        using_plugin = kwargs.get("amp_backend") or kwargs.get("strategy")
         out = []
-        on_before_optimizer_step = [
-            dict(name="Callback.on_before_optimizer_step", args=(trainer, model, ANY, 0)),
-            dict(name="on_before_optimizer_step", args=(ANY, 0)),
-        ]
         for i in range(batches):
             out.extend(
                 [
@@ -291,8 +286,6 @@ def _auto_train_batch(trainer, model, batches, device=torch.device("cpu"), curre
                     dict(name="Callback.on_batch_start", args=(trainer, model)),
                     dict(name="Callback.on_train_batch_start", args=(trainer, model, ANY, i)),
                     dict(name="on_train_batch_start", args=(ANY, i)),
-                    # without a precision plugin, we execute the closure inside the `optimizer.step`
-                    *([] if using_plugin else on_before_optimizer_step),
                     dict(name="forward", args=(ANY,)),
                     dict(name="training_step", args=(ANY, i)),
                     dict(name="training_step_end", args=(dict(loss=ANY),)),
@@ -305,7 +298,9 @@ def _auto_train_batch(trainer, model, batches, device=torch.device("cpu"), curre
                     *([dict(name="backward", args=(ANY, ANY, 0))] if not using_deepspeed else []),
                     dict(name="Callback.on_after_backward", args=(trainer, model)),
                     dict(name="on_after_backward"),
-                    *(on_before_optimizer_step if using_plugin else []),
+                    # note: unscaling happens here in the case of AMP
+                    dict(name="Callback.on_before_optimizer_step", args=(trainer, model, ANY, 0)),
+                    dict(name="on_before_optimizer_step", args=(ANY, 0)),
                     *([dict(name="log_grad_norm", args=ANY)] if not using_deepspeed else []),
                     dict(
                         name="clip_gradients",
@@ -334,7 +329,6 @@ def _auto_train_batch(trainer, model, batches, device=torch.device("cpu"), curre
     @staticmethod
     def _manual_train_batch(trainer, model, batches, device=torch.device("cpu"), **kwargs):
         using_deepspeed = kwargs.get("strategy") == "deepspeed"
-        using_plugin = kwargs.get("amp_backend") or kwargs.get("strategy")
         out = []
         for i in range(batches):
             out.extend(
@@ -355,11 +349,9 @@ def _manual_train_batch(trainer, model, batches, device=torch.device("cpu"), **k
                     dict(name="on_after_backward"),
                     # `manual_backward` calls the previous 3
                     dict(name="manual_backward", args=(ANY,)),
-                    *([dict(name="closure")] if using_plugin else []),
+                    dict(name="closure"),
                     dict(name="Callback.on_before_optimizer_step", args=(trainer, model, ANY, 0)),
                     dict(name="on_before_optimizer_step", args=(ANY, 0)),
-                    # without a precision plugin, we execute the closure inside the `optimizer.step`
-                    *([] if using_plugin else [dict(name="closure")]),
                     *([dict(name="log_grad_norm", args=ANY)] if not using_deepspeed else []),
                     dict(name="training_step", args=(ANY, i)),
                     dict(name="training_step_end", args=(dict(loss=ANY),)),