Repository: NVIDIA/nvidia-resiliency-ext
Pull request #40: Synchronize changes in MCore to NVRx for Metadata caching and newer changes in PyT DCP routines
Author: login: sbak5
Labels: name: ci-approved
description: Approved to run CI
Created: 2025-05-06T23:38:47Z | Merged: 2025-05-09T04:46:57Z

This PR includes changes made for refactoring of metadata caching in Mcore and PyT DCP routine signature changes


Diff:
diff --git a/examples/checkpointing/async_writer.py b/examples/checkpointing/async_writer.py
index 83e39dbe..9e900fbc 100644
--- a/examples/checkpointing/async_writer.py
+++ b/examples/checkpointing/async_writer.py
@@ -5,8 +5,11 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
+import torch.optim as optim
 from torch.distributed import checkpoint
 from torch.distributed.checkpoint import DefaultLoadPlanner, DefaultSavePlanner, FileSystemReader
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
+from torch.utils.data import DataLoader, DistributedSampler
 
 from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncCallsQueue, AsyncRequest
 from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import FileSystemWriterAsync
@@ -18,6 +21,18 @@
 # Set up basic logging configuration
 logging.basicConfig(level=logging.INFO)
 
+FEAT_SIZE = 4096
+DNN_OUT_SIZE = 128
+BATCH_SIZE = 100
+NUM_EPOCHS = 10
+DATASET_LEN = 100000
+CKPT_INTERVAL = 100
+
+
+def print_on_rank0(msg):
+    if dist.get_rank() == 0:
+        print(msg)
+
 
 def parse_args():
     parser = argparse.ArgumentParser(description="Async Checkpointing Example")
@@ -41,18 +56,30 @@ def parse_args():
     return parser.parse_args()
 
 
-class SimpleModel(nn.Module):
-    """A simple feedforward neural network for demonstration purposes."""
+class SimpleDataset(torch.utils.data.Dataset):
+    def __init__(self, size):
+        self.size = size
+
+    def __len__(self):
+        return self.size
+
+    def __getitem__(self, idx):
+        x = torch.rand((FEAT_SIZE,), dtype=torch.float32, device='cuda')
+        y = torch.rand((DNN_OUT_SIZE,), dtype=torch.float32, device='cuda')
+        return x, y
 
+
+class SimpleModel(nn.Module):
     def __init__(self):
-        super(SimpleModel, self).__init__()
-        self.fc1 = nn.Linear(10, 5)
-        self.fc2 = nn.Linear(5, 2)
-        self.activation = nn.ReLU()
+        super().__init__()
+        self.fc1 = nn.Linear(FEAT_SIZE, FEAT_SIZE)
+        self.fc2 = nn.Linear(FEAT_SIZE, DNN_OUT_SIZE)
 
     def forward(self, x):
-        x = self.activation(self.fc1(x))
-        return self.fc2(x)
+        x = self.fc1(x)
+        x = nn.functional.relu(x)
+        x = self.fc2(x)
+        return x
 
 
 def init_distributed_backend(backend="nccl"):
@@ -62,6 +89,7 @@ def init_distributed_backend(backend="nccl"):
         rank = dist.get_rank()
         torch.cuda.set_device(rank)
         logging.info(f"Process {rank} initialized with {backend} backend.")
+        return rank, torch.distributed.get_world_size()
     except Exception as e:
         logging.error(f"Failed to initialize distributed backend: {e}")
         raise
@@ -86,8 +114,8 @@ def save_checkpoint(checkpoint_dir, async_queue, model, thread_count):
     writer = FileSystemWriterAsync(checkpoint_dir, thread_count=thread_count)
     coordinator_rank = 0
 
-    save_state_dict_ret, *_ = save_state_dict_async_plan(
-        state_dict, writer, None, coordinator_rank, planner=planner
+    save_state_dict_ret = save_state_dict_async_plan(
+        state_dict, writer, None, coordinator_rank, planner=planner, enable_cache=True
     )
     save_request = get_save_and_finalize_callbacks(writer, save_state_dict_ret)
     async_queue.schedule_async_request(save_request)
@@ -109,20 +137,46 @@ def main():
     logging.info(f"Arguments: {args}")
 
     # Initialize distributed training
-    init_distributed_backend(backend="nccl")
+    rank, world_size = init_distributed_backend(backend="nccl")
+
+    # Define checkpoint directory based on iteration number
+    dataset = SimpleDataset(size=DATASET_LEN)
+    sampler = DistributedSampler(dataset, num_replicas=world_size, rank=rank)
+    dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, sampler=sampler)
 
-    # Initialize model and move to GPU
+    # Model, optimizer, and DDP
     model = SimpleModel().to("cuda")
+    fsdp_model = FSDP(model)
+    optimizer = optim.SGD(fsdp_model.parameters(), lr=0.01)
+    loss_fn = nn.MSELoss()
 
     # Create an async queue for handling asynchronous operations
     async_queue = AsyncCallsQueue(persistent=args.persistent_queue)
 
-    # Define checkpoint directory based on iteration number
-    iteration = 123
-    checkpoint_dir = f"{args.ckpt_dir}/iter_{iteration:07d}"
-
-    # Save the model asynchronously
-    save_checkpoint(checkpoint_dir, async_queue, model, args.thread_count)
+    iteration = 0
+    num_iters_in_epoch = len(dataloader)
+    print_on_rank0(f"num_iters_in_epoch: {num_iters_in_epoch}")
+
+    num_iters_for_10pct = num_iters_in_epoch // 10  # iters for 1/10 of epoch
+    checkpoint_dir = None
+    sampler.set_epoch(0)
+    for batch_idx, (data, target) in enumerate(dataloader):
+        async_queue.maybe_finalize_async_calls(blocking=False, no_dist=False)
+        if (batch_idx % num_iters_for_10pct) == 0 and rank == 0:
+            print(f"Epoch 0 progress: {100 * batch_idx / num_iters_in_epoch:.2f}%")
+        optimizer.zero_grad()
+        output = fsdp_model(data)
+        loss = loss_fn(output, target)
+        loss.backward()
+        optimizer.step()
+        if batch_idx % num_iters_for_10pct == 0:
+            iteration = batch_idx
+            checkpoint_dir = f"{args.ckpt_dir}/iter_{iteration:07d}"
+            # Save the model asynchronously
+            save_checkpoint(checkpoint_dir, async_queue, fsdp_model, args.thread_count)
+            print_on_rank0(f"Checkpoint Save triggered: {checkpoint_dir}, iteration: {iteration}")
+            iteration += batch_idx
+    print_on_rank0(f"Epoch 0 complete. Loss: {loss.item()}")
 
     logging.info("Finalizing checkpoint save...")
     async_queue.maybe_finalize_async_calls(blocking=True, no_dist=False)
