How can I extract and log grad norm for individual layers

I want to use gradients to monitor if the model is training properly, like this

I change the transformers.Trainer https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L1857 ,by add a function to get the parm.grad to log grad norm front the zero_grad.

But it not work for the model when use the DeepSpeed, parm.grad is None for all layers, I read the code and found comments said there is no way to return gradients when use DeepSpeed.

 elif self.distributed_type == DistributedType.DEEPSPEED:
    # `accelerator.backward(loss)` is doing that automatically. Therefore, its implementation is not needed
    # We cannot return the gradient norm because DeepSpeed does it.
    return None

https://github.com/huggingface/accelerate/blob/69e4c3c54da3201eda288b500d138761e7a5221c/src/accelerate/accelerator.py#L1499

So how can I get the grad norm for track training state when I am using the DeepSpeed?

I found a similar problem in megatron-deepspeed, but there seems to be no follow-up

https://github.com/bigscience-workshop/Megatron-DeepSpeed/issues/149

Start with the linked transformers Trainer location and accelerate.py around line 1499 to trace how DeepSpeed handles gradients, then compare the related Megatron-DeepSpeed issue. A complete resolution would need an agreed, tested way to obtain per-layer gradient norms during DeepSpeed training, but the issue does not specify the intended API or implementation location.