pytorch
e5503ace - [FSDP] Re-support model dtype change after FSDP init (#91192)

Commit
2 years ago
[FSDP] Re-support model dtype change after FSDP init (#91192) Closes https://github.com/pytorch/pytorch/issues/90838. To make mixed precision precise internally, https://github.com/pytorch/pytorch/pull/90660 changed the implementation to save `_orig_param_dtype`, `_low_prec_param_dtype`, and `_reduce_dtype` explicitly. However, these are computed at FSDP construction time, so it does not allow the user to change the model dtype after FSDP construction time but before lazy initialization. This PR recomputes those dtype attributes as needed if the model dtype changes in that window. Note that any mixed precision settings specified by the user take precedence over the model dtype. Pull Request resolved: https://github.com/pytorch/pytorch/pull/91192 Approved by: https://github.com/zhaojuanmao
Author
Committer
Parents
Loading