@@ -130,9 +184,9 @@ def main():
 
     # Synchronize processes before loading
     dist.barrier()
-
+    print_on_rank0(f"loading from {checkpoint_dir}")
     # Load the checkpoint
-    loaded_sd = load_checkpoint(checkpoint_dir, model)
+    loaded_sd = load_checkpoint(checkpoint_dir, fsdp_model)
 
     # Synchronize again to ensure all ranks have completed loading
     dist.barrier()
diff --git a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/cached_metadata_filesystem_reader.py b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/cached_metadata_filesystem_reader.py
new file mode 100644
index 00000000..11c3b3e3
--- /dev/null
+++ b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/cached_metadata_filesystem_reader.py
@@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" FS Reader with metadata cached support. """
+
+import os
+from typing import Union
+
+from torch.distributed.checkpoint import FileSystemReader, Metadata
+
+
+class CachedMetadataFileSystemReader(FileSystemReader):
+    """
+    Extends FileSystemReader to cache metadata for improved performance.
+
+    Attributes:
+        _cached_metadata (Metadata or None): Cached metadata from the file system.
+    """
+
+    def __init__(self, path: Union[str, os.PathLike]) -> None:
+        """
+        Initialize with file system path.
+
+        Args:
+            path (Union[str, os.PathLike]): Path to the checkpoint directory or file.
+        """
+        super().__init__(path=path)
+        self._cached_metadata = None
+
+    def read_metadata(self) -> Metadata:
+        """
+        Read metadata from file system, caching for subsequent calls.
+
+        Returns:
+            Metadata: Checkpoint metadata.
+        """
+        if self._cached_metadata is None:
+            self._cached_metadata = super().read_metadata()
+        return self._cached_metadata
diff --git a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py
index 42545043..abdb5ac2 100644
--- a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py
+++ b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py
@@ -274,7 +274,8 @@ class PersistentAsyncCaller(AsyncCaller):
     """
 
     def __init__(self):
-        super().__init__()
+        self.process: mp.Process = None
+        self.start_time: Optional[float] = None
         ctx = mp.get_context('spawn')
         # main queue to deliver `AsyncRequest` from host to the ckpt worker
         self.queue: mp.JoinableQueue = ctx.JoinableQueue()
@@ -392,6 +393,10 @@ def is_current_async_call_done(self, blocking: bool = False, no_dist: bool = Fal
         return is_done
 
     def close(self):
+        """Wait on the left async requests and terminate the PersistentAsyncCaller
+
+        Signals the PersistentAsyncCaller by sending a 'DONE' message to make it terminated
+        """
         logger.info(f"PersistentAsyncCaller: {self.rank}, Destroying Async Caller")
         if self.process:
             self.queue.put('DONE')
diff --git a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/filesystem_async.py b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/filesystem_async.py
index 7e688924..f733317b 100644
--- a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/filesystem_async.py
+++ b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/filesystem_async.py
@@ -15,6 +15,7 @@
 
 """ Storage writer for PyT Distributed format allowing asynchronous save. """
 import dataclasses
+import inspect
 import logging
 import os
 import pickle
@@ -55,8 +56,9 @@
 def _get_write_results_queue():
     global _results_queue
     if _results_queue is None:
-        ctx = mp.get_context('spawn')
-        _results_queue = ctx.Manager().Queue()
+        ctx = mp.get_context('fork')
+        with _disable_gc():
+            _results_queue = ctx.Manager().Queue()
     return _results_queue
 
 
@@ -360,6 +362,17 @@ def write_preloaded_data(
         local_results = []
         try:
             file_name, storage_key, (bytes_data, tensor_data) = write_bucket
+            extra_kwargs = {}
+            write_fn = _write_item
+            if "serialization_format" in inspect.signature(_write_item).parameters:
+                from torch.distributed.checkpoint.filesystem import SerializationFormat
+
+                extra_kwargs['serialization_format'] = SerializationFormat.TORCH_SAVE
+
+            if "transforms" in inspect.signature(_write_item).parameters:
+                assert len(transform_list) <= 1
+                write_fn = partial(_write_item, *transform_list)
+
             if use_msc:
                 import multistorageclient as msc
 
@@ -369,13 +382,13 @@ def write_preloaded_data(
             with open_file(file_name, "wb") as stream:
                 for write_item, data in bytes_data:
                     local_results.append(
-                        _write_item(*transform_list, stream, data, write_item, storage_key)
+                        write_fn(stream, data, write_item, storage_key, **extra_kwargs)
                     )
 
                 for write_item, tensor in tensor_data:
                     assert tensor.is_cpu
                     local_results.append(
-                        _write_item(*transform_list, stream, tensor, write_item, storage_key)
+                        write_fn(stream, tensor, write_item, storage_key, **extra_kwargs)
                     )
 
                 if use_fsync:
diff --git a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py
index 5c644253..2ba32ab4 100644
--- a/src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py
+++ b/src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py
@@ -49,53 +49,249 @@ def _compare_dataclasses(obj1, obj2):
     return differences if differences else "All fields are equal"
 
 
+class CheckpointMetadataCache:
+    """
+    Cache of metadata for checkpoint saving.
+
+    This class maintains a cache of metadata used during distributed checkpoint saving operations.
+    It stores various components of the save plan and metadata to optimize subsequent checkpoint
+    saves by avoiding redundant planning and metadata generation when the checkpoint structure
+    remains consistent across iterations.
+
+    The cache stores:
+    - cached_central_plan: The aggregated global save plan from all ranks
+    - cached_local_plan: The local save plan describing how the local state_dict is written
+    - cached_global_metadata: The global metadata (only held by the coordinator rank)
+    - validated_cache_reuse: Flag indicating if checkpoint structures are consistent
+    - validated_loaded_metadata_reuse: Flag indicating the metadata loaded from the prev checkpoint
+                                       is validated to reuse, which skips all metadata communications
+    - loaded_all_plans: Cached local plans from the previous checkpoint's metadata file
+
+    This caching mechanism helps optimize checkpoint saving by:
+    1. Avoiding redundant planning when checkpoint structures are consistent
+    2. Reusing global metadata when possible
+    3. Enabling decentralized planning when supported by the planner and storage writer
+    """
+
+    def __init__(self):
+        # Cached SavePlans to skip plan in `save_state_dict_async_plan`
+        # cached outcome of `SavePlan.prepare_global_plan`,
+        # which aggregates local plans from all ranks
+        self.cached_central_plan: SavePlan = None
+        # cached outcome of `SavePlan.prepare_local_plan` describes how local state_dict is written
+        self.cached_local_plan: SavePlan = None
+        # Cached global metadata, only `coordinator` for dist-ckpt holds
+        # if central plans are consistent over iters
+        self.cached_global_metadata: Metadata = None
+        # This variable records if the ckpt structures are consistent
+        # so the following checkpoint savings reuse `cached_global_metadata`
+        self.validated_cache_reuse: bool = False
+        # The knob to enable cached metadata communication in saving
+        self.validated_loaded_metadata_reuse = False
+        # The cached all_local_plans from the loaded metadata file of the previous checkpoint
+        self.loaded_all_plans: List[SavePlan] = None
+
+    def set_cached_global_metadata(self, cached_global_metadata):
+        """
+        Sets the cached global metadata and extracts local plans from it.
+
+        This method stores the global metadata from a previous checkpoint and attempts to extract
+        the local plans from it. The local plans are used to verify if the global metadata can
+        be reused in subsequent checkpoint saves.
+
+        Args:
+            cached_global_metadata (Metadata): The global metadata from a previous checkpoint
+                that contains information about the checkpoint structure and local plans.
+
+        Note:
+            If the metadata does not contain local plans, a debug message is logged indicating
+            that global metadata reuse verification will not be possible.
+        """
+        self.cached_global_metadata = cached_global_metadata
+        self.loaded_all_plans = getattr(self.cached_global_metadata, "all_local_plans", None)
+        if self.loaded_all_plans is None:
+            logger.debug("no all_local_plans in metadata - can't verify global metadata reuse...")
+
+    def set_cache_metadata(
+        self, central_plan: SavePlan, local_plan: SavePlan, global_md_verify_reuse: bool
+    ):
+        """
+        Sets the cached metadata and updates the cache flags.
+
+        This method updates the cache with the latest central plan, local plan, and metadata reuse
+        validation results. It also checks if the central plan is consistent with the cached plan.
+
+        Args:
+            central_plan (SavePlan): The latest central plan
+            local_plan (SavePlan): The latest local plan
+            global_md_verify_reuse (bool): Flag indicating if global metadata reuse is valid
+        """
+        self.validated_loaded_metadata_reuse = global_md_verify_reuse
+        self.validated_cache_reuse = bool(central_plan == self.cached_central_plan)
+        logger.debug(f"validated: {self.validated_cache_reuse}")
+        self.cached_central_plan = central_plan
+        self.cached_local_plan = local_plan
+
+    def prepare_save_state_dict_ret(
+        self,
+        rank: int,
+        coordinator: int,
+        save_state_dict_ret: Tuple['FileSystemWriterAsync', Union[Metadata, None]],
+    ) -> Tuple['FileSystemWriterAsync', Union[Metadata, None]]:
+        """
+        Prepares the save state dict return value based on the cached metadata.
+
+        This method checks if the global metadata can be reused from the previous checkpoint.
+        If so, it updates the save state dict return value with the cached global metadata.
+
+        Args:
+            rank (int): The rank of the current process
+            coordinator (int): The coordinator rank
+            save_state_dict_ret (Tuple[FileSystemWriterAsync, Union[Metadata, None]]):
+                                The return value of the save state dict
+
+        Returns:
+            Tuple[FileSystemWriterAsync, Union[Metadata, None]]:
+            The updated save state dict return value with the cached global metadata
+            if it can be reused.
+        """
+        if (
+            self.loaded_all_plans
+            and self.cached_global_metadata
+            and self.validated_loaded_metadata_reuse
+        ):
+            if coordinator == rank:
+                logger.debug(
+                    f"rank: {rank}, reuse global metadata from loaded"
+                    f" .metadata, {save_state_dict_ret[1]}"
+                )
+                save_state_dict_ret = list(save_state_dict_ret)
+                save_state_dict_ret[1] = self.cached_global_metadata
+
+        elif self.validated_cache_reuse:
+            logger.debug(f"rank: {rank}, cache validated")
+            if save_state_dict_ret[1]:  # when global_metadata is not cached
+                self.cached_global_metadata = save_state_dict_ret[1]  # Cache Metadata
+            # Only Coordinator rank holds cached global_metadata
+            # (None is returned for global_metadata)
+            elif coordinator == rank:
+                logger.debug(
+                    f"rank: {rank}, reuse global metadata cached from previous"
+                    f" save iteration, {save_state_dict_ret[1]}"
+                )
+                save_state_dict_ret = list(save_state_dict_ret)
+                save_state_dict_ret[1] = self.cached_global_metadata
+        return save_state_dict_ret
+
+    def get_cache_metadata(self) -> Optional[Tuple[SavePlan, SavePlan, bool, List[SavePlan]]]:
+        """
+        Retrieves the cached metadata components.
+
+        This method returns a tuple containing the cached central plan, local plan, cache reuse
+        validation, and all local plans from the previous checkpoint's metadata file.
+        """
+        return (
+            self.cached_central_plan,
+            self.cached_local_plan,
+            self.validated_cache_reuse,
+            self.loaded_all_plans,
+        )
+
+    def get_metadata_caching_status(self):
+        """
+        Retrieves the current caching status
+
+        This function returns the current caching status of the checkpoint metadata
+        """
+        return self.validated_cache_reuse, self.validated_loaded_metadata_reuse
+
+
+_checkpoint_metadata_cache = None
+
+
+def init_checkpoint_metadata_cache(cached_global_metadata: Metadata):
+    """
+    Initializes the checkpoint metadata cache.
+
+    This function creates a new CheckpointMetadataCache instance and
+    sets the cached global metadata from the previous checkpoint
+    """
+    global _checkpoint_metadata_cache
+    if _checkpoint_metadata_cache is None:
+        _checkpoint_metadata_cache = CheckpointMetadataCache()
+    _checkpoint_metadata_cache.set_cached_global_metadata(cached_global_metadata)
+
+
+def get_metadata_caching_status():
+    """
+    Retrieves the current caching status
+
+    This function returns the current caching status of the checkpoint metadata
+    """
+    global _checkpoint_metadata_cache
+    if _checkpoint_metadata_cache is not None:
+        return _checkpoint_metadata_cache.get_metadata_caching_status()
+
+
 def save_state_dict_async_plan(
     state_dict: STATE_DICT_TYPE,
     storage_writer: 'FileSystemWriterAsync',
     process_group: Optional[dist.ProcessGroup] = None,
     coordinator_rank: int = 0,
     planner: Optional[Union[SavePlanner, DefaultSavePlanner]] = None,
-    cached_ckpt_structure: Optional[Tuple[SavePlan, SavePlan, bool]] = None,
-    loaded_all_plans: Optional[List[SavePlan]] = None,
-) -> Tuple[Tuple['FileSystemWriterAsync', Union[Metadata, None], _DistWrapper], SavePlan, bool]:
+    enable_cache: bool = False,
+    metadata_cache: Optional[CheckpointMetadataCache] = None,
+) -> Tuple['FileSystemWriterAsync', Union[Metadata, None], _DistWrapper]:
     """
