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':

pytorch_lightning/trainer/trainer.py

@@ -91,6 +91,7 @@ def __init__(
         callbacks: Optional[Union[List[Callback], Callback]] = None,
         default_root_dir: Optional[str] = None,
         gradient_clip_val: float = 0,
+        gradient_clip_algorithm: str = 'norm',
         process_position: int = 0,
         num_nodes: int = 1,
         num_processes: int = 1,
@@ -197,6 +198,8 @@ def __init__(
 
             gradient_clip_val: 0 means don't clip.
 
+            gradient_clip_algorithm: 'value' means clip_by_value, 'norm' means clip_by_norm. Default: 'norm'
+
             limit_train_batches: How much of training dataset to check (float = fraction, int = num_batches)
 
             limit_val_batches: How much of validation dataset to check (float = fraction, int = num_batches)
@@ -347,7 +350,12 @@ def __init__(
 
         # init training tricks
         self.training_tricks_connector.on_trainer_init(
-            gradient_clip_val, track_grad_norm, accumulate_grad_batches, truncated_bptt_steps, terminate_on_nan
+            gradient_clip_val,
+            gradient_clip_algorithm,
+            track_grad_norm,
+            accumulate_grad_batches,
+            truncated_bptt_steps,
+            terminate_on_nan,
         )
         self.train_loop.on_trainer_init(
             max_epochs,

pytorch_lightning/utilities/__init__.py

@@ -23,7 +23,13 @@
     rank_zero_only,
     rank_zero_warn,
 )
-from pytorch_lightning.utilities.enums import AMPType, DeviceType, DistributedType, LightningEnum  # noqa: F401
+from pytorch_lightning.utilities.enums import (  # noqa: F401
+    AMPType,
+    DeviceType,
+    DistributedType,
+    GradClipAlgorithmType,
+    LightningEnum,
+)
 from pytorch_lightning.utilities.imports import (  # noqa: F401
     _APEX_AVAILABLE,
     _BOLTS_AVAILABLE,

pytorch_lightning/utilities/enums.py

@@ -97,3 +97,16 @@ class DeviceType(LightningEnum):
     CPU = 'CPU'
     GPU = 'GPU'
     TPU = 'TPU'
+
+
+class GradClipAlgorithmType(LightningEnum):
+    """ Define gradient_clip_algorithm types - training-tricks.
+    NORM type means "clipping gradients by norm". This computed over all model parameters together.
+    VALUE tpye means "clipping gradients by value". This will clip the gradient value for each parameter.
+
+    References:
+        clip_by_norm: https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_norm
+        clip_by_value: https://pytorch.org/docs/stable/nn.html#torch.nn.utils.clip_grad_value
+    """
+    VALUE = 'value'
+    NORM = 'norm'

tests/models/test_horovod.py

@@ -50,7 +50,7 @@ def _run_horovod(trainer_options, on_gpu=False):
     trainer_options.update(gpus=1 if on_gpu else None)
     tutils.reset_seed()
     # todo: Find why coverage breaks CI.
-    # append = '-a' if '.coverage' in os.listdir(_PROJECT_ROOT) else ''     # noqa E265
+    # append = '-a' if '.coverage' in os.listdir(_PROJECT_ROOT) else ''
     # str(num_processes), sys.executable, '-m', 'coverage', 'run', '--source', 'pytorch_lightning', append,   # noqa E265
     cmdline = [
         'horovodrun', '-np',
@@ -80,6 +80,24 @@ def test_horovod_cpu(tmpdir):
     _run_horovod(trainer_options)
 
 
+@RunIf(skip_windows=True, horovod=True)
+def test_horovod_cpu_clip_grad_by_value(tmpdir):
+    """Test Horovod running multi-process on CPU."""
+    trainer_options = dict(
+        default_root_dir=str(tmpdir),
+        weights_save_path=str(tmpdir),
+        gradient_clip_val=1.0,
+        gradient_clip_algorithm='value',
+        progress_bar_refresh_rate=0,
+        max_epochs=1,
+        limit_train_batches=0.4,
+        limit_val_batches=0.2,
+        accelerator='horovod',
+        deterministic=True,
+    )
+    _run_horovod(trainer_options)
+
+
 @RunIf(skip_windows=True, horovod=True)
 def test_horovod_cpu_implicit(tmpdir):
     """Test Horovod without specifying a backend, inferring from env set by `horovodrun`."""
@@ -114,6 +132,25 @@ def test_horovod_multi_gpu(tmpdir):
     _run_horovod(trainer_options, on_gpu=True)
 
 
+@RunIf(min_gpus=2, skip_windows=True, horovod_nccl=True)
+def test_horovod_multi_gpu_grad_by_value(tmpdir):
+    """Test Horovod with multi-GPU support."""
+    trainer_options = dict(
+        default_root_dir=str(tmpdir),
+        weights_save_path=str(tmpdir),
+        gradient_clip_val=1.0,
+        gradient_clip_algorithm='value',
+        progress_bar_refresh_rate=0,
+        max_epochs=1,
+        limit_train_batches=0.4,
+        limit_val_batches=0.2,
+        gpus=2,
+        deterministic=True,
+        accelerator='horovod',
+    )
+    _run_horovod(trainer_options, on_gpu=True)
+
+
 # https://discuss.pytorch.org/t/torch-cuda-amp-vs-nvidia-apex/74994
 # Check with (tgaddair) on Horovod issues if this feature is needed
 @pytest.mark.skip(reason="Horovod currently doesn't work with Apex")  # todo

tests/models/test_tpu.py

@@ -219,6 +219,26 @@ def test_tpu_grad_norm(tmpdir):
     tpipes.run_model_test(trainer_options, model, on_gpu=False, with_hpc=False)
 
 
+@RunIf(tpu=True)
+@pl_multi_process_test
+def test_tpu_clip_grad_by_value(tmpdir):
+    """Test if clip_gradients by value works on TPU."""
+    tutils.reset_seed()
+    trainer_options = dict(
+        default_root_dir=tmpdir,
+        progress_bar_refresh_rate=0,
+        max_epochs=4,
+        tpu_cores=1,
+        limit_train_batches=4,
+        limit_val_batches=4,
+        gradient_clip_val=0.5,
+        gradient_clip_algorithm='value'
+    )
+
+    model = BoringModel()
+    tpipes.run_model_test(trainer_options, model, on_gpu=False, with_hpc=False)
+
+
 @RunIf(tpu=True)
 @pl_multi_process_test
 def test_dataloaders_passed_to_fit(tmpdir):