Config to allow partial loading of the checkpoint (#855)

Summary:
Pull Request resolved: #855

Enable TNT restore option to allow partial loading of the checkpoints for DCP

Reviewed By: anshulverma, JKSenthil

Differential Revision: D59074582

fbshipit-source-id: 6275e582665611f84b0ac594789dd4bc408bd77c

Author: saumishr

tests/framework/callbacks/test_dcp_saver.py

@@ -392,6 +392,33 @@ def test_restore_planner_storage_components(self, mock_dist_cp: MagicMock) -> No
         self.assertIsInstance(planner, DummyLoadPlanner)
         self.assertIsInstance(storage_reader, DummyStorageReader)
 
+    @patch("torchtnt.framework.callbacks.dcp_saver.dcp")
+    def test_restore_allow_partial_loading(self, mock_dist_cp: MagicMock) -> None:
+        my_unit = DummyTrainUnit(input_dim=2)
+        restore_options = RestoreOptions(strict=False)
+        DistributedCheckpointSaver.restore(
+            path="path/to/snapshot",
+            unit=my_unit,
+            restore_options=restore_options,
+        )
+
+        allow_partial_load = mock_dist_cp.load.call_args[1][
+            "planner"
+        ].allow_partial_load
+        self.assertTrue(allow_partial_load)
+
+        restore_options = RestoreOptions(strict=True)
+        DistributedCheckpointSaver.restore(
+            path="path/to/snapshot",
+            unit=my_unit,
+            restore_options=restore_options,
+        )
+
+        allow_partial_load = mock_dist_cp.load.call_args[1][
+            "planner"
+        ].allow_partial_load
+        self.assertFalse(allow_partial_load)
+
 
 class DummyStatefulDataLoader:
     def __init__(self, dataloader: DataLoader) -> None:

torchtnt/framework/callbacks/dcp_saver.py

@@ -283,14 +283,16 @@ def restore(
             storage_reader: Instance of StorageReader used to perform reads. If this is not specified, it will automatically infer
                             the reader based on the checkpoint_id. If checkpoint_id is also None, an exception will be raised. (Default: ``None``)
         """
-        if planner is None:
-            planner = DefaultLoadPlanner()
+
+        restore_options = restore_options or RestoreOptions()
+        app_state = _prepare_app_state_for_restore(unit, restore_options)
 
         if storage_reader is None:
             storage_reader = Reader(path)
 
-        restore_options = restore_options or RestoreOptions()
-        app_state = _prepare_app_state_for_restore(unit, restore_options)
+        if planner is None:
+            allow_partial_load = not restore_options.strict
+            planner = DefaultLoadPlanner(allow_partial_load=allow_partial_load)
 
         if train_dataloader is not None:
             if not isinstance(train_dataloader, Stateful):