Fix gradient norm tracking and gradient clipping

pytorch_lightning/accelerators/accelerator.py

@@ -27,7 +27,7 @@
 from pytorch_lightning.trainer.states import TrainerFn
 from pytorch_lightning.utilities import _NATIVE_AMP_AVAILABLE
 from pytorch_lightning.utilities.apply_func import apply_to_collection, move_data_to_device
-from pytorch_lightning.utilities.enums import AMPType, GradClipAlgorithmType, LightningEnum
+from pytorch_lightning.utilities.enums import AMPType, LightningEnum
 from pytorch_lightning.utilities.types import STEP_OUTPUT
 
 if _NATIVE_AMP_AVAILABLE:
@@ -304,17 +304,6 @@ def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Opt
         model_ref = self.lightning_module
         model_ref.optimizer_zero_grad(current_epoch, batch_idx, optimizer, opt_idx)
 
-    def clip_gradients(
-        self,
-        optimizer: Optimizer,
-        clip_val: Union[int, float],
-        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
-    ) -> None:
-        """clips all the optimizer parameters to the given value"""
-        self.precision_plugin.clip_gradients(
-            optimizer, clip_val, gradient_clip_algorithm=gradient_clip_algorithm, model=self.model
-        )
-
     def setup_optimizers(self, trainer: "pl.Trainer") -> None:
         """
         Creates optimizers and schedulers

pytorch_lightning/core/lightning.py

@@ -1414,11 +1414,8 @@ def training_step(...):
             *args: Additional positional arguments to be forwarded to :meth:`~torch.Tensor.backward`
             **kwargs: Additional keyword arguments to be forwarded to :meth:`~torch.Tensor.backward`
         """
-        # make sure we're using manual opt
         self._verify_is_manual_optimization("manual_backward")
-
-        # backward
-        self.trainer.fit_loop.epoch_loop.batch_loop.backward(loss, None, None, *args, **kwargs)
+        self.trainer.accelerator.backward(loss, None, None, *args, **kwargs)
 
     def backward(
         self, loss: Tensor, optimizer: Optional[Optimizer], optimizer_idx: Optional[int], *args, **kwargs

pytorch_lightning/loops/batch/training_batch_loop.py

@@ -16,7 +16,6 @@
 from typing import Any, Callable, List, Optional, Tuple
 
 import numpy as np
-import torch
 from deprecate import void
 from torch import Tensor
 from torch.optim import Optimizer
@@ -246,25 +245,6 @@ def _tbptt_split_batch(self, batch: Any) -> List[Any]:
             splits = model_ref.tbptt_split_batch(batch, tbptt_steps)
         return splits
 
-    # TODO: remove this method and update tests
-    def backward(
-        self,
-        loss: Tensor,
-        optimizer: Optional[torch.optim.Optimizer],
-        opt_idx: Optional[int] = None,
-        *args: Any,
-        **kwargs: Any,
-    ) -> Tensor:
-        """Performs the backward step.
-
-        Args:
-            loss: The loss value to back-propagate on
-            optimizer: Current optimizer being used. ``None`` if using manual optimization.
-            opt_idx: Index of the current optimizer being used. ``None`` if using manual optimization.
-        """
-        self.trainer.accelerator.backward(loss, optimizer, opt_idx, *args, **kwargs)
-        return loss
-
     def _update_running_loss(self, current_loss: Tensor) -> None:
         """Updates the running loss value with the current value"""
         if self.trainer.lightning_module.automatic_optimization:

pytorch_lightning/loops/optimizer/optimizer_loop.py

@@ -14,7 +14,7 @@
 
 from copy import deepcopy
 from functools import partial
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import torch
 from torch import Tensor
@@ -32,7 +32,7 @@
 )
 from pytorch_lightning.trainer.connectors.logger_connector.result import ResultCollection
 from pytorch_lightning.trainer.progress import OptimizationProgress
-from pytorch_lightning.utilities import AMPType, AttributeDict, DeviceType, grad_norm
+from pytorch_lightning.utilities import AMPType, AttributeDict, DeviceType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _TPU_AVAILABLE
 
@@ -91,31 +91,6 @@ def on_run_end(self) -> Tuple[_OUTPUTS_TYPE, Optional[Any]]:
         self._hiddens = None
         return outputs, hiddens
 
