[FSDP][optim_state_dict][10/N] Make optim_state_dict and optim_state_dict_to_load public

fegin · fegin · commit ba23ecbc59e0 · 2023-01-12T16:12:33.000-08:00
Make optim_state_dict and optim_state_dict_to_load public APIs and consolidate them with state_dict by using the same state_dict_type to decide how to perform the optimizer state_dict save and load. Differential Revision: [D42488022](https://our.internmc.facebook.com/intern/diff/D42488022/) ghstack-source-id: 177584342 Pull Request resolved: #92118
diff --git a/test/distributed/fsdp/test_fsdp_optim_state.py b/test/distributed/fsdp/test_fsdp_optim_state.py
@@ -783,9 +783,8 @@ def test_flatten_sharded_optim_state_dict_transformer(self) -> None:
             num_iters=3,
         )
 
-    @unittest.skip("The test currently fails on CI.")
     @skip_if_lt_x_gpu(2)
-    def test_use_orig_params(self) -> None:
+    def ftest_use_orig_params(self) -> None:
         """Tests :meth:`optim_state_dict` for an FSDP-root nested model."""
         self._test_load_optim_state(
             _ModelClass.NESTED,
@@ -824,7 +823,7 @@ def _test_load_optim_state(
         """
         initializer = self._model_class[model_class]
         if osd_comm_method == _OSDCommMethod.OPTIM_STATE_DICT:
-            osd_method = FSDP._optim_state_dict
+            osd_method = FSDP.optim_state_dict
         elif osd_comm_method == _OSDCommMethod.FLATTEN_SHARDED_OSD:
             osd_method = FSDP.sharded_optim_state_dict
         else:
diff --git a/torch/distributed/fsdp/_common_utils.py b/torch/distributed/fsdp/_common_utils.py
@@ -24,7 +24,7 @@
     _CHECKPOINT_PREFIX,
 )
 
-from .api import FullStateDictConfig, ShardingStrategy, StateDictConfig, StateDictType
+from .api import FullStateDictConfig, ShardingStrategy, StateDictConfig, StateDictType, OptimStateDictConfig, FullOptimStateDictConfig
 
 FSDP_WRAPPED_MODULE = "_fsdp_wrapped_module"
 FSDP_PREFIX = FSDP_WRAPPED_MODULE + "."
@@ -39,6 +39,7 @@ def __init__(self) -> None:
         self._unshard_params_ctx: Dict[nn.Module, Generator] = {}
         self._state_dict_type: StateDictType = StateDictType.FULL_STATE_DICT
         self._state_dict_config: StateDictConfig = FullStateDictConfig()
+        self._state_dict_config: OptimStateDictConfig = FullOptimStateDictConfig()
         self._is_root: Optional[bool] = None
         self._handles: List[flat_param_file.FlatParamHandle] = []
         self._ignored_modules: Set[nn.Module] = set()
diff --git a/torch/distributed/fsdp/_init_utils.py b/torch/distributed/fsdp/_init_utils.py
@@ -37,6 +37,7 @@
     BackwardPrefetch,
     CPUOffload,
     FullStateDictConfig,
+    FullOptimStateDictConfig,
     MixedPrecision,
     ShardingStrategy,
     StateDictConfig,
@@ -374,6 +375,7 @@ def _init_prefetching_state(
 def _init_state_dict_state(state: _FSDPState) -> _FSDPState:
     state._state_dict_type = StateDictType.FULL_STATE_DICT
     state_dict_config: StateDictConfig = FullStateDictConfig()
+    state._optim_state_dict_config = FullOptimStateDictConfig()
     state._state_dict_config = state_dict_config
     unshard_params_ctx: Dict[nn.Module, Generator] = {}
     state._unshard_params_ctx = unshard_params_ctx
diff --git a/torch/distributed/fsdp/api.py b/torch/distributed/fsdp/api.py
@@ -284,3 +284,35 @@ class LocalStateDictConfig(StateDictConfig):
 @dataclass
 class ShardedStateDictConfig(StateDictConfig):
     pass
+
+
+@dataclass
+class OptimStateDictConfig:
+    """
+    ``OptimStateDictConfig`` is the base class for all optimizer state_dict
+    configuration classes.  Users should instantiate a child version
+    (i.e. ``FullOptimStateDictConfig``) in order to configure settings for the
+    particular type of ``optim_state_dict`` implementation FSDP will use.
+    """
+    # TODO: actually use this flag in the _optim_utils.py
+    offload_to_cpu: bool = True
+
+
+@dataclass
+class FullOptimStateDictConfig(OptimStateDictConfig):
+    rank0_only: bool = False
+
+@dataclass
+class LocalOptimStateDictConfig(OptimStateDictConfig):
+    offload_to_cpu: bool = False
+
+@dataclass
+class ShardedOptimStateDictConfig(OptimStateDictConfig):
+    pass
+
+
+@dataclass
+class StateDictSettings:
+    state_dict_type: StateDictType
+    state_dict_config: StateDictConfig
+    optim_state_dict_config: OptimStateDictConfig
diff --git a/torch/distributed/fsdp/fully_sharded_data_parallel.py b/torch/distributed/fsdp/fully_sharded_data_parallel.py