-    First stage of asynchronously saving a state dict to storage.
+    First stage of saving a state dict to storage.
 
-    This is an async adaptation of `torch.distributed.checkpoint.state_dict_saver`.
-    To support async save, the process is split into three stages:
+    This is an async adjustment of torch.distributed.checkpoint.state_dict_saver.
+    In order to support async save, saving should be split into three parts:
+    1. Planning
+    2. Actual saving
+    3. Finalization
 
-    1. **Planning**
-    2. **Actual saving** (must be asynchronous)
-    3. **Finalization**
+    Out of these, step (2) *must* happen asynchronously.
+    The first step is realized with this function.
 
-    The planning step is handled by this function and follows several steps as described in
-    the [PyTorch documentation](https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner).
+    The planning part consists of several steps, described here:
+    https://pytorch.org/docs/stable/distributed.checkpoint.html#torch.distributed.checkpoint.SavePlanner
 
     Args:
-        state_dict (STATE_DICT_TYPE): The state dict to save.
-        storage_writer (FileSystemWriterAsync): The storage writer.
-            Currently, only an instance of `FileSystemWriterAsync` is supported.
-        process_group (dist.ProcessGroup, optional): The process group used for save planning.
-        coordinator_rank (int, optional): The coordinator rank for planning. Defaults to 0.
-        planner (SavePlanner, optional): The save planner for the `torch.distributed.checkpoint` format.
-        cached_ckpt_structure (Tuple[SavePlan, SavePlan, bool], optional):
-            A tuple containing:
-            - `cached_central_plan` (SavePlan): A globally coordinated save plan cached in the previous iteration.
-            - `cached_local_plan` (SavePlan): A local plan cached in the previous iteration.
-            - `validated_cache_reuse` (bool): Whether global metadata and the planning dict are consistent over iterations.
-
-    Returns:
-        Tuple:
-        - The storage writer (same as input).
-        - Metadata from planning (or `None` if cached global metadata is reused).
-        - The distributed wrapper used for planning.
-
-    The return value of this function should be passed as input to
-    `save_state_dict_async_finalize`, along with `cached_plan`, to skip `reduce_scatter` during planning.
+        state_dict (STATE_DICT_TYPE): state dict to save
+        storage_writer (FileSystemWriterAsync): in current version only an instance of
+            FileSystemWriterAsync
+        process_group (dist.ProcessGroup, optional): process group used for save planning
+        coordinator_rank (int, optional): coordinator rank for planning. Defaults to 0.
+        planner (SavePlanner, optional): save planner for torch.distributed.checkpoint format
+        cached_ckpt_structure (Tuple[SavePlan, SavePlan, bool], Optional):
+            Each object of this tuple will be used in the order as following
+            cached_central_plan (SavePlan): a globally coordinated save plan
+                cached in the previous iteration
+            cached_local_plan (SavePlan): a local plan
+                cached in the previous iteration
+            validated_cache_reuse (bool): boolean value to tell global_metadata and planning dict
+                is consistent over iterations
+
+    Returns: Tuple of:
+        - storage writer (the one passed as input)
+        - metadata from planning (or None if we reuse cached global metadata)
+        - distributed wrapper used for planning
+    The return value of this function should be passed as an input to
+    `save_state_dict_async_finalize` and cached_plan to skip `reduce_scatter` at planning.
     """
-    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
+    cached_central_plan, cached_local_plan, validated_cache_reuse, loaded_all_plans = (
+        None,
+        None,
+        False,
+        None,
+    )
+    global _checkpoint_metadata_cache
+    metadata_cache = metadata_cache if metadata_cache is not None else _checkpoint_metadata_cache
+    if enable_cache and metadata_cache:
+        cached_central_plan, cached_local_plan, validated_cache_reuse, loaded_all_plans = (
+            metadata_cache.get_cache_metadata()
+        )
 
     rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0
     dist_wrapper = _DistWrapper(process_group, True, coordinator_rank)
@@ -146,18 +342,20 @@ def global_step(all_local_plans):
         )
 
         if not loaded_all_plans or not global_md_verify_reuse:
+            logger.debug(f"rank: {rank}, Passed cache non-reusable")
             all_local_plans = dist_wrapper.gather_object(local_plan)
             if dist_wrapper.is_coordinator:
                 _, global_metadata = planner.create_global_plan(all_local_plans)
                 global_metadata.all_local_plans = all_local_plans
         else:
-            logger.debug(f"rank: {rank}, Passed cached global metadata")
+            logger.debug(f"rank: {rank}, Passed cached global metadata, {global_md_verify_reuse}")
             global_metadata = None
         local_plan = planner.create_decentralized_global_plan(local_plan)
         local_plan = storage_writer.prepare_decentralized_global_plan(local_plan)
         central_plan = local_plan
     else:
         central_plan = dist_wrapper.reduce_scatter("plan", local_step, global_step)
+
     central_plan = planner.finish_plan(central_plan)
     end_plan = time()
     logger.debug(f"rank: {rank}, plan time: {end_plan - start_plan}")
@@ -167,13 +365,14 @@ def global_step(all_local_plans):
     storage_writer.prepare_write_data(central_plan, planner)
     end = time()
     logger.debug(f"{time()} rank: {rank}, write(async) time: {end - start}")
-    return (
-        (storage_writer, global_metadata, dist_wrapper),
-        central_plan,
-        local_plan,
-        cached_central_plan == central_plan,
-        global_md_verify_reuse,
-    )
+    save_state_dict_ret = (storage_writer, global_metadata, dist_wrapper)
+    if enable_cache and metadata_cache:
+        logger.debug(f"{time()} rank: {rank}, setting metadata caching")
+        metadata_cache.set_cache_metadata(central_plan, local_plan, global_md_verify_reuse)
+        save_state_dict_ret = metadata_cache.prepare_save_state_dict_ret(
+            rank, coordinator_rank, save_state_dict_ret
+        )
+    return save_state_dict_ret
 
 
 def verify_global_md_reuse(
diff --git a/src/nvidia_resiliency_ext/checkpointing/utils.py b/src/nvidia_resiliency_ext/checkpointing/utils.py
index da20ad4a..fe3f85ad 100644
--- a/src/nvidia_resiliency_ext/checkpointing/utils.py
+++ b/src/nvidia_resiliency_ext/checkpointing/utils.py
@@ -17,8 +17,9 @@
 import logging
 from contextlib import contextmanager
 from time import time
-from typing import Callable, Dict, List, Optional, TypeVar, Union
+from typing import Any, Callable, Dict, List, Optional, Tuple, TypeVar, Union
 
+import numpy as np
 import torch
 
 U = TypeVar("U")
@@ -81,16 +82,6 @@ def debug_msg(msg: str):
         last_logger.debug(f"{stacked_name} {msg}")
 
 
-def dict_list_map_outplace(f: Callable[[U], V], x: Union[Dict, List, U]) -> Union[Dict, List, V]:
-    """Maps dicts and lists *out-of-place* with a given function."""
-    if isinstance(x, dict):
-        return {k: dict_list_map_outplace(f, v) for k, v in x.items()}
-    elif isinstance(x, list):
-        return [dict_list_map_outplace(f, v) for v in x]
-    else:
-        return f(x)
-
-
 def preload_tensors(state_dict: Dict, non_blocking=True):
     """Preload tensors in state_dict to host memory through CPU memory
     Args:
