cosmos-framework/cosmos_framework/checkpoint/dcp.py

# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: OpenMDW-1.1

"""
Distributed checkpoint (DCP) directory structure and storage backends.

The checkpointer saves model state in a sharded format across multiple processes:

self.save_dirname/
├── iter_000000005/                    # Checkpoint at iteration 5
│   ├── model/                         # Model state shards
│   │   ├── __0_0.distcp              # Shard 0 from rank 0
│   │   └── __1_0.distcp              # Shard 1 from rank 1
│   ├── optim/                        # Optimizer state shards
│   │   ├── __0_0.distcp              # Shard 0 from rank 0
│   │   └── __1_0.distcp              # Shard 1 from rank 1
│   ├── scheduler/                    # Learning rate scheduler state
│   │   ├── __0_0.distcp              # Shard 0 from rank 0
│   │   └── __1_0.distcp              # Shard 1 from rank 1
│   └── trainer/                      # Additional training state
│       ├── __0_0.distcp              # Shard 0 from rank 0
│       └── __1_0.distcp              # Shard 1 from rank 1
│   └── dataloader/                   # Optional per-rank dataloader state
│       ├── rank_0.pkl
│       └── rank_1.pkl
└── latest_checkpoint.txt             # Points to most recent checkpoint folder, e.g. iter_000000005

Storage backends:
- Local filesystem:
  self.save_dirname = "{config_job.path_local}/checkpoints"

- S3 object store:
  self.save_dirname = "s3://{bucket}/{config_job.path}/checkpoints"
  where bucket = self.config_checkpoint.save_to_object_store.bucket

The sharded format enables efficient distributed saving/loading by:
1. Parallelizing I/O across processes
2. Reducing memory usage per process
3. Supporting both local and cloud storage backends
"""

import enum
import multiprocessing
import os
import re
import time
from multiprocessing import get_context
from typing import Any, Dict, List, Optional, Protocol, Tuple, Union, runtime_checkable

import torch
import torch.distributed as dist
import torch.distributed.checkpoint as dcp
from torch import nn
from torch.distributed.checkpoint.filesystem import FileSystemReader, FileSystemWriter
from torch.distributed.checkpoint.metadata import (
    STATE_DICT_TYPE,
    Metadata,
    StorageMeta,
)
from torch.distributed.checkpoint.state_dict import (
    StateDictOptions,
    get_model_state_dict,
    set_model_state_dict,
)
from torch.distributed.checkpoint.stateful import Stateful

from cosmos_framework.checkpoint.base import AbstractCheckpointer
from cosmos_framework.checkpoint.s3_filesystem import S3StorageReader, S3StorageWriter
from cosmos_framework.utils.config import CheckpointConfig, JobConfig
from cosmos_framework.model._base import ImaginaireModel
from cosmos_framework.utils import callback, distributed, log, misc
from cosmos_framework.utils.easy_io import easy_io
from cosmos_framework.utils.vfm.rand_state import get_rand_state_dict, set_rand_state_dict


class ModelWrapper(Stateful):
    """
    Wrapper for model state dict handling. Strips away the _orig_mod. prefix
    among other things from the state dict keys.
    """

    def __init__(self, model: nn.Module) -> None:
        self.model = model

    def state_dict(self) -> dict[str, Any]:
        return get_model_state_dict(self.model)

    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
        set_model_state_dict(
            self.model,
            model_state_dict=state_dict,
            options=StateDictOptions(strict=True),
        )


@runtime_checkable
class _DataloaderStateHandler(Protocol):
    """Structural contract for callbacks that participate in dataloader-state checkpointing."""

    checkpoint_component: str

    def has_checkpoint_state(self) -> bool: ...
    def state_dict(self) -> dict[Any, Any]: ...
    def load_state_dict(self, state_dict: dict[Any, Any]) -> None: ...


