Zijian Hu
Zijian Hu
@awgu thank you so much for the follow up! I guess some naive implementations like the below example should work but would appreciate your feedback. cc @wconstab @H-Huang @tianyu-l @XilunWu...
> I believe [`local_map`](https://pytorch.org/docs/main/distributed.tensor.html#torch.distributed.tensor.experimental.local_map) is a good fit for this case, to implement a custom `clip_grad_norm_` for DTensor. @zijian-hu let me draft a PR based on your sample so that...
@H-Huang in case Xilun is busy with other work items, I am more than happy draft this PR later today
@XilunWu @H-Huang the PR has been created https://github.com/pytorch/torchtitan/pull/649