@@ -127,3 +118,75 @@ def wrapped(state_dict, *args, **kwargs):
             fn(state_dict, *args, **kwargs)
 
     return wrapped
+
+
+# Dict utils copied from megatron.core.dist_checkpointing.dict_utils.py
+def diff(x1: Any, x2: Any, prefix: Tuple = ()) -> Tuple[list, list, list]:
+    """Recursive diff of dicts.
+
+    Args:
+        x1 (object): left dict
+        x2 (object): right dict
+        prefix (tuple): tracks recursive calls. Used for reporting differing keys.
+
+    Returns:
+        Tuple[list, list, list]: tuple of:
+            - only_left: Prefixes present only in left dict
+            - only_right: Prefixes present only in right dict
+            - mismatch: values present in both dicts but not equal across dicts.
+                For tensors equality of all elems is checked.
+                Each element is a tuple (prefix, type of left value, type of right value).
+    """
+    mismatch = []
+    if isinstance(x1, dict) and isinstance(x2, dict):
+        only_left = [prefix + (k,) for k in x1.keys() - x2.keys()]
+        only_right = [prefix + (k,) for k in x2.keys() - x1.keys()]
+        for k in x2.keys() & x1.keys():
+            _left, _right, _mismatch = diff(x1[k], x2[k], prefix + (k,))
+            only_left.extend(_left)
+            only_right.extend(_right)
+            mismatch.extend(_mismatch)
+    elif isinstance(x1, list) or isinstance(x1, tuple) or isinstance(x1, np.ndarray):
+        assert isinstance(x1, type(x2))
+        only_left = list(range(len(x1) - 1, len(x2) - 1, -1))
+        only_right = list(range(len(x1) - 1, len(x2) - 1, -1))
+        for i, (v1, v2) in enumerate(zip(x1, x2)):
+            _left, _right, _mismatch = diff(v1, v2, prefix + (i,))
+            only_left.extend(_left)
+            only_right.extend(_right)
+            mismatch.extend(_mismatch)
+    else:
+        only_left = []
+        only_right = []
+        if isinstance(x1, torch.Tensor) and isinstance(x2, torch.Tensor):
+            if x1.device != x2.device:
+                _is_mismatch = not torch.all(x1.cpu() == x2.cpu())
+            else:
+                _is_mismatch = not torch.all(x1 == x2)
+        # TODO: change with concrete type that has both replica_id and data attrs
+        elif hasattr(x1, 'replica_id') and hasattr(x2, 'replica_id'):
+            assert isinstance(x1, type(x2))
+            only_left, only_right, mismatch = diff(
+                x1.data, x2.data, prefix + (type(x1),)
+            )  # type: ignore
+            _is_mismatch = False
+        else:
+            try:
+                _is_mismatch = bool(x1 != x2)
+            except RuntimeError:
+                _is_mismatch = True
+
+        if _is_mismatch:
+            mismatch.append((prefix, type(x1), type(x2)))
+
+    return only_left, only_right, mismatch
+
+
+def dict_list_map_outplace(f: Callable[[U], V], x: Union[Dict, List, U]) -> Union[Dict, List, V]:
+    """Maps dicts and lists *out-of-place* with a given function."""
+    if isinstance(x, dict):
+        return {k: dict_list_map_outplace(f, v) for k, v in x.items()}
+    elif isinstance(x, list):
+        return [dict_list_map_outplace(f, v) for v in x]
+    else:
+        return f(x)
diff --git a/tests/checkpointing/unit/test_async_writer.py b/tests/checkpointing/unit/test_async_writer.py
index 5849e336..722240b8 100644
--- a/tests/checkpointing/unit/test_async_writer.py
+++ b/tests/checkpointing/unit/test_async_writer.py
@@ -13,6 +13,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import filecmp
+import pickle
+from copy import deepcopy
+from dataclasses import fields
 
 import torch
 from torch.distributed.checkpoint import (
@@ -21,8 +24,9 @@
     FileSystemReader,
     FileSystemWriter,
     load,
-    state_dict_saver,
+    save,
 )
