Clip grad norm support for dtensors (#4219)
* feat: clip grad norm support for dtensors
* feat: clip_grad_norm supports dtensors
* feat: add support for multiple device meshes
* refactor: separate method for mixed grads
* refactor: simplify the logic
* Apply batched suggestions from code review
Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>
* Update src/accelerate/accelerator.py
* Update src/accelerate/accelerator.py
---------
Co-authored-by: Marc Sun <57196510+SunMarc@users.noreply.github.com>