[FSDP] Add fast path for NO_SHARD clip_grad_norm_()
#89137
Conversation
[ghstack-poisoned]
awgu
requested review from
H-Huang,
kwen2501,
mrshenli,
pritamdamania87,
rohan-varma and
zhaojuanmao
as code owners
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/89137
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit ab9f9de: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
[ghstack-poisoned]
rohan-varma
left a comment
rohan-varma
left a comment
There was a problem hiding this comment.
LGTM, thanks for tracking the follow up issues around improved testing!
| # TODO: Gradient computation on CPU and GPU differ slightly causing | ||
| # drift unrelated to `clip_grad_norm_()`. | ||
| # https://github.com/pytorch/pytorch/issues/89133 | ||
| return |
There was a problem hiding this comment.
A bit hacky, but would it help if we just ran DDP on CPU for this case if having this coverage is important?
There was a problem hiding this comment.
Does DDP work on CPU with NCCL backend?
[ghstack-poisoned]
|
@pytorchbot merge |
Merge startedYour change will be merged once all checks pass (ETA 0-4 Hours). Learn more about merging in the wiki. Questions? Feedback? Please reach out to the PyTorch DevX Team |
Pull Request resolved: pytorch#89137 Approved by: https://github.com/rohan-varma
Stack from ghstack:
NO_SHARDclip_grad_norm_()#89137 [FSDP] Add fast path forNO_SHARDclip_grad_norm_()