TorchSnapshotSaver.restore non-strict option

Summary:
`strict` flag added to TorchSnapshotSaver.restore for when unit contains new stateful items that do not exist in previous snapshots

by default, strict is True to keep default behavior. To enable this feature, can set strict to false

Reviewed By: jiaxuzhu92

Differential Revision: D55501641

fbshipit-source-id: c475f82701f6b28490224b87f45a9de345de7023

Author: JKSenthil

tests/framework/callbacks/test_torchsnapshot_saver.py

@@ -227,6 +227,39 @@ def test_save_restore_no_lr_scheduler_restore(
         app_state = mock_torchsnapshot.Snapshot().restore.call_args.args[0]
         self.assertIn("lr_scheduler", app_state)
 
+    def test_restore_strict(self) -> None:
+        my_unit = DummyAutoUnit(module=torch.nn.Linear(2, 2))
+        with tempfile.TemporaryDirectory() as temp_dir:
+            state = get_dummy_train_state()
+            snapshot_cb = TorchSnapshotSaver(
+                temp_dir, save_every_n_train_steps=1, async_checkpoint=False
+            )
+            snapshot_cb.on_train_step_end(state, my_unit)
+
+            # add a new parameter to the module
+            my_unit.module2 = torch.nn.Linear(2, 2)
+
+            with self.assertRaisesRegex(
+                AssertionError,
+                "module2 is absent in both manifest and flattened.",
+            ):
+                TorchSnapshotSaver.restore(
+                    path=os.path.join(temp_dir, "epoch_0_step_0"),
+                    unit=my_unit,
+                    strict=True,
+                )
+
+            with self.assertLogs(level="WARNING") as log:
+                TorchSnapshotSaver.restore(
+                    path=os.path.join(temp_dir, "epoch_0_step_0"),
+                    unit=my_unit,
+                    strict=False,
+                )
+                self.assertEqual(
+                    log.output[0],
+                    "WARNING:torchtnt.utils.rank_zero_log:module2 was passed to `restore` but does not exists in the snapshot",
+                )
+
     @skip_if_not_distributed
     def test_save_restore_ddp(self) -> None:
         spawn_multi_process(

torchtnt/framework/callbacks/torchsnapshot_saver.py

@@ -267,6 +267,7 @@ def restore(
         restore_options: Optional[RestoreOptions] = None,
         storage_options: Optional[Dict[str, Any]] = None,
         knob_options: Optional[KnobOptions] = None,
+        strict: bool = True,
     ) -> None:
         """Utility method to restore snapshot state from a path.
 
@@ -281,6 +282,7 @@ def restore(
             restore_options: Controls what to  filter when restoring the state.
             storage_options: Additional keyword options for the storage plugin to use, to be passed to `torchsnapshot.Snapshot <https://pytorch.org/torchsnapshot/stable/api_reference.html#torchsnapshot.Snapshot>`_. See each storage plugin's documentation for customizations.
             knob_options: Additional keyword options for the snapshot knobs
+            strict: If ``False``, allows loading a snapshot even if not all keys exist in the unit's app_state.
         """
 
         _validate_snapshot_available()
@@ -319,9 +321,25 @@ def restore(
                         "train_dataloader was passed to `restore` but no train dataloader exists in the Snapshot"
                     )
 
+        if not strict:
+            # if app_state keys not in torchsnapshot checkpoint,
+            # remove them from app_state prior to checkpoint load
+            missing_stateful_keys = []
+            manifest = snapshot.get_manifest()
+            for stateful_key in app_state:
+                found = any((f"/{stateful_key}/" in key for key in manifest.keys()))
+                if not found:
+                    missing_stateful_keys.append(stateful_key)
+
+            for key in missing_stateful_keys:
+                rank_zero_warn(
+                    f"{key} was passed to `restore` but does not exists in the snapshot"
+                )
+                app_state.pop(key)
+
         knob_options = knob_options or KnobOptions()
         with _override_knobs(knob_options):
-            snapshot.restore(app_state)
+            snapshot.restore(app_state, strict=strict)
         rank_zero_info(f"Restored snapshot from path: {path}", logger=logger)