[DCP] Update DCP to use the updated FSDP optim state_dict APIs (#95303)

Pull Request resolved: pytorch/pytorch#95303 Approved by: https://github.com/fegin
cyyever · Feb 25, 2023 · dd8df34 · dd8df34
1 parent 88fb1cc
commit dd8df34
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/test/distributed/checkpoint/test_2d_fsdp_dt_checkpoint.py b/test/distributed/checkpoint/test_2d_fsdp_dt_checkpoint.py
@@ -139,7 +139,7 @@ def _test_fsdp_dt_checkpoint(self, fsdp_pg=None) -> None:
         with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
             state_dict = {
                 "model": model.state_dict(),
-                "optim": FSDP.sharded_optim_state_dict(model, optim),
+                "optim": FSDP.optim_state_dict(model, optim),
             }
 
             dist_cp.save_state_dict(
@@ -181,7 +181,7 @@ def _test_fsdp_dt_checkpoint(self, fsdp_pg=None) -> None:
                 optimizer_key="optim",
                 storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
             )
-            flattened_osd = FSDP.flatten_sharded_optim_state_dict(
+            flattened_osd = FSDP.optim_state_dict_to_load(
                 optim_state["optim"], model_2, optim_2
             )
             optim_2.load_state_dict(flattened_osd)

diff --git a/test/distributed/checkpoint/test_fsdp_optim_state.py b/test/distributed/checkpoint/test_fsdp_optim_state.py
@@ -40,7 +40,7 @@ def test_distributed_tensor_planner(self) -> None:
         with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
             state_dict = {
                 "model": model.state_dict(),
-                "optim": FSDP.sharded_optim_state_dict(model, optim),
+                "optim": FSDP.optim_state_dict(model, optim),
             }
 
             dist_cp.save_state_dict(
@@ -80,7 +80,7 @@ def test_distributed_tensor_planner(self) -> None:
                 storage_reader=dist_cp.FileSystemReader(CHECKPOINT_DIR),
             )
 
-            flattened_osd = FSDP.flatten_sharded_optim_state_dict(
+            flattened_osd = FSDP.optim_state_dict_to_load(
                 optim_state["optim"], model_2, optim_2
             )
             optim_2.load_state_dict(flattened_osd)

diff --git a/torch/distributed/checkpoint/optimizer.py b/torch/distributed/checkpoint/optimizer.py
@@ -213,7 +213,7 @@ def load_sharded_optimizer_state_dict(
     >>> # Save
     >>> with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
     >>>     state_dict = {
-    >>>         "optimizer": FSDP.sharded_optim_state_dict(model, optim, optim_params),
+    >>>         "optimizer": FSDP.optim_state_dict(model, optim),
     >>>         "model": model.state_dict()
     >>>     }
     >>>     dist_cp.save_state_dict(
@@ -241,7 +241,7 @@ def load_sharded_optimizer_state_dict(
     >>>         storage_reader=dist_cp.FileSystemReader("checkpoint"),
     >>>     )
     >>>
-    >>>     flattened_osd = FSDP.flatten_sharded_optim_state_dict(
+    >>>     flattened_osd = FSDP.optim_state_dict_to_load(
     >>>        optim_state["optimizer"], model, optim
     >>>     )
     >>>