-    def backward(
-        self,
-        loss: Tensor,
-        optimizer: Optional[torch.optim.Optimizer],
-        opt_idx: Optional[int] = None,
-        *args: Any,
-        **kwargs: Any,
-    ) -> Tensor:
-        """Performs the backward step.
-
-        Args:
-            loss: The loss value to back-propagate on
-            optimizer: Current optimizer being used. ``None`` if using manual optimization.
-            opt_idx: Index of the current optimizer being used. ``None`` if using manual optimization.
-        """
-        self.trainer.accelerator.backward(loss, optimizer, opt_idx, *args, **kwargs)
-
-        if not self.trainer.fit_loop.should_accumulate():
-            # track gradients
-            grad_norm_dict = self._track_and_norm_grad(optimizer=optimizer)
-            if grad_norm_dict:
-                self.trainer.lightning_module._current_fx_name = "on_after_backward"
-                self.trainer.lightning_module.log_grad_norm(grad_norm_dict)
-        return loss
-
     def _run_optimization(
         self,
         batch_idx: int,
@@ -217,7 +192,7 @@ def _make_backward_fn(
         """
 
         def backward_fn(loss: Tensor):
-            self.backward(loss, optimizer, opt_idx)
+            self.trainer.accelerator.backward(loss, optimizer, opt_idx)
 
             # check if loss or model weights are nan
             if self.trainer.terminate_on_nan:
@@ -355,22 +330,3 @@ def _training_step(
         # the loss will get scaled for amp. avoid any modifications to it
         loss = closure_loss.detach().clone()
         return AttributeDict(closure_loss=closure_loss, loss=loss, result_collection=result_collection)
-
-    def _track_and_norm_grad(self, optimizer: torch.optim.Optimizer) -> Dict[str, float]:
-        """Tracks gradient norms and clips the gradients of all parameters optimized by the current optimizer.
-
-        Args:
-            optimizer: the current optimizer
-        """
-        # track gradient norms
-        grad_norm_dict = {}
-        can_log = (self.trainer.global_step + 1) % self.trainer.log_every_n_steps == 0
-        should_track = float(self.trainer.track_grad_norm) > 0
-        if should_track and can_log:
-            grad_norm_dict = grad_norm(self.trainer.lightning_module, self.trainer.track_grad_norm)
-
-        # clip gradients
-        self.trainer.accelerator.clip_gradients(
-            optimizer, self.trainer.gradient_clip_val, gradient_clip_algorithm=self.trainer.gradient_clip_algorithm
-        )
-        return grad_norm_dict

pytorch_lightning/plugins/precision/apex_amp.py

@@ -97,9 +97,9 @@ def pre_optimizer_step(
         **kwargs: Any,
     ) -> bool:
         """Hook to do something before each optimizer step."""
+        lambda_closure()  # APEX amp does not support closures
         super().pre_optimizer_step(model, optimizer, optimizer_idx, lambda_closure, **kwargs)
         # the following should be in a `optimizer_step` hook but we don't have one in the precision plugin.
-        lambda_closure()  # APEX amp does not support closures
         optimizer.step(**kwargs)
         return False
 

pytorch_lightning/plugins/precision/deepspeed_precision.py

@@ -42,9 +42,9 @@ def pre_optimizer_step(
         **kwargs: Any,
     ) -> bool:
         """Hook to do something before each optimizer step."""
+        lambda_closure()  # DeepSpeed does not support closures
         super().pre_optimizer_step(model, optimizer, optimizer_idx, lambda_closure, **kwargs)
         # the following should be in a `optimizer_step` hook but we don't have one in the precision plugin.
-        lambda_closure()  # DeepSpeed does not support closures
         deepspeed_engine = model.trainer.model
         deepspeed_engine.step()
         return False

pytorch_lightning/plugins/precision/native_amp.py

@@ -97,6 +97,7 @@ def pre_optimizer_step(
                 " To request, please file a Github issue in PyTorch and tag @mcarilli"
             )
         result = True
+        # FIXME: is this correct for manual?
         if model.automatic_optimization:
             result = lambda_closure()
         self.scaler.unscale_(optimizer)

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -21,7 +21,7 @@
 
 import pytorch_lightning as pl
 from pytorch_lightning.core.hooks import CheckpointHooks
-from pytorch_lightning.utilities import GradClipAlgorithmType
+from pytorch_lightning.utilities import grad_norm, GradClipAlgorithmType
 from pytorch_lightning.utilities.types import _PARAMETERS
 
 
@@ -100,9 +100,25 @@ def pre_optimizer_step(
         **kwargs: Any,
     ) -> bool:
         """Hook to do something before each optimizer step."""
+        trainer = model.trainer
+        self._track_and_norm_grad(trainer)
+        self.clip_gradients(
+            optimizer, trainer.gradient_clip_val, gradient_clip_algorithm=trainer.gradient_clip_algorithm, model=model
+        )
         model.trainer.call_hook("on_before_optimizer_step", optimizer, optimizer_idx)
         return True
 
+    def _track_and_norm_grad(self, trainer: "pl.Trainer") -> None:
+        """Tracks the model's gradient norms."""
+        if float(trainer.track_grad_norm) < 0:
+            return
+        grad_norm_dict = grad_norm(trainer.lightning_module, trainer.track_grad_norm)
+        if grad_norm_dict:
+            prev_fx = trainer.lightning_module._current_fx_name
+            trainer.lightning_module._current_fx_name = "on_before_optimizer_step"
+            trainer.lightning_module.log_grad_norm(grad_norm_dict)
+            trainer.lightning_module._current_fx_name = prev_fx
+
     def post_optimizer_step(self, optimizer: Optimizer, optimizer_idx: int) -> None:
         """Hook to do something after each optimizer step."""
 

tests/models/test_hooks.py

@@ -236,7 +236,7 @@ def __init__(self, called):
         super().__init__()
         pl_module_hooks = get_members(LightningModule)
         # remove non-hooks
-        pl_module_hooks.difference_update({"optimizers"})
+        pl_module_hooks.difference_update({"optimizers", "log", "log_dict"})
         # remove most `nn.Module` hooks
         module_hooks = get_members(torch.nn.Module)
         module_hooks.difference_update({"forward", "zero_grad", "train"})
@@ -277,8 +277,10 @@ 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("plugins") == "deepspeed"
+        using_plugin = kwargs.get("amp_backend") or kwargs.get("plugins")
         out = []
         on_before_optimizer_step = [
+            dict(name="log_grad_norm", args=ANY),
             dict(name="Callback.on_before_optimizer_step", args=(trainer, model, ANY, 0)),
             dict(name="on_before_optimizer_step", args=(ANY, 0)),
         ]
@@ -292,10 +294,8 @@ 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, 0)),
                     dict(name="on_train_batch_start", args=(ANY, i, 0)),