class _DataloaderWrapper:
    """Adapter that surfaces a dataloader-state callback's checkpoint API.

    Walks the registered callbacks at construction time and binds to the
    first callback that:

    1. Declares ``checkpoint_component == "dataloader"``, AND
    2. Returns ``True`` from ``has_checkpoint_state()``.

    The bound callback's ``state_dict`` / ``load_state_dict`` methods are
    re-exposed via :meth:`state_dict` / :meth:`load_state_dict`.  Callers
    must gate those on :meth:`has_state` — invoking them when nothing was
    bound raises :class:`RuntimeError`.

    Note: only the first callback tagged ``checkpoint_component=="dataloader"``
    is considered; if it does not currently want its state checkpointed,
    no further callbacks are searched.  In practice there is at most one
    such callback (see ``DataLoaderStateCallback``).
    """

    def __init__(self, callbacks: callback.CallBackGroup | None) -> None:
        self._callback: _DataloaderStateHandler | None = None
        if callbacks is None:
            return
        for current_callback in callbacks._callbacks:
            if getattr(current_callback, "checkpoint_component", None) != "dataloader":
                continue
            if current_callback.has_checkpoint_state():
                self._callback = current_callback
            return

    def has_state(self) -> bool:
        return self._callback is not None

    def state_dict(self) -> dict[Any, Any]:
        if self._callback is None:
            raise RuntimeError("No dataloader state handler is registered, cannot save dataloader state.")
        return self._callback.state_dict()

    def load_state_dict(self, state_dict: dict[Any, Any]) -> None:
        if self._callback is None:
            raise RuntimeError("No dataloader state handler is registered, cannot load dataloader state.")
        self._callback.load_state_dict(state_dict)


class AsyncMode(str, enum.Enum):
    DISABLED = "disabled"
    ASYNC_WITH_PINNED_MEM = "async_with_pinned_mem"


class Terminate:
    pass


class SaveDone:
    def __init__(self, iteration: int, elapsed_time: float, succeeded: bool):
        self.iteration = iteration
        self.elapsed_time = elapsed_time
        self.succeeded = succeeded

    def __str__(self):
        return f"SaveDone(iteration={self.iteration}, elapsed_time={self.elapsed_time}, succeeded={self.succeeded})"


def save_checkpoint_in_background(
    receiver_queue: multiprocessing.Queue,
    sender_queue: multiprocessing.Queue,
    config_checkpoint: CheckpointConfig,
    config_job: JobConfig,
) -> None:
    """
    Handles model checkpoint saving in a separate background process using PyTorch's distributed functionality.
    This function runs in a dedicated process to avoid blocking the main training loop.

    Args:
        receiver_queue: Queue to receive state dictionaries and commands from the main process
        sender_queue: Queue to send completion signals back to the main process
        config_checkpoint: Configuration settings for checkpoint saving behavior
        config_job: Configuration settings for the training job

    Flow:
        1. Initializes distributed processing environment
        2. Continuously waits for state dictionaries to save
        3. Saves checkpoints asynchronously
        4. Signals completion back to main process
        5. Terminates when receiving a Terminate signal

    Raises:
        AssertionError: If received object is neither Terminate signal nor valid state dict tuple

    Note:
        - Uses a different port than the main process to avoid conflicts
        - Disables TorchElastic agent store for checkpoint operations
        - Automatically cleans up distributed process group on exit
    """
    # Configure distributed environment
    os.environ["MASTER_PORT"] = str(int(os.environ["MASTER_PORT"]) + 2)
    os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False"

    # Set up GPU device and distributed processing
    torch.cuda.set_device(int(os.environ["LOCAL_RANK"]))
    if dist.is_initialized():
        dist.destroy_process_group()
    dist.init_process_group(backend="gloo")

    # Initialize checkpointing mechanism
    checkpoint_handler = DistributedCheckpointer(
        config_checkpoint=config_checkpoint,
        config_job=config_job,
        callbacks=None,
        disable_async=True,
    )

    while True:
        log.info(f"Checkpoint background process is ready for next task, waiting for new state_dict")
        received_data = receiver_queue.get()
        log.info(f"Checkpoint background process received new state_dict")

        if isinstance(received_data, Terminate):
            log.info(f"Checkpoint background process received termination signal, closing sender queue")
            break

        assert isinstance(received_data, tuple), "Received data must be a tuple of (state_dict, checkpoint_path)"
        state_dict, checkpoint_path = received_data

        # Save checkpoint and measure time taken.
        start_time = time.monotonic()
        iteration = state_dict["trainer"][0]["iteration"]
        succeeded = False

        try:
            log.info(f"Saving checkpoint to {checkpoint_path}")
            checkpoint_handler.save_state_dict_worker(state_dict, checkpoint_path)
            succeeded = True
        except Exception as e:
            log.error(f"Error saving checkpoint to {checkpoint_path}: {e}")
            # continue because if the thread exits, the main thread keeps on adding to the queue
        finally:
            elapsed_time = time.monotonic() - start_time
            log.info(
                f"Checkpoint save completed in background process. "
                f"Time taken: {elapsed_time:.2f} seconds, iteration: {iteration}, "
                f"status: {'SUCCESS' if succeeded else 'FAILURE'}"
            )
            sender_queue.put(SaveDone(iteration, elapsed_time, succeeded))

    log.info("Cleaning up: destroying distributed process group")
    dist.destroy_process_group()


