pytorch
1ab6ac46 - [FSDP][optim_state_dict][6/N] Refactor the optim_state_dict APIs to support hooks (#90798)

Commit
2 years ago
[FSDP][optim_state_dict][6/N] Refactor the optim_state_dict APIs to support hooks (#90798) **What does this PR do?** This PR splits the FSDP optim_state_dict APIs into common implementation parts that are shared for different frontend APIs (we have many now and will consolidate them gradually). This PR also add `_optim_state_dict_post_hook` and `_load_optim_state_dict_pre_hook` for the integration with `NamedOptimzer`. Pull Request resolved: https://github.com/pytorch/pytorch/pull/90798 Approved by: https://github.com/rohan-varma, https://github.com/awgu
Author
Committer
Parents
Loading