extract and log grad norm for individual layers

the paper NormFormer: Improved Transformer Pretraining with Extra Normalization https://arxiv.org/abs/2110.09456 suggests that under preLN:
> gradients at earlier layers tend to be larger than gradients at later layers

so we want to verify that this is so in our case before acting on it and potentially integrating NormFormer.

So we need to expand the tensorboard and logs to log grad norm for individual layers, perhaps as in the paper we can log 5 layers: 0, 1, int(n_layers/2), -2, -1

For reference see p6 in the paper (graphs and discussion).

Currently only L2 average of all layer grad norms is calculated and logged (a single number).

------------------

Performance-wise this will require doing some extra calculations but not much - it's really just figuring out how to broadcast this info to TB so that multiple-points can be logged at once.

Additionally, we could activate this tool on demand - e.g. after encountering a spike we could roll back to the last good checkpoint and run a cycle with this debug feature enabled.