Update clip gradients signature for precision plugins

CHANGELOG.md

@@ -79,6 +79,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `artifact_location` argument to `MLFlowLogger` which will be passed to the `MlflowClient.create_experiment` call ([#6677](https://github.com/PyTorchLightning/pytorch-lightning/pull/6677))
 
 
+- Added `model` parameter to precision plugins' `clip_gradients` signature ([#6764](https://github.com/PyTorchLightning/pytorch-lightning/pull/6764))
+
+
 ### Changed
 
 - Renamed `pytorch_lightning.callbacks.swa` to `pytorch_lightning.callbacks.stochastic_weight_avg` ([#6259](https://github.com/PyTorchLightning/pytorch-lightning/pull/6259))

pytorch_lightning/accelerators/accelerator.py

@@ -318,7 +318,7 @@ def optimizer_zero_grad(self, current_epoch: int, batch_idx: int, optimizer: Opt
     def clip_gradients(self, optimizer: Optimizer, clip_val: Union[int, float]) -> None:
         """clips all the optimizer parameters to the given value"""
 
-        self.precision_plugin.clip_gradients(optimizer, clip_val)
+        self.precision_plugin.clip_gradients(self.model, optimizer, clip_val)
 
     def on_train_epoch_end(self, outputs: Sequence[_STEP_OUTPUT_TYPE]) -> None:
         """Hook to do something on the end of an training epoch

pytorch_lightning/plugins/precision/deepspeed_precision.py

@@ -75,7 +75,13 @@ def backward(
 
         return closure_loss
 
-    def clip_gradients(self, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0) -> None:
+    def clip_gradients(
+        self,
+        model: 'LightningModule',
+        optimizer: 'Optimizer',
+        clip_val: Union[int, float],
+        norm_type: float = 2.0
+    ) -> None:
         """
         DeepSpeed handles clipping gradients via the training type plugin.
         """

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -98,7 +98,13 @@ def pre_optimizer_step(
     def post_optimizer_step(self, optimizer: 'Optimizer', optimizer_idx: int) -> None:
         """Hook to do something after each optimizer step."""
 
-    def clip_gradients(self, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0) -> None:
+    def clip_gradients(
+        self,
+        model: 'LightningModule',
+        optimizer: 'Optimizer',
+        clip_val: Union[int, float],
+        norm_type: float = 2.0
+    ) -> None:
         """Clips the gradients to a specific value"""
         if clip_val is None:
             return

pytorch_lightning/plugins/precision/sharded_native_amp.py

@@ -23,6 +23,8 @@
 if TYPE_CHECKING:
     from torch.optim import Optimizer
 
+    from pytorch_lightning.core import LightningModule
+
 
 class ShardedNativeMixedPrecisionPlugin(NativeMixedPrecisionPlugin):
     """Mixed Precision for Sharded Training
@@ -32,7 +34,13 @@ def __init__(self) -> None:
         super().__init__()
         self.scaler = ShardedGradScaler()
 
-    def clip_gradients(self, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0) -> None:
+    def clip_gradients(
+        self,
+        model: 'LightningModule',
+        optimizer: 'Optimizer',
+        clip_val: Union[int, float],
+        norm_type: float = 2.0
+    ) -> None:
         if clip_val <= 0:
             return
 

tests/plugins/test_precision_plugin.py

@@ -0,0 +1,26 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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 inspect import signature
+
+from pytorch_lightning.plugins.precision import PrecisionPlugin
+
+
+def test_precision_clip_gradients_signature():
+
+    expected_params_list = ['self', 'model', 'optimizer', 'clip_val', 'norm_type']
+
+    params = signature(PrecisionPlugin.clip_gradients).parameters
+    params_list = [param.name for param in params.values()]
+
+    assert params_list == expected_params_list