Repository: Lightning-AI/pytorch-lightning

Issue:
Gradient norms are not logged unless row_log_interval==1

## 🐛 Bug
In version 0.9 the guards to calculate the gradient norms and then log the metrics can't be satisfied in the same batch unless the row_log_interval  is 1. In most places the guard seems to be `(batch_idx + 1) % self.row_log_interval == 0` such as here:

https://github.com/PyTorchLightning/pytorch-lightning/blob/b40de5464a953ff5866a255f4670d318bd8fd65a/pytorch_lightning/trainer/training_loop.py#L749-L757

However in `run_batch_backward_pass` it is `batch_idx % self.row_log_interval == 0`

https://github.com/PyTorchLightning/pytorch-lightning/blob/b40de5464a953ff5866a255f4670d318bd8fd65a/pytorch_lightning/trainer/training_loop.py#L929-L939


### To Reproduce

Steps to reproduce the behavior:

1. Run the code sample below (taken from #1527 ).
2. Confirm that gradients are not being logged in tensorboard.
3. Change row_log_interval to 1 and rerun the code.
4 Confirm that gradients are now being logged.

#### Code sample

```
import pytorch_lightning as pl
from torch.utils.data import TensorDataset, DataLoader
from pytorch_lightning.loggers.tensorboard import TensorBoardLogger
from torch.optim import SGD
import torch.nn as nn
import torch


class MWENet(pl.LightningModule):
    def __init__(self):
        super(MWENet, self).__init__()

        self.first = nn.Conv2d(1, 1, 3)
        self.second = nn.Conv2d(1, 1, 3)
        self.loss = nn.L1Loss()

    def train_dataloader(self):
        xs, ys = torch.zeros(16, 1, 10, 10), torch.ones(16, 1, 6, 6)
        ds = TensorDataset(xs, ys)
        return DataLoader(ds)

    def forward(self, xs):
        out = self.first(xs)
        out = self.second(out)
        return out

    def configure_optimizers(self):
        first = SGD(self.first.parameters(), lr=0.01)
        second = SGD(self.second.parameters(), lr=0.01)
        return [second, first]

    def training_step(self, batch, batch_idx, optimizer_idx):
        xs, ys = batch
        out = self.forward(xs)
        return {'loss': self.loss(out, ys)}


net = MWENet()
logger = TensorBoardLogger('tb_logs', name='testing')
trainer = pl.Trainer(
    track_grad_norm=2,
    row_log_interval=2,
    max_epochs=50,
    logger=logger)
trainer.fit(net)

```
### Expected behavior

Gradients should be logged if track_grad_norm is True

### Environment

* CUDA:
        - GPU:
                - GeForce GTX 1080 Ti
        - available:         True
        - version:           10.2
* Packages:
        - numpy:             1.19.1
        - pyTorch_debug:     False
        - pyTorch_version:   1.6.0
        - pytorch-lightning: 0.9.0
        - tensorboard:       2.2.1
        - tqdm:              4.48.2
* System:
        - OS:                Windows
        - architecture:
                - 64bit
                - WindowsPE
        - processor:         AMD64 Family 23 Model 113 Stepping 0, AuthenticAMD
        - python:            3.7.9
        - version:           10.0.19041

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi! thanks for your contribution!, great first issue!
yeah, it's a bug, mind send a PR?

Resolving patch:
diff --git a/pytorch_lightning/trainer/training_loop.py b/pytorch_lightning/trainer/training_loop.py
--- a/pytorch_lightning/trainer/training_loop.py
+++ b/pytorch_lightning/trainer/training_loop.py
@@ -376,13 +376,12 @@ def on_before_backward(self, batch_idx, optimizer):
         return grad_norm_dic
 
     def _track_gradient_norm(self, batch_idx):
-        grad_norm_dic = {}
-        if batch_idx % self.trainer.row_log_interval == 0:
+        grad_norm_dict = {}
+        if (batch_idx + 1) % self.trainer.row_log_interval == 0:
             if float(self.trainer.track_grad_norm) > 0:
                 model = self.trainer.get_model()
-                grad_norm_dic = model.grad_norm(
-                    self.trainer.track_grad_norm)
-        return grad_norm_dic
+                grad_norm_dict = model.grad_norm(self.trainer.track_grad_norm)
+        return grad_norm_dict
 
     def log_training_step_metrics(self, opt_closure_result, batch_callback_metrics, batch_log_metrics):
         # track callback metrics


Test patch:
diff --git a/tests/models/test_grad_norm.py b/tests/models/test_grad_norm.py
--- a/tests/models/test_grad_norm.py
+++ b/tests/models/test_grad_norm.py
@@ -1,4 +1,5 @@
 import os
+from unittest.mock import patch
 
 import numpy as np
 import pytest
@@ -73,3 +74,28 @@ def test_grad_tracking(tmpdir, norm_type, rtol=5e-3):
         log, mod = [log[k] for k in common], [mod[k] for k in common]
 
         assert np.allclose(log, mod, rtol=rtol)
+
+
+@pytest.mark.parametrize("row_log_interval", [1, 2, 3])
+def test_grad_tracking_interval(tmpdir, row_log_interval):
+    """ Test that gradient norms get tracked in the right interval and that everytime the same keys get logged. """
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        track_grad_norm=2,
+        row_log_interval=row_log_interval,
+        max_steps=10,
+    )
+
+    with patch.object(trainer.logger, "log_metrics") as mocked:
+        model = EvalModelTemplate()
+        trainer.fit(model)
+        expected = trainer.global_step // row_log_interval
+        grad_norm_dicts = []
+        for _, kwargs in mocked.call_args_list:
+            metrics = kwargs.get("metrics", {})
+            grad_norm_dict = {k: v for k, v in metrics.items() if k.startswith("grad_")}
+            if grad_norm_dict:
+                grad_norm_dicts.append(grad_norm_dict)
+
+        assert len(grad_norm_dicts) == expected
+        assert all(grad_norm_dicts[0].keys() == g.keys() for g in grad_norm_dicts)