pytorch
0e8565d1 - [FSDP][optim_state_dict][8/N] Enable fully_shard optim state_dict save and load (#91234)

Commit
2 years ago
[FSDP][optim_state_dict][8/N] Enable fully_shard optim state_dict save and load (#91234) **What does this PR do?** This PR refactor `_optim_utils.py` to use `_FSDPState` instead of `FullyShardedDataParallel` class. This change enables the support of optim state_dict for `fully_shard`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/91234 Approved by: https://github.com/rohan-varma
Author
Committer
Parents
Loading