+from torch.distributed.fsdp import FullyShardedDataParallel as FSDP
 
 from nvidia_resiliency_ext.checkpointing.async_ckpt.core import AsyncCallsQueue, AsyncRequest
 from nvidia_resiliency_ext.checkpointing.async_ckpt.filesystem_async import FileSystemWriterAsync
@@ -30,9 +34,9 @@
     save_state_dict_async_finalize,
     save_state_dict_async_plan,
 )
-
-from . import TempNamedDir
-from .test_utilities import TestModel, Utils
+from nvidia_resiliency_ext.checkpointing.utils import diff
+from tests.checkpointing.unit import TempNamedDir
+from tests.checkpointing.unit.test_utilities import TestModel, Utils
 
 
 class TestAsyncSave:
@@ -48,21 +52,21 @@ def finalize_fn():
         return AsyncRequest(save_fn, save_args, [finalize_fn], preload_fn=preload_fn)
 
     def async_save_checkpoint(
-        self, checkpoint_dir, state_dict, planner, async_queue, thread_count=1
+        self, checkpoint_dir, state_dict, planner, async_queue, thread_count=1, caching=False
     ):
         """Performs an asynchronous model checkpoint save."""
         writer = FileSystemWriterAsync(checkpoint_dir, thread_count=thread_count)
         coordinator_rank = 0
 
