[FSDP] Fix input grad propagation when using param mixed precision #90921
Conversation
[ghstack-poisoned]
awgu
requested review from
mrshenli,
pritamdamania87,
zhaojuanmao,
rohan-varma,
H-Huang,
kwen2501 and
wanchaol
as code owners
pytorch-bot
bot
added
the
release notes: distributed (fsdp)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/90921
Note: Links to docs will display an error until the docs builds have been completed. ✅ No FailuresAs of commit cd938ce: This comment was automatically generated by Dr. CI and updates every 15 minutes. |
ghstack-source-id: 9f301e7252ae7e38cc4eecfcbe73ad3afb0dc459 Pull Request resolved: #90921
…recision" [ghstack-poisoned]
ghstack-source-id: 30fbaa288f4cf524360f4e8db459d1fe09ae22b0 Pull Request resolved: #90921
…recision" [ghstack-poisoned]
ghstack-source-id: 7b779515038cb08910611f663e7da589e98d0935 Pull Request resolved: #90921
…recision" For parameter mixed precision, we cast the inputs to the low precision parameter dtype. If the input has tensors that require gradient, then we must cast them in place in order for them to receive a gradient. Otherwise, the tensor that resulted from the out-of-place cast receives the gradient and is not in scope to the user. To preserve BC as much as possible, this PR only does the in-place cast if the tensor requires gradient. [ghstack-poisoned]
ghstack-source-id: a6b8861faa3fcc0ee326c35ab0bcdc25508d688b Pull Request resolved: #90921
|
|
||
| with torch.no_grad(): | ||
| return (_apply_to_tensors(cast_fn, args), _apply_to_tensors(cast_fn, kwargs)) | ||
| return x.to(dtype) |
There was a problem hiding this comment.
hmm, I am a bit concerned about unforseen BC issues, but also our testing surface is quite solid. Did we consider doing this only for inputs that require grad to err on the safe side?
There was a problem hiding this comment.
I think if x does not require gradient, then x.to(dtype) will not be tracked by autograd. if x does require gradient, then x.to(dtype) will be tracked by autograd. This should be exactly the behavior we want. In other words, the casing should be handled naturally already if I understand correctly.
…recision" For parameter mixed precision, we cast the inputs to the low precision parameter dtype. If the input has tensors that require gradient, then we must cast them in place in order for them to receive a gradient. The cast should be tracked by autograd (e.g. with `grad_fn` equal to `ToCopyBackward0`). This removes the `torch.no_grad` context when calling `_apply_to_tensors`. [ghstack-poisoned]
ghstack-source-id: ad234ed096a8c598818af2dfe6b77ae2ef1ebd54 Pull Request resolved: #90921
|
@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 |
Stack from ghstack:
_module_to_handles,HandleConfig; use term "fqn"; clarify docs #90840 [FSDP][BE] Remove_module_to_handles,HandleConfig; use term "fqn"; clarify docsFor parameter mixed precision, we cast the inputs to the low precision parameter dtype. If the input has tensors that require gradient, then we must cast them in place in order for them to receive a gradient. The cast should be tracked by autograd (e.g. with
grad_fnequal toToCopyBackward0). This removes thetorch.no_gradcontext when calling_apply_to_tensors.