Use scale and clip grads utilties from Automodel for DTensor v2
Open
enhancement
t-pytdensor
- Dominant language
- Python
- Stars
- 2k
- Forks
- 561
- Avg merge
- 4d 5h
- Merged PRs (30d)
- 145
Description
**Describe the solution you'd like**
Use utilities from https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo_automodel/components/training/utils.py#L203 for scaling and clipping grads for DTensor policy v2.
Contributor guide
Assessment
This issue has not been assessed yet.