Use megatron method to cache plan

Author: Leahlijuan

src/ml_flashpoint/adapter/megatron/save_strategies.py

@@ -28,6 +28,8 @@
     _replace_state_dict_keys_with_sharded_keys,
     mcore_to_pyt_state_dict,
 )
+from torch.distributed.checkpoint.metadata import Metadata
+from torch.distributed.checkpoint.planner import SavePlan
 from torch.distributed.checkpoint.utils import _DistWrapper
 from typing_extensions import override
 
@@ -94,6 +96,13 @@ def __init__(
         self._storage_writer: MemoryStorageWriter = storage_writer
         self._checkpoint_saver: MLFlashpointCheckpointSaver = storage_writer.checkpoint_saver
 
+        # Cache for state dict saving
+        self.cached_central_plan: SavePlan | None = None
+        self.cached_local_plan: SavePlan | None = None
+        self.cached_global_metadata: Metadata | None = None
+        self.validated_cache_reuse: bool = False
+        self.use_cached_ckpt_structure: bool = True
+
     @override
     def can_handle_sharded_objects(self) -> bool:
         # Not currently used, but in case it is, ensure this strategy is used for ShardedObjects as well.
@@ -157,14 +166,45 @@ def async_save(self, sharded_state_dict: ShardedStateDict, checkpoint_dir: Union
         # we also use Megatron's SavePlanner during saving for compatibility.
         planner: MCoreSavePlanner = MCoreSavePlanner(can_run_decentralized_global_plan=False)
         world_dist_wrapper = _DistWrapper(group=None, use_dist=not disable_dist, coordinator_rank=0)
-        plan, write_buckets, global_metadata = statedictsaver.generate_plan(
+        # Try twice to validate the generated `central_plan` is the same across iterations
+        # If so, reuse `cached_central_plan` and `cached_global_metadata`
+        # From the 3rd iteration, `save_state_dict_async_plan` will not generate `global_metadata`
+        # (return None) so `self.cached_global_metadata` is reused
+        args_cached_plans = None
+        loaded_all_plans = None
+        if self.use_cached_ckpt_structure:
+            if self.cached_global_metadata:
+                loaded_all_plans = getattr(self.cached_global_metadata, "all_local_plans", None)
+
+            args_cached_plans = (
+                self.cached_central_plan,
+                self.cached_local_plan,
+                self.validated_cache_reuse,
+            )
+
+        (
+            (plan, write_buckets, global_metadata),
+            self.cached_central_plan,
+            self.cached_local_plan,
+            self.validated_cache_reuse,
+        ) = statedictsaver.generate_plan(
             checkpoint_id=checkpoint_id,
             state_dict=pyt_state_dict,
             storage_writer=self._storage_writer,
             planner=planner,
             world_dist_wrapper=world_dist_wrapper,
+            cached_ckpt_structure=args_cached_plans,
+            loaded_all_plans=loaded_all_plans,
         )
 
+        if self.validated_cache_reuse:
+            if global_metadata is None and self.cached_global_metadata:
+                global_metadata = self.cached_global_metadata
+
+        # If we have a valid global_metadata (either new or reused), cache it for next time
+        if global_metadata is not None:
+            self.cached_global_metadata = global_metadata
+
         # 5. Stage to CPU.
         staged_write_buckets = self._storage_writer.stage_write_data_buckets(
             checkpoint_id, write_buckets, non_blocking=True

src/ml_flashpoint/adapter/pytorch/custom_state_dict_saver.py

@@ -22,7 +22,6 @@
 
 import torch.cuda
 from torch import distributed as torchdist
-from torch.distributed.checkpoint import Metadata
 from torch.distributed.checkpoint import state_dict_saver as torchdistsaver
 from torch.distributed.checkpoint.logger import _dcp_method_logger
 from torch.distributed.checkpoint.planner import SavePlan
@@ -46,7 +45,14 @@ def generate_plan(
     storage_writer: MemoryStorageWriter,
     planner: torchdistsaver.SavePlanner,
     world_dist_wrapper: _DistWrapper,
-) -> tuple[SavePlan, list[ObjectWriteBucket], Metadata]:
+    cached_ckpt_structure: tuple[SavePlan, SavePlan, bool] | None = None,
+    loaded_all_plans: list[SavePlan] | None = None,
+) -> tuple[
+    tuple[SavePlan, list[ObjectWriteBucket], torchdistsaver.Metadata | None],
+    SavePlan,
+    SavePlan,
+    bool,
+]:
     """Performs the planning phase of checkpointing.
 
     This function is similar to PyTorch's `state_dict_saver.save` but only
@@ -62,15 +68,28 @@ def generate_plan(
         planner: The SavePlanner to use for the save.
         world_dist_wrapper: The distributed wrapper for world (all ranks) communication.
             Typically created as `_DistWrapper(process_group, not no_dist, coordinator_rank)`.
+        cached_ckpt_structure: Tuple of (cached_central_plan, cached_local_plan, validated_cache_reuse).
+        loaded_all_plans: List of all local plans from the previous checkpoint (for validation).
+
     Returns:
-        A tuple containing the updated local plan, write buckets, and global metadata.
+        A tuple containing:
+            - (final_local_plan, write_buckets, global_metadata)
+            - final_local_plan (for caching)
+            - local_plan (for caching)
+            - validated_cache_reuse (bool)
     """
+    cached_central_plan, cached_local_plan, validated_cache_reuse = (None, None, False)
+    if cached_ckpt_structure:
+        cached_central_plan, cached_local_plan, validated_cache_reuse = cached_ckpt_structure
+
     global_metadata: torchdistsaver.Metadata | None = None
 
     ckpt_kwargs = {"checkpoint_id": storage_writer.current_checkpoint_id, "process_group": world_dist_wrapper.group}
+    local_plan = cached_local_plan
 
     @_dcp_method_logger(**ckpt_kwargs)
     def local_step() -> SavePlan:
+        nonlocal local_plan
         storage_meta = storage_writer.storage_meta()
         planner.set_up_planner(
             state_dict=state_dict,
@@ -79,7 +98,9 @@ def local_step() -> SavePlan:
         )
         storage_writer.set_up_storage_writer(world_dist_wrapper.is_coordinator)
 
-        local_plan = planner.create_local_plan()
+        if not validated_cache_reuse:
+            local_plan = planner.create_local_plan()
+
         local_plan = storage_writer.prepare_local_plan(local_plan)
         return local_plan
 
@@ -91,19 +112,29 @@ def global_step(all_local_plans: list[SavePlan]) -> list[SavePlan]:
         all_local_plans = storage_writer.prepare_global_plan(all_local_plans)
         return all_local_plans
 
-    with log_execution_time(logger=_LOGGER, name="generate_plan__reduce_scatter_plan"):
-        _LOGGER.debug("Executing plan reduce_scatter to get updated_local_plan...")
-        updated_local_plan = world_dist_wrapper.reduce_scatter("plan", local_step, global_step)
-
-    with log_execution_time(logger=_LOGGER, name="generate_plan__broadcast_metadata"):
-        _LOGGER.debug("Executing global_metadata broadcast...")
-        # TODO(perf): - can broadcast only to local rank 0 to reduce comms
-        global_metadata = world_dist_wrapper.broadcast_object(global_metadata)
-
-    final_local_plan = planner.finish_plan(updated_local_plan)
+    central_plan = None
+    if validated_cache_reuse and cached_central_plan:
+        _LOGGER.debug("Passed cache reusable")
+        local_step()
+        central_plan = cached_central_plan
+    else:
+        with log_execution_time(logger=_LOGGER, name="generate_plan__reduce_scatter_plan"):
+            _LOGGER.debug("Executing plan reduce_scatter to get central_plan...")
+            central_plan = world_dist_wrapper.reduce_scatter("plan", local_step, global_step)
+
+        with log_execution_time(logger=_LOGGER, name="generate_plan__broadcast_metadata"):
+            _LOGGER.debug("Executing global_metadata broadcast...")
+            global_metadata = world_dist_wrapper.broadcast_object(global_metadata)
+
+    final_local_plan = planner.finish_plan(central_plan)
     write_buckets = storage_writer.prepare_write_data_buckets(checkpoint_id, final_local_plan, planner)
 
-    return final_local_plan, write_buckets, global_metadata
+    return (
+        (final_local_plan, write_buckets, global_metadata),
+        central_plan,  # cached_central_plan
+        local_plan,  # cached_local_plan
+        cached_central_plan == central_plan,  # validated_cache_reuse
+    )
 
 
 @log_execution_time(logger=_LOGGER, name="write_data", level=logging.INFO)

tests/adapter/megatron/test_save_strategies.py

@@ -169,9 +169,10 @@ def test_async_save_initialization_calls_success(
                 _,
             ) = async_save_setup
             mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), dummy_write_buckets, mocker.MagicMock()),
                 mocker.MagicMock(),
-                dummy_write_buckets,
                 mocker.MagicMock(),
+                False,
             )
 
             mock_memory_storage_writer_cls = mocker.patch(
@@ -211,9 +212,10 @@ def test_async_save_reinitializes_storage_writer_with_thread_count(
                 _,
             ) = async_save_setup
             mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), dummy_write_buckets, mocker.MagicMock()),
                 mocker.MagicMock(),
