add init_optim_state flag to restore options (#992)

Summary: Pull Request resolved: #992

Reviewed By: diego-urgell

Differential Revision: D73678205

fbshipit-source-id: 28e1b748d2b6f9d347f4c78c049113f25c6a8457

Author: JKSenthil

tests/framework/callbacks/test_dcp_saver.py

@@ -255,6 +255,27 @@ def test_save_restore_no_lr_scheduler_restore(
         app_state = mock_dist_cp.load.call_args.args[0]["app_state"].state_dict()
         self.assertIn("lr_scheduler", app_state)
 
+    @patch("torchtnt.framework.callbacks.dcp_saver._init_optim_state")
+    @patch("torchtnt.framework.callbacks.dcp_saver.dcp")
+    def test_save_restore_no_init_optim_state(
+        self, _: MagicMock, mock_init_optim_state: MagicMock
+    ) -> None:
+        my_unit = DummyTrainUnit(input_dim=2)
+        restore_options = RestoreOptions(init_optim_states=False)
+        DistributedCheckpointSaver.restore(
+            path="path/to/snapshot",
+            unit=my_unit,
+            restore_options=restore_options,
+        )
+        mock_init_optim_state.assert_not_called()
+
+        DistributedCheckpointSaver.restore(
+            path="path/to/snapshot",
+            unit=my_unit,
+            restore_options=RestoreOptions(),
+        )
+        mock_init_optim_state.assert_called()
+
     @skip_if_not_distributed
     def test_save_restore_ddp(self) -> None:
         spawn_multi_process(

torchtnt/framework/callbacks/checkpointer_types.py

@@ -43,6 +43,9 @@ class RestoreOptions:
         restore_optimizers: Whether to restore the optimizer states.
         restore_lr_schedulers: Whether to restore the lr scheduler states.
         strict: Whether to strictly restore app state and the module state dict.
+        init_optim_states: Whether to initialize the optimizer state. Defaults to True. Toggle off
+            if running into issues with loading optimizer state. This will reset optimizer state,
+            which may affect training in some cases.
     """
 
     restore_modules: bool = True
@@ -52,3 +55,4 @@ class RestoreOptions:
     restore_optimizers: bool = True
     restore_lr_schedulers: bool = True
     strict: bool = True
+    init_optim_states: bool = True

torchtnt/framework/callbacks/dcp_saver.py

@@ -370,14 +370,15 @@ def restore_with_id(
             predict_dataloader,
         )
 
-        # necessary for loading optimizers since states are initialized lazy
-        for obj in app_state.values():
-            # sometimes optimizers are actually held in a wrapper which handles calling
-            # state_dict and load_state_dict, sa is the case for
-            # `torchtnt.utils.prepare_module.FSDPOptimizerWrapper`, this handles that case.
-            optimizer = getattr(obj, "optimizer", obj)
-            if isinstance(optimizer, torch.optim.Optimizer):
-                _init_optim_state(optimizer)
+        if restore_options.init_optim_states:
+            # if optimizers states are initialized lazy
+            for obj in app_state.values():
+                # sometimes optimizers are actually held in a wrapper which handles calling
+                # state_dict and load_state_dict, sa is the case for
+                # `torchtnt.utils.prepare_module.FSDPOptimizerWrapper`, this handles that case.
+                optimizer = getattr(obj, "optimizer", obj)
+                if isinstance(optimizer, torch.optim.Optimizer):
+                    _init_optim_state(optimizer)
 
         with get_or_create_gloo_pg(candidate_pg=process_group) as pg:
             dcp.load(