pytorch/torch/distributed/checkpoint
Iris b8b7480065 [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212)
This is the last PR for integrating 2D into core distributed.

This PR does the following:
1. Add optimizer.py: this adds ability to load a state_dict in conjunction with FSDP sharded optimzer state.
2. Update default_planner.py to support 2D checkpoint.
3. Add test_fsdp_optim_state.py as a unit test for No. 1.
4. Fix bug in torch/testing/_internal/distributed/checkpoint_utils.py
5. Rename the filename for the APIs that should be private. Will organize and cleanup further in following PRs. #90328

Docstring and integration test will be added in the following PRs.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/90212
Approved by: https://github.com/wanchaol
2022-12-08 02:53:29 +00:00
..
__init__.py
_dedup_tensors.py [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212) 2022-12-08 02:53:29 +00:00
_nested_dict.py [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212) 2022-12-08 02:53:29 +00:00
_nested_tensor.py [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212) 2022-12-08 02:53:29 +00:00
_traverse.py [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212) 2022-12-08 02:53:29 +00:00
api.py
default_planner.py [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212) 2022-12-08 02:53:29 +00:00
filesystem.py [PT-D][Checkpoint]Add MultiThreaded FileSystemWriter for distributed checkpointing and Update tests (#87987) 2022-11-30 08:19:41 +00:00
metadata.py
optimizer.py [Checkpoint][2D][6/N] Add optimizer and update default_planner to core distributed (#90212) 2022-12-08 02:53:29 +00:00
planner_helpers.py
planner.py
resharding.py
state_dict_loader.py
state_dict_saver.py
storage.py
utils.py [Checkpoint][2D][3/N] Add nested_tensors for distributed checkpoint to core distributed (#89501) 2022-11-28 23:21:38 +00:00