def _replace_keys_with_ema_keys(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE:
    """
    Renames model parameters from "net." to "net_ema.".
    """
    if not all(k.startswith("net.") for k in state_dict.keys()):
        raise ValueError("State dict must start with net. keys when load_ema_to_reg is True")
    return {k.replace("net.", "net_ema."): v for k, v in state_dict.items()}


class CustomLoadPlanner(dcp.DefaultLoadPlanner):
    """
    CustomLoadPlanner that supports ignoring keys during checkpoint load.
    This is useful when the checkpoint is saved with a different component
    architecture, e.g. different RoPE embeddings than the current model.
    """

    def __init__(
        self,
        flatten_state_dict: bool = True,
        flatten_sharded_tensors: bool = True,
        allow_partial_load: bool = False,
        keys_to_skip_loading: List[str] = [],
        load_ema_to_reg: bool = False,
    ) -> None:
        super().__init__(
            flatten_state_dict=flatten_state_dict,
            flatten_sharded_tensors=flatten_sharded_tensors,
            allow_partial_load=allow_partial_load,
        )
        self.keys_to_skip_loading = keys_to_skip_loading
        self.load_ema_to_reg = load_ema_to_reg
        if len(keys_to_skip_loading) > 0:
            log.info(f"Skipping loading of keys that match the following patterns: {keys_to_skip_loading}")

    def set_up_planner(
        self,
        state_dict: STATE_DICT_TYPE,
        metadata: Metadata | None = None,
        is_coordinator: bool = False,
    ) -> None:
        state_dict = self._skip_keys_if_found(state_dict)

        if self.load_ema_to_reg:
            state_dict = _replace_keys_with_ema_keys(state_dict)

        super().set_up_planner(
            state_dict=state_dict,
            metadata=metadata,
            is_coordinator=is_coordinator,
        )

    def _skip_keys_if_found(
        self,
        state_dict: STATE_DICT_TYPE,
    ) -> Dict[str, Any]:
        """
        While loading the checkpoint, skip the weight loading for the keys
        that contain any element of `self.keys_to_skip_loading` as a substring.
        """
        if len(self.keys_to_skip_loading) == 0:
            return state_dict

        new_state_dict = {}
        for fqn, obj in state_dict.items():
            if any(skip_key in fqn for skip_key in self.keys_to_skip_loading):
                log.warning(f"Skipping loading of key: {fqn}")
                continue
            new_state_dict[fqn] = obj
        return new_state_dict


class CustomSavePlanner(dcp.DefaultSavePlanner):
    """
    Custom save planner that enables an override for cache_plans_key when
    caching of save plans is enabled. Caching of save plans reduces checkpointing
    time by reusing the same save plan across checkpoints. This reduces the
    checkpointing time by ~60% (benchmarked using the 235B-A22B Qwen3-VL model
    on 64 GB200 nodes).
    """

    def __init__(
        self,
        flatten_state_dict: bool = True,
        flatten_sharded_tensors: bool = True,
        dedup_save_to_lowest_rank: bool = False,
        save_reg_to_ema: bool = False,
        enable_plan_caching: bool = False,
        cache_plans_key: str | None = None,
    ) -> None:
        super().__init__(
            flatten_state_dict=flatten_state_dict,
            flatten_sharded_tensors=flatten_sharded_tensors,
            dedup_save_to_lowest_rank=dedup_save_to_lowest_rank,
            enable_plan_caching=enable_plan_caching,
        )
        if cache_plans_key is not None:
            self._cached_plans_key = cache_plans_key

        self.save_reg_to_ema = save_reg_to_ema

    def set_up_planner(
        self,
        state_dict: STATE_DICT_TYPE,
        storage_meta: StorageMeta | None = None,
        is_coordinator: bool = False,
    ) -> None:
        if self.save_reg_to_ema:
            state_dict = _replace_keys_with_ema_keys(state_dict)

        super().set_up_planner(
            state_dict=state_dict,
            storage_meta=storage_meta,
            is_coordinator=is_coordinator,
        )


class DistributedCheckpointer(AbstractCheckpointer):
    CHECKPOINT_KEYS = ["model", "optim", "scheduler", "trainer", "dataloader"]

    def __init__(
        self,
        config_checkpoint: CheckpointConfig,
        config_job: JobConfig,
        callbacks: Optional[callback.CallBackGroup] = None,
        disable_async: bool = False,
    ):
        super().__init__(config_checkpoint, config_job, callbacks)
        self.config_checkpoint = config_checkpoint
        if config_checkpoint.dcp_async_mode_enabled and not disable_async:
            self.async_mode = AsyncMode.ASYNC_WITH_PINNED_MEM
        else:
            self.async_mode = AsyncMode.DISABLED

        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
            ctx = get_context("spawn")
            self.mp_queue_send = ctx.Queue()
            self.mp_queue_recv = ctx.Queue()
            self.mp = ctx.Process(
                target=save_checkpoint_in_background,
                args=(
                    self.mp_queue_send,
                    self.mp_queue_recv,
                    config_checkpoint,
                    config_job,
                ),
                daemon=True,
            )
            self.mp.start()
            self.cpu_offload_state_dict = None
            self.staging_ckpt_file = None
            self.staging_stream = torch.cuda.Stream()
            self.checkpoint_in_progress = False

    def keys_to_resume_during_load(self) -> tuple[set[str], str | None, bool | None]:
        """
        Determines the keys to resume from the checkpoint and the checkpoint path.
        If the checkpoint is the latest checkpoint of the same model, then it is a
        normal resume. If the checkpoint is a different model's checkpoint, then it is
        a warm start.

        Args:
            None

        Returns:
            resume_keys: The keys to resume from the checkpoint.
            checkpoint_path: The path to the checkpoint. If the checkpoint is a different
            warm_start: Whether to warm start the training from a different model's checkpoint.
                If the checkpoint is a different model's checkpoint, then this is True.
                If the checkpoint is the latest checkpoint of the same model, then this is False.
        """
        latest_checkpoint_file = self._read_latest_checkpoint_file()

        resume_keys = []
        warm_start = None

        if latest_checkpoint_file is not None:
            # 1. Resume training from the latest checkpoint of the same model.
            warm_start = False
            checkpoint_path = os.path.join(self.load_dirname, latest_checkpoint_file)
            resume_keys.extend(self.CHECKPOINT_KEYS)

        else:
            if self.load_path and not str(self.load_path).endswith(".pt"):
                # 2. Warm Start: Resume training from a different model's checkpoint
                # specified by `load_path`.
                warm_start = True
                checkpoint_path = self.load_path

                if self.load_s3_backend_key:
                    checkpoint_path = f"s3://{self.config_checkpoint.load_from_object_store.bucket}/{checkpoint_path}"

                    # If the path doesn't end with specific checkpoint, read the latest
                    # checkpoint file to determine the most recent checkpoint iteration.
                    if not re.search(r"/checkpoints/iter_\d{9}/?$", checkpoint_path):
                        old_ckpt_path = checkpoint_path
                        latest_ckpt_path = os.path.join(checkpoint_path, "checkpoints/latest_checkpoint.txt")

                        # If the latest checkpoint file exists, use it to determine the
                        # checkpoint path. Otherwise, use the original path.
                        if easy_io.exists(latest_ckpt_path, backend_key=self.load_s3_backend_key):
                            checkpoint_file = easy_io.load(
                                latest_ckpt_path, backend_key=self.load_s3_backend_key
                            ).strip()
                            checkpoint_path = f"{checkpoint_path}/checkpoints/{checkpoint_file}"
                        else:
                            log.warning(
                                f"Latest checkpoint file {latest_ckpt_path} not found, load from {old_ckpt_path}"
                            )
                            checkpoint_path = old_ckpt_path

                if self.load_training_state:
                    resume_keys.extend(self.CHECKPOINT_KEYS)
                else:
                    resume_keys.append("model")
                    if self.only_load_scheduler_state:
                        resume_keys.append("scheduler")
            else:
                checkpoint_path = None

        if len(self.keys_not_to_resume) > 0:
            for key in self.keys_not_to_resume:
                assert key in self.CHECKPOINT_KEYS, f"Invalid key to resume: {key} not in {self.CHECKPOINT_KEYS}"
            resume_keys = [key for key in resume_keys if key not in self.keys_not_to_resume]

        return set(resume_keys), checkpoint_path, warm_start

    @misc.timer("checkpoint loading")
    def load(
        self,
        model: ImaginaireModel,
        optimizer: torch.optim.Optimizer | None = None,
        scheduler: torch.optim.lr_scheduler.LRScheduler | None = None,
        grad_scaler: torch.amp.GradScaler | None = None,
    ) -> int:
        if self.callbacks is not None:
            self.callbacks.on_load_checkpoint_start(model)

        resume_keys, checkpoint_path, warm_start = self.keys_to_resume_during_load()
        resume_keys = sorted(resume_keys)
        log.critical(f"Resuming ckpt {checkpoint_path} with keys: {resume_keys}")

        iteration = 0

        if checkpoint_path is not None:
            self._check_checkpoint_exists(checkpoint_path)

            for key in resume_keys:
                dist.barrier()

                cur_key_ckpt_full_path = os.path.join(checkpoint_path, key)
                log.critical(f"Start loading checkpoint from {cur_key_ckpt_full_path}")

                storage_reader = self.get_storage_reader(cur_key_ckpt_full_path)
                strict_resume = self.config_checkpoint.strict_resume

                # Note that we only allow skipping loading of keys during warm start. If the checkpoint is
                # the latest checkpoint of the same model, then we don't need to skip any keys.
                keys_to_skip_loading = self.config_checkpoint.keys_to_skip_loading if warm_start else []

                load_planner = CustomLoadPlanner(
                    allow_partial_load=not strict_resume,
                    keys_to_skip_loading=keys_to_skip_loading,
                )

                if key == "model":
                    log.info("- Loading the model...")
                    _model_wrapper = ModelWrapper(model)
                    _state_dict = _model_wrapper.state_dict()
                    dcp.load(
                        _state_dict,
                        storage_reader=storage_reader,
                        planner=load_planner,
                    )
                    if self.config_checkpoint.load_ema_to_reg:
                        # The model has both net.* and net_ema.* submodules, so _state_dict
                        # contains both sets of keys after dcp.load(). Copy EMA weights into
                        # regular model weights so we can resume from EMA and reset EMA.
                        for sd_key in list(_state_dict.keys()):
                            if sd_key.startswith("net."):
                                key_ema = "net_ema." + sd_key.removeprefix("net.")
                                assert key_ema in _state_dict, (
                                    f"EMA key {key_ema} not found in state_dict. "
                                    "Ensure the model has net_ema submodule."
                                )
                                _state_dict[sd_key] = _state_dict[key_ema]
                    elif warm_start and any(str(s).startswith("net_ema") for s in self.keys_to_skip_loading):
                        # Only when net_ema.* is explicitly skipped on load (e.g. an HF->DCP
                        # init from convert_model_to_dcp that has only net.*): the skipped
                        # net_ema.* keep build_net() construction values (random init when
                        # vlm_config.pretrained_weights.enabled=False), which would seed EMA
                        # from random weights -> copy net.* -> net_ema.* so EMA starts from the
                        # freshly-loaded init. When net_ema.* IS loaded (e.g. a training DCP
                        # that carries a trained EMA), do NOT clobber it.
                        log.info("Warm start: net_ema. skipped on load -> resetting net_ema = net.")
                        for sd_key in list(_state_dict.keys()):
                            if sd_key.startswith("net."):
                                key_ema = "net_ema." + sd_key.removeprefix("net.")
                                if key_ema in _state_dict:
                                    _state_dict[key_ema] = _state_dict[sd_key]
                    results = _model_wrapper.load_state_dict(_state_dict)
                    if results is not None:
                        if len(results.missing_keys) > 0:
                            raise ValueError(f"Missing keys (not found in checkpoint): {results.missing_keys}")
                        if len(results.unexpected_keys) > 0:
                            raise ValueError(
                                f"Unexpected keys (found in checkpoint but not in model): {results.unexpected_keys}"
                            )

                elif key == "optim":
                    log.info("- Loading the optimizer...")
                    _state_dict = optimizer.state_dict()
                    dcp.load(
                        _state_dict,
                        storage_reader=storage_reader,
                        planner=load_planner,
                    )
                    optimizer.load_state_dict(_state_dict)

                elif key == "scheduler":
                    log.info("- Loading the scheduler...")
                    _state_dict = scheduler.state_dict()
                    dcp.load(
                        _state_dict,
                        storage_reader=storage_reader,
                        planner=load_planner,
                    )
                    scheduler.load_state_dict(_state_dict)

                elif key == "trainer":
                    log.info("- Loading the trainer...")

                    # Use rank-specific key for RNG state to support correct per-rank restoration
                    rng_key = f"rng_state_{dist.get_rank()}"
                    current_rng_state = get_rand_state_dict()
                    _state_dict = {
                        "grad_scaler": grad_scaler.state_dict(),
                        "iteration": iteration,
                    }
                    # Check if rng_key exists in checkpoint metadata to avoid failure with strict_resume=True
                    metadata = storage_reader.read_metadata()
                    rng_key_exists = any(
                        k.startswith(f"{rng_key}.") or k == rng_key for k in metadata.state_dict_metadata.keys()
                    )
                    if rng_key_exists:
                        _state_dict[rng_key] = current_rng_state

                    dcp.load(
                        _state_dict,
                        storage_reader=storage_reader,
                        planner=load_planner,
                    )
                    grad_scaler.load_state_dict(_state_dict["grad_scaler"])
                    iteration = _state_dict["iteration"]
                    set_rand_state_dict(_state_dict.get(rng_key, current_rng_state))

                elif key == "dataloader":
                    if not easy_io.exists(cur_key_ckpt_full_path, backend_key=self.load_s3_backend_key):
                        log.info(
                            f"Checkpoint {cur_key_ckpt_full_path} does not exist, skip loading dataloader.",
                            rank0_only=False,
                        )
                        continue

                    rank = dist.get_rank()
                    dataloader_pkl_path = os.path.join(cur_key_ckpt_full_path, f"rank_{rank}.pkl")
                    if not easy_io.exists(dataloader_pkl_path, backend_key=self.load_s3_backend_key):
                        log.info(f"No dataloader checkpoint found at {dataloader_pkl_path}", rank0_only=False)
                        continue

                    log.info(f"- Loading the dataloader {cur_key_ckpt_full_path}...", rank0_only=False)
                    _state_dict = easy_io.load(
                        dataloader_pkl_path,
                        file_format="pkl",
                        backend_key=self.load_s3_backend_key,
                    )
                    dataloader_wrapper = _DataloaderWrapper(self.callbacks)
                    if dataloader_wrapper.has_state():
                        dataloader_wrapper.load_state_dict(_state_dict)

                else:
                    raise ValueError(f"Invalid key: {key}. not support to resume.")

            if self.callbacks is not None and resume_keys:
                # Note that this callback is never used in the codebase.
                self.callbacks.on_load_checkpoint(model, state_dict={})
            log.info(f"Loaded checkpoint from {checkpoint_path} in iteration {iteration}")

        else:
            log.info("Training from scratch.")

        torch.cuda.empty_cache()

        if self.callbacks is not None:
            self.callbacks.on_load_checkpoint_end(model, iteration=iteration, checkpoint_path=checkpoint_path)
        return iteration

    def _checkpoint_async_with_pinned_memory(
        self, checkpoint_file: str, state_dict: Dict[str, Tuple[Any, str]]
    ) -> None:
        assert self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM, "Async mode must be AsyncMode.ASYNC_WITH_PINNED_MEM"

        from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict

        if self.cpu_offload_state_dict is None:
            log.info(f"Preparing the CPU memory for staging")
            self.cpu_offload_state_dict = _create_cpu_state_dict(state_dict, pin_memory=True, share_memory=True)

        log.info(f"Staging the state_dict in CPU memory")
        with torch.cuda.stream(self.staging_stream):
            self.cpu_offload_state_dict = _copy_state_dict(
                state_dict,
                self.cpu_offload_state_dict,
                non_blocking=True,
            )
            self.staging_ckpt_file = checkpoint_file

        self.staging_stream.synchronize()
        log.info(f"Staging the state_dict in CPU memory completed")

        self.mp_queue_send.put_nowait((self.cpu_offload_state_dict, self.staging_ckpt_file))
        self.checkpoint_in_progress = True
        log.info(f"Submitted checkpoint to background process")

    def _wait_for_previous_async_checkpoint(self) -> None:
        """
        Gets the results of previously submitted checkpoints.
        Pass them to callbacks if checkpoint succeeded.
        """
        assert self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM, "Async mode must be AsyncMode.ASYNC_WITH_PINNED_MEM"

        if not self.checkpoint_in_progress:
            return

        success = False
        try:
            log.info(f"Waiting for checkpoint save result")

            # Note that we set a timeout of 1 hour to avoid blocking the main process
            # indefinitely. Gloo and NCCL timeouts are ~30 minutes, so this timeout
            # should typically be sufficient.
            save_done: SaveDone = self.mp_queue_recv.get(timeout=3600)

            log.info(f"Received checkpoint save result: {save_done}")

            if self.callbacks is not None and save_done.succeeded:
                self.callbacks.on_save_checkpoint_success(
                    iteration=save_done.iteration, elapsed_time=save_done.elapsed_time
                )
            self.checkpoint_in_progress = False
            success = save_done.succeeded

        except Exception as e:
            log.error(f"Error waiting for checkpoint save result: {e}")

        if not success:
            # Terminate training execution upon a failed checkpoint save attempt.
            # A failure at this stage typically indicates a non-recoverable system error.
            # Continuing execution would result in subsequent persistent failures and
            # unnecessary waste of GPU resources.
            raise RuntimeError("Previous checkpoint save failed. Exiting...")

    def get_storage_writer(self, checkpoint_path: str) -> Union[S3StorageWriter, FileSystemWriter]:
        if self.save_to_object_store:
            return S3StorageWriter(
                credential_path=self.config_checkpoint.save_to_object_store.credentials,
                path=checkpoint_path,
                enable_gcs_patch_in_boto3=self.config_checkpoint.enable_gcs_patch_in_boto3,
            )
        return FileSystemWriter(path=checkpoint_path)

    def get_storage_reader(self, checkpoint_path: str) -> Union[S3StorageReader, FileSystemReader]:
        if self.load_from_object_store:
            return S3StorageReader(
                credential_path=self.config_checkpoint.load_from_object_store.credentials,
                path=checkpoint_path,
                enable_gcs_patch_in_boto3=self.config_checkpoint.enable_gcs_patch_in_boto3,
            )
        return FileSystemReader(checkpoint_path)

    def _save_as_pkl(self, obj: Any, output_dir: str) -> None:
        """Save per-rank Python checkpoint state such as no-replace dataloader progress."""
        rank = dist.get_rank()
        path = os.path.join(output_dir, f"rank_{rank}.pkl")
        easy_io.dump(
            obj,
            path,
            file_format="pkl",
            backend_key=self.save_s3_backend_key,
        )
        log.info(f"Saved state to {path}")

    def save_state_dict_worker(self, to_save_dict: Dict[str, Tuple[Any, str]], checkpoint_file: str) -> None:
        for key, (v, full_checkpoint_path) in to_save_dict.items():
            if key == "dataloader":
                self._save_as_pkl(v, full_checkpoint_path)
            else:
                storage_writer = self.get_storage_writer(full_checkpoint_path)
                # Note that it is ok to create a new CustomSavePlanner object
                # for each checkpoint save since the save plans are cached in a
                # class dictionary.
                save_planner = CustomSavePlanner(
                    dedup_save_to_lowest_rank=True,
                    enable_plan_caching=True,
                    cache_plans_key=f"custom_planner_{key}",
                )
                dcp.save(
                    v,
                    storage_writer=storage_writer,
                    planner=save_planner,
                )

        if distributed.is_rank0():
            log.info(f"Saving last checkpoint file {checkpoint_file}")
            self._write_latest_checkpoint_file(checkpoint_file)

        log.info(f"Saved checkpoint to {os.path.join(self.save_dirname, checkpoint_file)}")

    def save(
        self,
        model: ImaginaireModel,
        optimizer: torch.optim.Optimizer,
        scheduler: torch.optim.lr_scheduler.LRScheduler,
        grad_scaler: torch.amp.GradScaler,
        iteration: int,
    ) -> None:
        """Save network weights, optimizer parameters, scheduler parameters to a checkpoint.

        Args:
            model (ImaginaireModel): The PyTorch model.
            optimizer (torch.optim.Optimizer): The model optimizer.
            scheduler (torch.optim.lr_scheduler.LRScheduler): The optimization scheduler.
            grad_scaler (torch.amp.GradScaler): The gradient scaler (for mixed precision training).
            iteration (int): Current iteration number.
        """
        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
            self._wait_for_previous_async_checkpoint()

        if self.callbacks is not None:
            self.callbacks.on_save_checkpoint_start(model, iteration)

        checkpoint_file = f"iter_{iteration:09}"

        # Use rank-specific key for RNG state to ensure each rank saves its own state
        rng_key = f"rng_state_{dist.get_rank()}"

        to_save_dict = {
            "model": ModelWrapper(model).state_dict(),
            "optim": optimizer.state_dict(),
            "scheduler": scheduler.state_dict(),
            "trainer": {
                "grad_scaler": grad_scaler.state_dict(),
                "iteration": iteration,
                rng_key: get_rand_state_dict(),
            },
        }
        dataloader_wrapper = _DataloaderWrapper(self.callbacks)
        if dataloader_wrapper.has_state():
            to_save_dict["dataloader"] = dataloader_wrapper.state_dict()

        if self.callbacks is not None:
            self.callbacks.on_save_checkpoint(model, state_dict=to_save_dict)

        for k in to_save_dict.keys():
            output_dirname = os.path.join(self.save_dirname, f"iter_{iteration:09}/{k}")
            to_save_dict[k] = (to_save_dict[k], output_dirname)

        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
            dataloader_entry = to_save_dict.pop("dataloader", None)
            if dataloader_entry is not None:
                dataloader_state, dataloader_save_dir = dataloader_entry
                self._save_as_pkl(dataloader_state, dataloader_save_dir)
            self._checkpoint_async_with_pinned_memory(checkpoint_file, to_save_dict)
        else:
            start_time = time.monotonic()
            self.save_state_dict_worker(to_save_dict, checkpoint_file)
            elapsed_time = time.monotonic() - start_time
            log.info(f"Checkpoint save completed: Time taken: {elapsed_time:.2f} seconds")

            if self.callbacks is not None:
                self.callbacks.on_save_checkpoint_success(iteration=iteration, elapsed_time=elapsed_time)

        # This measures exposed (synchronous) checkpoint time, on_save_checkpoint_success()
        # is instead called to measure the entire duration for asynchronous checkpoint for the async case too.
        if self.callbacks is not None:
            self.callbacks.on_save_checkpoint_end(model=None, iteration=iteration)

    def finalize(self) -> None:
        super().finalize()
        if self.async_mode == AsyncMode.ASYNC_WITH_PINNED_MEM:
            if self.mp and self.mp.is_alive():
                # Wait for the previous checkpoint to complete.
                self._wait_for_previous_async_checkpoint()

                self.mp_queue_send.put(Terminate())
                self.mp.join()