Minor fixes related to clipping (Lightning-AI#10130)

Co-authored-by: Adrian Wälchli <[email protected]>
Co-authored-by: Rohit Gupta <[email protected]>

Author: 3 people

pytorch_lightning/accelerators/accelerator.py

@@ -329,6 +329,7 @@ def optimizer_step(
             opt_idx: index of the current optimizer
             lambda_closure: closure calculating the loss value
             model: reference to the model, optionally defining optimizer step related hooks
+            **kwargs: Any extra arguments to ``optimizer.step``
         """
         model = model or self.lightning_module
         make_optimizer_step = self.precision_plugin.pre_optimizer_step(
@@ -349,9 +350,7 @@ def clip_gradients(
         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
-        )
+        self.precision_plugin.clip_gradients(optimizer, clip_val, gradient_clip_algorithm=gradient_clip_algorithm)
 
     def setup_optimizers(self, trainer: "pl.Trainer") -> None:
         """Creates optimizers and schedulers.

pytorch_lightning/core/lightning.py

@@ -1411,10 +1411,7 @@ 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.accelerator.backward(loss, None, None, *args, **kwargs)
 
     def backward(
@@ -1487,7 +1484,7 @@ def clip_gradients(
         self,
         optimizer: Optimizer,
         gradient_clip_val: Optional[Union[int, float]] = None,
-        gradient_clip_algorithm: Optional[Union[str, GradClipAlgorithmType]] = None,
+        gradient_clip_algorithm: Optional[str] = None,
     ):
         """Handles gradient clipping internally.
 
@@ -1505,8 +1502,9 @@ def clip_gradients(
             gradient_clip_val = self.trainer.gradient_clip_val or 0.0
         elif self.trainer.gradient_clip_val is not None and self.trainer.gradient_clip_val != gradient_clip_val:
             raise MisconfigurationException(
-                "You have set `Trainer(gradient_clip_val)` and have passed"
-                " `gradient_clip_val` inside `clip_gradients`. Please use only one of them."
+                f"You have set `Trainer(gradient_clip_val={self.trainer.gradient_clip_val!r})`"
+                f" and have passed `clip_gradients(gradient_clip_val={gradient_clip_val!r})`."
+                " Please use only one of them."
             )
 
         if gradient_clip_algorithm is None:
@@ -1518,8 +1516,9 @@ def clip_gradients(
                 and self.trainer.gradient_clip_algorithm != gradient_clip_algorithm
             ):
                 raise MisconfigurationException(
-                    "You have set `Trainer(gradient_clip_algorithm)` and have passed"
-                    " `gradient_clip_algorithm` inside `clip_gradients`. Please use only one of them."
+                    f"You have set `Trainer(gradient_clip_algorithm={self.trainer.gradient_clip_algorithm.value!r})`"
+                    f" and have passed `clip_gradients(gradient_clip_algorithm={gradient_clip_algorithm!r})"
+                    " Please use only one of them."
                 )
 
         if not isinstance(gradient_clip_val, (int, float)):
@@ -1543,10 +1542,6 @@ def configure_gradient_clipping(
     ):
         """Perform gradient clipping for the optimizer parameters. Called before :meth:`optimizer_step`.
 
-        Note:
-            This hook won't be called when using deepspeed since it handles gradient clipping internally.
-            Consider setting ``gradient_clip_val`` and ``gradient_clip_algorithm`` inside ``Trainer``."
-
         Args:
             optimizer: Current optimizer being used.
             optimizer_idx: Index of the current optimizer being used.

pytorch_lightning/loops/optimization/optimizer_loop.py

@@ -488,11 +488,10 @@ def _track_and_norm_grad(self, optimizer: torch.optim.Optimizer, opt_idx: int) -
             )
 
         # clip gradients
-        if not self.trainer.accelerator_connector.use_deepspeed:
-            self.trainer.lightning_module.configure_gradient_clipping(
-                optimizer,
-                opt_idx,
-                gradient_clip_val=self.trainer.gradient_clip_val,
-                gradient_clip_algorithm=self.trainer.gradient_clip_algorithm,
-            )
+        self.trainer.lightning_module.configure_gradient_clipping(
+            optimizer,
+            opt_idx,
+            gradient_clip_val=self.trainer.gradient_clip_val,
+            gradient_clip_algorithm=self.trainer.gradient_clip_algorithm,
+        )
         return grad_norm_dict

pytorch_lightning/plugins/precision/apex_amp.py

@@ -101,7 +101,7 @@ def pre_optimizer_step(
         model: Union["pl.LightningModule", Module],
         optimizer: Optimizer,
         optimizer_idx: int,
-        lambda_closure: Callable,
+        lambda_closure: Callable[[], Any],
         **kwargs: Any,
     ) -> bool:
         """Hook to do something before each optimizer step."""

pytorch_lightning/plugins/precision/deepspeed_precision.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 typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Union
 
 from torch import Tensor
 from torch.nn import Module
@@ -51,7 +51,7 @@ def pre_optimizer_step(
         model: Union["pl.LightningModule", Module],
         optimizer: Optimizer,
         optimizer_idx: int,
-        lambda_closure: Callable,
+        lambda_closure: Callable[[], Any],
         **kwargs: Any,
     ) -> bool:
         """Hook to do something before each optimizer step."""
@@ -69,14 +69,13 @@ def pre_optimizer_step(
             )
         # DeepSpeed handles the optimizer step internally
         deepspeed_engine = model.trainer.model if isinstance(model, pl.LightningModule) else model
-        deepspeed_engine.step()
+        deepspeed_engine.step(**kwargs)
         return False
 
     def clip_gradients(
         self,
         optimizer: Optimizer,
-        clip_val: Union[int, float],
+        clip_val: Union[int, float] = 0.0,
         gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
-        model: Optional[Module] = None,
     ) -> None:
         """DeepSpeed handles gradient clipping internally."""

pytorch_lightning/plugins/precision/fully_sharded_native_amp.py

@@ -11,36 +11,23 @@
 # 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 typing import Optional, Union
-
-from torch.nn import Module
-from torch.optim import Optimizer
+from typing import Any
 
 from pytorch_lightning.plugins.precision.sharded_native_amp import ShardedNativeMixedPrecisionPlugin
-from pytorch_lightning.utilities import GradClipAlgorithmType
+from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
 class FullyShardedNativeMixedPrecisionPlugin(ShardedNativeMixedPrecisionPlugin):
     """Mixed Precision for Full Sharded Training."""
 
     precision = "mixed"
 
-    def clip_gradients(
-        self,
-        optimizer: Optimizer,
-        clip_val: Union[int, float],
-        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.VALUE,
-        model: Optional[Module] = None,
-    ) -> None:
-        clip_val = float(clip_val)
-        if clip_val <= 0:
-            return
+    def clip_grad_by_norm(self, *_: Any, **__: Any) -> None:
         # see https://fairscale.readthedocs.io/en/latest/api/nn/fsdp_tips.html
         # section `Gradient Clipping`, using `torch.nn.utils.clip_grad_norm_` is incorrect
         # for FSDP module. To overcome this, needs to call sharded_module.clip_grad_norm(clip_val)
         # however we rely on LightningModule's configure_sharded_model to wrap FSDP, it would be hard to
         # trace back the root FSDP. Now we only support clip by value.
-        assert (
-            gradient_clip_algorithm == GradClipAlgorithmType.VALUE
-        ), "`gradient_clip_algorithm`: `norm` is currently not supported for `FullyShardedNativeMixedPrecisionPlugin`"
-        self.clip_grad_by_value(optimizer, clip_val)
+        raise MisconfigurationException(
+            f"`gradient_clip_algorithm='norm'` is currently not supported for `{self.__class__.__name__}`"
+        )

pytorch_lightning/plugins/precision/ipu_precision.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 typing import Any, Callable, Optional, Union
+from typing import Any, Callable, Union
 
 from torch.nn import Module
 from torch.optim import LBFGS, Optimizer
@@ -67,12 +67,9 @@ def pre_optimizer_step(
     def clip_gradients(
         self,
         optimizer: Optimizer,
-        clip_val: Union[int, float],
+        clip_val: Union[int, float] = 0.0,
         gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
-        model: Optional[Module] = None,
     ) -> None:
-        """Clips the gradients."""
-        if clip_val is None or float(clip_val) <= 0:
+        if clip_val <= 0:
             return
-
         raise MisconfigurationException("IPUs currently do not support clipping gradients.")

pytorch_lightning/plugins/precision/native_amp.py

@@ -69,7 +69,7 @@ def pre_optimizer_step(
         model: Union["pl.LightningModule", Module],
         optimizer: Optimizer,
         optimizer_idx: int,
-        lambda_closure: Callable,
+        lambda_closure: Callable[[], Any],
         **kwargs: Any,
     ) -> bool:
         if self.is_bfloat16:
@@ -86,7 +86,7 @@ def pre_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)
+            self.scaler.step(optimizer, **kwargs)
             self.scaler.update()
         return False
 

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -113,22 +113,15 @@ def pre_optimizer_step(
     def clip_gradients(
         self,
         optimizer: Optimizer,
-        clip_val: Union[int, float],
+        clip_val: Union[int, float] = 0.0,
         gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
-        model: Optional[Module] = None,
     ) -> None:
         """Clips the gradients."""
-        if clip_val is None:
-            return
-
-        clip_val = float(clip_val)
         if clip_val <= 0:
             return
-
         if gradient_clip_algorithm == GradClipAlgorithmType.VALUE:
             self.clip_grad_by_value(optimizer, clip_val)
         elif gradient_clip_algorithm == GradClipAlgorithmType.NORM:
-            # TODO: there should be a mechanism to set `norm_type`
             self.clip_grad_by_norm(optimizer, clip_val)
 
     def clip_grad_by_value(self, optimizer: Optimizer, clip_val: Union[int, float]) -> None:

pytorch_lightning/plugins/precision/sharded_native_amp.py

@@ -29,7 +29,5 @@ def __init__(self, precision: Union[int, str] = 16, use_cpu: bool = False) -> No
         if not self.use_cpu:
             self.scaler = ShardedGradScaler()
 
-    def clip_grad_by_norm(
-        self, optimizer: "OSS", clip_val: Union[int, float], norm_type: float = 2.0, eps: float = 1e-6
-    ) -> None:
-        optimizer.clip_grad_norm(clip_val, norm_type=norm_type)
+    def clip_grad_by_norm(self, optimizer: "OSS", clip_val: Union[int, float]) -> None:
+        optimizer.clip_grad_norm(clip_val)

pytorch_lightning/plugins/training_type/deepspeed.py

@@ -33,10 +33,9 @@
 from pytorch_lightning.plugins.training_type.ddp import DDPPlugin
 from pytorch_lightning.trainer.optimizers import _get_default_scheduler_config
 from pytorch_lightning.trainer.states import TrainerFn
-from pytorch_lightning.utilities import AMPType
+from pytorch_lightning.utilities import AMPType, GradClipAlgorithmType
 from pytorch_lightning.utilities.apply_func import apply_to_collection
 from pytorch_lightning.utilities.distributed import log, rank_zero_info, rank_zero_only
-from pytorch_lightning.utilities.enums import GradClipAlgorithmType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.imports import _DEEPSPEED_AVAILABLE
 from pytorch_lightning.utilities.model_helpers import is_overridden
@@ -422,17 +421,16 @@ def _setup_model_and_optimizer(
         return deepspeed_engine, deepspeed_optimizer
 
     def init_deepspeed(self):
-        # check that `configure_gradient_clipping` hook isn't overriden since deepspeed handles
-        # gradient clipping internally
+        # deepspeed handles gradient clipping internally
         if is_overridden("configure_gradient_clipping", self.lightning_module, pl.LightningModule):
             rank_zero_warn(
-                "Since deepspeed handles gradient clipping internally, this hook will"
-                " be ignored. Consider setting `gradient_clip_val` and `gradient_clip_algorithm`"
-                " inside `Trainer`."
+                "Since deepspeed handles gradient clipping internally, `LightningModule.configure_gradient_clipping`"
+                " will be ignored. Consider setting `Trainer(gradient_clip_val=..., gradient_clip_algorithm='norm')`"
+                " which will use the internal mechanism."
             )
 
         if self.lightning_module.trainer.gradient_clip_algorithm == GradClipAlgorithmType.VALUE:
-            raise MisconfigurationException("Deepspeed does not support clipping gradients by value.")
+            raise MisconfigurationException("DeepSpeed does not support clipping gradients by value.")
 
         accumulation_scheduler = self.lightning_module.trainer.accumulation_scheduler
 
@@ -480,7 +478,7 @@ def _initialize_deepspeed_train(self, model):
         else:
             rank_zero_info(
                 "You have not specified an optimizer or scheduler within the DeepSpeed config."
-                "Using `configure_optimizers` to define optimizer and scheduler."
+                " Using `configure_optimizers` to define optimizer and scheduler."
             )
             optimizer, lr_scheduler, _ = self._init_optimizers()
 
@@ -534,7 +532,7 @@ def _initialize_deepspeed_inference(self, model):
         if "optimizer" not in self.config:
             rank_zero_info(
                 "You have not specified an optimizer or scheduler within the DeepSpeed config."
-                "Using `configure_optimizers` to define optimizer and scheduler."
+                " Using `configure_optimizers` to define optimizer and scheduler."
             )
             optimizer, lr_scheduler, _ = self._init_optimizers()
             scheduler = lr_scheduler["scheduler"]

pytorch_lightning/trainer/trainer.py

@@ -257,7 +257,7 @@ def __init__(
             gpus: Number of GPUs to train on (int) or which GPUs to train on (list or str) applied per node
 
             gradient_clip_val: The value at which to clip gradients. Passing ``gradient_clip_val=None`` disables
-                gradient clipping.
+                gradient clipping. If using Automatic Mixed Precision (AMP), the gradients will be unscaled before.
 
             gradient_clip_algorithm: The gradient clipping algorithm to use. Pass ``gradient_clip_algorithm="value"``
                 to clip by value, and ``gradient_clip_algorithm="norm"`` to clip by norm. By default it will
@@ -380,7 +380,8 @@ def __init__(
 
             ipus: How many IPUs to train on.
 
-            track_grad_norm: -1 no tracking. Otherwise tracks that p-norm. May be set to 'inf' infinity-norm.
+            track_grad_norm: -1 no tracking. Otherwise tracks that p-norm. May be set to 'inf' infinity-norm. If using
+                Automatic Mixed Precision (AMP), the gradients will be unscaled before logging them.
 
             val_check_interval: How often to check the validation set. Use float to check within a training epoch,
                 use int to check every n steps (batches).

tests/core/test_lightning_module.py

@@ -386,11 +386,14 @@ def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_va
     trainer = Trainer(
         default_root_dir=tmpdir, max_epochs=1, limit_train_batches=2, limit_val_batches=0, gradient_clip_val=1e-4
     )
-    with pytest.raises(MisconfigurationException, match=r".*have set `Trainer\(gradient_clip_val\)` and have passed.*"):
+    with pytest.raises(
+        MisconfigurationException,
+        match=r"gradient_clip_val=0.0001\)` and have passed `clip_gradients\(gradient_clip_val=0.01",
+    ):
         trainer.fit(model)
 
     class TestModel(BoringModel):
-        custom_gradient_clip_algorithm = "value"
+        custom_gradient_clip_algorithm = "foo"
 
         def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_val, gradient_clip_algorithm):
             self.clip_gradients(optimizer, gradient_clip_algorithm=self.custom_gradient_clip_algorithm)
@@ -404,6 +407,7 @@ def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_va
         gradient_clip_algorithm="norm",
     )
     with pytest.raises(
-        MisconfigurationException, match=r".*have set `Trainer\(gradient_clip_algorithm\)` and have passed.*"
+        MisconfigurationException,
+        match=r"gradient_clip_algorithm='norm'\)` and have passed `clip_gradients\(gradient_clip_algorithm='foo'",
     ):
         trainer.fit(model)

tests/models/test_hooks.py

@@ -281,24 +281,6 @@ def _auto_train_batch(trainer, model, batches, device=torch.device("cpu"), curre
             dict(name="Callback.on_before_optimizer_step", args=(trainer, model, ANY, 0)),
             dict(name="on_before_optimizer_step", args=(ANY, 0)),
         ]
-
-        # deepspeed handles gradient clipping internally
-        configure_gradient_clipping = (
-            []
-            if using_deepspeed
-            else [
-                dict(
-                    name="clip_gradients",
-                    args=(ANY,),
-                    kwargs=dict(gradient_clip_val=None, gradient_clip_algorithm=None),
-                ),
-                dict(
-                    name="configure_gradient_clipping",
-                    args=(ANY, 0),
-                    kwargs=dict(gradient_clip_val=None, gradient_clip_algorithm=None),
-                ),
-            ]
-        )
         for i in range(batches):
             out.extend(
                 [
@@ -323,7 +305,16 @@ 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"),
-                    *configure_gradient_clipping,
+                    dict(
+                        name="clip_gradients",
+                        args=(ANY,),
+                        kwargs=dict(gradient_clip_val=None, gradient_clip_algorithm=None),
+                    ),
+                    dict(
+                        name="configure_gradient_clipping",
+                        args=(ANY, 0),
+                        kwargs=dict(gradient_clip_val=None, gradient_clip_algorithm=None),
+                    ),
                     *(on_before_optimizer_step if using_plugin else []),
                     dict(
                         name="optimizer_step",