OSS Zero Overhead Checkpointing Implementation (#1010)

Summary:

X-link: pytorch/pytorch#156207

This diff updates DCP driver code/APIs to support Zero Overhead Checkpointing

Differential Revision: D72391401

Author: Saiteja64

torchtnt/framework/callbacks/dcp_saver.py

@@ -25,6 +25,7 @@
     DefaultLoadPlanner,
     DefaultSavePlanner,
 )
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
 from torch.distributed.checkpoint.planner import LoadPlanner, SavePlanner
 from torchtnt.framework.state import EntryPoint
 
@@ -193,13 +194,15 @@ def _checkpoint_impl(
             with get_timing_context(state, f"{self.__class__.__name__}.async_save"):
                 # Redundant check for safety
                 self._wait(log_warning=True)
-                self._prev_snapshot = dcp.async_save(
+                prev_snapshot = dcp.async_save(
                     state_dict={"app_state": MultiStateful(app_state)},
                     checkpoint_id=checkpoint_id,
                     process_group=self._process_group,
                     storage_writer=storage_writer,
                     planner=planner,
                 )
+                assert isinstance(prev_snapshot, STATE_DICT_TYPE)
+                self._prev_snapshot = prev_snapshot
                 if curr_snapshot_wait:
                     self._wait(log_warning=False)
         else:

torchtnt/utils/prepare_module.py

@@ -13,7 +13,6 @@
 from typing import (
     Any,
     Callable,
-    cast,
     Collection,
     ContextManager,
     Dict,