Repository: Lightning-AI/pytorch-lightning

Issue:
Speed up gradient norm clipping.
## 🚀 Feature
Improve the speed of gradient clipping.

### Motivation
I've noticed that clipping the gradients is quite slow, e.g., for a ResNet50 model (with some minor extensions) gradient clipping takes ~207 ms. The major bottlenecks are the `torch.where` computation within a loop and the manual norm computation.

### Pitch
With slight modifications that don't hurt readability and work well on TPU,
grad_norm can be computed in ~75ms, i.e.,  2.7x as fast. For implementation details see the PR I'm going to link in a second.

Reproduce: Train any model you've got lying around with gradient_clip_val=0.1 and take a look at the trace timeline ;)

As a reference, here's a before and after comparison for my model (trace after 40 warmup steps):
Current:
![image](https://user-images.githubusercontent.com/16101605/88942052-09cd7100-d28a-11ea-93db-73c72e350ac4.png)

With speed improvements:
![image](https://user-images.githubusercontent.com/16101605/88941994-f4584700-d289-11ea-9c72-668a8b8d227e.png)



Discussion:

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_tricks.py b/pytorch_lightning/trainer/training_tricks.py
--- a/pytorch_lightning/trainer/training_tricks.py
+++ b/pytorch_lightning/trainer/training_tricks.py
@@ -64,27 +64,29 @@ def clip_gradients(self):
 
         # this code is a modification of torch.nn.utils.clip_grad_norm_
         # with TPU support based on https://github.com/pytorch/xla/blob/master/TROUBLESHOOTING.md
-        if self.gradient_clip_val > 0:
-            model = self.get_model()
-            parameters = model.parameters()
-            max_norm = float(self.gradient_clip_val)
-            norm_type = float(2.0)
-            if isinstance(parameters, torch.Tensor):
-                parameters = [parameters]
-            parameters = list(filter(lambda p: p.grad is not None, parameters))
-            if norm_type == math.inf:
-                total_norm = max(p.grad.data.abs().max() for p in parameters)
-            else:
-                device = parameters[0].device
-                total_norm = torch.zeros([], device=device if parameters else None)
-                for p in parameters:
-                    param_norm = p.grad.data.pow(norm_type).sum()
-                    total_norm.add_(param_norm)
-                total_norm = (total_norm ** (1. / norm_type))
-            eps = EPSILON_FP16 if self.precision == 16 else EPSILON
-            clip_coef = torch.tensor(max_norm, device=device) / (total_norm + eps)
-            for p in parameters:
-                p.grad.data.mul_(torch.where(clip_coef < 1, clip_coef, torch.tensor(1., device=device)))
+        if self.gradient_clip_val <= 0:
+            return
+        model = self.get_model()
+        parameters = model.parameters()
+        max_norm = float(self.gradient_clip_val)
+        norm_type = float(2.0)
+        if isinstance(parameters, torch.Tensor):
+            parameters = [parameters]
+        parameters = list(filter(lambda p: p.grad is not None, parameters))
+        if norm_type == math.inf:
+            total_norm = max(p.grad.data.abs().max() for p in parameters)
+        else:
+            device = parameters[0].device
+            out = torch.empty(len(parameters), device=device)
+            for i, p in enumerate(parameters):
+                torch.norm(p.grad.data.to(device), norm_type, out=out[i])
+            total_norm = torch.norm(out, norm_type)
+
+        eps = EPSILON_FP16 if self.precision == 16 else EPSILON
+        clip_coef = torch.tensor(max_norm, device=device) / (total_norm + eps)
+        clip_coef = torch.min(clip_coef, torch.ones_like(clip_coef))
+        for p in parameters:
+            p.grad.data.mul_(clip_coef.to(p.grad.data.device))
 
     def print_nan_gradients(self) -> None:
         model = self.get_model()


Test patch:
diff --git a/tests/models/test_tpu.py b/tests/models/test_tpu.py
--- a/tests/models/test_tpu.py
+++ b/tests/models/test_tpu.py
@@ -238,6 +238,25 @@ def test_early_stop_checkpoints_on_tpu(tmpdir):
     assert torch_xla._XLAC._xla_get_default_device() == 'xla:5'
 
 
+@pytest.mark.skipif(not TPU_AVAILABLE, reason="test requires TPU machine")
+@pl_multi_process_test
+def test_tpu_grad_norm(tmpdir):
+    """Test if grad_norm works on TPU."""
+    trainer_options = dict(
+        default_root_dir=tmpdir,
+        progress_bar_refresh_rate=0,
+        max_epochs=1,
+        distributed_backend='tpu',
+        tpu_cores=1,
+        limit_train_batches=0.4,
+        limit_val_batches=0.4,
+        gradient_clip_val=0.1,
+    )
+
+    model = EvalModelTemplate()
+    tpipes.run_model_test(trainer_options, model, on_gpu=False, with_hpc=False)
+
+
 @pytest.mark.skipif(not TPU_AVAILABLE, reason="test requires TPU machine")
 @pl_multi_process_test
 def test_dataloaders_passed_to_fit(tmpdir):