on_after_backward is called before gradient is unscale_ when using mixed precision
🐛 Bug
From my understanding, one of the main purpose of callback on_after_backward, is to check and log gradients. However, when AMP is being enabled, the gradients you are accessing are not unscaled. I.E. all the numbers and norms you look at will be super large and not really useful.
To Reproduce
Use following code snippet to log some layer's l2 norm.
    def on_after_backward(self):
            with torch.no_grad():
                if (self.model.*****.weight.grad is not None):
                    norm_value = self.model.*****.weight.grad.detach().norm(2).item()
                    self.log("norm2", norm_value)

The value registered will be wildly different when using 32bit precision vs. when you enable Native AMP using:
                     amp_backend='native',
                     precision=16,

Expected behavior
The weights accessible inside the on_after_backward callback should ideally be consistent, with or without AMP.
More specifically, lightning should consider calling the self.trainer.scaler.unscale_ before on_after_backward. Or maybe provide the option to do so.
Edit: Sorry I just realized this might not be possible since unscale_ should be called after gradient accumulation but on_after_backward is called for every step. This lead to another issue: when conducting gradient accumulation i sometimes need to manually check if current step is The Step to do weight update (the last step of accumulation). One possible solution for both of these issues is perhaps another separate callback that will be called every N accumulation steps. Right after _unscale but right before gradient clipping and updating? Although this will definitely make things a bit complicated.
Environment

CUDA:

GPU:

TITAN RTX
GeForce RTX 2080 SUPER


available:         True
version:           10.2


Packages:

numpy:             1.19.1
pyTorch_debug:     False
pyTorch_version:   1.6.0
pytorch-lightning: 1.0.1
tqdm:              4.49.0


System:

OS:                Linux
architecture:

64bit



processor:         x86_64
python:            3.7.9
version:           #52-Ubuntu SMP Thu Sep 10 10:58:49 UTC 2020



Additional context
Currently Lightning calls self.trainer.scaler.unscale_(optimizer) if gradient_clip_val is enabled. Hence if i want to use norm clipping currently i can't manually do un-scale inside the callback.