accelerate
0f7e35fe - Clip grad norm support for dtensors (#4219)

Commit
5 days ago
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>
Parents
Loading