-                dummy_write_buckets,
                 mocker.MagicMock(),
+                False,
             )
 
             # Set a specific thread_count on the original storage_writer
@@ -256,12 +258,21 @@ def test_async_save_generate_plan_call_success(self, mocker, async_save_setup, s
             ) = async_save_setup
             mock_planner = MockMCoreSavePlanner.return_value
             mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()),
                 mocker.MagicMock(),
                 mocker.MagicMock(),
-                mocker.MagicMock(),
+                False,
             )
 
-            expected_kwarg_keys = {"checkpoint_id", "state_dict", "storage_writer", "planner", "world_dist_wrapper"}
+            expected_kwarg_keys = {
+                "checkpoint_id",
+                "state_dict",
+                "storage_writer",
+                "planner",
+                "world_dist_wrapper",
+                "cached_ckpt_structure",
+                "loaded_all_plans",
+            }
 
             # When
             strategy.async_save(sharded_state_dict, checkpoint_id.data)
@@ -281,6 +292,9 @@ def test_async_save_generate_plan_call_success(self, mocker, async_save_setup, s
             assert kwargs["planner"] is mock_planner
             assert "world_dist_wrapper" in kwargs
             assert kwargs["world_dist_wrapper"].use_dist is False
+            assert "cached_ckpt_structure" in kwargs
+            assert "loaded_all_plans" in kwargs
+            assert "cached_global_metadata" not in kwargs
 
         def test_generate_plan_failure(self, mocker, async_save_setup):
             """Tests that an exception in generate_plan is propagated."""
@@ -303,7 +317,12 @@ def test_async_save_async_fn_call_success(
 
             mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
             strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
-            mock_statedictsaver.generate_plan.return_value = (dummy_save_plan, dummy_write_buckets, dummy_metadata)
+            mock_statedictsaver.generate_plan.return_value = (
+                (dummy_save_plan, dummy_write_buckets, dummy_metadata),
+                mocker.MagicMock(),
+                mocker.MagicMock(),
+                False,
+            )
             staged_write_buckets = [
                 ObjectWriteBucket(
                     object_id=CheckpointObjectId(f"/test_checkpoint/staged_obj_{i}"),
@@ -339,9 +358,10 @@ def test_async_save_async_fn_failure(self, mocker, async_save_setup, checkpoint_
             mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
             strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
             mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()),
                 mocker.MagicMock(),
                 mocker.MagicMock(),
-                mocker.MagicMock(),
+                False,
             )
             mock_statedictsaver.write_data.side_effect = Exception("Test Exception")
 
@@ -368,7 +388,12 @@ def test_async_save_finalize_fns_calls(
             finalize_checkpoint_spy = mocker.spy(checkpoint_saver, "finalize_checkpoint")
             mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
             strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
-            mock_statedictsaver.generate_plan.return_value = (dummy_save_plan, dummy_write_buckets, dummy_metadata)
+            mock_statedictsaver.generate_plan.return_value = (
+                (dummy_save_plan, dummy_write_buckets, dummy_metadata),
+                mocker.MagicMock(),
+                mocker.MagicMock(),
+                False,
+            )
 
             mock_memory_storage_writer_cls = mocker.patch(
                 "ml_flashpoint.adapter.megatron.save_strategies.MemoryStorageWriter"
@@ -427,7 +452,12 @@ def test_finalize_fns_failure(
             finalize_checkpoint_spy = mocker.spy(checkpoint_saver, "finalize_checkpoint")
             mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
             strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
-            mock_statedictsaver.generate_plan.return_value = (dummy_save_plan, mocker.MagicMock(), dummy_metadata)
+            mock_statedictsaver.generate_plan.return_value = (
+                (dummy_save_plan, mocker.MagicMock(), dummy_metadata),
+                mocker.MagicMock(),
+                mocker.MagicMock(),
+                False,
+            )
             mock_statedictsaver.finish_write.side_effect = ValueError("Finish Write Failed")
 
             # When
@@ -465,13 +495,73 @@ def test_async_save_rank_determination(
             # Mock dependencies to ensure success path
             mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
             mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), mocker.MagicMock(), mocker.MagicMock()),
                 mocker.MagicMock(),
                 mocker.MagicMock(),
-                mocker.MagicMock(),
+                False,
             )
 
             # When
             actual_async_request = strategy.async_save(sharded_state_dict, checkpoint_id.data)
 
             # Then
             assert actual_async_request.async_fn_kwargs["rank"] == expected_rank
+
+        def test_async_save_caching_flow(self, mocker, async_save_setup, storage_writer):
+            """Tests the caching flow across multiple async_save calls."""
+            # Given
+            mock_statedictsaver = mocker.patch("ml_flashpoint.adapter.megatron.save_strategies.statedictsaver")
+            strategy, checkpoint_id, sharded_state_dict, _ = async_save_setup
+
+            cached_plan = mocker.MagicMock()
+            cached_metadata = mocker.MagicMock()
+
+            # First call: No cache
+            mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), [], mocker.MagicMock()),
+                cached_plan,  # cached_central_plan returned
+                mocker.MagicMock(),
+                False,
+            )
+
+            # When 1
+            strategy.async_save(sharded_state_dict, checkpoint_id.data)
+
+            # Then 1
+            assert strategy.cached_central_plan == cached_plan
+            assert strategy.validated_cache_reuse is False
+
+            # Second call: Cache validation success
+            mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), [], cached_metadata),
+                cached_plan,
+                mocker.MagicMock(),
+                True,  # validated_cache_reuse
+            )
+
+            # When 2
+            strategy.async_save(sharded_state_dict, checkpoint_id.data)
+
+            # Then 2
+            assert strategy.validated_cache_reuse is True
+            assert strategy.cached_global_metadata == cached_metadata
+
+            # Third call: Reuse cache
+            mock_statedictsaver.generate_plan.return_value = (
+                (mocker.MagicMock(), [], None),  # Returns None for metadata
+                cached_plan,
+                mocker.MagicMock(),
+                True,
+            )
+
+            # During third call, async_save should use self.cached_global_metadata
+
+            # When 3
+            strategy.async_save(sharded_state_dict, checkpoint_id.data)
+
+            # Then 3
+            # Ensure generate_plan was called without cached_global_metadata
+            _, kwargs = mock_statedictsaver.generate_plan.call_args
+            assert "cached_global_metadata" not in kwargs
+            # And cached_global_metadata in strategy should still be the same
+            assert strategy.cached_global_metadata == cached_metadata