-                    # these are before the training step because
-                    # they are not part of the `training_step_and_backward` closure, however,
-                    # with native amp, the closure is run first and then the optimizer step.
-                    *(on_before_optimizer_step if not using_native_amp else []),
+                    # without a precision plugin, we execute the closure inside the `optimizer.step`
+                    *(on_before_optimizer_step if not using_plugin else []),
                     dict(name="forward", args=(ANY,)),
                     dict(name="training_step", args=(ANY, i)),
                     dict(name="training_step_end", args=(dict(loss=ANY),)),
@@ -308,7 +308,7 @@ 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_native_amp else []),
+                    *(on_before_optimizer_step if using_plugin else []),
                     dict(
                         name="optimizer_step",
                         args=(current_epoch, i, ANY, 0, ANY),
@@ -344,6 +344,7 @@ 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="log_grad_norm", args=ANY),
                     dict(name="Callback.on_before_optimizer_step", args=(trainer, model, ANY, 0)),
                     dict(name="on_before_optimizer_step", args=(ANY, 0)),
                     dict(name="training_step", args=(ANY, i)),
@@ -456,6 +457,7 @@ def training_step(self, batch, batch_idx):
         progress_bar_refresh_rate=0,
         weights_summary=None,
         callbacks=[callback],
+        track_grad_norm=1,
         **kwargs,
     )
 
@@ -554,7 +556,9 @@ def training_step(self, batch, batch_idx):
         dict(name="Callback.teardown", args=(trainer, model), kwargs=dict(stage="fit")),
         dict(name="teardown", kwargs=dict(stage="fit")),
     ]
-    assert called == expected
+    assert [c["name"] for c in called] == [e["name"] for e in expected]
+    for c, e in zip(called, expected):
+        assert c["name"] == e["name"]
 
 
 def test_trainer_model_hook_system_fit_no_val_and_resume(tmpdir):