-        save_state_dict_ret, *_ = save_state_dict_async_plan(
-            state_dict, writer, None, coordinator_rank, planner=planner
+        save_state_dict_ret = save_state_dict_async_plan(
+            state_dict, writer, None, coordinator_rank, planner=planner, enable_cache=caching
         )
         async_request = self.get_async_save_request(writer, save_state_dict_ret)
         async_queue.schedule_async_request(async_request)
 
     def sync_save_checkpoint(self, checkpoint_dir, state_dict, planner):
         """Performs a synchronous model checkpoint save using FileSystemWriter."""
-        state_dict_saver.save(
+        save(
             state_dict=state_dict,
             storage_writer=FileSystemWriter(checkpoint_dir),
             planner=planner,
@@ -80,9 +84,8 @@ def load_checkpoint(self, checkpoint_dir, state_dict):
     def test_async_is_equivalent_to_sync(self, tmp_path_dist_ckpt):
         """Verifies that async checkpointing produces the same results as sync checkpointing."""
         Utils.initialize_distributed()
-        model = TestModel((1024, 1024), 10)
+        model = FSDP(TestModel((1024, 1024), 8))
         async_queue = AsyncCallsQueue()
-
         with (
             TempNamedDir(tmp_path_dist_ckpt / 'async_checkpoint', sync=True) as async_ckpt_dir,
             TempNamedDir(tmp_path_dist_ckpt / 'sync_checkpoint', sync=True) as sync_ckpt_dir,
@@ -110,14 +113,68 @@ def test_async_is_equivalent_to_sync(self, tmp_path_dist_ckpt):
             ], f"Differences found in saved files: {comparison.diff_files}"
 
             # Load and compare state dicts
-            loaded_async_state_dict = self.load_checkpoint(async_ckpt_dir, state_dict.copy())
-            loaded_sync_state_dict = self.load_checkpoint(sync_ckpt_dir, state_dict.copy())
-
+            loaded_async_state_dict = self.load_checkpoint(async_ckpt_dir, deepcopy(state_dict))
+            loaded_sync_state_dict = self.load_checkpoint(sync_ckpt_dir, deepcopy(state_dict))
+            assert loaded_sync_state_dict.keys() == state_dict.keys()
             for key in loaded_sync_state_dict.keys():
                 assert key in loaded_async_state_dict, f"Missing key in async checkpoint: {key}"
                 assert torch.equal(
                     loaded_async_state_dict[key], loaded_sync_state_dict[key]
                 ), f"Mismatch for key '{key}' between async and sync checkpoints."
                 assert torch.equal(
-                    loaded_async_state_dict[key], state_dict[key]
+                    loaded_sync_state_dict[key], state_dict[key]
                 ), f"Mismatch for key '{key}' between async checkpoint and original state_dict."
+
+    def test_cached_metadata(self, tmp_path_dist_ckpt):
+        Utils.initialize_distributed()
+        async_queue = AsyncCallsQueue()
+
+        model = FSDP(TestModel((1024, 1024), 8))
+        state_dict_non_cached = model.state_dict()
+        state_dict_cached = deepcopy(state_dict_non_cached)
+        loaded_non_cached, loaded_cached = None, None
+        md_non_cached, md_cached = None, None
+        planner = DefaultSavePlanner()
+
+        with TempNamedDir(tmp_path_dist_ckpt / 'ckpt_dir', sync=True) as ckpt_dir:
+            self.async_save_checkpoint(
+                ckpt_dir, state_dict_non_cached, planner, async_queue, caching=True
+            )
+            async_queue.maybe_finalize_async_calls(blocking=True, no_dist=False)
+            loaded_non_cached = self.load_checkpoint(ckpt_dir, state_dict_non_cached)
+            md_path = ckpt_dir.__enter__() / '.metadata'
+            with md_path.open('rb') as f:
+                md_non_cached = pickle.load(f)
+
+        # Run over 3 iterations with cached metadata enabled
+        # The 3rd iteration will run with cached metadata
+        # `ckpt_dir` at the 3rd iteration 2 will be maintained for comparison
+        ckpt_dir = None
+        for i in range(3):
+            ckpt_dir = TempNamedDir(tmp_path_dist_ckpt / f'ckpt_dir_{i}_cached', sync=True)
+            self.async_save_checkpoint(
+                ckpt_dir.__enter__(), state_dict_cached, planner, async_queue, caching=True
+            )
+            async_queue.maybe_finalize_async_calls(blocking=True, no_dist=False)
+            if i < 2:
+                ckpt_dir.cleanup()
+        loaded_cached = self.load_checkpoint(ckpt_dir.__enter__(), state_dict_cached)
+        md_path = ckpt_dir.__enter__() / '.metadata'
+
+        with md_path.open('rb') as f:
+            md_cached = pickle.load(f)
+
+        # Check loaded state dict
+        diffs = diff(loaded_non_cached, loaded_cached)
+
+        assert not any(
+            len(x) for x in diffs
+        ), 'Cached metadata doesn\'t produce the same state_dict in loading'
+        # Check metadata recorded in .metadata, torch.distributed.metadata.Metadata
+        for field in fields(md_non_cached):
+            if field.name not in ['storage_data', 'storage_meta']:
+                diffs = diff(getattr(md_non_cached, field.name), getattr(md_cached, field.name))
+                assert not any(
+                    len(x) for x in diffs
+                ), f'{field.name} is different in metadata from non-cached, cached metadata impls'
+        ckpt_dir.cleanup()
diff --git a/tests/checkpointing/unit/test_async_writer_msc.py b/tests/checkpointing/unit/test_async_writer_msc.py
index f686d02a..3a5d83c1 100644
--- a/tests/checkpointing/unit/test_async_writer_msc.py
+++ b/tests/checkpointing/unit/test_async_writer_msc.py
@@ -55,7 +55,7 @@ def async_save_checkpoint(
         writer = FileSystemWriterAsync(checkpoint_dir, thread_count=thread_count, use_msc=True)
         coordinator_rank = 0
 
-        save_state_dict_ret, *_ = save_state_dict_async_plan(
+        save_state_dict_ret = save_state_dict_async_plan(
             state_dict, writer, None, coordinator_rank, planner=planner
         )
         async_request = self.get_async_save_request(writer, save_state_dict_ret)


Reviews:
user: login: rhewett-nv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T02:10:32Z
user: login: rhewett-nv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T15:30:06Z
user: login: rhewett-nv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T15:43:24Z
user: login: rhewett-nv
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T16:30:51Z
user: login: sbak5
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T21:15:07Z
user: login: sbak5
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T21:15:54Z
user: login: sbak5
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T21:16:17Z
user: login: sbak5
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-07T21:17:11Z
user: login: hexinw-nvidia
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-09T04:46:46Z

Review comments:
diff_hunk: @@ -160,7 +160,6 @@ def close(self):
         """Terminate the async caller at exit of an application or some termination conditions"""
         logger.info(f"AsyncCaller: {torch.distributed.get_rank()}, Destroying Async Caller")
 
-    @abstractmethod
path: src/nvidia_resiliency_ext/checkpointing/async_ckpt/core.py
user: login: rhewett-nv
body: Keep this
created_at: 2025-05-07T02:08:22Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" FS Reader with metadata cached support. """
+
+import os
+from typing import Union
+
+from torch.distributed.checkpoint import FileSystemReader, Metadata
+
+
+class CachedMetadataFileSystemReader(FileSystemReader):
+    """
+    Extends FileSystemReader to cache metadata for improved performance.
+
+    Attributes:
+        _cached_metadata (Metadata or None): Cached metadata from the file system.
+    """
+
+    def __init__(self, path: Union[str, os.PathLike]) -> None:
+        """
+        Initialize with file system path.
+
+        Args:
+            path (Union[str, os.PathLike]): Path to the checkpoint directory or file.
+        """
+        super().__init__(path=path)
+        self._cached_metadata = None
+
+    def read_metadata(self) -> Metadata:
path: src/nvidia_resiliency_ext/checkpointing/async_ckpt/cached_metadata_filesystem_reader.py
user: login: rhewett-nv
body: Is there ever a usecase where we want to flush the cache and reread it?
created_at: 2025-05-07T15:30:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -50,53 +50,248 @@ def _compare_dataclasses(obj1, obj2):
     return differences if differences else "All fields are equal"
 
 
+class CheckpointMetadataCache:
+    """
+    Cache of metadata for checkpoint saving.
+
+    This class maintains a cache of metadata used during distributed checkpoint saving operations.
+    It stores various components of the save plan and metadata to optimize subsequent checkpoint
+    saves by avoiding redundant planning and metadata generation when the checkpoint structure
+    remains consistent across iterations.
+
+    The cache stores:
+    - cached_central_plan: The aggregated global save plan from all ranks
+    - cached_local_plan: The local save plan describing how the local state_dict is written
+    - cached_global_metadata: The global metadata (only held by the coordinator rank)
+    - validated_cache_reuse: Flag indicating if checkpoint structures are consistent
+    - validated_loaded_metadata_reuse: Flag enabling cached metadata communication in saving
+    - loaded_all_plans: Cached local plans from the previous checkpoint's metadata file
+
+    This caching mechanism helps optimize checkpoint saving by:
+    1. Avoiding redundant planning when checkpoint structures are consistent
+    2. Reusing global metadata when possible
+    3. Enabling decentralized planning when supported by the planner and storage writer
+    """
+
+    def __init__(self):
+        # Cached SavePlans to skip plan in `save_state_dict_async_plan`
+        # cached outcome of `SavePlan.prepare_global_plan`,
+        # which aggregates local plans from all ranks
+        self.cached_central_plan: SavePlan = None
+        # cached outcome of `SavePlan.prepare_local_plan` describes how local state_dict is written
+        self.cached_local_plan: SavePlan = None
+        # Cached global metadata, only `coordinator` for dist-ckpt holds
+        # if central plans are consistent over iters
+        self.cached_global_metadata: Metadata = None
+        # This variable records if the ckpt structures are consistent
+        # so the following checkpoint savings reuse `cached_global_metadata`
+        self.validated_cache_reuse: bool = False
path: src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py
user: login: rhewett-nv
body: I don't understand the use of the word `validated` in this var and the next.  Can we make them (and the comments) more clear?
created_at: 2025-05-07T15:43:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -127,3 +117,75 @@ def wrapped(state_dict, *args, **kwargs):
             fn(state_dict, *args, **kwargs)
 
     return wrapped
+
+
+# Dict utils copied from megatron.core.dist_checkpointing.dict_utils.py
+def diff(x1: Any, x2: Any, prefix: Tuple = ()) -> Tuple[list, list, list]:
+    """Recursive diff of dicts.
path: src/nvidia_resiliency_ext/checkpointing/utils.py
user: login: rhewett-nv
body: This function works for types that are not `dict`.  
created_at: 2025-05-07T16:30:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,51 @@
+# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+""" FS Reader with metadata cached support. """
+
+import os
+from typing import Union
+
+from torch.distributed.checkpoint import FileSystemReader, Metadata
+
+
+class CachedMetadataFileSystemReader(FileSystemReader):
+    """
+    Extends FileSystemReader to cache metadata for improved performance.
+
+    Attributes:
+        _cached_metadata (Metadata or None): Cached metadata from the file system.
+    """
+
+    def __init__(self, path: Union[str, os.PathLike]) -> None:
+        """
+        Initialize with file system path.
+
+        Args:
+            path (Union[str, os.PathLike]): Path to the checkpoint directory or file.
+        """
+        super().__init__(path=path)
+        self._cached_metadata = None
+
+    def read_metadata(self) -> Metadata:
path: src/nvidia_resiliency_ext/checkpointing/async_ckpt/cached_metadata_filesystem_reader.py
user: login: sbak5
body: This is used to bypass any collective comm for metadata when the job resumes from the previous checkpoint. 
We load a checkpoint and keep it in some variable, which is used to fill out the metadata needed for the upcoming saving.
created_at: 2025-05-07T21:15:06Z
author_association: CONTRIBUTOR
diff_hunk: @@ -50,53 +50,248 @@ def _compare_dataclasses(obj1, obj2):
     return differences if differences else "All fields are equal"
 
 
+class CheckpointMetadataCache:
+    """
+    Cache of metadata for checkpoint saving.
+
+    This class maintains a cache of metadata used during distributed checkpoint saving operations.
+    It stores various components of the save plan and metadata to optimize subsequent checkpoint
+    saves by avoiding redundant planning and metadata generation when the checkpoint structure
+    remains consistent across iterations.
+
+    The cache stores:
+    - cached_central_plan: The aggregated global save plan from all ranks
+    - cached_local_plan: The local save plan describing how the local state_dict is written
+    - cached_global_metadata: The global metadata (only held by the coordinator rank)
+    - validated_cache_reuse: Flag indicating if checkpoint structures are consistent
+    - validated_loaded_metadata_reuse: Flag enabling cached metadata communication in saving
+    - loaded_all_plans: Cached local plans from the previous checkpoint's metadata file
+
+    This caching mechanism helps optimize checkpoint saving by:
+    1. Avoiding redundant planning when checkpoint structures are consistent
+    2. Reusing global metadata when possible
+    3. Enabling decentralized planning when supported by the planner and storage writer
+    """
+
+    def __init__(self):
+        # Cached SavePlans to skip plan in `save_state_dict_async_plan`
+        # cached outcome of `SavePlan.prepare_global_plan`,
+        # which aggregates local plans from all ranks
+        self.cached_central_plan: SavePlan = None
+        # cached outcome of `SavePlan.prepare_local_plan` describes how local state_dict is written
+        self.cached_local_plan: SavePlan = None
+        # Cached global metadata, only `coordinator` for dist-ckpt holds
+        # if central plans are consistent over iters
+        self.cached_global_metadata: Metadata = None
+        # This variable records if the ckpt structures are consistent
+        # so the following checkpoint savings reuse `cached_global_metadata`
+        self.validated_cache_reuse: bool = False
path: src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py
user: login: sbak5
body: validated means the cached metadata aligns with state dict over iterations so it's safe to reuse.
created_at: 2025-05-07T21:15:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -50,53 +50,248 @@ def _compare_dataclasses(obj1, obj2):
     return differences if differences else "All fields are equal"
 
 
+class CheckpointMetadataCache:
+    """
+    Cache of metadata for checkpoint saving.
+
+    This class maintains a cache of metadata used during distributed checkpoint saving operations.
+    It stores various components of the save plan and metadata to optimize subsequent checkpoint
+    saves by avoiding redundant planning and metadata generation when the checkpoint structure
+    remains consistent across iterations.
+
+    The cache stores:
+    - cached_central_plan: The aggregated global save plan from all ranks
+    - cached_local_plan: The local save plan describing how the local state_dict is written
+    - cached_global_metadata: The global metadata (only held by the coordinator rank)
+    - validated_cache_reuse: Flag indicating if checkpoint structures are consistent
+    - validated_loaded_metadata_reuse: Flag enabling cached metadata communication in saving
+    - loaded_all_plans: Cached local plans from the previous checkpoint's metadata file
+
+    This caching mechanism helps optimize checkpoint saving by:
+    1. Avoiding redundant planning when checkpoint structures are consistent
+    2. Reusing global metadata when possible
+    3. Enabling decentralized planning when supported by the planner and storage writer
+    """
+
+    def __init__(self):
+        # Cached SavePlans to skip plan in `save_state_dict_async_plan`
+        # cached outcome of `SavePlan.prepare_global_plan`,
+        # which aggregates local plans from all ranks
+        self.cached_central_plan: SavePlan = None
+        # cached outcome of `SavePlan.prepare_local_plan` describes how local state_dict is written
+        self.cached_local_plan: SavePlan = None
+        # Cached global metadata, only `coordinator` for dist-ckpt holds
+        # if central plans are consistent over iters
+        self.cached_global_metadata: Metadata = None
+        # This variable records if the ckpt structures are consistent
+        # so the following checkpoint savings reuse `cached_global_metadata`
+        self.validated_cache_reuse: bool = False
path: src/nvidia_resiliency_ext/checkpointing/async_ckpt/state_dict_saver.py
user: login: sbak5
body: I can add more words to the comment.
created_at: 2025-05-07T21:16:17Z
author_association: CONTRIBUTOR
diff_hunk: @@ -127,3 +117,75 @@ def wrapped(state_dict, *args, **kwargs):
             fn(state_dict, *args, **kwargs)
 
     return wrapped
+
+
+# Dict utils copied from megatron.core.dist_checkpointing.dict_utils.py
+def diff(x1: Any, x2: Any, prefix: Tuple = ()) -> Tuple[list, list, list]:
+    """Recursive diff of dicts.
path: src/nvidia_resiliency_ext/checkpointing/utils.py
user: login: sbak5
body: It's intended. It's intended to compare elements recursively over iteratable objects. 
created_at: 2025-05-07T21:17:11Z
author_association: CONTRIBUTOR

PR comments:
user: login: sbak5
created_at: 2025-05-06T23:54:25Z
body: a new test is needed to verify metadata caching.
author_association: CONTRIBUTOR
user: login: apaithankar
created_at: 2025-05-07T02:05:14Z
body: Will the tests be added in a separate pull request?
author_association: CONTRIBUTOR
user: login: sbak5
created_at: 2025-05-07T07:00:53Z
body: > Will the tests be added in a separate pull request?

Will move two unit tests from Mcore for metadata caching and add an example for metadata caching.
author_association: CONTRIBUTOR