OSS Zero Overhead Checkpointing Implementation (#1010)

Summary:
Pull Request resolved: #1010

X-link: pytorch/pytorch#156207

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

Reviewed By: diego-urgell

Differential Revision: D72391401

fbshipit-source-id: 3e872e0faa9558a61e6bcf9026dab001f4f52f0b

Author: Saiteja64

torchtnt/framework/callbacks/dcp_saver.py

@@ -193,13 +193,14 @@ 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,
                 )
+                self._prev_snapshot = cast(Future, 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,