pytorch
08245060 - [FSDP] Allow nested FSDP wrapper to use different mixed precision (#90523)

Commit
2 years ago
[FSDP] Allow nested FSDP wrapper to use different mixed precision (#90523) The main change is to move `args` and `kwargs` dtype convertion from `_root_pre_forward` to `_pre_forward`, so that every FSDP has a chance to apply its own precision. Pull Request resolved: https://github.com/pytorch/pytorch/pull/90523 Approved by: https://github.com/awgu, https://github.com/rohan-varma
Author
Committer
Parents
Loading