Add Trainer(gradient_clip_algorithm='value'|'norm') (#6123)

* add changelog

* add clip by value

* fix bug in training tricks.rst

* fix bug in trainer.rst

* Update trainer.rst

* Update trainer.rst

* Update CHANGELOG.md

Co-authored-by: Jirka Borovec <[email protected]>

* Update pytorch_lightning/plugins/precision/deepspeed_precision.py

Co-authored-by: Jirka Borovec <[email protected]>

* Update pytorch_lightning/utilities/enums.py

Co-authored-by: Jirka Borovec <[email protected]>

* yapf formatting

* update training tricks

* update based on comment

* update based on comment

* Update pytorch_lightning/trainer/trainer.py

Co-authored-by: ananthsub <[email protected]>

* update based on comment

* pep8

* mypy

* mypy

* Update docs/source/advanced/training_tricks.rst

Co-authored-by: thomas chaton <[email protected]>

* Update sharded_native_amp.py

* Update test_sharded_parity.py

* update test codes

* Update test_tpu.py

* Update pytorch_lightning/trainer/connectors/training_trick_connector.py

Co-authored-by: Carlos Mocholí <[email protected]>

* Update test_trainer.py

* Update enums.py

* Update enums.py

* add super-class initialization to precision plugins.

* add clip_grad horovod cpu test

* add clip_grad horovod cpu test

* use subprocess check_call

* change order of horovod tests

* set max_epochs 2 in horovod test

* remove clip_grad_val test from horovod-cpu

* remove "type: ignore"

* divide clip grad val test in horovod

* update based on comments

* add super-class initialization to precision plugins.

* bugfix

* bugfix

* revert some changes

* revert some changes

* Update tests/models/test_horovod.py

* merge master

* Delete signature test

No point in testing a signature

Co-authored-by: Jirka Borovec <[email protected]>
Co-authored-by: thomas chaton <[email protected]>
Co-authored-by: ananthsub <[email protected]>
Co-authored-by: Carlos Mocholí <[email protected]>
Co-authored-by: Jirka Borovec <[email protected]>

Author: 5 people

CHANGELOG.md

@@ -19,6 +19,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Trigger warning when non-metric logged value with multi processes hasn't been reduced ([#6417](https://github.com/PyTorchLightning/pytorch-lightning/pull/6417))
 
 
+- Added `gradient_clip_algorithm` argument to Trainer for gradient clipping by value ([#6123](https://github.com/PyTorchLightning/pytorch-lightning/pull/6123)).
+
+
 - Added a way to print to terminal without breaking up the progress bar ([#5470](https://github.com/PyTorchLightning/pytorch-lightning/pull/5470))
 
 

docs/source/advanced/training_tricks.rst

@@ -26,8 +26,10 @@ The effect is a large effective batch size of size KxN.
 
 Gradient Clipping
 -----------------
-Gradient clipping may be enabled to avoid exploding gradients. Specifically, this will `clip the gradient
-norm <https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_norm_>`_ computed over all model parameters together.
+Gradient clipping may be enabled to avoid exploding gradients. By default, this will `clip the gradient norm
+<https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_norm_>`_ computed over all model parameters together.
+If ``gradient_clip_algorithm`` option is set to ``value``, which is ``norm`` by default, this will
+`clip the gradient value <https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_value_>`_ for each parameter instead.
 
 .. seealso:: :class:`~pytorch_lightning.trainer.trainer.Trainer`
 
@@ -39,6 +41,10 @@ norm <https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_norm_>`_
     # clip gradients with norm above 0.5
     trainer = Trainer(gradient_clip_val=0.5)
 
+    # clip gradients with value above 0.5
+    # gradient_clip_algorithm types => :class:`~pytorch_lightning.utilities.enums.GradClipAlgorithmType`
+    trainer = Trainer(gradient_clip_val=0.5, gradient_clip_algorithm='value')
+
 ----------
 
 Stochastic Weight Averaging

pytorch_lightning/accelerators/accelerator.py

@@ -24,7 +24,7 @@
 from pytorch_lightning.trainer.states import TrainerState
 from pytorch_lightning.utilities import rank_zero_warn
 from pytorch_lightning.utilities.apply_func import move_data_to_device
-from pytorch_lightning.utilities.enums import AMPType, LightningEnum
+from pytorch_lightning.utilities.enums import AMPType, GradClipAlgorithmType, LightningEnum
 
 if TYPE_CHECKING:
     from torch.cuda.amp import GradScaler
@@ -315,10 +315,14 @@ 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]) -> None:
+    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(self.model, optimizer, clip_val)
+        self.precision_plugin.clip_gradients(self.model, optimizer, clip_val, gradient_clip_algorithm)
 
     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/apex_amp.py

@@ -30,6 +30,7 @@ class ApexMixedPrecisionPlugin(MixedPrecisionPlugin):
     """Mixed Precision Plugin based on Nvidia/Apex (https://github.com/NVIDIA/apex)"""
 
     def __init__(self, amp_level: str = "O2") -> None:
+        super().__init__()
         self.backend = AMPType.APEX
         self.amp_level = amp_level
 

pytorch_lightning/plugins/precision/deepspeed_precision.py

@@ -16,6 +16,7 @@
 import torch
 
 from pytorch_lightning.plugins.precision.precision_plugin import PrecisionPlugin
+from pytorch_lightning.utilities import GradClipAlgorithmType
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.warnings import WarningCache
 
@@ -80,7 +81,7 @@ def clip_gradients(
         model: 'LightningModule',
         optimizer: 'Optimizer',
         clip_val: Union[int, float],
-        norm_type: float = 2.0
+        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
     ) -> None:
         """
         DeepSpeed handles clipping gradients via the training type plugin.

pytorch_lightning/plugins/precision/double.py

@@ -67,6 +67,7 @@ class DoublePrecisionPlugin(PrecisionPlugin):
     precision: int = 64
 
     def __init__(self) -> None:
+        super().__init__()
         self.patches: List[_DoublePrecisionPatch] = []
 
     def connect(

pytorch_lightning/plugins/precision/native_amp.py

@@ -30,6 +30,7 @@
 class NativeMixedPrecisionPlugin(MixedPrecisionPlugin):
 
     def __init__(self) -> None:
+        super().__init__()
         if not _NATIVE_AMP_AVAILABLE:
             raise MisconfigurationException(
                 "You have asked for native AMP but your PyTorch version does not support it."

pytorch_lightning/plugins/precision/precision_plugin.py

@@ -17,6 +17,7 @@
 import torch
 
 from pytorch_lightning.plugins.base_plugin import Plugin
+from pytorch_lightning.utilities import GradClipAlgorithmType
 
 if TYPE_CHECKING:
     from torch.nn import Module
@@ -33,6 +34,13 @@ class PrecisionPlugin(Plugin):
     EPSILON: float = 1e-6
     precision: Union[str, int] = 32
 
+    def __init__(self) -> None:
+        super().__init__()
+        self.clip_grad_funcs = {
+            GradClipAlgorithmType.VALUE: self.clip_grad_by_value,
+            GradClipAlgorithmType.NORM: self.clip_grad_by_norm,
+        }
+
     def master_params(self, optimizer: 'Optimizer') -> Generator[torch.Tensor, None, None]:
         """The master params of the model. Returns the plain model params here.
         Maybe different in other precision plugins.
@@ -103,20 +111,29 @@ def clip_gradients(
         model: 'LightningModule',
         optimizer: 'Optimizer',
         clip_val: Union[int, float],
-        norm_type: float = 2.0
+        gradient_clip_algorithm: GradClipAlgorithmType = GradClipAlgorithmType.NORM,
     ) -> None:
-        """Clips the gradients to a specific value"""
+        """Clips the gradients"""
         if clip_val is None:
             return
 
-        grad_clip_val = float(clip_val)
-
-        if grad_clip_val <= 0:
+        clip_val = float(clip_val)
+        if clip_val <= 0:
             return
 
+        clip_grad_func = self.clip_grad_funcs[gradient_clip_algorithm]
+        clip_grad_func(optimizer, clip_val)  # type: ignore
+
+    def clip_grad_by_value(self, optimizer: 'Optimizer', clip_val: Union[int, float]) -> None:
+        """Clip gradients by value"""
         parameters = list(self.master_params(optimizer))
+        torch.nn.utils.clip_grad_value_(parameters, clip_value=clip_val)
 
-        max_norm = grad_clip_val
+    def clip_grad_by_norm(self, optimizer: 'Optimizer', clip_val: Union[int, float], norm_type: float = 2.0) -> None:
+        """Clip gradients by norm"""
+        # TODO: separate TPU case from here
+        parameters = list(self.master_params(optimizer))
+        max_norm = clip_val
 
         if isinstance(parameters, torch.Tensor):
             parameters = [parameters]

pytorch_lightning/plugins/precision/sharded_native_amp.py

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

pytorch_lightning/trainer/connectors/training_trick_connector.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from pytorch_lightning.callbacks import GradientAccumulationScheduler
+from pytorch_lightning.utilities import GradClipAlgorithmType
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 
 
@@ -23,6 +24,7 @@ def __init__(self, trainer):
     def on_trainer_init(
         self,
         gradient_clip_val,
+        gradient_clip_algorithm,
         track_grad_norm,
         accumulate_grad_batches,
         truncated_bptt_steps,
@@ -32,7 +34,12 @@ def on_trainer_init(
         self.trainer.terminate_on_nan = terminate_on_nan
 
         # gradient clipping
+        if gradient_clip_algorithm not in list(GradClipAlgorithmType):
+            raise MisconfigurationException(
+                f"gradient_clip_algorithm should be in {list(GradClipAlgorithmType)}"
+            )
         self.trainer.gradient_clip_val = gradient_clip_val
+        self.trainer.gradient_clip_algorithm = gradient_clip_algorithm
 
         # gradient norm tracking
         if not isinstance(track_grad_norm, (int, float)) and track_grad_norm != 'inf':