Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/__init__.py

@@ -0,0 +1,21 @@
+from . import _extension
+from .api import CheckpointException
+from .default_planner import DefaultLoadPlanner, DefaultSavePlanner
+from .filesystem import FileSystemReader, FileSystemWriter
+from .hf_storage import HuggingFaceStorageReader, HuggingFaceStorageWriter
+from .metadata import (
+    BytesStorageMetadata,
+    ChunkStorageMetadata,
+    Metadata,
+    TensorStorageMetadata,
+)
+from .optimizer import load_sharded_optimizer_state_dict
+from .planner import LoadPlan, LoadPlanner, ReadItem, SavePlan, SavePlanner, WriteItem
+from .quantized_hf_storage import QuantizedHuggingFaceStorageReader
+
+# pyrefly: ignore [deprecated]
+from .state_dict_loader import load, load_state_dict
+
+# pyrefly: ignore [deprecated]
+from .state_dict_saver import async_save, save, save_state_dict
+from .storage import StorageReader, StorageWriter

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_async_executor.py

@@ -0,0 +1,33 @@
+# pyre-strict
+# mypy: allow-untyped-defs
+import abc
+import os
+from concurrent.futures import Future
+
+import torch.distributed as dist
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+from torch.distributed.checkpoint.planner import SavePlanner
+from torch.distributed.checkpoint.storage import StorageWriter
+
+
+class _AsyncCheckpointExecutor(abc.ABC):
+    @abc.abstractmethod
+    def execute_save(
+        self,
+        staging_future_or_state_dict: STATE_DICT_TYPE | Future[STATE_DICT_TYPE],
+        *,
+        checkpoint_id: str | os.PathLike | None = None,
+        storage_writer: StorageWriter | None = None,
+        planner: SavePlanner | None = None,
+        process_group: dist.ProcessGroup | None = None,
+        no_dist: bool = False,
+        use_collectives: bool = True,
+    ) -> Future:
+        """
+        Execute the checkpoint save request asynchronously.
+
+        This method is intended to be used as an abstraction for
+        implementing async checkpointing. The actual checkpoint save
+        operation is executed in a separate thread or process depending
+        on the implementation of this interface.
+        """

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_async_process_executor.py

@@ -0,0 +1,453 @@
+# pyre-strict
+# mypy: allow-untyped-defs
+import gc
+import logging
+import os
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from enum import Enum
+from typing import Any
+from uuid import uuid4
+
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from torch.distributed import PrefixStore, TCPStore
+from torch.distributed.checkpoint._async_executor import _AsyncCheckpointExecutor
+from torch.distributed.checkpoint.logger import _dcp_method_logger, _init_logger
+from torch.distributed.checkpoint.metadata import Metadata, STATE_DICT_TYPE
+from torch.distributed.checkpoint.planner import SavePlanner
+from torch.distributed.checkpoint.storage import StorageWriter
+from torch.distributed.checkpoint.utils import _DistWrapper
+from torch.distributed.elastic.agent.server.api import _get_fq_hostname
+from torch.distributed.elastic.utils.distributed import get_free_port
+
+
+logger = logging.getLogger()
+
+
+class _CheckpointSaveProcessControlOpts(Enum):
+    INIT_COMPLETE = "init_complete"
+    TERMINATE = "terminate"
+
+
+@dataclass(init=False, unsafe_hash=True)
+class _CheckpointRequestIdentifier:
+    checkpoint_id: str | os.PathLike | None
+    uuid: str
+
+    def __init__(self, checkpoint_id: str | os.PathLike | None):
+        self.checkpoint_id = checkpoint_id
+        self.uuid = str(uuid4())
+
+
+@dataclass
+class _AsyncCheckpointRequest:
+    staged_state_dict: STATE_DICT_TYPE
+    checkpoint_request_id: _CheckpointRequestIdentifier
+    storage_writer: StorageWriter | None = None
+    planner: SavePlanner | None = None
+    no_dist: bool = False
+    use_collectives: bool = True
+
+
+@dataclass(init=False)
+class _ProcessGroupInitInfo:
+    local_rank: int
+    global_rank: int
+    world_size: int
+    tcp_store_master_addr: str
+    tcp_store_master_port: int
+    use_prefix_store: bool
+    disable_automatic_gc: bool
+    disable_manual_gc: bool
+
+    def __init__(self, process_group: dist.ProcessGroup | None = None):
+        self.local_rank = dist.get_node_local_rank(fallback_rank=0)
+        self.global_rank = dist.get_rank(process_group)
+        self.world_size = dist.get_world_size(process_group)
+        self.use_prefix_store = os.environ.get("DCP_USE_PREFIX_STORE", "0") == "1"
+        self.disable_automatic_gc = (
+            os.environ.get("DCP_DISABLE_AUTOMATIC_GC", "0") == "1"
+        )
+        self.disable_manual_gc = os.environ.get("DCP_DISABLE_MANUAL_GC", "0") == "1"
+
+        # Let coordinator rank find a port on the localhost.
+        # Broadcast the (master_addr, port) to all ranks; each rank in the
+        # checkpoint daemon process will use TCPStore (master_addr, port)
+        # for collective communication.
+        dist_wrapper: _DistWrapper = _DistWrapper(
+            group=process_group,
+            use_dist=True,
+            coordinator_rank=0,
+        )
+
+        def get_master_addr_and_port() -> tuple[str, int]:
+            if self.use_prefix_store:
+                master_addr = os.environ.get("MASTER_ADDR")
+                master_port = os.environ.get("MASTER_PORT")
+                if master_addr is None:
+                    raise AssertionError("DCP needs MASTER_ADDR to use prefix store")
+                if master_port is None:
+                    raise AssertionError("DCP needs MASTER_PORT to use prefix store")
+                master_port = int(master_port)
+            else:
+                master_addr = os.environ.get("MASTER_ADDR")
+                if master_addr is None:
+                    master_addr = _get_fq_hostname()
+                master_port = get_free_port()
+
+            return master_addr, master_port
+
+        self.tcp_store_master_addr, self.tcp_store_master_port = dist_wrapper.broadcast(
+            step="get_master_addr_and_port",
+            map_fun=get_master_addr_and_port,
+        )
+
+
+class _AsyncCheckpointProcess:
+    def __init__(
+        self,
+        pg_init_info: _ProcessGroupInitInfo,
+    ):
+        self.ctx = mp.get_context("spawn")
+        self._process_pipe, child_end = self.ctx.Pipe()
+
+        self._save_process = self.ctx.Process(
+            target=self._checkpointing_subprocess,
+            args=(
+                pg_init_info,
+                child_end,
+            ),
+            daemon=True,
+        )
+
+        self._save_process.start()
+
+        # Close the parent's copy of child end after we pass it into the child,
+        # so the recv()s on it will fail-fast if the child process dies.
+        child_end.close()
+
+        # Wait for the checkpoint background process to initialize.
+        # Using default GLOO init timeout.
+        response = self._wait_for_response(timeout=1800)
+        if not response == _CheckpointSaveProcessControlOpts.INIT_COMPLETE:
+            raise AssertionError(f"Expected INIT_COMPLETE response, got {response}")
+
+    def __del__(self) -> None:
+        if self._save_process.is_alive():
+            try:
+                logger.info("Terminating the checkpoint background process.")
+                self._send(_CheckpointSaveProcessControlOpts.TERMINATE)
+                self._save_process.join(timeout=5)
+            finally:
+                if self._save_process.is_alive():
+                    logger.warning(
+                        "Checkpoint background process is still alive after termination request. Sending SIGTERM."
+                    )
+                    self._save_process.terminate()
+
+    def _send(self, data: Any) -> None:
+        self._process_pipe.send(data)
+
+    def _wait_for_response(self, timeout: float | None = None) -> Any:
+        if not self._save_process.is_alive():
+            logger.info("Checkpoint background process is dead calling join()...")
+            self._save_process.join()
+            raise RuntimeError(
+                f"Checkpoint background process is dead. Exit code: {self._save_process.exitcode}"
+            )
+
+        if timeout is not None and not self._process_pipe.poll(timeout=timeout):
+            raise RuntimeError(
+                f"Timed out after {timeout}s while waiting for response from checkpointer process pid: {self._save_process.pid}"
+            )
+
+        try:
+            response = self._process_pipe.recv()
+        except EOFError:
+            raise RuntimeError(  # noqa: B904
+                f"Checkpoint background process is dead. Exit code: {self._save_process.exitcode}"
+            )
+
+        if isinstance(response, BaseException):
+            raise response
+
+        return response
+
+    def save(
+        self,
+        staged_state_dict: STATE_DICT_TYPE,
+        *,
+        checkpoint_id: str | os.PathLike | None = None,
+        storage_writer: StorageWriter | None = None,
+        planner: SavePlanner | None = None,
+        no_dist: bool = False,
+        use_collectives: bool = True,
+    ) -> Metadata:
+        # Create a unique identifier to locate requests/responses
+        # from the checkpoint daemon process.
+        checkpoint_request_id = _CheckpointRequestIdentifier(checkpoint_id)
+        async_cp_request = _AsyncCheckpointRequest(
+            staged_state_dict=staged_state_dict,
+            checkpoint_request_id=checkpoint_request_id,
+            storage_writer=storage_writer,
+            planner=planner,
+            no_dist=no_dist,
+            use_collectives=use_collectives,
+        )
+        self._send(async_cp_request)
+        result = self._wait_for_response()
+        if not isinstance(result, Metadata):
+            raise AssertionError(f"Expected Metadata response, got {type(result)}")
+        return result
+
+    @staticmethod
+    def _execute_save(
+        state_dict: STATE_DICT_TYPE,
+        *,
+        checkpoint_request_id: _CheckpointRequestIdentifier,
+        storage_writer: StorageWriter | None = None,
+        planner: SavePlanner | None = None,
+        no_dist: bool = False,
+        use_collectives: bool = True,
+    ) -> Metadata:
+        from torch.distributed.checkpoint.state_dict_saver import save
+
+        metadata = save(
+            state_dict,
+            checkpoint_id=checkpoint_request_id.checkpoint_id,
+            storage_writer=storage_writer,
+            planner=planner,
+            no_dist=no_dist,
+            use_collectives=use_collectives,
+        )
+        return metadata
+
+    @staticmethod
+    def _checkpointing_subprocess(
+        pg_init_info: _ProcessGroupInitInfo,
+        parent_conn,
+    ) -> None:
+        # Phase 1: Process Group Initialization
+        # Only needs to execute once during the lifetime of the checkpoint background process.
+        try:
+            _init_logger(pg_init_info.global_rank)
+
+            # Setup environment variables for process group initialization.
+            os.environ["TORCHELASTIC_USE_AGENT_STORE"] = "False"
+            os.environ["MASTER_ADDR"] = pg_init_info.tcp_store_master_addr
+            os.environ["MASTER_PORT"] = str(pg_init_info.tcp_store_master_port)
+            os.environ["LOCAL_RANK"] = str(pg_init_info.local_rank)
+            os.environ["RANK"] = str(pg_init_info.global_rank)
+            os.environ["WORLD_SIZE"] = str(pg_init_info.world_size)
+
+            logger.info(
+                "Initializing dist.ProcessGroup in checkpoint background process on port %s",
+                pg_init_info.tcp_store_master_port,
+            )
+            # NOTE: GLOO backend is enforced here.
+            if pg_init_info.use_prefix_store:
+                logger.info(
+                    "Initializing dist.ProcessGroup in checkpoint background process with prefix store"
+                )
+                store = PrefixStore(
+                    "AsyncCheckpointProcess/",
+                    TCPStore(
+                        pg_init_info.tcp_store_master_addr,
+                        pg_init_info.tcp_store_master_port,
+                    ),
+                )
+                dist.init_process_group(
+                    backend=dist.Backend.GLOO,
+                    store=store,
+                    world_size=pg_init_info.world_size,
+                    rank=pg_init_info.global_rank,
+                )
+            else:
+                dist.init_process_group(backend=dist.Backend.GLOO)
+            dist.barrier()
+
+            logger.info("Checkpoint background process is running...")
+            parent_conn.send(_CheckpointSaveProcessControlOpts.INIT_COMPLETE)
+
+            if pg_init_info.disable_automatic_gc:
+                # Disable automatic garbage collection
+                # GC can optionally be called manually after each checkpoint
+                gc.disable()
+                logger.info("Disabled automatic garbage collection")
+        except BaseException as e:  # noqa: B036
+            logger.error(
+                f"Checkpoint background process failed during initialization: {e}"  # noqa: G004
+            )
+            parent_conn.send(e)
+            return
+
+        # Phase 2: Serving Loop
+        try:
+            first_request = True
+            while True:
+                logger.info("Waiting for checkpoint save request...")
+                obj = parent_conn.recv()
+                if (
+                    isinstance(obj, _CheckpointSaveProcessControlOpts)
+                    and obj == _CheckpointSaveProcessControlOpts.TERMINATE
+                ):
+                    logger.info("Terminating the checkpoint background process.")
+                    return
+                if not isinstance(obj, _AsyncCheckpointRequest):
+                    raise AssertionError(
+                        f"Expected _AsyncCheckpointRequest, got {type(obj)}"
+                    )
+                logger.info(
+                    f"Received async checkpoint request with id={obj.checkpoint_request_id.checkpoint_id}"  # noqa: G004
+                )
+
+                try:
+                    response = _AsyncCheckpointProcess._execute_save(
+                        obj.staged_state_dict,
+                        checkpoint_request_id=obj.checkpoint_request_id,
+                        storage_writer=obj.storage_writer,
+                        planner=obj.planner,
+                        no_dist=obj.no_dist,
+                        use_collectives=obj.use_collectives,
+                    )
+                    parent_conn.send(response)
+                    logger.info(
+                        f"Completed checkpoint save request for checkpoint_id={obj.checkpoint_request_id}"  # noqa: G004
+                    )
+
+                    # in theory this manual gc should not be needed as we shouldn't be leaking anything from checkpointing process
+                    if (
+                        pg_init_info.disable_automatic_gc
+                        and not pg_init_info.disable_manual_gc
+                    ):
+                        del obj
+
+                        collected_objects = gc.collect()
+
+                        logger.info(
+                            f"Manual garbage collection completed - collected {collected_objects} objects."  # noqa: G004
+                        )
+                        if first_request:
+                            # Freeze GC to not check GC for large checkpoint save plans
+                            # After freezing, subsequent gc.collect() calls will only scan
+                            # NEW objects created after this point, not the frozen save plan
+                            logger.info(
+                                "First checkpoint request completed - freezing gc"
+                            )
+                            gc.freeze()
+                    first_request = False
+                except BaseException as e:  # noqa: B036
+                    logger.error(
+                        f"Checkpoint save failed for checkpoint_id={obj.checkpoint_request_id.checkpoint_id}: {e}"  # noqa: G004
+                    )
+                    parent_conn.send(e)
+                    # Continue serving loop - don't exit process
+        finally:
+            logger.info("Checkpoint background process is shutting down...")
+            dist.destroy_process_group()
+            parent_conn.close()
+
+
+_CHECKPOINT_PROCESS: _AsyncCheckpointProcess | None = None
+
+
+class _ProcessBasedAsyncCheckpointExecutor(_AsyncCheckpointExecutor):
+    def __init__(self) -> None:
+        self._executor = ThreadPoolExecutor(max_workers=1)
+
+    @staticmethod
+    def _execute_save_impl(
+        *,
+        pg_init_info: _ProcessGroupInitInfo | None,
+        staging_future_or_state_dict: Future[STATE_DICT_TYPE] | STATE_DICT_TYPE,
+        checkpoint_id: str | os.PathLike | None = None,
+        storage_writer: StorageWriter | None = None,
+        planner: SavePlanner | None = None,
+        process_group: dist.ProcessGroup | None = None,
+        no_dist: bool = False,
+        use_collectives: bool = True,
+    ) -> Metadata:
+        global _CHECKPOINT_PROCESS
+        if _CHECKPOINT_PROCESS is None:
+            if pg_init_info is None:
+                raise AssertionError(
+                    "pg_init_info must not be None when _CHECKPOINT_PROCESS is None"
+                )
+            ckpt_kwargs = {}
+            if (ckpt_id := getattr(storage_writer, "checkpoint_id", None)) is not None:
+                ckpt_kwargs["checkpoint_id"] = ckpt_id
+                ckpt_kwargs["process_group"] = process_group
+
+            @_dcp_method_logger(**ckpt_kwargs)
+            def create_checkpoint_daemon_process() -> None:
+                global _CHECKPOINT_PROCESS
+                # pyrefly: ignore [bad-argument-type]
+                _CHECKPOINT_PROCESS = _AsyncCheckpointProcess(pg_init_info=pg_init_info)
+
+            create_checkpoint_daemon_process()
+
+        if _CHECKPOINT_PROCESS is None:
+            raise AssertionError(
+                "_CHECKPOINT_PROCESS must not be None after initialization"
+            )
+        staged_state_dict = (
+            staging_future_or_state_dict.result()
+            if isinstance(staging_future_or_state_dict, Future)
+            else staging_future_or_state_dict
+        )
+        return _CHECKPOINT_PROCESS.save(
+            staged_state_dict=staged_state_dict,
+            checkpoint_id=checkpoint_id,
+            storage_writer=storage_writer,
+            planner=planner,
+            no_dist=no_dist,
+            use_collectives=use_collectives,
+        )
+
+    def execute_save(
+        self,
+        staging_future_or_state_dict: Future[STATE_DICT_TYPE] | STATE_DICT_TYPE,
+        *,
+        checkpoint_id: str | os.PathLike | None = None,
+        storage_writer: StorageWriter | None = None,
+        planner: SavePlanner | None = None,
+        process_group: dist.ProcessGroup | None = None,
+        no_dist: bool = False,
+        use_collectives: bool = True,
+    ) -> Future:
+        """
+        NOTE:
+
+        - Checkpoint process is implemented as a daemon process.
+        The AsyncCheckpointProcess' lifetime is tied to the lifetime of the
+        main process (e.g. trainer process).
+
+        - The first call to execute_save_in_process() will initialize the checkpoint
+        daemon process. Subsequent async checkpoint requests will not need process
+        initialization. Therefore, the first async checkpoint request will take longer to complete.
+
+        - Process initialization can have significant overhead, dominated by latency for all ranks to spawn
+        a background process + process group initialization in the background process.
+        """
+
+        global _CHECKPOINT_PROCESS
+        pg_init_info: _ProcessGroupInitInfo | None = None
+        if _CHECKPOINT_PROCESS is None:
+            # Find a port on coordinator rank and broadcast
+            # to all ranks.
+            pg_init_info = _ProcessGroupInitInfo(process_group)
+
+        f: Future = self._executor.submit(
+            self._execute_save_impl,
+            pg_init_info=pg_init_info,
+            staging_future_or_state_dict=staging_future_or_state_dict,
+            checkpoint_id=checkpoint_id,
+            storage_writer=storage_writer,
+            planner=planner,
+            no_dist=no_dist,
+            use_collectives=use_collectives,
+        )
+        f.add_done_callback(lambda f: self._executor.shutdown(wait=False))
+
+        return f

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_async_thread_executor.py

@@ -0,0 +1,70 @@
+# pyre-strict
+# mypy: allow-untyped-defs
+import os
+from concurrent.futures import Future, ThreadPoolExecutor
+
+import torch.distributed as dist
+from torch.distributed.checkpoint._async_executor import _AsyncCheckpointExecutor
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+from torch.distributed.checkpoint.planner import SavePlanner
+from torch.distributed.checkpoint.storage import StorageWriter
+
+
+def save_wrapper(
+    staging_future_or_state_dict: Future[STATE_DICT_TYPE] | STATE_DICT_TYPE,
+    *,
+    checkpoint_id: str | os.PathLike | None = None,
+    storage_writer: StorageWriter | None = None,
+    planner: SavePlanner | None = None,
+    process_group: dist.ProcessGroup | None = None,
+    no_dist: bool = False,
+    use_collectives: bool = True,
+) -> Future:
+    from torch.distributed.checkpoint.state_dict_saver import save
+
+    staged_dict = (
+        staging_future_or_state_dict.result()
+        if isinstance(staging_future_or_state_dict, Future)
+        else staging_future_or_state_dict
+    )
+    return save(
+        staged_dict,
+        checkpoint_id=checkpoint_id,
+        storage_writer=storage_writer,
+        planner=planner,
+        process_group=process_group,
+        no_dist=no_dist,
+        use_collectives=use_collectives,
+    )
+
+
+class _ThreadBasedAsyncCheckpointExecutor(_AsyncCheckpointExecutor):
+    def __init__(self) -> None:
+        self._executor = ThreadPoolExecutor(
+            max_workers=1, thread_name_prefix="AsyncCheckpointExecutor"
+        )
+
+    def execute_save(
+        self,
+        staging_future_or_state_dict: Future[STATE_DICT_TYPE] | STATE_DICT_TYPE,
+        *,
+        checkpoint_id: str | os.PathLike | None = None,
+        storage_writer: StorageWriter | None = None,
+        planner: SavePlanner | None = None,
+        process_group: dist.ProcessGroup | None = None,
+        no_dist: bool = False,
+        use_collectives: bool = True,
+    ) -> Future:
+        f: Future = self._executor.submit(
+            save_wrapper,
+            staging_future_or_state_dict=staging_future_or_state_dict,
+            checkpoint_id=checkpoint_id,
+            storage_writer=storage_writer,
+            planner=planner,
+            process_group=process_group,
+            no_dist=no_dist,
+            use_collectives=use_collectives,
+        )
+        f.add_done_callback(lambda f: self._executor.shutdown(wait=False))
+
+        return f

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_checkpointer.py

@@ -0,0 +1,103 @@
+from concurrent.futures import Future
+from typing import Any
+
+import torch.distributed as dist
+import torch.distributed.checkpoint.state_dict_loader as loader
+import torch.distributed.checkpoint.state_dict_saver as saver
+from torch.distributed.checkpoint.metadata import Metadata, STATE_DICT_TYPE
+from torch.distributed.checkpoint.storage import (
+    LoadPlanner,
+    SavePlanner,
+    StorageReader,
+    StorageWriter,
+)
+
+
+__all__: list[str] = []
+
+
+class _Checkpointer:
+    """This base class specifies a high level API for saving and loading
+    distributed `state_dict` 's. It provides an abstraction over the low-level APIs
+    provided by :py:mod:`torch.distributed.checkpoint.storage`, essentially calling
+    :py:meth: `torch.distributed.state_dict_saver.save` and
+    :py:meth: `torch.distributed.state_dict_loader.load` with the provided storage
+    readers and writers.
+
+    .. warning::
+        This feature is experimental and subject to removal/change.
+
+    """
+
+    def __init__(
+        self,
+        storage_writer: StorageWriter,
+        storage_reader: StorageReader,
+        *,
+        process_group: dist.ProcessGroup | None = None,
+        coordinator_rank: int = 0,
+        no_dist: bool = False,
+        load_planner: LoadPlanner | None = None,
+        save_planner: SavePlanner | None = None,
+    ):
+        """Initializes the Checkpointer instance.
+
+        Args:
+            storage_writer: Instance of StorageWrite use to perform writes.
+            storage_reader: StorageReader used to load data from.
+            process_group: ProcessGroup to be used for cross-rank synchronization.
+            coordinator_rank: Rank to use to coordinate the checkpoint. rank0 is used by default.
+            no_dist: If ``True``, distributed checkpoint will not load in SPMD style. (Default: ``False``)
+            loader_planner: Instance of LoadPlanner to use when loading.
+            save_planner: Instance of SavePlanner to use when saving.
+        """
+        self.storage_writer = storage_writer
+        self.storage_reader = storage_reader
+        self.process_group = process_group
+        self.coordinator_rank = coordinator_rank
+        self.no_dist = no_dist
+        self.load_planner = load_planner
+        self.save_planner = save_planner
+
+    def save(
+        self,
+        state_dict: STATE_DICT_TYPE,
+    ) -> Metadata:
+        """Calls :py:meth: `torch.distributed.state_dict_saver.save`. Utilizing values passed during initialization."""
+        return saver.save(
+            state_dict,
+            self.storage_writer,
+            process_group=self.process_group,
+            coordinator_rank=self.coordinator_rank,
+            no_dist=self.no_dist,
+            planner=self.save_planner,
+        )
+
+    def async_save(
+        self,
+        state_dict: STATE_DICT_TYPE,
+    ) -> Future:
+        """
+        Calls :py:meth: `torch.distributed.state_dict_saver._async_save`. Utilizing values passed during initialization.
+
+        Returns:
+            Future: A future holding the resultant Metadata object from `save`.
+        """
+        response = saver.async_save(
+            state_dict,
+            storage_writer=self.storage_writer,
+            process_group=self.process_group,
+            planner=self.save_planner,
+        )
+        if not isinstance(response, Future):
+            raise AssertionError("response should be a Future instance")
+        return response
+
+    def load(self, state_dict: dict[str, Any]) -> None:
+        """Calls :py:meth: `torch.distributed.state_dict_loader.load`. Utilizing values passed during initialization."""
+        loader.load(
+            state_dict,
+            storage_reader=self.storage_reader,
+            process_group=self.process_group,
+            planner=self.load_planner,
+        )

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_consolidate_hf_safetensors.py

@@ -0,0 +1,749 @@
+# pyre-strict
+
+import concurrent.futures
+import glob
+import json
+import logging
+import math
+import os
+import struct
+import time
+from dataclasses import dataclass, field
+from typing import Any
+
+import torch
+from torch import distributed as dist
+from torch.distributed.checkpoint._hf_utils import (
+    _gen_file_name,
+    _get_dcp_custom_metadata,
+    _get_safetensors_file_metadata,
+    _metadata_fn,
+    DATA_OFFSETS_KEY,
+    DEFAULT_EXTRA_METADATA_KEY,
+    DTYPE_KEY,
+    SAVED_OFFSETS_KEY,
+    SHAPE_KEY,
+    SUFFIX,
+)
+
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+@dataclass
+class _FqnData:
+    """
+    Dataclass to store information about a tensor (identified by its fully qualified name).
+
+    Attributes:
+        offset_in_file: Byte offset where this tensor's data begins in the output file
+        shape_in_file: Shape of the tensor in the output file
+        dtype_size: Size of the tensor's data type in bytes
+        dtype_str: String representation of the tensor's data type
+    """
+
+    offset_in_file: int = 0
+    shape_in_file: list[int] = field(default_factory=list)
+    dtype_size: int = 0
+    dtype_str: str = ""
+
+
+@dataclass
+class _OutputFileData:
+    """
+    Dataclass to store information about an output safetensors file.
+
+    Attributes:
+        metadata_size: Size of the metadata section in bytes
+        fqn_data: Dictionary mapping tensor names to their metadata
+    """
+
+    metadata_size: int = 0
+    fqn_data: dict[str, _FqnData] = field(default_factory=dict)
+
+
+@dataclass
+class _InputFileData:
+    """
+    Dataclass to store information about an input safetensors file.
+
+    Attributes:
+        metadata_size: Size of the metadata section in bytes
+        metadata: Json metadata from the safetensors file
+    """
+
+    metadata_size: int = 0
+    metadata: Any = None
+
+
+def _parse_input_metadata(
+    input_files_data: dict[str, _InputFileData],
+    output_files_data: dict[str, _OutputFileData],
+) -> None:
+    """
+    Parse metadata from input safetensors files to determine the full tensor shapes and types.
+
+    This function analyzes the metadata from all input files to determine the complete shape
+    of each tensor after consolidation. It updates the output_files_data with this information.
+
+    Args:
+        input_files_data: dict of metadata from input safetensors files
+        output_files_data: Dictionary mapping output file paths to their metadata
+
+    Raises:
+        ValueError: If no DCP custom metadata is found in a safetensors file
+    """
+
+    from safetensors.torch import _getdtype  # type: ignore[import]
+
+    # Dictionary to track the full size of each tensor across all shards
+    fqn_to_size_mapping: dict[str, tuple[list[int], str]] = {}
+
+    for file_data in input_files_data.values():
+        safetensors_metadata = file_data.metadata
+        dcp_sharding_info = _get_dcp_custom_metadata(safetensors_metadata)
+        if not dcp_sharding_info:
+            raise ValueError(
+                "No DCP custom metadata found in safetensors file. The file must be saved with DCP to be consolidated."
+            )
+
+        for key, val in safetensors_metadata.items():
+            if key == DEFAULT_EXTRA_METADATA_KEY:
+                continue
+
+            # Get the shape of this tensor shard and its offset in the full tensor
+            sizes = val[SHAPE_KEY]
+            offsets = dcp_sharding_info[key][SAVED_OFFSETS_KEY]
+
+            if key not in fqn_to_size_mapping:
+                # First time seeing this tensor - calculate its full size by adding offsets to dimensions
+                cur_size = [size + offset for size, offset in zip(sizes, offsets)]
+                fqn_to_size_mapping[key] = (cur_size, val[DTYPE_KEY])
+            else:
+                # We've seen this tensor before - update its size if this shard extends beyond current known dimensions
+                cur_size = fqn_to_size_mapping[key][0]
+                for i in range(len(sizes)):
+                    cur_size[i] = max(cur_size[i], sizes[i] + offsets[i])
+
+    # Now that we know the full size of each tensor, populate the output file data
+    for fqn, tensor_info in fqn_to_size_mapping.items():
+        tensor_size = tensor_info[0]
+        dtype_str = tensor_info[1]
+        for output_data in output_files_data.values():
+            # Add this tensor to the output file if it's already assigned there
+            if fqn in output_data.fqn_data:
+                output_data.fqn_data[fqn] = _FqnData(
+                    shape_in_file=tensor_size,
+                    dtype_size=torch.finfo(_getdtype(dtype_str)).bits
+                    // 8,  # Convert bits to bytes
+                    dtype_str=dtype_str,
+                )
+
+
+def _write_metadata(
+    output_files_data: dict[str, _OutputFileData],
+) -> None:
+    """
+    Write metadata to the beginning of each output safetensors file.
+
+    This function writes the metadata section to each output file, including information
+    about tensor shapes, data types, and offsets. It also updates the offset_in_file
+    field for each tensor in the output_files_data.
+
+    Args:
+        output_files_data: Dictionary mapping output file paths to their metadata
+    """
+    # Process each output file
+    for file_path, output_data in output_files_data.items():
+        with open(file_path, "wb") as f:
+            metadata = {}
+            curr_offset = 0
+
+            # Calculate offsets for each tensor in the file
+            for fqn, fqn_data in output_data.fqn_data.items():
+                # Calculate the end offset by multiplying all dimensions and the data type size
+                end_offset = (
+                    curr_offset
+                    + math.prod(fqn_data.shape_in_file) * fqn_data.dtype_size
+                )
+
+                # Store metadata for this tensor
+                metadata[fqn] = {
+                    SHAPE_KEY: fqn_data.shape_in_file,
+                    DTYPE_KEY: fqn_data.dtype_str,
+                    DATA_OFFSETS_KEY: [
+                        curr_offset,
+                        end_offset,
+                    ],  # Start and end byte offsets
+                }
+                # Store the offset for later use when writing the actual tensor data
+                fqn_data.offset_in_file = curr_offset
+
+                # Update current offset for the next tensor
+                curr_offset = end_offset
+
+            # Convert metadata to JSON and encode as bytes
+            json_metadata = json.dumps(metadata)
+            json_bytes = json_metadata.encode("utf-8")
+
+            # Write the metadata size as an 8-byte unsigned integer (little-endian)
+            size_in_bytes = len(json_bytes)
+            header_len = struct.pack("<Q", size_in_bytes)
+
+            # Write the header length and metadata to the file
+            f.write(header_len)
+            f.write(json_bytes)
+
+            # Store the total metadata size (header + JSON) for later use
+            output_data.metadata_size = f.tell()
+
+
+def _read_tensor_data(
+    f,
+    start_offset: int,
+    end_offset: int,
+    metadata_size: int,
+) -> bytes:
+    """
+    Read a specific byte range of tensor data from an open safetensors file.
+
+    Args:
+        f: An open file object (handle) for the safetensors file
+        start_offset: Start offset of tensor data within the data section
+        end_offset: End offset of tensor data within the data section
+        metadata_size: Size of the metadata header
+
+    Returns:
+        Raw tensor data as bytes
+    """
+    absolute_start = metadata_size + start_offset
+    length = end_offset - start_offset
+
+    f.seek(absolute_start)
+    return f.read(length)
+
+
+def _process_output_file(
+    output_file: str,
+    output_data: _OutputFileData,
+    input_files_data: dict[str, _InputFileData],
+) -> None:
+    """
+    Process a single output file by writing tensor data from input files using direct reads.
+
+    This function is designed to be run in parallel for different output files.
+
+    Args:
+        output_file: Path to the output file
+        output_data: Metadata for the output file
+        input_files_data: Dictionary mapping input file paths to their metadata
+    """
+
+    sorted_tensors = sorted(
+        output_data.fqn_data.items(), key=lambda x: x[1].offset_in_file
+    )
+
+    file_handles = {}
+    dcp_metadata = {}
+    for safetensors_file, file_data in input_files_data.items():
+        dcp_metadata[safetensors_file] = _get_dcp_custom_metadata(file_data.metadata)
+
+    try:
+        # Open all input files for reading
+        for safetensors_file in input_files_data:
+            file_handles[safetensors_file] = open(safetensors_file, "rb")  # noqa: SIM115
+
+        with open(output_file, "r+b") as output_stream:
+            output_stream.seek(0, os.SEEK_END)
+            # Process each tensor in sequential output order
+            for tensor_fqn, tensor_fqn_data in sorted_tensors:
+                full_tensor_mv = memoryview(
+                    bytearray(
+                        math.prod(tensor_fqn_data.shape_in_file)
+                        * tensor_fqn_data.dtype_size
+                    )
+                )
+
+                # Process each input safetensors file
+                for safetensors_file in input_files_data:
+                    file_metadata = input_files_data[safetensors_file].metadata
+                    input_metadata_size = input_files_data[
+                        safetensors_file
+                    ].metadata_size
+
+                    if tensor_fqn not in file_metadata:
+                        continue
+
+                    metadata = file_metadata[tensor_fqn]
+
+                    data_offsets = metadata[DATA_OFFSETS_KEY]
+
+                    # Use explicit reads to fetch tensor data efficiently
+                    data_to_write = _read_tensor_data(
+                        file_handles[safetensors_file],
+                        data_offsets[0],
+                        data_offsets[1],
+                        input_metadata_size,
+                    )
+
+                    # Get the offsets of this tensor shard within the full tensor
+                    # pyrefly: ignore [unsupported-operation]
+                    fqn_custom_metadata = dcp_metadata[safetensors_file][tensor_fqn]  # type: ignore[index]
+                    offsets_of_tensor_being_read = fqn_custom_metadata[
+                        SAVED_OFFSETS_KEY
+                    ]  # type: ignore[index]
+
+                    # Write this tensor shard to the appropriate position in the output file
+                    _write_sub_tensor_to_file_optimized(
+                        full_tensor_mv,
+                        data_to_write,
+                        tensor_fqn_data.dtype_size,  # Size of each element in bytes
+                        tensor_fqn_data.shape_in_file,  # Full tensor shape
+                        offsets_of_tensor_being_read,  # Where this shard belongs in the full tensor
+                        metadata[SHAPE_KEY],  # Shape of this shard
+                    )
+
+                output_stream.write(full_tensor_mv)
+
+    finally:
+        for f in file_handles.values():
+            f.close()
+
+
+def _write_data(
+    input_files_data: dict[str, _InputFileData],
+    output_files_data: dict[str, _OutputFileData],
+    num_threads: int = 1,
+) -> None:
+    """
+    Write tensor data from input files to the output files using memory mapping.
+
+    This function reads tensor data from each input file and writes it to the appropriate
+    position in the output files based on the tensor's offsets. When num_threads > 1,
+    the work is split across threads with each thread handling a different output file.
+
+    Args:
+        input_files_data: Dictionary mapping input file paths to their metadata
+        output_files_data: Dictionary mapping output file paths to their metadata
+        num_threads: Number of threads to use for parallel processing
+    """
+    if num_threads <= 1 or len(output_files_data) <= 1:
+        # Sequential processing
+        for output_file, output_data in output_files_data.items():
+            _process_output_file(output_file, output_data, input_files_data)
+    else:
+        # Parallel processing with ThreadPoolExecutor
+        with concurrent.futures.ThreadPoolExecutor(
+            max_workers=min(num_threads, len(output_files_data))
+        ) as executor:
+            futures = []
+            for output_file, output_data in output_files_data.items():
+                futures.append(
+                    executor.submit(
+                        _process_output_file,
+                        output_file,
+                        output_data,
+                        input_files_data,
+                    )
+                )
+
+            # Wait for all futures to complete
+            for future in concurrent.futures.as_completed(futures):
+                # Handle any exceptions that might have occurred
+                try:
+                    future.result()
+                except Exception as e:
+                    print(f"Error processing output file: {e}")
+                    raise
+
+
+def _write_sub_tensor_to_file_optimized(
+    full_tensor_mv: memoryview,
+    sub_tensor_bytes: bytes,
+    element_size: int,
+    tensor_shape: list[int],
+    sub_tensor_offsets: list[int],
+    sub_tensor_shape: list[int],
+) -> None:
+    """
+    Optimized version that writes the maximum number of contiguous bytes possible.
+
+    Uses a unified algorithm that calculates the maximum contiguous bytes that can be
+    written in each iteration and continues until the entire subtensor is written.
+    Handles all sharding patterns efficiently:
+    - Full sub-tensor at once for row-wise sharding
+    - Row-by-row for column-wise sharding
+    - Optimized chunks for other patterns
+
+    Args:
+        full_tensor_mv: Buffer to write the full tensor to
+        sub_tensor_bytes: Raw tensor data as bytes
+        element_size: Size of each element in bytes
+        tensor_shape: Shape of the full tensor
+        sub_tensor_offsets: Starting offsets of the sub-tensor within the full tensor
+        sub_tensor_shape: Shape of the sub-tensor
+    """
+    # Handle empty tensors
+    if not tensor_shape or not sub_tensor_shape:
+        return
+
+    # Calculate tensor strides for efficient indexing
+    tensor_strides = [1]
+    for i in range(len(tensor_shape) - 1, 0, -1):
+        tensor_strides.insert(0, tensor_strides[0] * tensor_shape[i])
+
+    sub_tensor_strides = [1]
+    for i in range(len(sub_tensor_shape) - 1, 0, -1):
+        sub_tensor_strides.insert(0, sub_tensor_strides[0] * sub_tensor_shape[i])
+
+    total_elements = math.prod(sub_tensor_shape)
+
+    elements_written = 0
+    while elements_written < total_elements:
+        # Convert linear index to multi-dimensional indices
+        temp_idx = elements_written
+        indices = []
+        for dim_size in reversed(sub_tensor_shape):
+            indices.append(temp_idx % dim_size)
+            temp_idx //= dim_size
+        indices.reverse()
+
+        # Calculate maximum contiguous elements we can write from this position
+        max_contiguous = _calculate_max_contiguous_elements(
+            indices, sub_tensor_shape, tensor_shape
+        )
+
+        # Calculate source position in bytes
+        src_pos = sum(idx * stride for idx, stride in zip(indices, sub_tensor_strides))
+        src_byte_offset = src_pos * element_size
+
+        # Calculate destination position in bytes
+        dest_indices = [
+            idx + offset for idx, offset in zip(indices, sub_tensor_offsets)
+        ]
+        dest_pos = sum(
+            idx * stride for idx, stride in zip(dest_indices, tensor_strides)
+        )
+        dest_byte_offset = dest_pos * element_size
+
+        # Write the contiguous chunk
+        bytes_to_write = max_contiguous * element_size
+        chunk_data = sub_tensor_bytes[
+            src_byte_offset : src_byte_offset + bytes_to_write
+        ]
+        full_tensor_mv[dest_byte_offset : dest_byte_offset + bytes_to_write] = (
+            chunk_data
+        )
+
+        elements_written += max_contiguous
+
+
+def _calculate_max_contiguous_elements(
+    indices: list[int],
+    sub_tensor_shape: list[int],
+    tensor_shape: list[int],
+) -> int:
+    """
+    Calculate the maximum number of contiguous elements that can be written from current position.
+
+    This determines the largest chunk by checking how elements are laid out in memory
+    and finding natural boundaries where contiguity breaks.
+
+    Args:
+        indices: Current position indices in the sub-tensor
+        sub_tensor_shape: Shape of the sub-tensor being written
+        tensor_shape: Shape of the full tensor
+
+    Raises:
+        ValueError: If input lists are empty, have mismatched lengths, or contain invalid values
+    """
+    # Validate input lists are not empty
+    if not indices or not sub_tensor_shape or not tensor_shape:
+        raise ValueError("Input lists cannot be empty")
+
+    # Validate all lists have the same length (same number of dimensions)
+    if not (len(indices) == len(sub_tensor_shape) == len(tensor_shape)):
+        raise ValueError(
+            f"All input lists must have the same length. Got indices: {len(indices)}, "
+            f"sub_tensor_shape: {len(sub_tensor_shape)}, tensor_shape: {len(tensor_shape)}"
+        )
+
+    # Validate indices are within bounds of sub_tensor_shape
+    for i, (idx, sub_dim) in enumerate(zip(indices, sub_tensor_shape)):
+        if idx >= sub_dim:
+            raise ValueError(
+                f"Index {idx} at dimension {i} is out of bounds for sub-tensor shape {sub_tensor_shape}"
+            )
+
+    # Validate sub_tensor dimensions don't exceed tensor dimensions
+    for i, (sub_dim, tensor_dim) in enumerate(zip(sub_tensor_shape, tensor_shape)):
+        if sub_dim > tensor_dim:
+            raise ValueError(
+                f"Sub-tensor dimension {sub_dim} at position {i} exceeds tensor dimension {tensor_dim}"
+            )
+
+    # Start with elements remaining in the last dimension
+    max_contiguous = sub_tensor_shape[-1] - indices[-1]
+
+    # Check if we can extend across multiple dimensions
+    # We can write across dimension boundaries if we're writing complete "rows"
+    # and the layout in destination tensor maintains contiguity
+
+    # For 2D case: check if we can write multiple complete rows
+    if len(sub_tensor_shape) >= 2:
+        # If we're at the start of a row and can write complete rows
+        if indices[-1] == 0:  # At start of last dimension (column)
+            rows_remaining = sub_tensor_shape[-2] - indices[-2]  # Rows left to write
+
+            # Check if writing complete rows maintains contiguity in destination
+            # This is true for row-wise sharding or when sub-tensor spans full width
+            if sub_tensor_shape[-1] == tensor_shape[-1]:  # Full width
+                max_contiguous = rows_remaining * sub_tensor_shape[-1]
+
+            # For higher dimensions, check if we can extend further
+            if len(sub_tensor_shape) >= 3 and indices[-2] == 0:
+                # Check if we can write complete 2D slices
+                remaining_in_dim = sub_tensor_shape[-3] - indices[-3]
+                if (
+                    sub_tensor_shape[-1] == tensor_shape[-1]
+                    and sub_tensor_shape[-2] == tensor_shape[-2]
+                ):
+                    max_contiguous = (
+                        remaining_in_dim * sub_tensor_shape[-2] * sub_tensor_shape[-1]
+                    )
+
+    return max_contiguous
+
+
+def _write_overall_metadata_file(
+    output_dir: str,
+    output_files_data: dict[str, _OutputFileData],
+) -> None:
+    """
+    Write the overall metadata file that maps tensor names to their file locations.
+
+    This creates a model.safetensors.index.json file that HuggingFace models use
+    to locate tensors across multiple files.
+
+    Args:
+        output_dir: Directory where the metadata file will be written
+        output_files_data: Dictionary mapping output file paths to their metadata
+    """
+    total_size = 0
+    weight_map = {}
+    for output_path, value in output_files_data.items():
+        for fqn, fqn_data in value.fqn_data.items():
+            total_size += math.prod(fqn_data.shape_in_file) * fqn_data.dtype_size
+            weight_map[fqn] = os.path.basename(output_path)
+
+    metadata_to_write: dict[str, Any] = {}
+    metadata_to_write["metadata"] = {"total_size": total_size}
+    metadata_to_write["weight_map"] = weight_map
+
+    metadata_path = os.path.join(output_dir, f"{_metadata_fn}")
+    with open(metadata_path, "w") as metadata_file:
+        json.dump(metadata_to_write, metadata_file, indent=2)
+
+
+def _consolidate_safetensors_files(
+    input_dir: str,
+    output_dir: str,
+    fqn_to_file_mapping: dict[str, str],
+    num_threads: int,
+) -> dict[str, _OutputFileData]:
+    output_files_data: dict[str, _OutputFileData] = {}
+    # Create multiple output files based on the provided mapping
+    for fqn, filename in fqn_to_file_mapping.items():
+        output_path = os.path.join(output_dir, filename)
+
+        if output_path not in output_files_data:
+            output_files_data[output_path] = _OutputFileData(fqn_data={fqn: _FqnData()})
+        else:
+            output_files_data[output_path].fqn_data[fqn] = _FqnData()
+
+    # Find all safetensors files in the input directory
+    safetensors_files = glob.glob(os.path.join(input_dir, f"*{SUFFIX}"))
+
+    # Read metadata from all input files
+    input_files_data: dict[str, _InputFileData] = {}
+    for safetensor_file in safetensors_files:
+        with open(safetensor_file, "rb") as f:
+            metadata, size = _get_safetensors_file_metadata(f)
+            input_files_data[safetensor_file] = _InputFileData(
+                metadata_size=size, metadata=metadata
+            )
+    # Step 1: Parse metadata to determine tensor shapes and types
+    _parse_input_metadata(input_files_data, output_files_data)
+
+    # Step 2: Write metadata headers to output files
+    _write_metadata(output_files_data)
+    # Step 3: Write actual tensor data from input files to output files
+    _write_data(input_files_data, output_files_data, num_threads)
+
+    return output_files_data
+
+
+def consolidate_safetensors_files(
+    input_dir: str,
+    output_dir: str,
+    fqn_to_index_mapping: dict[str, int],
+    num_threads: int = 1,
+) -> None:
+    """
+    Main function to consolidate sharded safetensors files into one or more output files.
+
+    This function orchestrates the entire consolidation process:
+    1. Sets up the output file structure based on the fqn_to_index_mapping
+    2. Finds all safetensors files in the input directory
+    3. Parses metadata from all input files
+    4. Writes metadata to the output files
+    5. Writes tensor data from input files to output files
+    6. Writes overall model.index.safetensors.json file with weight map
+
+    Args:
+        input_dir: Directory containing sharded safetensors files
+        output_dir: Directory where consolidated files will be written
+        fqn_to_index_mapping: Optional mapping of tensor names to output file indices.
+                             If None, all tensors will be consolidated into a single file.
+        num_threads: Number of threads to use for parallel processing of saving data to output files.
+    """
+    start_time = time.time()
+    logger.info(
+        "Consolidating safetensors files from %s to %s. Beginning at time %f",
+        input_dir,
+        output_dir,
+        start_time,
+    )
+
+    max_index = max(fqn_to_index_mapping.values())
+    fqn_to_file_mapping = {
+        fqn: _gen_file_name(idx, max_index) for fqn, idx in fqn_to_index_mapping.items()
+    }
+
+    output_files_data = _consolidate_safetensors_files(
+        input_dir, output_dir, fqn_to_file_mapping, num_threads
+    )
+
+    # Step 4: Write overall model.index.safetensors.json file with weight map
+    _write_overall_metadata_file(output_dir, output_files_data)
+
+    logger.info("Done consolidating. Took %.2f secs.", time.time() - start_time)
+
+
+def consolidate_safetensors_files_on_every_rank(
+    input_dir: str,
+    output_dir: str,
+    fqn_to_index_mapping: dict[str, int],
+    num_threads: int = 1,
+    process_group: dist.ProcessGroup | None = None,
+) -> None:
+    """
+    Consolidate sharded safetensors files across multiple ranks, with each rank handling a subset of output files.
+
+    This function distributes the consolidation work by assigning output files to different ranks.
+    All tensors with the same index in fqn_to_index_mapping are processed by the same rank,
+    as they belong to the same output file.
+
+    If process_group is provided, rank and world_size will be derived from it. Otherwise,
+    they will be automatically detected from the distributed environment if available.
+
+    Args:
+        input_dir: Directory containing sharded safetensors files
+        output_dir: Directory where consolidated files will be written
+        fqn_to_index_mapping: Mapping of tensor names to output file indices
+        num_threads: Number of threads to use for parallel processing on each rank
+        process_group: PyTorch distributed process group (default: None, will use default group)
+    """
+
+    start_time = time.time()
+    # Derive rank and world_size from process_group or default distributed environment
+    if dist.is_available() and dist.is_initialized():
+        rank = dist.get_rank(group=process_group)
+        world_size = dist.get_world_size(group=process_group)
+    else:
+        # Default to single process mode if distributed is not initialized
+        rank = 0
+        world_size = 1
+        logger.warning(
+            "Distributed environment not initialized. Running in single process mode."
+        )
+    logger.info(
+        "Rank %d/%d: Consolidating safetensors files from %s to %s",
+        rank,
+        world_size,
+        input_dir,
+        output_dir,
+    )
+
+    # Find all unique indices in the mapping
+    unique_indices = set(fqn_to_index_mapping.values())
+
+    # Distribute indices across ranks
+    indices_for_this_rank = []
+    for idx in unique_indices:
+        # Simple distribution: index % world_size == rank
+        if idx % world_size == rank:
+            indices_for_this_rank.append(idx)
+
+    logger.info(
+        "Rank %d: Assigned %d output files out of %d total files",
+        rank,
+        len(indices_for_this_rank),
+        len(unique_indices),
+    )
+
+    # Filter the fqn_to_index_mapping to only include tensors for this rank
+    filtered_mapping = {
+        fqn: idx
+        for fqn, idx in fqn_to_index_mapping.items()
+        if idx in indices_for_this_rank
+    }
+
+    output_files_data: dict[str, _OutputFileData] = {}
+    if filtered_mapping:
+        # Convert index mapping to filename mapping
+        max_index = max(unique_indices)
+        filtered_filename_mapping = {}
+        for fqn, idx in filtered_mapping.items():
+            filename = _gen_file_name(idx, max_index)
+            filtered_filename_mapping[fqn] = filename
+
+        # Call the existing consolidation function with the filtered mapping
+        output_files_data = _consolidate_safetensors_files(
+            input_dir=input_dir,
+            output_dir=output_dir,
+            fqn_to_file_mapping=filtered_filename_mapping,
+            num_threads=num_threads,
+        )
+
+    logger.info(
+        "Rank %d: Done consolidating. Processed %d unique indices in %.2f secs.",
+        rank,
+        len(indices_for_this_rank),
+        time.time() - start_time,
+    )
+
+    # Wait for all ranks to complete and gather output_files_data on rank 0
+    if dist.is_available() and dist.is_initialized():
+        gathered_output_files_data: list[dict[str, _OutputFileData]] | None = (
+            [{} for _ in range(world_size)] if rank == 0 else None
+        )
+        dist.gather_object(
+            output_files_data,
+            gathered_output_files_data,
+            dst=0,
+            group=process_group,
+        )
+
+        if rank == 0:
+            # Merge all output_files_data from all ranks
+            all_output_files_data: dict[str, _OutputFileData] = {}
+            if gathered_output_files_data is None:
+                raise AssertionError
+            for rank_data in gathered_output_files_data:
+                all_output_files_data.update(rank_data)
+
+            _write_overall_metadata_file(output_dir, all_output_files_data)
+            logger.info("Rank 0: Wrote overall metadata file.")
+            logger.info("Total time taken: %.2f secs.", time.time() - start_time)
+        dist.barrier(group=process_group)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_save_plans.py

@@ -0,0 +1,65 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+import dataclasses
+from collections import defaultdict
+from typing import TYPE_CHECKING
+
+from torch.distributed.checkpoint.planner import SavePlan, WriteItem
+
+
+if TYPE_CHECKING:
+    from torch.distributed.checkpoint.metadata import MetadataIndex
+
+__all__ = ["dedup_save_plans"]
+
+
+def dedup_save_plans(
+    all_plans: list[SavePlan],
+    save_to_lowest_rank: bool = False,
+) -> list[SavePlan]:
+    """
+    Removes duplicate entries from appearing on multiple SavePlans. For each duplicate across
+    a set of SavePlans, only the smallest SavePlan in terms of planned storage keeps the entry.
+
+    Please note that this function does not modify the original SavePlans, but rather returns
+    """
+
+    # Map to query the plan indices that a write item is duplicated in
+    write_item_to_plan_indices: dict[MetadataIndex, set[int]] = defaultdict(set)
+    # Map to query the write item from its index
+    write_item_idx_to_write_item: dict[MetadataIndex, WriteItem] = {}
+    # Set of write item indices that are present in each plan
+    # After deduplication, this will be the set of write item indices that are present in the final plans
+    plan_to_item_indices: list[set[MetadataIndex]] = [
+        {item.index for item in plan.items} for plan in all_plans
+    ]
+
+    for plan_idx, plan in enumerate(all_plans):
+        for write_item in plan.items:
+            # map each write item to its plan
+            write_item_to_plan_indices[write_item.index].add(plan_idx)
+            write_item_idx_to_write_item[write_item.index] = write_item
+    plan_to_size = [0] * len(all_plans)
+    for write_item_idx, plan_indices in write_item_to_plan_indices.items():
+        if save_to_lowest_rank:
+            select_plan_idx = min(plan_indices)
+        else:
+            select_plan_idx = min(
+                plan_indices, key=lambda plan_idx: plan_to_size[plan_idx]
+            )
+
+        write_item = write_item_idx_to_write_item[write_item_idx]
+        # Ignore the storage size of anything that is not a tensor, since
+        # we don't know how much storage they represent
+        plan_to_size[select_plan_idx] += write_item.tensor_storage_size() or 1
+        for plan_idx in plan_indices - {select_plan_idx}:
+            plan_to_item_indices[plan_idx].discard(write_item_idx)
+    # Sanity check
+    if len(all_plans) != len(plan_to_item_indices):
+        raise AssertionError("len(all_plans) != len(plan_to_item_indices)")
+    # Create new plans with the updated write items post deduplication
+    return [
+        dataclasses.replace(
+            plan, items=[item for item in plan.items if item.index in item_indexes]
+        )
+        for plan, item_indexes in zip(all_plans, plan_to_item_indices)
+    ]

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_dedup_tensors.py

@@ -0,0 +1,62 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+import dataclasses
+import logging
+from typing import TYPE_CHECKING
+
+from torch.distributed.checkpoint.planner import SavePlan
+
+
+if TYPE_CHECKING:
+    from torch.distributed.checkpoint.metadata import MetadataIndex
+
+__all__ = ["dedup_tensors"]
+
+
+def init_logger() -> logging.Logger:
+    logger = logging.getLogger(__name__)
+    level = logging.INFO
+    logger.setLevel(level)
+    console = logging.StreamHandler()
+    formatter = logging.Formatter(
+        "%(asctime)s %(filename)s:%(lineno)s %(levelname)s p:%(processName)s t:%(threadName)s: %(message)s"
+    )
+    console.setFormatter(formatter)
+    console.setLevel(level)
+    logger.addHandler(console)
+    logger.propagate = False
+    return logger
+
+
+logger = init_logger()
+
+
+# TODO add docstring for dedup_tensors
+def dedup_tensors(all_plans: list[SavePlan]) -> list[SavePlan]:
+    all_plans = list(all_plans)
+    key_to_plan: dict[MetadataIndex, list[int]] = {}
+    for plan_idx, plan in enumerate(all_plans):
+        for write_item in plan.items:
+            key_to_plan.setdefault(write_item.index, []).append(plan_idx)
+
+    replicated_items = {k: v for k, v in key_to_plan.items() if len(v) > 1}
+
+    # Remove duplicates by always keeping the first entry.
+    # Compute the per-rank remove set.
+    plan_to_keys: dict[int, list[MetadataIndex]] = {}
+    for key, plans in replicated_items.items():
+        for plan_idx in plans[1:]:
+            plan_to_keys.setdefault(plan_idx, []).append(key)
+    if len(plan_to_keys) > 0:
+        logger.info("Duplicate keys to remove: %s", plan_to_keys)
+
+    for plan_idx, keys in plan_to_keys.items():
+        key_set = set(keys)
+        # rewrite items and remove elements
+        new_items = [
+            write_item
+            for write_item in all_plans[plan_idx].items
+            if write_item.index not in key_set
+        ]
+        all_plans[plan_idx] = dataclasses.replace(all_plans[plan_idx], items=new_items)
+
+    return all_plans

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/__init__.py

@@ -0,0 +1,53 @@
+"""
+Checkpoint functionality for machine learning models.
+
+This module provides classes for saving and loading model checkpoints in a distributed
+training environment. It includes functionality for coordinating checkpoint operations
+across multiple processes and customizing the checkpoint process through hooks.
+
+Key components:
+- Checkpointer: Main class for orchestrating checkpoint operations (save, load)
+- CheckpointWriter: Handles writing state dictionaries to storage
+- CheckpointReader: Handles reading state dictionaries from storage read
+- Barrier: Synchronization mechanism for distributed checkpointing
+- RankInfo: Information about the current rank in a distributed environment
+"""
+
+from .barriers import (
+    Barrier,
+    BarrierConfig,
+    create_barrier_from_config,
+    TCPStoreBarrier,
+)
+from .builder import make_async_checkpointer, make_sync_checkpointer
+from .checkpoint_reader import CheckpointReader
+from .checkpoint_writer import CheckpointWriter, CheckpointWriterConfig, WriterHook
+from .checkpointer import AsyncCheckpointer, Checkpointer, SyncCheckpointer
+from .config import CheckpointerConfig
+from .staging import CheckpointStager, CheckpointStagerConfig, DefaultStager
+from .types import RankInfo, STATE_DICT
+from .utils import wrap_future
+
+
+__all__ = [
+    "Barrier",
+    "TCPStoreBarrier",
+    "CheckpointReader",
+    "CheckpointWriter",
+    "CheckpointWriterConfig",
+    "WriterHook",
+    "Checkpointer",
+    "SyncCheckpointer",
+    "AsyncCheckpointer",
+    "CheckpointerConfig",
+    "BarrierConfig",
+    "create_barrier_from_config",
+    "CheckpointStager",
+    "CheckpointStagerConfig",
+    "DefaultStager",
+    "RankInfo",
+    "STATE_DICT",
+    "wrap_future",
+    "make_sync_checkpointer",
+    "make_async_checkpointer",
+]

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/barriers.py

@@ -0,0 +1,267 @@
+"""
+Barrier implementations for synchronizing distributed checkpoint operations.
+
+This module provides abstract and concrete barrier implementations that ensure
+all ranks in a distributed training environment complete their checkpoint operations
+before proceeding, which is essential for data consistency.
+"""
+
+import abc
+import logging
+from collections import Counter
+from dataclasses import dataclass, field
+from datetime import timedelta
+from typing import Any, Optional
+
+import torch.distributed as dist
+import torch.distributed.elastic.utils.store as store_util
+
+
+logger = logging.getLogger()
+
+
+# Registry of barrier types
+BARRIER_REGISTRY: dict[str, type] = {}
+
+
+def register_barrier(barrier_class: type) -> type:
+    """Register a barrier class in the global registry."""
+    if hasattr(barrier_class, "barrier_type"):
+        BARRIER_REGISTRY[barrier_class.barrier_type] = barrier_class
+    return barrier_class
+
+
+@dataclass
+class BarrierConfig:
+    """
+    Configuration for barrier construction.
+
+    This class provides a flexible way to configure different barrier implementations
+    with their specific constructor arguments. The barrier type will be looked up
+    from a registry and instantiated with rank_info and barrier_args.
+
+    Attributes:
+        barrier_type: A string identifying the barrier type (e.g., "tcp_store").
+                     If None, no barrier will be used.
+        barrier_args: Dictionary of arguments to pass to the barrier constructor.
+                     rank_info will be automatically injected as the first argument.
+
+    Examples:
+        # No barrier
+        BarrierConfig()
+
+        # TCPStore barrier
+        BarrierConfig(
+            barrier_type="tcp_store",
+            barrier_args={
+                'timeout_barrier_init_secs': 30,
+                'barrier_prefix_list': ['checkpoint'],
+                'use_checkpoint_barrier_tcpstore_libuv': False,
+                'tcpstore_port': 12345,
+                'master_address': 'localhost'
+            }
+        )
+    """
+
+    barrier_type: str | None = None
+    barrier_args: dict[str, Any] = field(default_factory=dict)
+
+
+def create_barrier_from_config(
+    barrier_config: BarrierConfig,
+) -> Optional["Barrier"]:
+    """
+    Create a barrier instance from BarrierConfig.
+
+    Args:
+        barrier_config: Configuration for barrier construction.
+
+    Returns:
+        Barrier instance or None if no barrier type is configured.
+
+    Raises:
+        ValueError: If the barrier_type is not found in the registry.
+    """
+    if barrier_config.barrier_type is None:
+        return None
+
+    if barrier_config.barrier_type not in BARRIER_REGISTRY:
+        raise ValueError(
+            f"Unknown barrier type: {barrier_config.barrier_type}. "
+            f"Available types: {list(BARRIER_REGISTRY.keys())}"
+        )
+
+    barrier_class = BARRIER_REGISTRY[barrier_config.barrier_type]
+    return barrier_class(**barrier_config.barrier_args)
+
+
+class Barrier(abc.ABC):
+    """
+    Abstract base class for synchronization barriers.
+
+    A barrier ensures that all ranks in a distributed environment reach a certain
+    point in execution before any rank proceeds further, which is essential for
+    coordinating operations like checkpointing across multiple processes.
+    """
+
+    @abc.abstractmethod
+    def __init__(self, **kwargs: dict[str, Any]):
+        """
+        Initialize a barrier.
+
+        Args:
+            **kwargs: Keyword arguments for specific barrier implementations.
+                     Common arguments may include rank information, barrier prefixes,
+                     timeout settings, and other barrier-specific configuration.
+        """
+        # No implementation needed in the abstract base class
+
+    @abc.abstractmethod
+    def execute_barrier(self) -> None:
+        """
+        Execute a synchronization barrier.
+
+        This method uses the barrier_prefix provided during initialization to
+        coordinate synchronization across processes.
+        """
+
+
+@register_barrier
+class DistBarrier(Barrier):
+    """
+    A barrier implementation using PyTorch's distributed barrier for synchronization.
+
+    This barrier uses the built-in torch.distributed.barrier() function to coordinate
+    synchronization across multiple processes. It's simpler than TCPStoreBarrier but
+    requires an initialized process group.
+    """
+
+    barrier_type = "dist_barrier"
+
+    def __init__(
+        self,
+    ) -> None:
+        """
+        Initialize a DistBarrier.
+
+        This barrier requires an initialized PyTorch distributed process group.
+        No additional arguments are needed as it uses the current process group.
+
+        Raises:
+            AssertionError: If the distributed process group is not initialized.
+        """
+        if not dist.is_initialized():
+            raise AssertionError("DistBarrier requires an initialized process group.")
+
+    def execute_barrier(self) -> None:
+        """
+        Execute a synchronization barrier using the prefix provided during initialization.
+        """
+        # Note: dist.barrier() doesn't support explicit timeouts
+        # The timeout is handled by the underlying implementation
+        dist.barrier()
+
+
+@register_barrier
+class TCPStoreBarrier(Barrier):
+    """
+    A barrier implementation using PyTorch's TCPStore for synchronization.
+
+    This barrier uses a TCP-based distributed key-value store to coordinate
+    synchronization across multiple processes. It uses a single TCP store
+    for all barrier operations, with different prefixes to distinguish between
+    different barrier types.
+    """
+
+    barrier_type = "tcp_store"
+
+    def __init__(
+        self,
+        global_rank: int,
+        global_world_size: int,
+        barrier_prefix: str,
+        timeout_barrier_init_secs: int,
+        use_checkpoint_barrier_tcpstore_libuv: bool,
+        tcpstore_port: int,
+        master_address: str,
+        timeout_secs: int,
+    ):
+        """
+        Initialize a TCPStoreBarrier.
+
+        Args:
+            global_rank: The rank of the current process in the distributed environment.
+            global_world_size: The total number of processes in the distributed environment.
+            barrier_prefix: A string prefix to identify this specific barrier.
+            timeout_barrier_init_secs: Timeout in seconds for initializing the TCPStore.
+            use_checkpoint_barrier_tcpstore_libuv: Whether to use libuv for the TCPStore.
+            tcpstore_port: Port number for the TCPStore.
+            master_address: Address of the master node for the TCPStore.
+            timeout_secs: Maximum time in seconds to wait for all ranks to reach the barrier.
+        """
+        logger.info(
+            "Initializing TCPStore master_address=%s tcpstore_port=%s rank=%s "
+            "world_size=%s barrier_prefix=%s timeout_barrier_init_secs=%s "
+            "use_checkpoint_barrier_tcpstore_libuv=%s timeout_secs=%s",
+            master_address,
+            tcpstore_port,
+            global_rank,
+            global_world_size,
+            barrier_prefix,
+            timeout_barrier_init_secs,
+            use_checkpoint_barrier_tcpstore_libuv,
+            timeout_secs,
+        )
+
+        # Counter collection to track barrier seq on a per barrier prefix basis.
+        self._tcp_store_barrier_seq: Counter = Counter()
+        self._barrier_prefix = barrier_prefix
+
+        # Store rank and world size for barrier operations
+        self._global_rank = global_rank
+        self._global_world_size = global_world_size
+        self._timeout_secs = timeout_secs
+
+        # Create a single TCP store for all barrier operations
+        self._tcp_store = dist.TCPStore(
+            master_address,
+            int(tcpstore_port),
+            world_size=self._global_world_size,
+            timeout=timedelta(seconds=timeout_barrier_init_secs),
+            is_master=(self._global_rank == 0),
+        )
+
+    def execute_barrier(self) -> None:
+        """
+        Execute a synchronization barrier using the prefix provided during initialization.
+
+        The implementation uses a sequence number that is incremented every time
+        a barrier is reached. The sequence number is per barrier prefix to allow
+        different barriers to operate concurrently.
+        """
+        barrier_prefix = self._barrier_prefix
+
+        logger.info(
+            "Executing barrier barrier_prefix=%s timeout_secs=%s",
+            barrier_prefix,
+            self._timeout_secs,
+        )
+
+        def _rank_key(rank: int) -> str:
+            return f"rank{rank}"
+
+        # Track which barrier sequence this rank is joining.
+        self._tcp_store.set(
+            _rank_key(self._global_rank),
+            str(self._tcp_store_barrier_seq[barrier_prefix]),
+        )
+
+        # Execute barrier for that sequence number (for the specific prefix).
+        store_util.barrier(
+            store=self._tcp_store,
+            world_size=self._global_world_size,
+            key_prefix=(
+                barrier_prefix + str(self._tcp_store_barrier_seq[barrier_prefix])
+            ),
+        )
+        self._tcp_store_barrier_seq[barrier_prefix] += 1

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/builder.py

@@ -0,0 +1,174 @@
+"""
+Factory functions for creating checkpointer instances with sensible defaults.
+
+This module provides high-level factory functions that simplify the creation
+of checkpointer instances by automatically handling component initialization
+and configuration with reasonable defaults.
+"""
+
+from collections.abc import Callable
+from typing import Any
+
+import torch.distributed as dist
+
+from .barriers import create_barrier_from_config
+from .checkpoint_process import CheckpointProcess
+from .checkpoint_reader import CheckpointReader
+from .checkpoint_writer import CheckpointWriter, CheckpointWriterConfig, WriterHook
+from .checkpointer import AsyncCheckpointer, SyncCheckpointer
+from .config import CheckpointerConfig
+from .staging import DefaultStager
+from .types import RankInfo
+
+
+def _get_default_rank_info() -> RankInfo:
+    """
+    Get default rank information from the current distributed environment.
+
+    Returns:
+        RankInfo: Rank information from the default process group if initialized,
+                 otherwise single-rank fallback.
+    """
+    if dist.is_initialized():
+        return RankInfo(
+            global_world_size=dist.get_world_size(),
+            global_rank=dist.get_rank(),
+        )
+    else:
+        # Single-rank fallback
+        return RankInfo(global_world_size=1, global_rank=0)
+
+
+def default_subprocess_init_fn(*_: Any) -> None:
+    """Default subprocess initialization function (no-op)."""
+
+
+def default_writer_init_fn(rank_info: RankInfo) -> CheckpointWriter:
+    """Default checkpoint writer initialization function."""
+    return CheckpointWriter(
+        config=CheckpointWriterConfig(),
+        rank_info=rank_info,
+    )
+
+
+def make_sync_checkpointer(
+    config: CheckpointerConfig = CheckpointerConfig(),
+    rank_info: RankInfo | None = None,
+    commit_hook: WriterHook | None = None,
+) -> SyncCheckpointer:
+    """
+    Factory function to create a SyncCheckpointer instance with sensible defaults.
+
+    This function creates a synchronous checkpointer with default components, automatically
+    detecting rank information from the default process group if available, and using the
+    provided component configurations.
+
+    Args:
+        config: CheckpointerConfig containing component-specific configurations
+               (writer_config, staging_config, process_config). Defaults to CheckpointerConfig().
+        rank_info: RankInfo for distributed training. Defaults to auto-detection from
+                  the default PyTorch distributed process group if initialized, otherwise
+                  falls back to single-rank (world_size=1, rank=0).
+        commit_hook: Optional hook for custom actions before and after checkpoint commits.
+
+    Returns:
+        SyncCheckpointer: A configured synchronous checkpointer instance.
+
+    Examples:
+        # Simplest usage - auto-detect rank, default config
+        checkpointer = make_sync_checkpointer()
+
+        # Explicit rank configuration
+        checkpointer = make_sync_checkpointer(
+            rank_info=RankInfo(global_world_size=4, global_rank=0)
+        )
+
+        # Disable barrier
+        from .barriers import BarrierConfig
+        config = CheckpointerConfig(barrier_config=BarrierConfig(barrier_type=None))
+        checkpointer = make_sync_checkpointer(config=config)
+    """
+    if rank_info is None:
+        rank_info = _get_default_rank_info()
+
+    reader = CheckpointReader(
+        rank_info=rank_info,
+    )
+
+    barrier = create_barrier_from_config(config.barrier_config)
+
+    writer = CheckpointWriter(
+        config=config.writer_config,
+        rank_info=rank_info,
+        barrier=barrier,
+        commit_hook=commit_hook,
+    )
+
+    return SyncCheckpointer(
+        writer=writer,
+        reader=reader,
+    )
+
+
+def make_async_checkpointer(
+    config: CheckpointerConfig = CheckpointerConfig(),
+    rank_info: RankInfo | None = None,
+    subprocess_init_fn: Callable[..., None] = default_subprocess_init_fn,
+    subprocess_init_args: tuple[Any, ...] = (),
+    checkpoint_writer_init_fn: Callable[..., CheckpointWriter] = default_writer_init_fn,
+    checkpoint_writer_init_args: dict[str, Any] | None = None,
+) -> AsyncCheckpointer:
+    """
+    Factory function to create an AsyncCheckpointer instance with sensible defaults.
+
+    This function creates an asynchronous checkpointer using the provided configuration,
+    automatically detecting rank information if not provided.
+
+    Args:
+        config: CheckpointerConfig containing component-specific configurations.
+        rank_info: RankInfo for distributed training. Defaults to auto-detection.
+        subprocess_init_fn: Function to initialize the subprocess. Defaults to no-op.
+        subprocess_init_args: Arguments to pass to subprocess_init_fn.
+        checkpoint_writer_init_fn: Function to create CheckpointWriter instance.
+        checkpoint_writer_init_args: Arguments to pass to checkpoint_writer_init_fn.
+
+    Returns:
+        AsyncCheckpointer: A configured asynchronous checkpointer instance.
+
+    Examples:
+        # Create with default config
+        checkpointer = make_async_checkpointer()
+
+        # Create with custom init functions
+        checkpointer = make_async_checkpointer(
+            subprocess_init_fn=my_subprocess_init_fn,
+            checkpoint_writer_init_fn=my_writer_init_fn
+        )
+    """
+    if rank_info is None:
+        rank_info = _get_default_rank_info()
+
+    reader = CheckpointReader(
+        rank_info=rank_info,
+    )
+
+    checkpoint_stager = DefaultStager(
+        config=config.staging_config,
+    )
+
+    checkpoint_writer_init_args = checkpoint_writer_init_args or {}
+
+    checkpoint_process = CheckpointProcess(
+        rank_info=rank_info,
+        config=config.process_config,
+        subprocess_init_fn=subprocess_init_fn,
+        subprocess_init_args=subprocess_init_args,
+        checkpoint_writer_init_fn=checkpoint_writer_init_fn,
+        checkpoint_writer_init_args=checkpoint_writer_init_args,
+    )
+
+    return AsyncCheckpointer(
+        checkpoint_stager=checkpoint_stager,
+        checkpoint_process=checkpoint_process,
+        reader=reader,
+    )

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/checkpoint_process.py

@@ -0,0 +1,361 @@
+import logging
+import os
+import traceback
+from collections.abc import Callable
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from enum import Enum
+from multiprocessing.connection import Connection
+from typing import Any
+
+import torch.multiprocessing as mp
+from torch.multiprocessing.spawn import ProcessExitedException
+
+from .checkpoint_writer import CheckpointWriter
+from .types import RankInfo, STATE_DICT
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class CheckpointProcessConfig:
+    """
+    Configuration options for the CheckpointProcess.
+
+    This class provides configuration options for the checkpoint process,
+    including initialization functions, timeouts, and writer configuration.
+
+    Attributes:
+        subprocess_init_timeout_secs: Maximum time in seconds to wait for subprocess initialization.
+        subprocess_shutdown_timeout_secs: Maximum time in seconds to wait for subprocess shutdown.
+    """
+
+    subprocess_init_timeout_secs: int = 30
+    subprocess_shutdown_timeout_secs: int = 60
+
+
+class RequestType(Enum):
+    PING = "ping"
+    WRITE_CHECKPOINT = "write_checkpoint"
+    TERMINATE_PROCESS = "exit"
+
+
+@dataclass
+class WorkerRequest:
+    """
+    A dataclass for storing the command to be sent to the worker process.
+    Note: This relies on pickling to send the command to the worker process. Handle
+    backward compatibility accordingly.
+    """
+
+    request_type: RequestType
+    payload: dict[str, Any]
+
+
+@dataclass
+class WorkerResponse:
+    request_type: RequestType
+    success: bool
+    error_msg: str | None = None
+    payload: dict[str, Any] | None = None
+
+
+class CheckpointProcess:
+    """
+    A checkpoint writer that writes checkpoints to a remote process.
+    """
+
+    def __init__(
+        self,
+        rank_info: RankInfo,
+        config: CheckpointProcessConfig,
+        subprocess_init_fn: Callable[[Any], None],
+        subprocess_init_args: tuple[Any, ...],
+        checkpoint_writer_init_fn: Callable[..., CheckpointWriter],
+        checkpoint_writer_init_args: dict[str, Any],
+    ):
+        self._executor = ThreadPoolExecutor(max_workers=1)
+        self._rank_info = rank_info
+        self._config = config
+        self._subprocess_init_fn = subprocess_init_fn
+        self._subprocess_init_args = subprocess_init_args
+        self._checkpoint_writer_init_fn = checkpoint_writer_init_fn
+        self._checkpoint_writer_init_args = checkpoint_writer_init_args
+        self.process = None
+        self._parent_end: Connection | None = None
+        self._child_end: Connection | None = None
+
+        self.process_creation_future = self._executor.submit(
+            self._create_subprocess,
+            config,
+        )
+
+    def _create_subprocess(
+        self,
+        config: CheckpointProcessConfig,
+    ) -> None:
+        logger.info(
+            "Creating checkpoint subprocess for rank %d", self._rank_info.global_rank
+        )
+
+        spawn_context = mp.get_context("spawn")
+        self._parent_end, child_end = spawn_context.Pipe()
+
+        # Known workaround for https://github.com/pytorch/pytorch/issues/37377
+        os.environ["MKL_SERVICE_FORCE_INTEL"] = "GNU"
+
+        logger.debug("Spawning subprocess for rank_info=%s", self._rank_info)
+        self.process = mp.spawn(
+            fn=CheckpointProcess._subprocess,
+            args=(
+                self._rank_info,
+                child_end,
+                self._subprocess_init_fn,
+                self._subprocess_init_args,
+                self._checkpoint_writer_init_fn,
+                self._checkpoint_writer_init_args,
+            ),
+            nprocs=1,
+            join=False,
+            daemon=True,
+        )
+
+        # close the child end of the pipe so recv on it will fail
+        # fast when the child process is terminated unexpectedly.
+        child_end.close()
+        self._send(
+            request_type=RequestType.PING,
+            payload={},
+        )
+
+        logger.debug(
+            "Waiting for checkpoint subprocess to initialize (timeout: %ds)",
+            config.subprocess_init_timeout_secs,
+        )
+
+        # wait for the timeout or a response from subprocess
+        if self._parent_end is None:
+            raise AssertionError("Parent end of pipe should be initialized")
+        if not self._parent_end.poll(timeout=config.subprocess_init_timeout_secs):
+            msg = f"Timed out after {config.subprocess_init_timeout_secs}s waiting for checkpoint subprocess to initialize"
+            logger.error(msg)
+            raise TimeoutError(msg)
+
+        self._recv()
+        logger.info("Checkpoint subprocess initialized successfully")
+
+    @staticmethod
+    def _subprocess(
+        sub_rank: int,
+        rank_info: RankInfo,
+        parent_pipe: Connection,
+        subprocess_init_fn: Callable[[Any], None],
+        subprocess_init_args: tuple[Any, ...],
+        checkpoint_writer_init_fn: Callable[..., CheckpointWriter],
+        checkpoint_writer_init_args: dict[str, Any],
+    ) -> None:
+        logger.debug(
+            "Checkpoint subprocess started for rank %d/%d (PID: %d)",
+            rank_info.global_rank,
+            rank_info.global_world_size,
+            os.getpid(),
+        )
+
+        if sub_rank != 0:
+            raise AssertionError("We need only one checkpointer per parent training")
+        request = WorkerRequest(request_type=RequestType.PING, payload={})
+
+        try:
+            # Calling initialize callback, so we can perform app-specific initialization of the subprocess.
+            subprocess_init_fn(*subprocess_init_args)
+
+            # Initialize checkpoint writer - automatically include rank_info in init_args
+            writer_init_args = dict(checkpoint_writer_init_args)
+            if "rank_info" not in writer_init_args:
+                writer_init_args["rank_info"] = rank_info
+            checkpoint_writer = checkpoint_writer_init_fn(**writer_init_args)
+
+            while True:
+                request = parent_pipe.recv()
+
+                if request.request_type == RequestType.PING:
+                    parent_pipe.send(
+                        WorkerResponse(request_type=RequestType.PING, success=True)
+                    )
+                elif request.request_type == RequestType.WRITE_CHECKPOINT:
+                    path = request.payload["path"]
+                    logger.info("Writing checkpoint to %s", path)
+
+                    checkpoint_writer.write(
+                        path=path,
+                        state_dict=request.payload["state_dict"],
+                        **request.payload["kwargs"],
+                    )
+
+                    logger.info("Checkpoint written successfully to %s", path)
+                    parent_pipe.send(
+                        WorkerResponse(RequestType.WRITE_CHECKPOINT, success=True)
+                    )
+                elif request.request_type == RequestType.TERMINATE_PROCESS:
+                    logger.debug("Received termination request.")
+                    parent_pipe.send(
+                        WorkerResponse(RequestType.TERMINATE_PROCESS, success=True)
+                    )
+                    logger.info("Subprocess terminated gracefully")
+                    break
+                else:
+                    error_msg = f"Unknown request type: {request.request_type}"
+                    logger.error(error_msg)
+                    raise ValueError(error_msg)
+
+        except Exception as e:
+            error_text = traceback.format_exc()
+            logger.error(
+                "Exception in subprocess  (%s): %s", type(e).__name__, error_text
+            )
+
+            # Communicating exception via the queue to the main process
+            parent_pipe.send(
+                WorkerResponse(
+                    request_type=request.request_type,
+                    success=False,
+                    error_msg=error_text,
+                )
+            )
+            parent_pipe.close()
+            logger.exception("Subprocess terminated due to exception")
+
+    def _send(self, request_type: RequestType, payload: dict[str, Any]) -> None:
+        try:
+            if self._parent_end is None:
+                raise AssertionError("Parent end of pipe should be initialized")
+            self._parent_end.send(
+                WorkerRequest(
+                    request_type=request_type,
+                    payload=payload,
+                )
+            )
+        except OSError as e:
+            error_msg = "Child process terminated unexpectedly"
+            logger.exception(
+                "Communication failed during %s request", request_type.value
+            )
+            raise RuntimeError(error_msg) from e
+
+    def _recv(self) -> dict[str, Any] | None:
+        try:
+            if self._parent_end is None:
+                raise AssertionError("Parent end of pipe should be initialized")
+            response = self._parent_end.recv()
+            if response.success is False:
+                error_msg = (
+                    f"Unexpected response from worker process: {response.error_msg}"
+                )
+                logger.error(error_msg)
+                raise RuntimeError(error_msg)
+            return response.payload
+        except (EOFError, BrokenPipeError, ConnectionResetError) as e:
+            error_msg = f"Child process terminated unexpectedly: {e}"
+            logger.error(error_msg)
+            raise RuntimeError(error_msg) from e
+
+    def write(
+        self,
+        state_dict: STATE_DICT | Future[STATE_DICT],
+        path: str,
+        **kwargs: Any,
+    ) -> Future[None] | None:
+        logger.debug("Waiting for subprocess initialization to complete")
+
+        # wait until the process is started
+        self.process_creation_future.result()
+
+        return self._executor.submit(
+            self._write,
+            state_dict,
+            path,
+            **kwargs,
+        )
+
+    def _write(
+        self,
+        state_dict: STATE_DICT | Future[STATE_DICT],
+        path: str,
+        **kwargs: Any,
+    ) -> None:
+        logger.debug("Starting checkpoint write to %s", path)
+
+        # wait for staging state_dict to be available
+        if isinstance(state_dict, Future):
+            logger.debug("Waiting for state_dict Future to resolve")
+            sd = state_dict.result()
+        else:
+            sd = state_dict
+
+        # Log state_dict info only if debug logging is enabled (performance-conscious)
+        if logger.isEnabledFor(logging.DEBUG):
+            if hasattr(sd, "keys"):
+                logger.debug("State_dict contains %d keys", len(sd.keys()))
+
+        self._send(
+            request_type=RequestType.WRITE_CHECKPOINT,
+            payload={
+                "state_dict": sd,
+                "path": path,
+                "kwargs": kwargs,
+            },
+        )
+
+        logger.debug("Waiting for write completion response")
+        # wait for response
+        self._recv()
+        logger.debug("Checkpoint write to %s completed successfully", path)
+
+    def close(self) -> None:
+        logger.debug(
+            "Closing CheckpointProcess for rank %d", self._rank_info.global_rank
+        )
+        self._executor.shutdown(wait=True, cancel_futures=True)
+
+        if self.process and self.process.processes[0].is_alive():
+            subprocess_pid = self.process.processes[0].pid
+            # send graceful termination to sub process
+            try:
+                # pyrefly: ignore [missing-attribute]
+                self._parent_end.send(
+                    WorkerRequest(
+                        request_type=RequestType.TERMINATE_PROCESS,
+                        payload={},
+                    )
+                )
+            except BrokenPipeError:
+                logger.warning(
+                    "BrokenPipeError when sending termination request - subprocess (PID: %d) may have already terminated",
+                    subprocess_pid,
+                )
+                # subprocess terminated unexpectedly and below code will raise a
+                # ProcessExitedException.
+
+            logger.debug(
+                "Waiting for subprocess to terminate gracefully (timeout: %ds)",
+                self._config.subprocess_shutdown_timeout_secs,
+            )
+
+            try:
+                if not self.process.join(
+                    timeout=self._config.subprocess_shutdown_timeout_secs
+                ):
+                    # graceful shutdown failed, kill the process.
+                    logger.warning(
+                        "Subprocess (PID: %d) did not terminate gracefully within %ds, killing it",
+                        subprocess_pid,
+                        self._config.subprocess_shutdown_timeout_secs,
+                    )
+                    self.process.processes[0].kill()
+                    logger.info("Subprocess killed forcefully")
+            except ProcessExitedException:
+                logger.exception("ProcessExitedException during subprocess termination")
+                raise
+
+        logger.debug("CheckpointProcess closed successfully")

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/checkpoint_reader.py

@@ -0,0 +1,223 @@
+"""
+Checkpoint reader functionality for machine learning models.
+
+This module provides classes for reading checkpoints from storage, including
+determining checkpoint layout and configuring the reader.
+"""
+
+import logging
+import os
+from itertools import zip_longest
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch._subclasses.fake_tensor import FakeTensorMode
+
+from .types import RankInfo, STATE_DICT
+
+
+logger = logging.getLogger(__name__)
+
+
+class CheckpointReader:
+    """
+    Handles reading state dictionaries from storage.
+
+    This class is responsible for reading model state dictionaries from storage according
+    to the specified checkpoint layout. It supports synchronization barriers to ensure
+    all ranks in a distributed setting complete their checkpoint operations.
+    """
+
+    def __init__(
+        self,
+        rank_info: RankInfo,
+    ):
+        """
+        Initialize a CheckpointReader.
+
+        Args:
+            rank_info: Information about the current rank in a distributed setting.
+        """
+
+        self._rank_info = rank_info
+
+    def read(
+        self,
+        path: str,
+        state_dict: STATE_DICT | None = None,
+        *,
+        map_location: Any = None,
+        **kwargs: dict[str, Any],
+    ) -> tuple[STATE_DICT, list[str]]:
+        """
+        Reads a state dictionary from storage.
+
+        Args:
+            path (str): The path from which to read the checkpoint.
+            map_location (Any): Device mapping function or device name for relocating tensors.
+            **kwargs: Additional keyword arguments passed to torch.load.
+
+        Returns:
+            STATE_DICT: The loaded state dictionary.
+            list[str]: List of missing keys.
+        """
+        logger.debug(
+            "Reading checkpoint from %s for rank %s",
+            path,
+            self._rank_info.global_rank,
+        )
+
+        dir_path = Path(path)
+        file_path = dir_path / f"checkpoint_{self._rank_info.global_rank}.pt"
+
+        # Check if the file exists
+        if not os.path.exists(file_path):
+            logger.error("Checkpoint file not found at %s", file_path)
+            raise FileNotFoundError(f"Checkpoint file not found at {file_path}")
+
+        if state_dict is None:
+            result: tuple[STATE_DICT, list[str]] = (
+                torch.load(file_path, map_location=map_location),
+                [],
+            )
+        else:
+            result = self._partial_read(
+                file_path, state_dict, map_location=map_location, **kwargs
+            )
+        logger.debug("Successfully read checkpoint file from %s", file_path)
+        return result
+
+    def _partial_read(
+        self,
+        file_path: Path,
+        state_dict: STATE_DICT,
+        *,
+        map_location: Any = None,
+        **kwargs: dict[str, Any],
+    ) -> tuple[STATE_DICT, list[str]]:
+        """
+        Reads only the keys present in state_dict from the checkpoint file.
+
+        This method optimizes checkpoint loading by only loading the tensors that
+        are actually needed, based on the keys present in the input state_dict.
+        This can significantly reduce memory usage and loading time for large checkpoints
+        when only a subset of the model needs to be loaded.
+
+        Args:
+            file_path (str): The path to the checkpoint file.
+            state_dict (STATE_DICT): The state dictionary containing keys to load.
+            map_location (Any): Device mapping function or device name for relocating tensors.
+            **kwargs: Additional keyword arguments passed to torch.load.
+
+        Returns:
+            tuple[STATE_DICT, list[str]]: The updated state dictionary with loaded values and a list of missing keys.
+        """
+
+        with FakeTensorMode():
+            metadata_dict = torch.load(file_path, map_location=map_location)
+
+        missing_keys = []
+
+        with open(file_path, "rb") as file:
+            # Helper function to load tensor data from file
+            def load_tensor(
+                target: torch.Tensor | None, source: torch.Tensor, full_key: str
+            ) -> torch.Tensor:
+                if target is not None and (
+                    target.size() != source.size() or target.dtype != source.dtype
+                ):
+                    raise RuntimeError(
+                        f"Target tensor size={target.size()} dtype={target.dtype} does not match "
+                        f"source tensor size={source.size()} dtype={source.dtype} for key {full_key}"
+                    )
+
+                tensor_offset = source.untyped_storage()._checkpoint_offset
+
+                if tensor_offset is None:
+                    raise AssertionError(
+                        "checkpoint_offset for tensor in torch serialized file is not set. This could "
+                        "happen if the checkpoint was saved with a older version of Pytorch. "
+                        "Please make sure that the checkpoint was saved with Pytorch 2.7 or later."
+                    )
+
+                tensor_len = source.nelement() * source.element_size()
+                file.seek(
+                    tensor_offset + source.element_size() * int(source.storage_offset())
+                )
+                if target is None:
+                    target = torch.empty(
+                        source.size(), dtype=source.dtype, device=source.device
+                    )
+
+                buffer = file.read(tensor_len)
+                cpu_tensor = torch.frombuffer(buffer, dtype=source.dtype)
+                tensor = cpu_tensor.view(source.size())
+                target.copy_(tensor)
+                return target
+
+            # Helper function to recursively process nested structures
+            def process_value(
+                target_value: Any, source_value: Any, key_path: str
+            ) -> Any:
+                source_type = type(source_value)
+                if source_type is torch._subclasses.fake_tensor.FakeTensor:
+                    source_type = torch.Tensor
+                if target_value is not None and not isinstance(
+                    target_value, source_type
+                ):
+                    raise RuntimeError(
+                        f"Target value {key_path} is set to {type(target_value)}, but source value is {type(source_value)}"
+                    )
+                if isinstance(source_value, torch.Tensor):
+                    return load_tensor(target_value, source_value, key_path)
+                elif isinstance(source_value, dict):
+                    if target_value is None:
+                        # create a new map with all the keys present in source_value
+                        target_value = dict.fromkeys(source_value.keys())
+
+                    # pyrefly: ignore [missing-attribute]
+                    for key in list(target_value.keys()):
+                        current_path = f"{key_path}.{key}" if key_path else key
+                        if key in source_value:
+                            target_value[key] = process_value(
+                                target_value[key], source_value[key], current_path
+                            )
+                        else:
+                            missing_keys.append(current_path)
+
+                    return target_value
+                elif isinstance(source_value, list):
+                    if target_value is None:
+                        target_value = [None] * len(source_value)
+                    result = []
+                    for i, (target_item, source_item) in enumerate(
+                        zip_longest(target_value, source_value, fillvalue=None)
+                    ):
+                        current_path = f"{key_path}[{i}]" if key_path else f"[{i}]"
+                        result.append(
+                            process_value(target_item, source_item, current_path)
+                        )
+                    return result
+                else:
+                    return source_value
+
+            # Start recursive processing from the root of the state dictionary
+            updated_state_dict = process_value(state_dict, metadata_dict, "")
+
+        if missing_keys:
+            if len(missing_keys) > 10:
+                logger.warning(
+                    "Missing %s keys from checkpoint: %s... (and %s more)",
+                    len(missing_keys),
+                    missing_keys[:10],
+                    len(missing_keys) - 10,
+                )
+            else:
+                logger.warning(
+                    "Missing %s keys from checkpoint: %s",
+                    len(missing_keys),
+                    missing_keys,
+                )
+
+        return updated_state_dict, missing_keys

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/checkpoint_writer.py

@@ -0,0 +1,163 @@
+"""
+Checkpoint writer functionality for machine learning models.
+
+This module provides classes for writing checkpoints to storage, including
+determining checkpoint layout, configuring the writer, and defining hooks
+for custom actions during the checkpoint writing process.
+"""
+
+import abc
+import logging
+import os
+from concurrent.futures import Future
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+
+from .barriers import Barrier
+from .types import RankInfo, STATE_DICT
+
+
+logger = logging.getLogger(__name__)
+
+
+class WriterHook(abc.ABC):
+    """
+    Abstract base class for checkpoint commit hooks.
+
+    A commit hook provides callbacks that are executed before and after a checkpoint
+    is committed to storage. This allows for custom actions to be performed at specific
+    points in the checkpoint writing process, such as metadata updates, cleanup operations,
+    or notifications.
+    """
+
+    @abc.abstractmethod
+    def pre_commit(self, path: str, **kwargs: dict[str, Any]) -> None:
+        """
+        Performs actions before committing the checkpoint.
+        """
+
+    @abc.abstractmethod
+    def post_commit(self, path: str, **kwargs: dict[str, Any]) -> None:
+        """
+        Performs actions after committing the checkpoint.
+        """
+
+
+@dataclass
+class CheckpointWriterConfig:
+    """
+    Configuration options for the CheckpointWriter.
+
+    Attributes:
+        write_barrier_timeout_secs: Maximum time in seconds to wait for all ranks
+            to reach the checkpoint barrier before timing out. Default is 600 seconds.
+    """
+
+    write_barrier_timeout_secs: int = 600
+
+
+class CheckpointWriter:
+    """
+    Handles writing state dictionaries to storage.
+
+    This class is responsible for writing model state dictionaries to storage according
+    to the specified checkpoint layout. It supports synchronization barriers to ensure
+    all ranks in a distributed setting complete their checkpoint operations.
+    """
+
+    def __init__(
+        self,
+        config: CheckpointWriterConfig,
+        rank_info: RankInfo,
+        barrier: Barrier | None = None,
+        commit_hook: WriterHook | None = None,
+    ):
+        """
+        Initialize a CheckpointWriter.
+
+        Args:
+            config: Configuration options for the checkpoint writer.
+            rank_info: Information about the current rank in a distributed setting.
+            barrier: Optional synchronization barrier for distributed checkpointing.
+                    Note: The barrier should be initialized with the appropriate barrier_prefix
+                    and timeout_secs parameters.
+            commit_hook: Optional hook for custom actions before and after checkpoint commits.
+        """
+
+        self._config = config
+        self._rank_info = rank_info
+        self._commit_hook = commit_hook
+        self._barrier = barrier
+
+    def write(
+        self,
+        path: str,
+        state_dict: STATE_DICT,
+        **kwargs: dict[str, Any],
+    ) -> Future[None] | None:
+        """
+        Writes the state_dict to storage.
+
+        Args:
+            path (str): The path to write the checkpoint to.
+            state_dict (STATE_DICT): The state_dict to write.
+            **kwargs: Additional keyword arguments passed to hooks.
+
+        Returns:
+            Optional[Future[None]]: A future for tracking the write operation, if applicable.
+        """
+        logger.debug(
+            "Writing checkpoint to %s for rank %s",
+            path,
+            self._rank_info.global_rank,
+        )
+        dir_path = Path(path)
+        full_path = dir_path / f"checkpoint_{self._rank_info.global_rank}.pt"
+        os.makedirs(
+            os.path.dirname(full_path),
+            exist_ok=True,
+        )
+        torch.save(state_dict, full_path)
+        logger.debug("Successfully saved checkpoint file to %s", full_path)
+
+        # Execute pre-commit hook if available
+        commit_hook = self._commit_hook
+        if commit_hook is not None:
+            logger.debug("Executing pre-commit hook for %s", path)
+            commit_hook.pre_commit(path, **kwargs)
+
+        # Wait for all ranks to finish writing if barrier is available
+        barrier = self._barrier
+        if barrier is not None:
+            logger.info(
+                "Waiting for all ranks at barrier with timeout %ss",
+                self._config.write_barrier_timeout_secs,
+            )
+            barrier.execute_barrier()
+            logger.info("All ranks passed barrier")
+        else:
+            logger.info("No barrier configured, skipping synchronization")
+
+        # Execute commit hook if available
+        if commit_hook is not None:
+            logger.debug("Executing commit hook for %s", path)
+            commit_hook.post_commit(path, **kwargs)
+
+        logger.info(
+            "Successfully wrote checkpoint to %s for rank %s",
+            path,
+            self._rank_info.global_rank,
+        )
+        return None
+
+    def close(self) -> None:
+        """
+        Close the writer and release any resources.
+
+        This is a no-op for the base CheckpointWriter but may be overridden
+        by subclasses that need to perform cleanup.
+        """
+        logger.debug("Closing checkpoint writer")

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/checkpointer.py

@@ -0,0 +1,341 @@
+import abc
+import logging
+from concurrent.futures import Future
+from typing import Any, TypeVar
+
+from .checkpoint_process import CheckpointProcess
+from .checkpoint_reader import CheckpointReader
+from .checkpoint_writer import CheckpointWriter
+from .staging import CheckpointStager
+from .types import STATE_DICT
+from .utils import wrap_future
+
+
+logger = logging.getLogger(__name__)
+
+LOG_INTERVAL = 60
+T = TypeVar("T")
+
+
+class Checkpointer(abc.ABC):
+    """
+    WARNING: This class is experimental, and is created to validate certain ideas,
+    and is subjected to change or deprecation and we strong discourage any usages at
+    this time.
+
+    Abstract base class that defines the API for checkpointing.
+
+    This class defines the interface for coordinating the writing and loading of model
+    state dictionaries to and from storage. It provides abstract methods to save and load model states
+    with support for both synchronous and asynchronous operations.
+
+    Concrete implementations of this class must implement all the abstract methods.
+    """
+
+    @abc.abstractmethod
+    def save(
+        self,
+        path: str,
+        state_dict: STATE_DICT,
+        **kwargs: dict[str, Any],
+    ) -> tuple[Future, Future] | None:
+        """
+        Save a state dictionary to storage.
+
+        Args:
+            path: The path where the checkpoint should be saved.
+            state_dict: The state dictionary to save.
+            **kwargs: Additional keyword arguments to pass to the writer.
+
+        Returns:
+            For synchronous implementations: None
+            For asynchronous implementations: tuple of (stage_future, write_future)
+                                            representing the staging and writing operations.
+        """
+
+    @abc.abstractmethod
+    def load(
+        self,
+        path: str,
+        state_dict: STATE_DICT | None = None,
+        *,
+        default_map_location: Any = None,
+        strict: bool = False,
+        **kwargs: dict[str, Any],
+    ) -> STATE_DICT:
+        """
+        Load a state dictionary from storage.
+
+        Args:
+            path: The path from which to load the checkpoint.
+            state_dict: Optional state dictionary to update with loaded values.
+                        If provided, only keys in this dictionary will be loaded.
+            default_map_location: Device mapping function or device name for relocating tensors.
+            strict: If True, raises an error when there are missing keys in the checkpoint.
+            **kwargs: Additional keyword arguments to pass to the reader.
+
+        Returns:
+            The loaded state dictionary.
+        """
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """
+        Close the checkpointer and release any resources.
+
+        This method should be called when the checkpointer is no longer needed to ensure
+        proper cleanup of resources.
+        """
+
+
+class SyncCheckpointer(Checkpointer):
+    """
+    Synchronous implementation of Checkpointer.
+
+    This class coordinates the writing and loading of model state dictionaries to and from storage
+    using only synchronous operations. It provides a simple, efficient interface for checkpoint
+    operations without async overhead.
+
+    Attributes:
+        _writer: CheckpointWriter for writing state dictionaries to storage.
+        _reader: CheckpointReader for reading state dictionaries from storage.
+
+    Example:
+        checkpointer = SyncCheckpointer(writer=writer, reader=reader)
+        checkpointer.save(state_dict, path)
+        loaded_state_dict = checkpointer.load(path)
+    """
+
+    def __init__(
+        self,
+        writer: CheckpointWriter,
+        reader: CheckpointReader,
+    ):
+        """
+        Initialize a synchronous checkpointer.
+
+        Args:
+            writer: CheckpointWriter for writing checkpoints to storage.
+            reader: CheckpointReader for reading checkpoints from storage.
+        """
+        self._writer = writer
+        self._reader = reader
+
+    def save(
+        self,
+        path: str,
+        state_dict: STATE_DICT,
+        **kwargs: dict[str, Any],
+    ) -> tuple[Future, Future] | None:
+        """
+        Save a state dictionary to storage synchronously.
+
+        Args:
+            path: The path where the checkpoint should be saved.
+            state_dict: The state dictionary to save.
+            **kwargs: Additional keyword arguments to pass to the writer.
+
+        Returns:
+            Always returns None as operations are synchronous.
+
+        Example:
+            checkpointer.save("/path/to/checkpoint", state_dict)
+        """
+        logger.debug("Saving checkpoint synchronously to %s", path)
+        self._writer.write(path, state_dict, **kwargs)
+        return None
+
+    def load(
+        self,
+        path: str,
+        state_dict: STATE_DICT | None = None,
+        *,
+        default_map_location: Any = None,
+        strict: bool = False,
+        **kwargs: dict[str, Any],
+    ) -> STATE_DICT:
+        """
+        Load a state dictionary from storage.
+
+        Args:
+            path: The path from which to load the checkpoint.
+            state_dict: Optional state dictionary to update with loaded values.
+                        If provided, only keys in this dictionary will be loaded.
+            default_map_location: Device mapping function or device name for relocating tensors.
+            strict: If True, raises an error when there are missing keys in the checkpoint.
+            **kwargs: Additional keyword arguments to pass to the reader.
+
+        Returns:
+            The loaded state dictionary.
+
+        Raises:
+            RuntimeError: If strict=True and there are missing keys in the checkpoint.
+            FileNotFoundError: If the checkpoint file is not found.
+        """
+        logger.info("Loading checkpoint from %s", path)
+
+        loaded_state_dict, missing_keys = self._reader.read(
+            path=path,
+            state_dict=state_dict,
+            map_location=default_map_location,
+            **kwargs,
+        )
+        if strict and missing_keys is not None and missing_keys != []:
+            raise RuntimeError(f"Checkpoint at {path} is missing keys: {missing_keys}")
+        return loaded_state_dict
+
+    def close(self) -> None:
+        """
+        Close the checkpointer and release any resources.
+
+        This method should be called when the checkpointer is no longer needed to ensure
+        proper cleanup of resources.
+        """
+        self._writer.close()
+        logger.info("SyncCheckpointer closed")
+
+
+class AsyncCheckpointer(Checkpointer):
+    """
+    Asynchronous implementation of Checkpointer.
+
+    This class coordinates the writing and loading of model state dictionaries to and from storage
+    using asynchronous operations for saving. It provides efficient async checkpoint operations
+    with staging and background writing capabilities.
+
+    Attributes:
+        _reader: CheckpointReader for reading state dictionaries from storage.
+        _checkpoint_stager: Stager for async operations.
+        _checkpoint_process: Process for async operations.
+        _write_future: Future representing the ongoing async write operation.
+
+    Example:
+        checkpointer = AsyncCheckpointer(
+            reader=reader,
+            checkpoint_stager=stager,
+            checkpoint_process=process
+        )
+        stage_future, write_future = checkpointer.save(state_dict, path)
+        # ... do other work ...
+        write_future.result()  # Wait for completion
+    """
+
+    def __init__(
+        self,
+        checkpoint_stager: CheckpointStager,
+        checkpoint_process: CheckpointProcess,
+        reader: CheckpointReader,
+    ):
+        """
+        Initialize an asynchronous checkpointer.
+
+        Args:
+            checkpoint_stager: Stager for async operations.
+            checkpoint_process: Process for async operations.
+            reader: CheckpointReader for reading checkpoints from storage.
+        """
+        self._reader = reader
+        self._checkpoint_stager = checkpoint_stager
+        self._checkpoint_process = checkpoint_process
+        self._write_future: Future[Any] | None = None
+
+    def save(
+        self,
+        path: str,
+        state_dict: STATE_DICT,
+        **kwargs: Any,
+    ) -> tuple[Future, Future] | None:
+        """
+        Save a state dictionary to storage asynchronously.
+
+        Args:
+            path: The path where the checkpoint should be saved.
+            state_dict: The state dictionary to save.
+            **kwargs: Additional keyword arguments to pass to the stager and writer.
+
+        Returns:
+            A tuple of (stage_future, write_future) representing the staging and writing operations.
+
+        Example:
+            stage_future, write_future = checkpointer.save("/path/to/checkpoint", state_dict)
+            # ... do other work ...
+            write_future.result()  # Wait for completion
+        """
+        logger.info(
+            "Initiating checkpoint save to %s. Will wait for prev checkpoints to complete.",
+            path,
+        )
+        # Wait for previous checkpoint ops to finish and verify they are successful
+        if self._write_future is not None:
+            self._write_future.result()
+
+        logger.debug("Starting state dictionary staging")
+        staging_result = self._checkpoint_stager.stage(
+            state_dict=state_dict,
+            **kwargs,
+        )
+
+        logger.debug("Starting checkpoint write to %s", path)
+        self._write_future = self._checkpoint_process.write(
+            staging_result, path, **kwargs
+        )
+        logger.info("Checkpoint save to %s initiated", path)
+
+        # Return futures for the staging and writing operations
+        if self._write_future is not None:
+            return wrap_future(staging_result), self._write_future
+        else:
+            # This should not happen since we just assigned _write_future above
+            raise RuntimeError("Write future is unexpectedly None")
+
+    def load(
+        self,
+        path: str,
+        state_dict: STATE_DICT | None = None,
+        *,
+        default_map_location: Any = None,
+        strict: bool = False,
+        **kwargs: Any,
+    ) -> STATE_DICT:
+        """
+        Load a state dictionary from storage.
+
+        Loading is always performed synchronously, even in AsyncCheckpointer.
+
+        Args:
+            path: The path from which to load the checkpoint.
+            state_dict: Optional state dictionary to update with loaded values.
+                        If provided, only keys in this dictionary will be loaded.
+            default_map_location: Device mapping function or device name for relocating tensors.
+            strict: If True, raises an error when there are missing keys in the checkpoint.
+            **kwargs: Additional keyword arguments to pass to the reader.
+
+        Returns:
+            The loaded state dictionary.
+
+        Raises:
+            RuntimeError: If strict=True and there are missing keys in the checkpoint.
+            FileNotFoundError: If the checkpoint file is not found.
+        """
+        logger.info("Loading checkpoint from %s", path)
+
+        loaded_state_dict, missing_keys = self._reader.read(
+            path=path,
+            state_dict=state_dict,
+            map_location=default_map_location,
+            **kwargs,
+        )
+        if strict and missing_keys is not None and missing_keys != []:
+            raise RuntimeError(f"Checkpoint at {path} is missing keys: {missing_keys}")
+        return loaded_state_dict
+
+    def close(self) -> None:
+        """
+        Close the checkpointer and release any resources.
+
+        This method should be called when the checkpointer is no longer needed to ensure
+        proper cleanup of async resources.
+        """
+        self._checkpoint_stager.close()
+        self._checkpoint_process.close()
+        logger.info("AsyncCheckpointer closed")

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/config.py

@@ -0,0 +1,44 @@
+"""
+Configuration classes for checkpointer construction.
+
+This module provides configuration dataclasses that consolidate all
+configuration options needed to construct checkpointers.
+"""
+
+from dataclasses import dataclass, field
+
+from .barriers import BarrierConfig
+from .checkpoint_process import CheckpointProcessConfig
+from .checkpoint_writer import CheckpointWriterConfig
+from .staging import CheckpointStagerConfig
+
+
+@dataclass
+class CheckpointerConfig:
+    """
+    Configuration class for checkpointer construction.
+
+    This class consolidates the core component configuration options needed to construct
+    a checkpointer, providing a clean separation of concerns where each component
+    manages its own configuration.
+
+    Attributes:
+        writer_config: Configuration options for the checkpoint writer component.
+        barrier_config: Configuration for barrier construction and arguments.
+        staging_config: Configuration options for the async staging component.
+        process_config: Configuration options for the async checkpoint process component.
+
+    """
+
+    writer_config: CheckpointWriterConfig = field(
+        default_factory=CheckpointWriterConfig
+    )
+    barrier_config: BarrierConfig = field(default_factory=BarrierConfig)
+
+    # Below configs are used for async checkpointing
+    staging_config: CheckpointStagerConfig = field(
+        default_factory=CheckpointStagerConfig
+    )
+    process_config: CheckpointProcessConfig = field(
+        default_factory=CheckpointProcessConfig
+    )

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/staging.py

@@ -0,0 +1,216 @@
+"""
+Experimental staging module for PyTorch Distributed Checkpointing.
+
+This module provides advanced staging capabilities for checkpoints including:
+- Asynchronous staging using ThreadPoolExecutor
+- Pinned memory allocation for faster CPU-GPU transfers
+- Shared memory support for multi-process scenarios
+- Non-blocking CUDA operations with stream synchronization
+- Caching of frequently used storages for efficient memory management
+- Automatic resource cleanup and memory management
+
+Classes:
+    CheckpointStager: Abstract base class defining the staging interface
+    StagingOptions: Configuration dataclass for staging behavior
+    DefaultStager: Default implementation with comprehensive staging features
+"""
+
+import abc
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass
+from typing import Any, TypeVar
+
+import torch
+from torch.distributed.checkpoint._state_dict_stager import StateDictStager
+
+from .types import STATE_DICT
+
+
+T = TypeVar("T")
+
+
+class CheckpointStager(abc.ABC):
+    """
+    Abstract base class for checkpoint staging implementations.
+
+    CheckpointStager defines the interface that all staging implementations
+    must follow. Staging is the process of offloading state dictionaries
+    for async checkpointing.
+    """
+
+    @abc.abstractmethod
+    def stage(
+        self,
+        state_dict: STATE_DICT,
+        **kwargs: Any,
+    ) -> STATE_DICT | Future[STATE_DICT]:
+        """
+        Stage a state dictionary for checkpointing.
+
+        Args:
+            state_dict: The state dictionary to stage
+            **kwargs: Additional staging parameters
+
+        Returns:
+            Either a staged state dictionary (synchronous) or a Future
+            that will resolve to the staged state dictionary (asynchronous)
+        """
+
+    @abc.abstractmethod
+    def close(self) -> None:
+        """
+        Clean up all resources used by the stager.
+        """
+
+
+@dataclass
+class CheckpointStagerConfig:
+    """
+    Configuration options for checkpoint staging behavior.
+
+    Attributes:
+        use_pinned_memory (bool): Enable pinned memory allocation for faster
+            CPU-GPU transfers. Requires CUDA to be available. Default: True
+        use_shared_memory (bool): Enable shared memory for multi-process
+            scenarios. Useful when multiple processes need access to the
+            same staged data. Default: True
+        use_async_staging (bool): Enable asynchronous staging using a
+            background thread pool. Allows overlapping computation with
+            staging operations. Requires CUDA. Default: True
+        use_non_blocking_copy (bool): Use non-blocking device memory
+            copies with stream synchronization. Improves performance by
+            allowing CPU work to continue during GPU transfers. Default: True
+
+    Note:
+        CUDA-dependent features will raise exception if CUDA is not available.
+    """
+
+    use_pinned_memory: bool = True
+    use_shared_memory: bool = True
+    use_async_staging: bool = True
+    use_non_blocking_copy: bool = True
+
+
+class DefaultStager(CheckpointStager):
+    """
+    DefaultStager provides a full-featured staging implementation that combines
+    multiple optimization techniques for efficient checkpoint preparation.
+
+    The staging process works as follows:
+    1. State dictionary is submitted for staging (sync or async)
+    2. Tensors are copied from GPU to optimized CPU storage
+    3. CUDA operations are synchronized if non-blocking copies are used
+    4. Staged state dictionary is returned or made available via Future
+
+    NOTE: state_dict should be deep-copyable object as staging will create a
+    copy of it.
+
+    Usage Patterns:
+        # Synchronous staging
+        stager = DefaultStager(CheckpointStagerConfig(use_async_staging=False))
+        staged_dict = stager.stage(state_dict)
+        stager.close()
+
+        # Asynchronous staging
+        stager = DefaultStager(CheckpointStagerConfig(use_async_staging=True))
+        future = stager.stage(state_dict)
+        # ... do other work ...
+        staged_dict = future.result()
+        stager.close()
+
+        # Context manager pattern (recommended)
+        with DefaultStager(config) as stager:
+            result = stager.stage(state_dict)
+            # Automatic cleanup on exit
+
+    Performance Considerations:
+        - Async staging provides best performance when model computation
+          can overlap with staging operations
+        - Pinned memory improves CPU-GPU transfer speeds but uses more memory
+        - Shared memory allows efficient IPC to checkpoint process
+        - Non-blocking copies reduce GPU idle time during memory transfers
+
+    Thread Safety:
+        DefaultStager is not thread-safe. Each thread should use its own
+        instance, or external synchronization should be provided.
+    """
+
+    def __init__(
+        self,
+        config: CheckpointStagerConfig = CheckpointStagerConfig(),
+    ):
+        self._config = config
+        self._state_dict_stager = StateDictStager(
+            pin_memory=config.use_pinned_memory, share_memory=config.use_shared_memory
+        )
+        self._staging_executor = None
+        self._staging_stream = None
+
+        if self._config.use_async_staging:
+            self._staging_executor = ThreadPoolExecutor(max_workers=1)
+            if torch.accelerator.is_available():
+                # Note: stream needs to be initialized on the main thread after default cuda
+                # stream is setup/used to avoid the risk of accidentally reusing the main
+                # compute stream or in other cases kernels actually launching from the
+                # main thread.
+                self._staging_stream = torch.Stream()
+
+        if self._config.use_non_blocking_copy:
+            if not torch.accelerator.is_available():
+                raise AssertionError(
+                    "Non-blocking copy requires that the current accelerator is available."
+                )
+
+    def stage(
+        self,
+        state_dict: STATE_DICT,
+        **kwargs: Any,
+    ) -> STATE_DICT | Future[STATE_DICT]:
+        if self._config.use_async_staging:
+            if self._staging_executor is None:
+                raise AssertionError(
+                    "Staging executor should be initialized for async staging"
+                )
+            return self._staging_executor.submit(
+                self._stage,
+                state_dict,
+                **kwargs,
+            )
+        else:
+            return self._stage(state_dict, **kwargs)
+
+    def _stage(self, state_dict: STATE_DICT, **kwargs: Any) -> STATE_DICT:
+        state_dict = self._state_dict_stager.stage(
+            state_dict, non_blocking=self._config.use_non_blocking_copy, **kwargs
+        )
+
+        if self._config.use_non_blocking_copy:
+            if not (self._staging_stream or not self._config.use_async_staging):
+                raise AssertionError(
+                    "Non-blocking copy in a background thread for async staging needs staging_stream to be initialized."
+                )
+
+            # waits for the enqued copy operations to finish.
+            self._staging_stream.synchronize() if self._staging_stream else torch.accelerator.synchronize()
+
+        return state_dict
+
+    def close(self) -> None:
+        """
+        Clean up all resources used by the DefaultStager. Shuts down the ThreadPoolExecutor
+        used for async staging operations and cleans up the underlying StateDictStager's
+        cached storages. Should be called when the stager is no longer needed to prevent
+        resource leaks, especially in long-running applications. After calling close(),
+        the stager should not be used for further staging operations.
+
+        state_dict should be deep-copyable object.
+
+        Example:
+            stager = DefaultStager(CheckpointStagerConfig(use_async_staging=True))
+            # ... do staging operations ...
+            stager.close()  # Clean up all resources
+        """
+        if self._staging_executor:
+            self._staging_executor.shutdown(wait=True)
+
+        self._state_dict_stager.close()

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/types.py

@@ -0,0 +1,28 @@
+"""
+Type definitions for distributed training and checkpointing.
+
+This module provides type definitions and classes for managing rank information
+in distributed training environments, which is essential for proper checkpoint
+saving and loading.
+"""
+
+from dataclasses import dataclass
+from typing import Any, TypeAlias
+
+
+# Type alias for state dictionaries used in checkpointing
+STATE_DICT: TypeAlias = dict[str, Any]
+
+
+@dataclass
+class RankInfo:
+    """
+    Information about the current rank in a distributed training environment.
+
+    Attributes:
+        global_rank: The global rank ID of the current process.
+        global_world_size: The total number of processes in the distributed environment.
+    """
+
+    global_rank: int
+    global_world_size: int

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_experimental/utils.py

@@ -0,0 +1,42 @@
+"""
+Utility functions for the experimental checkpoint module.
+
+This module contains helper functions and utilities used across the experimental
+checkpoint functionality.
+"""
+
+from concurrent.futures import Future
+from typing import Any
+
+
+def wrap_future(original_result: Any) -> Future[None]:
+    """
+    Wraps a result (Future or not) to return a Future with None result.
+
+    If the input is a Future, returns a new Future that completes with None when
+    the original Future completes successfully, or propagates any exception.
+    If the input is not a Future, returns a completed Future with None result.
+
+    Args:
+        original_result: The result to wrap (Future or any other value).
+
+    Returns:
+        A Future that completes with None on success or propagates exceptions.
+    """
+    masked_future: Future[None] = Future()
+
+    if isinstance(original_result, Future):
+
+        def on_complete(_: Future[Any]) -> None:
+            try:
+                original_result.result()
+                masked_future.set_result(None)
+            except Exception as e:
+                masked_future.set_exception(e)
+
+        original_result.add_done_callback(on_complete)
+    else:
+        # Return a completed future with None result
+        masked_future.set_result(None)
+
+    return masked_future

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_extension.py

@@ -0,0 +1,222 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import abc
+import io
+from collections.abc import Sequence
+from typing import cast, IO
+
+# introduced as collections.abc.Buffer in Python 3.12
+from typing_extensions import Buffer
+
+from torch._utils import try_import
+
+
+# NOTE: everything in this file is experimental, and subject to
+# change.  Feedback and bug fixes are always welcome.
+
+pyzstd_module_name = "pyzstd"
+pyzstd = try_import(pyzstd_module_name)
+zstandard_module_name = "zstandard"
+zstandard = try_import(zstandard_module_name)
+
+
+__all__ = [
+    "Extension",
+    "StreamTransformExtension",
+    "ZStandard",
+    "ExtensionRegistry",
+]
+
+
+class Extension(abc.ABC):
+    """
+    Extensions provide modular additions to functionality within distributed checkpointing,
+    which affect the layout or format of the written artifacts.  Extensions may be
+    built into pytorch, or provided externally.
+
+    When writing, the caller provides a list of extension instances of the appropriate
+    type.  Each extension can output a descriptor which is used to reconstitute the
+    extension at read-time.
+    """
+
+    @staticmethod
+    @abc.abstractmethod
+    def registry_name() -> str:
+        """
+        See ExtensionRegistry.from_descriptor_list
+        """
+
+    @staticmethod
+    @abc.abstractmethod
+    def from_descriptor(version: str) -> "Extension":
+        """
+        See ExtensionRegistry.from_descriptor_list
+        """
+
+    @abc.abstractmethod
+    def get_descriptor(self) -> str:
+        """
+        Return descriptor name to be included in metadata.  The form should be
+        "extension_name[@local-domain][/version]".
+        """
+
+
+class StreamTransformExtension(Extension):
+    """
+    An extension which performs transformation on a byte stream, such as compression
+    or encryption.
+
+    Implementations should try to be memory friendly and performant.  For example, don't
+    read the whole input, then transform it, and write it back.  If at all possible, do it in
+    chunks.  But, don't read/transform/write one byte at a time, either.
+    """
+
+    @abc.abstractmethod
+    def transform_to(self, output: IO[bytes]) -> IO[bytes]:
+        """
+        Takes a writeable output stream, and generates a new stream which implements the
+        output transform.  Input data written to the returned stream will be transformed
+        and written to the `output` argument stream.
+        """
+
+    @abc.abstractmethod
+    def transform_from(self, input: IO[bytes]) -> IO[bytes]:
+        """
+        Takes a readable input stream, and generates a new stream which implements the
+        input transform.  When the returned stream is read, data will be read from the
+        'input' stream, transformed, and returned.
+        """
+
+
+class ZStandard(StreamTransformExtension):
+    @staticmethod
+    def is_available() -> bool:
+        return zstandard is not None or pyzstd is not None
+
+    @staticmethod
+    def from_descriptor(version: str) -> "ZStandard":
+        if version.partition(".")[0] != "1":
+            raise ValueError(f"Unknown extension {version=}")
+        if not ZStandard.is_available():
+            raise ValueError(
+                f"Stream with ZStandard compression cannot be processed because "
+                f"no module named '{zstandard_module_name}' or '{pyzstd_module_name}'"
+            )
+        return ZStandard()
+
+    @staticmethod
+    def registry_name() -> str:
+        return "stream.zstd"
+
+    def __init__(self) -> None:
+        super().__init__()
+        if not ZStandard.is_available():
+            raise ValueError(
+                f"ZStandard extension is unavailable because no module named '{zstandard_module_name}' or '{pyzstd_module_name}'"
+            )
+
+    def get_descriptor(self) -> str:
+        return f"{self.registry_name()}/1"
+
+    def transform_to(self, output: IO[bytes]) -> IO[bytes]:
+        if zstandard is not None:
+            compressor = zstandard.ZstdCompressor()  # type: ignore[union-attr]
+            return compressor.stream_writer(output)
+
+        class Writer(io.RawIOBase):
+            def __init__(self, output: IO[bytes]) -> None:
+                self.output = output
+                self.compressor = pyzstd.ZstdCompressor()  # type: ignore[union-attr]
+
+            def writeable(self) -> bool:
+                return True
+
+            def write(self, b: Buffer) -> int | None:
+                outdata = self.compressor.compress(b)
+                if outdata:
+                    self.output.write(outdata)
+                return len(memoryview(b))
+
+            def flush(self) -> None:
+                outdata = self.compressor.flush()
+                if outdata:
+                    self.output.write(outdata)
+                self.output.flush()
+
+        return cast(IO[bytes], Writer(output))
+
+    def transform_from(self, input: IO[bytes]) -> IO[bytes]:
+        if zstandard is not None:
+            decompressor = zstandard.ZstdDecompressor()  # type: ignore[union-attr]
+            return decompressor.stream_reader(input)
+
+        class Reader(io.RawIOBase):
+            def __init__(self, input: IO[bytes]) -> None:
+                self.input = input
+                self.decompressor = pyzstd.EndlessZstdDecompressor()  # type: ignore[union-attr]
+
+            def readable(self) -> bool:
+                return True
+
+            def readinto(self, b: Buffer) -> int | None:
+                # This needs to read enough so it can decompress
+                # something so the output doesn't look like EOF.  This
+                # means reading at least one block.  The max block
+                # size is 128KB, so we read that plus some
+                # overhead to be sure.
+
+                if self.decompressor.needs_input:
+                    indata = self.input.read((128 + 6) * 1024)
+                else:
+                    indata = b""
+
+                bview = memoryview(b)
+                blen = len(bview)
+                outdata = self.decompressor.decompress(indata, blen)
+                if outdata is None:
+                    return None
+
+                count = len(outdata)
+                bview[:count] = outdata
+                return count
+
+            def seekable(self) -> bool:
+                return False
+
+        return cast(IO[bytes], Reader(input))
+
+
+class ExtensionRegistry:
+    def __init__(self) -> None:
+        # Populate default registry contents
+        self.extensions: dict[str, type[Extension]] = {
+            cls.registry_name(): cls for cls in (ZStandard,)
+        }
+
+    def register(self, cls: type[Extension]) -> None:
+        self.extensions[cls.registry_name()] = cls
+
+    def from_descriptor_list(self, descriptors: Sequence[str]) -> Sequence[Extension]:
+        """
+        Given a seuquence of descriptor strings as returned by
+        Extension.get_descriptor at save time, creates a sequence of
+        Extension instances.  The name[@local-domain] preceding the
+        version number is used to look up an implementation class in
+        the registry, and the version is passed to the class's
+        from_descriptor static method.  If the registry contains no
+        match, this will throw ValueError.  If the from_descriptor
+        method raises an exception, that will pass through to the
+        caller.
+        """
+
+        def from_descriptor(desc: str) -> Extension:
+            name, _, version = desc.partition("/")
+            if version is None:
+                version = 0
+            ext = self.extensions.get(name)
+            if not ext:
+                raise ValueError(f"Unknown extension {name=}")
+            # pyrefly: ignore [bad-argument-type]
+            return ext.from_descriptor(version)
+
+        return [from_descriptor(desc) for desc in descriptors]

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_fsspec_filesystem.py

@@ -0,0 +1,162 @@
+# Mypy will not try inferring the types of any 3rd party libraries installed.
+# mypy: ignore-errors
+
+import io
+import os
+from collections.abc import Generator, Sequence
+from contextlib import contextmanager
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from fsspec.core import url_to_fs
+
+from torch.distributed.checkpoint._extension import StreamTransformExtension
+from torch.distributed.checkpoint.filesystem import (
+    FileSystemBase,
+    FileSystemReader,
+    FileSystemWriter,
+    SerializationFormat,
+)
+
+
+if TYPE_CHECKING:
+    from fsspec import AbstractFileSystem
+
+
+__all__ = [
+    "FsspecWriter",
+    "FsspecReader",
+]
+
+
+class FileSystem(FileSystemBase):
+    def __init__(self) -> None:
+        self.fs: AbstractFileSystem | None = None
+
+    @contextmanager
+    def create_stream(
+        self, path: str | os.PathLike, mode: str
+    ) -> Generator[io.IOBase, None, None]:
+        if self.fs is None:
+            raise AssertionError("fs should not be None")
+        path = os.fspath(path)
+
+        # fsspec does not support concurrent transactions, and not all
+        # AbstractFileSystem have working rollback implementations, so
+        # just manually delete the file if necessary on errors.
+        with self.fs.open(path, mode) as stream:
+            try:
+                yield stream
+            except:  # noqa: B001,E722
+                if any(ch in mode for ch in "w+a"):  # cleanup file if not read-only
+                    try:
+                        self.rm_file(path)
+                    except:  # noqa: B001,E722
+                        pass
+                raise
+
+    def concat_path(self, path: str | os.PathLike, suffix: str) -> str | os.PathLike:
+        return os.path.join(path, suffix)
+
+    def init_path(self, path: str | os.PathLike, **kwargs) -> str | os.PathLike:
+        self.fs, _ = url_to_fs(path, **kwargs)
+        return path
+
+    def rename(self, path: str | os.PathLike, new_path: str | os.PathLike) -> None:
+        self.fs.rename(path, new_path)
+
+    def mkdir(self, path: str | os.PathLike) -> None:
+        self.fs.makedirs(path, exist_ok=True)
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
+        if isinstance(checkpoint_id, Path):
+            return False
+
+        try:
+            url_to_fs(checkpoint_id)
+        except ValueError:
+            return False
+
+        return True
+
+    def exists(self, path: str | os.PathLike) -> bool:
+        return self.fs.exists(path)
+
+    def rm_file(self, path: str | os.PathLike) -> None:
+        self.fs.rm(path)
+
+    def ls(self, path: str | os.PathLike) -> list[str]:
+        # setting detail to False explicitly to keep the list[str] return type,
+        # instead of the list[Dict] return type when detail=True
+        return self.fs.ls(path, detail=False)
+
+
+# TODO: add the dcp.async_save mixin
+class FsspecWriter(FileSystemWriter):
+    """
+    Basic implementation of StorageWriter using FFspec.
+
+    This implementation makes the following assumptions and simplifications:
+
+    * The checkpoint path is an empty or non-existing directory.
+    * File creation is atomic
+
+    The checkpoint consist of one file per write request plus
+    a `.metadata` file with the serialized metadata.
+
+    """
+
+    def __init__(
+        self,
+        path: str | os.PathLike,
+        single_file_per_rank: bool = True,
+        sync_files: bool = True,
+        thread_count: int = 1,
+        per_thread_copy_ahead: int = 10_000_000,
+        overwrite: bool = True,
+        _extensions: Sequence[StreamTransformExtension] | None = None,
+        serialization_format: SerializationFormat = SerializationFormat.TORCH_SAVE,
+        **kwargs,
+    ) -> None:
+        """
+        Initialize the writer pointing to `path`.
+
+        Args:
+            path: directory where the checkpoint will be written to.
+            single_file_per_rank: Produce one file per rank instead of one file per tensor/blob. Default to True.
+            sync_files : force files to be synced to permanent storage. Default to True.
+            thread_count: Number of IO threads to use to write. Default to 1.
+            per_thread_copy_ahead: How many bytes to copy from the GPU ahead of saving then. Default 10Mb.
+            overwrite: Whether to allow overwriting existing checkpoints. Defaults to True.
+            _extensions: Extensions to apply to output streams (EXPERIMENTAL)
+
+        N. B. If sync_files is disabled, there's no guarantee that the checkpoint will be consistent in the case of a failure.
+        """
+        super().__init__(
+            path,
+            single_file_per_rank,
+            sync_files,
+            thread_count,
+            per_thread_copy_ahead,
+            overwrite=overwrite,
+            _extensions=_extensions,
+            serialization_format=serialization_format,
+        )
+        self.fs = FileSystem()
+        self.path = self.fs.init_path(path, **kwargs)
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
+        return FileSystem.validate_checkpoint_id(checkpoint_id)
+
+
+class FsspecReader(FileSystemReader):
+    def __init__(self, path: str | os.PathLike, **kwargs) -> None:
+        super().__init__(path)
+        self.fs = FileSystem()
+        self.path = self.fs.init_path(path, **kwargs)
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
+        return FileSystem.validate_checkpoint_id(checkpoint_id)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_hf_utils.py

@@ -0,0 +1,106 @@
+import io
+import json
+import struct
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+
+_metadata_fn: str = "model.safetensors.index.json"
+
+FILE_NAME = "model-{cpt_idx}-of-{num_files}"
+SHARDED_FILE_NAME = "shard-{shard_idx}-model-{cpt_idx}-of-{num_files}"
+SUFFIX = ".safetensors"
+
+# metadata keys
+CUSTOM_METADATA_KEY = "DCP_SHARDING_INFO"
+DEFAULT_EXTRA_METADATA_KEY = "__metadata__"
+SAVED_OFFSETS_KEY = "saved_offsets"
+SHAPE_KEY = "shape"
+DATA_KEY = "data"
+DTYPE_KEY = "dtype"
+DATA_OFFSETS_KEY = "data_offsets"
+
+DTYPE_MAP = {
+    "F16": torch.float16,
+    "F32": torch.float32,
+    "F64": torch.float64,
+    "I8": torch.int8,
+    "U8": torch.uint8,
+    "I16": torch.int16,
+    "I32": torch.int32,
+    "I64": torch.int64,
+    "BF16": torch.bfloat16,
+}
+
+HF_DCP_VERSION: float = 1.0
+DCP_VERSION_KEY = "DCP_VERSION"
+DCP_SHARDING_INFO_KEY = "DCP_SHARDING_INFO"
+
+FORMAT_KEY = "format"
+FORMAT_VALUE = "pt"
+
+NUM_BYTES_FOR_HEADER_LEN = 8
+
+SHARDED_DIR_NAME = "sharded"
+
+
+@dataclass
+class _HFStorageInfo:
+    """This is the per entry storage info."""
+
+    relative_path: str
+    shape: torch.Size
+    dtype: torch.dtype
+
+
+def _gen_file_name(
+    index: int, largest_index: int, shard_index: int | None = None
+) -> str:
+    if shard_index is not None:
+        return (
+            SHARDED_FILE_NAME.format(
+                shard_idx=f"{shard_index}".zfill(5),
+                cpt_idx=f"{index}".zfill(5),
+                num_files=f"{largest_index}".zfill(5),
+            )
+            + SUFFIX
+        )
+    else:
+        return (
+            FILE_NAME.format(
+                cpt_idx=f"{index}".zfill(5), num_files=f"{largest_index}".zfill(5)
+            )
+            + SUFFIX
+        )
+
+
+def _get_safetensors_file_metadata(file_bytes: io.IOBase) -> tuple[Any, int]:
+    # this uses the same logic that's done in HF code base
+    # https://github.com/2404589803/huggingface_hub/blob/main/src/huggingface_hub/hf_api.py#L5308
+    # and follows their documentation on how their files are serialized
+    # https://huggingface.co/docs/safetensors/index#format
+
+    header_len_bytes = file_bytes.read(NUM_BYTES_FOR_HEADER_LEN)
+    header_len = struct.unpack("<Q", header_len_bytes)[0]
+    header_json = file_bytes.read(header_len)
+    metadata = json.loads(header_json)
+    return (metadata, header_len + NUM_BYTES_FOR_HEADER_LEN)
+
+
+def _get_dtype(dtype_str: str) -> torch.dtype:
+    try:
+        dtype = DTYPE_MAP[dtype_str]
+    except KeyError:
+        dtype = torch.get_default_dtype()
+
+    return dtype
+
+
+def _get_dcp_custom_metadata(metadata: Any) -> Any | None:
+    if DEFAULT_EXTRA_METADATA_KEY in metadata:
+        custom_metadata = metadata[DEFAULT_EXTRA_METADATA_KEY]
+        if CUSTOM_METADATA_KEY in custom_metadata:
+            return json.loads(custom_metadata[CUSTOM_METADATA_KEY])
+    return None

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_nested_dict.py

@@ -0,0 +1,69 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+
+from . import _version
+from ._traverse import (
+    OBJ_PATH,
+    set_element,
+    STATE_DICT_ITEM,
+    traverse_state_dict,
+    traverse_state_dict_v_2_3,
+)
+
+
+"""
+TODO:
+Need to add ability to handle tuple, OrderedDict, NamedTuple.
+Update mappings from dict to a class.
+Change set_element to recreate the right type for tuple, OrderedDict, and NamedTuple.
+"""
+
+
+FLATTEN_MAPPING = dict[str, OBJ_PATH]
+
+
+# TODO: Update Docstring for nested_dict.py
+def flatten_state_dict(
+    state_dict: STATE_DICT_TYPE,
+) -> tuple[STATE_DICT_TYPE, FLATTEN_MAPPING]:
+    """
+    Flatten ``state_dict`` made of nested dicts and lists into a top level dictionary.
+
+    Use ``unflatten_state_dict`` to revert this process.
+    Returns:
+        A tuple with the flatten state_dict and a mapping from original to new state_dict.
+    N.B. The new keys are derived from the object paths, joined by dot.
+        For example: ``{ 'a': {'b':...}}`` results in the key `a.b`.
+    """
+    flattened: STATE_DICT_TYPE = {}
+    mappings: FLATTEN_MAPPING = {}
+
+    def flat_copy(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None:
+        new_fqn = ".".join(map(str, path))
+        if new_fqn in flattened:
+            raise ValueError(f"duplicated flatten key {new_fqn}")
+        flattened[new_fqn] = value
+        mappings[new_fqn] = path
+
+    # We started to flatten dictionary since v2.4. But in order to not break
+    # the checkpoints that were saved before v2.4, we need to keep the old
+    # traversal so that we can reconstruct those checkpoints.
+    use_v_2_3 = (
+        _version._derived_version is not None and _version._derived_version == "2_3"
+    )
+    if use_v_2_3:
+        traverse_state_dict_v_2_3(state_dict, flat_copy)
+    else:
+        traverse_state_dict(state_dict, flat_copy)
+    return flattened, mappings
+
+
+def unflatten_state_dict(
+    state_dict: STATE_DICT_TYPE, mapping: FLATTEN_MAPPING
+) -> STATE_DICT_TYPE:
+    """Restore the original nested state_dict according to ``mapping`` and the flattened ``state_dict``."""
+    nested: STATE_DICT_TYPE = {}
+    for key, value in state_dict.items():
+        set_element(nested, mapping[key], value)
+    return nested

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_pg_transport.py

@@ -0,0 +1,384 @@
+import logging
+import pickle
+import time
+from collections.abc import Callable, Generator
+from contextlib import contextmanager
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import cast, TypeVar
+
+import torch
+from torch.distributed import ProcessGroup, Work
+from torch.distributed._shard.sharded_tensor import (
+    Shard as ShardedTensorShard,
+    ShardedTensor,
+    ShardMetadata,
+)
+from torch.distributed._shard.sharded_tensor.metadata import ShardedTensorMetadata
+from torch.distributed.tensor import _DTensorSpec, DTensor
+from torch.utils._pytree import (
+    KeyPath,
+    tree_flatten_with_path,
+    tree_unflatten,
+    TreeSpec,
+)
+
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+
+@dataclass
+class _TensorMeta:
+    """
+    This is the metadata for a tensor that is used to transfer checkpoints.
+    It contains the shape, the dtype, the storage offset and the stride of the
+    tensor.
+
+    This must be pickleable so that it can be sent over the wire.
+    """
+
+    shape: torch.Size
+    dtype: torch.dtype
+    storage_offset: int
+    stride: tuple[int, ...]
+    nbytes: int
+
+
+@dataclass
+class _DTensorMeta:
+    """
+    This is the metadata for a DTensor that is used to transfer checkpoints.
+    It contains the metadata for the local tensor and the spec of the DTensor.
+
+    This must be pickleable so that it can be sent over the wire.
+    """
+
+    local: _TensorMeta
+    spec: _DTensorSpec
+
+
+@dataclass
+class _ShardedTensorMeta:
+    """
+    This is the metadata for a ShardedTensor that is used to transfer checkpoints.
+    It contains the metadata for all local shards and the global tensor metadata.
+
+    This must be pickleable so that it can be sent over the wire.
+    """
+
+    local_shards_meta: list[_TensorMeta]
+    local_shards_shard_metadata: list[
+        ShardMetadata
+    ]  # Original shard metadata for each local shard
+    sharded_tensor_metadata: ShardedTensorMetadata
+
+
+@dataclass
+class _StateDictMeta:
+    """
+    This is the metadata for a state dict that is used to transfer checkpoints.
+    It contains the step, the pytree spec of the state dict and the metadata for
+    each tensor in the state dict.
+
+    This must be pickleable so that it can be sent over the wire.
+
+    Args:
+        step: the step of the checkpoint to verify consistency
+        treespec: the pytree spec of the state dict
+        paths: the path of each leaf in the state dict
+        non_tensor_leaves: the metadata for each tensor in the state dict and any
+            non-tensor leaves in the state dict
+    """
+
+    treespec: TreeSpec
+    paths: list[KeyPath]
+    non_tensor_leaves: list[object | _TensorMeta | _DTensorMeta | _ShardedTensorMeta]
+
+
+@contextmanager
+def _timeit(name: str) -> Generator[None, None, None]:
+    start = time.perf_counter()
+    yield
+    dur = time.perf_counter() - start
+    logger.info("%s took %ss", name, dur)
+
+
+def _prepare_tensor(tensor: torch.Tensor) -> tuple[torch.Tensor, _TensorMeta]:
+    return (
+        _cast_tensor(tensor, torch.uint8),
+        _TensorMeta(
+            shape=tensor.shape,
+            dtype=tensor.dtype,
+            storage_offset=cast(int, tensor.storage_offset()),
+            stride=tensor.stride(),
+            nbytes=tensor.untyped_storage().nbytes(),
+        ),
+    )
+
+
+def _prepare_state_dict(
+    state_dict: object,
+    device: torch.device,
+) -> tuple[_StateDictMeta, list[torch.Tensor]]:
+    leaves: list[tuple[KeyPath, object]]
+    leaves, treespec = tree_flatten_with_path(state_dict)
+
+    paths: list[KeyPath] = []
+    non_tensor_leaves: list[
+        object | _TensorMeta | _DTensorMeta | _ShardedTensorMeta
+    ] = []
+    tensors: list[torch.Tensor] = []
+    for key_path, v in leaves:
+        paths.append(key_path)
+
+        if isinstance(v, DTensor):
+            tensor, tensor_meta = _prepare_tensor(v._local_tensor)
+
+            tensors.append(tensor)
+
+            non_tensor_leaves.append(
+                _DTensorMeta(
+                    local=tensor_meta,
+                    spec=v._spec,
+                )
+            )
+        elif isinstance(v, ShardedTensor):
+            # Handle ShardedTensor by extracting all local shards
+            local_shards = v.local_shards()
+
+            # Prepare metadata for all local shards
+            local_shards_meta = []
+            local_shards_shard_metadata = []
+            for shard in local_shards:
+                tensor, tensor_meta = _prepare_tensor(shard.tensor)
+                tensors.append(tensor)
+                local_shards_meta.append(tensor_meta)
+                local_shards_shard_metadata.append(shard.metadata)
+
+            non_tensor_leaves.append(
+                _ShardedTensorMeta(
+                    local_shards_meta=local_shards_meta,
+                    local_shards_shard_metadata=local_shards_shard_metadata,
+                    sharded_tensor_metadata=v.metadata(),  # Complete metadata
+                )
+            )
+        elif isinstance(v, torch.Tensor):
+            tensor, tensor_meta = _prepare_tensor(v)
+            tensors.append(tensor)
+            non_tensor_leaves.append(tensor_meta)
+        else:
+            non_tensor_leaves.append(v)
+
+    return (
+        _StateDictMeta(
+            treespec=treespec,
+            paths=paths,
+            non_tensor_leaves=non_tensor_leaves,
+        ),
+        tensors,
+    )
+
+
+def _cast_tensor(tensor: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    """
+    Casts the underlying storage to a tensor of the given dtype.
+
+    The returned tensor will be of size ``storage.nbytes``.
+
+    This works for all datatypes and supports strided/offset tensors with the
+    caveat that the cast tensor may be larger than the original tensor due to
+    the differences in striding.
+    """
+    if type(tensor) is not torch.Tensor:
+        raise AssertionError(f"can only cast standard tensors not {type(tensor)}")
+    storage = tensor.untyped_storage()
+    ret = torch.tensor(storage, dtype=dtype, device=tensor.device)
+    if ret.untyped_storage() is not storage:
+        raise AssertionError("storage should be the same")
+    return ret
+
+
+class PGTransport:
+    """
+    This is a checkpoint transport that uses the process group to transfer checkpoints.
+    This allows for fast recovery of workers by fetching the current weights
+    from an existing worker.
+
+    Args:
+        pg: the process group to use for communication
+        timeout: the timeout for communication
+        device: the device to use for tensors
+        state_dict: if specified this function will be called to do an inplace
+            receive into the returned state_dict. This is much faster than
+            having to allocate new tensors and transferring them to the CPU.
+    """
+
+    def __init__(
+        self,
+        pg: ProcessGroup,
+        timeout: timedelta,
+        device: torch.device,
+        state_dict: Callable[[], object] | None = None,
+    ) -> None:
+        self._work: list[Work] = []
+        self._pg = pg
+        self._timeout = timeout
+        self._device = device
+        self._state_dict = state_dict
+
+    def send_checkpoint(self, dst_ranks: list[int], state_dict: object) -> None:
+        """
+        Send a checkpoint to multiple destination ranks.
+
+        The process:
+        1. Prepares the state dict by converting tensors to a serializable format
+        2. Sends metadata as pickled data
+        3. Sends each tensor sequentially to all destination ranks
+
+        Args:
+            dst_ranks: List of destination ranks to send the checkpoint to
+            state_dict: The state dictionary containing model parameters
+        """
+        with _timeit("preparing state_dict"):
+            meta, tensors = _prepare_state_dict(state_dict, device=self._device)
+
+        work = []
+
+        with _timeit("send meta"):
+            buf = pickle.dumps(meta)
+            len_t = torch.tensor([len(buf)], dtype=torch.int64, device=self._device)
+            buf_t = torch.frombuffer(buf, dtype=torch.uint8).to(self._device)
+            for dst_rank in dst_ranks:
+                work.append(self._pg.send([len_t], dst_rank, tag=1))
+                work.append(self._pg.send([buf_t], dst_rank, tag=2))
+
+        with _timeit("send tensors"):
+            for i, t in enumerate(tensors):
+                original_device = t.device
+                t = t.to(self._device)
+                for dst_rank in dst_ranks:
+                    work.append(self._pg.send([t], dst_rank, tag=3 + i))
+
+                # if we did a copy we should wait for the work to complete so we
+                # can free the memory to avoid OOMs
+                if original_device == torch.device("cpu"):
+                    for w in work:
+                        w.wait()
+                    work = []
+
+            for w in work:
+                w.wait()
+
+    def recv_checkpoint(self, src_rank: int) -> object:
+        """
+        Receive a checkpoint from a source rank.
+
+        The process:
+        1. Receives metadata about the checkpoint structure
+        2. Receives each tensor, potentially reusing existing tensors for in-place updates
+        3. Reconstructs the original state dict structure
+
+        Args:
+            src_rank: The source rank to receive the checkpoint from
+
+        Returns:
+            The reconstructed state dictionary with model parameters
+        """
+        state_dict = self._state_dict() if self._state_dict else {}
+        state_dict_leaves, _ = tree_flatten_with_path(state_dict)
+
+        dst_tensors: dict[KeyPath, object] = dict(state_dict_leaves)
+
+        len_t = torch.zeros(1, dtype=torch.int64, device=self._device)
+        self._pg.recv([len_t], src_rank, tag=1).wait()
+        length = cast(int, len_t.item())
+
+        buf = torch.empty(length, dtype=torch.uint8, device=self._device)
+        self._pg.recv([buf], src_rank, tag=2).wait()
+
+        meta: _StateDictMeta = pickle.loads(buf.cpu().numpy().tobytes())
+
+        i: int = 0
+        works: list[Work] = []
+
+        def recv(path: KeyPath, v: _TensorMeta) -> torch.Tensor:
+            nonlocal i
+
+            inplace = dst_tensors.get(path)
+            if (
+                isinstance(inplace, torch.Tensor)
+                and inplace.device.type == self._device.type
+            ):
+                if isinstance(inplace, DTensor):
+                    inplace = inplace._local_tensor
+                t = _cast_tensor(inplace, torch.uint8)
+                if t.nbytes != v.nbytes:
+                    raise AssertionError("inplace tensor storage must be the same size")
+            else:
+                t = torch.empty(v.nbytes, dtype=torch.uint8, device=self._device)
+
+            work = self._pg.recv([t], src_rank, tag=3 + i)
+            i += 1
+
+            if inplace is None:
+                # if not inplace we need to copy it to CPU to avoid OOMing
+                work.wait()
+                t = t.cpu()
+            else:
+                works.append(work)
+
+            return torch.as_strided(
+                t.view(v.dtype),
+                size=v.shape,
+                stride=v.stride,
+                storage_offset=v.storage_offset,
+            )
+
+        values: list[object] = []
+        for path, v in zip(meta.paths, meta.non_tensor_leaves):
+            if isinstance(v, _TensorMeta):
+                values.append(recv(path, v))
+            elif isinstance(v, _DTensorMeta):
+                tensor = recv(path, v.local)
+                # pyrefly: ignore [bad-argument-type, bad-argument-count, unexpected-keyword]
+                values.append(DTensor(tensor, v.spec, requires_grad=False))
+            elif isinstance(v, _ShardedTensorMeta):
+                # Receive all local shards that were sent to us
+                local_shards = []
+                current_rank = self._pg.rank()
+
+                # Receive tensors for each local shard that was sent
+                for j, shard_meta in enumerate(v.local_shards_meta):
+                    tensor = recv(path, shard_meta)
+
+                    # Use the original shard metadata that was stored during preparation
+                    # but update the placement to reflect the current rank/device
+                    original_shard_metadata = v.local_shards_shard_metadata[j]
+                    updated_shard_metadata = ShardMetadata(
+                        shard_offsets=original_shard_metadata.shard_offsets,
+                        shard_sizes=original_shard_metadata.shard_sizes,
+                        placement=f"rank:{current_rank}/{tensor.device.type}",
+                    )
+
+                    local_shard = ShardedTensorShard(
+                        tensor=tensor, metadata=updated_shard_metadata
+                    )
+                    local_shards.append(local_shard)
+
+                # Use complete metadata to reconstruct ShardedTensor
+                sharded_tensor = (
+                    ShardedTensor._init_from_local_shards_and_global_metadata(
+                        local_shards=local_shards,
+                        sharded_tensor_metadata=v.sharded_tensor_metadata,
+                    )
+                )
+                values.append(sharded_tensor)
+            else:
+                values.append(v)
+
+        for work in works:
+            work.wait()
+
+        return tree_unflatten(values, meta.treespec)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_sharded_tensor_utils.py

@@ -0,0 +1,107 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import copy
+from typing import TYPE_CHECKING
+
+import torch.distributed as dist
+from torch.distributed._shard.sharded_tensor import Shard, ShardedTensor, ShardMetadata
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+from torch.distributed.remote_device import _remote_device
+
+from ._traverse import OBJ_PATH, set_element, STATE_DICT_ITEM, traverse_state_dict
+from .utils import _element_wise_add, _normalize_device_info
+
+
+if TYPE_CHECKING:
+    from torch.distributed._shard.sharded_tensor.metadata import ShardedTensorMetadata
+
+
+# TODO: We need to refactor this code.
+def _flatten_sharded_tensors(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE:
+    r"""
+    Transform ``state_dict`` by flattening all nested ShardedTensor instances found.
+
+    The resulting ShardedTensor instances are only correct regarding the local shard and
+    MUST not be used for any other purpose but checkpointing, as no operator will work with them.
+
+    This function should be used in conjunction with a state_dict produced by FSDP's
+    StateDictType.SHARDED_STATE_DICT methods.
+    """
+    new_state_dict: STATE_DICT_TYPE = {}
+
+    def rewrite_dict(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None:
+        if not isinstance(value, ShardedTensor):
+            set_element(new_state_dict, path, value)
+            return
+        shards = value.local_shards()
+
+        if len(shards) == 0:
+            return
+        if len(shards) != 1:
+            set_element(new_state_dict, path, value)
+            return
+
+        outer_shard = shards[0]
+
+        inner_st = outer_shard.tensor
+        if not isinstance(inner_st, ShardedTensor):
+            set_element(new_state_dict, path, value)
+            return
+
+        if len(inner_st.local_shards()) != 1:
+            raise ValueError("Cannot handle inner tensor with more than 1 shard")
+        inner_shard = inner_st.local_shards()[0]
+
+        local_shards = [
+            Shard(
+                tensor=inner_shard.tensor,
+                metadata=ShardMetadata(
+                    shard_offsets=_element_wise_add(
+                        outer_shard.metadata.shard_offsets,
+                        inner_shard.metadata.shard_offsets,
+                    ),
+                    shard_sizes=inner_shard.metadata.shard_sizes,
+                    placement=f"rank:{dist.get_rank()}/{inner_shard.tensor.device}",
+                ),
+            )
+        ]
+
+        st_meta: ShardedTensorMetadata = copy.deepcopy(value.metadata())
+        other_rank = 0 if dist.get_rank() > 0 else 1
+        device_info = _normalize_device_info(inner_shard.tensor.device.type, 0)
+
+        # Remove the outer ST shard the inner ST covers
+        for i, shard_md in enumerate(st_meta.shards_metadata):
+            if shard_md.shard_offsets == outer_shard.metadata.shard_offsets:
+                st_meta.shards_metadata.pop(i)
+                break
+
+        # Attribute other rank for the other shards
+        for shard_md in st_meta.shards_metadata:
+            shard_md.placement = _remote_device(f"rank:{other_rank}/{device_info}")
+
+        # Add other inner shards from the inner tensor
+        for inner_md in inner_st.metadata().shards_metadata:
+            if inner_md.shard_offsets != inner_shard.metadata.shard_offsets:
+                st_meta.shards_metadata.append(
+                    ShardMetadata(
+                        shard_offsets=_element_wise_add(
+                            outer_shard.metadata.shard_offsets,
+                            inner_md.shard_offsets,
+                        ),
+                        shard_sizes=inner_md.shard_sizes,
+                        placement=f"rank:{other_rank}/{device_info}",
+                    )
+                )
+
+        # Finally add this shard
+        st_meta.shards_metadata.append(local_shards[0].metadata)
+
+        st = ShardedTensor._init_from_local_shards_and_global_metadata(
+            local_shards=local_shards,
+            sharded_tensor_metadata=st_meta,
+        )
+        set_element(new_state_dict, path, st)
+
+    traverse_state_dict(state_dict, rewrite_dict)
+    return new_state_dict

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_state_dict_stager.py

@@ -0,0 +1,473 @@
+# mypy: allow-untyped-defs
+import types
+import warnings
+import weakref
+from copyreg import dispatch_table
+from typing import Any
+
+import torch
+import torch.cuda._pin_memory_utils as pin_memory_utils
+from torch.storage import UntypedStorage
+from torch.utils.weak import WeakIdKeyDictionary
+
+
+class StateDictStager:
+    """
+    A class for optimizing storage objects during staging for async checkpointing.
+
+    StateDictStager stages the state_dict to CPU DRAM while applying optimizations
+    like memory sharing and pinning to improve performance. It caches storage objects
+    to avoid redundant copies and can be configured to automatically share memory
+    (for multi-process usage) and pin memory (for faster CPU-GPU transfers).
+
+    Attributes:
+        pin_memory (bool): Whether to pin CPU memory for faster CPU-GPU transfers
+        share_memory (bool): Whether to share memory across processes
+        pin_memory_min_bytes (int): Minimum tensor size in bytes to pin memory (default: 5)
+        _cached_storage_mapping (WeakIdKeyDictionary): Maps storage objects to optimized CPU storages using weak references
+    """
+
+    def __init__(
+        self,
+        pin_memory: bool = False,
+        share_memory: bool = False,
+        pin_memory_min_bytes: int = 5,
+    ):
+        if pin_memory and not torch.cuda.is_available():
+            warnings.warn(
+                "Ignoring pin_memory flag for checkpoint staging as pinning memory"
+                "requires CUDA, but CUDA is not available. ",
+                stacklevel=2,
+            )
+            self.pin_memory = False
+        else:
+            self.pin_memory = pin_memory
+        self.share_memory = share_memory
+        # Mapping from original storage objects to CPU storages using weak references
+        self._cached_storage_mapping = WeakIdKeyDictionary()
+        self.pin_memory_min_bytes = pin_memory_min_bytes
+
+        def _deepcopy_atomic(x, _):
+            return x
+
+        def _deepcopy_list(x, memo, non_blocking=False):
+            y: list = []
+            memo[id(x)] = y
+            append = y.append
+            for a in x:
+                append(
+                    self.deepcopy_with_tensor_offload(
+                        a, memo, non_blocking=non_blocking
+                    )
+                )
+            return y
+
+        def _deepcopy_tuple(x, memo, non_blocking=False):
+            y = [
+                self.deepcopy_with_tensor_offload(a, memo, non_blocking=non_blocking)
+                for a in x
+            ]
+            # We're not going to put the tuple in the memo, but it's still important we
+            # check for it, in case the tuple contains recursive mutable structures.
+            try:
+                return memo[id(x)]
+            except KeyError:
+                pass
+
+            # Check if any elements changed during deepcopy
+            for k, j in zip(x, y):
+                if k is not j:
+                    # At least one element changed, create new tuple
+                    return tuple(y)
+
+            # No elements changed, return original tuple
+            return x
+
+        def _deepcopy_dict(x, memo, non_blocking=False):
+            y: dict = {}
+            memo[id(x)] = y
+            for key, value in x.items():
+                y[
+                    self.deepcopy_with_tensor_offload(
+                        key, memo, non_blocking=non_blocking
+                    )
+                ] = self.deepcopy_with_tensor_offload(
+                    value, memo, non_blocking=non_blocking
+                )
+            return y
+
+        def _deepcopy_method(x, memo, non_blocking=False):  # Copy instance methods
+            return type(x)(
+                x.__func__,
+                self.deepcopy_with_tensor_offload(
+                    x.__self__, memo, non_blocking=non_blocking
+                ),
+            )
+
+        d: dict[Any, Any] = {}
+        self._deepcopy_dispatch = d
+        d[type(None)] = _deepcopy_atomic
+        d[int] = _deepcopy_atomic
+        d[float] = _deepcopy_atomic
+        d[bool] = _deepcopy_atomic
+        d[complex] = _deepcopy_atomic
+        d[bytes] = _deepcopy_atomic
+        d[str] = _deepcopy_atomic
+        d[types.CodeType] = _deepcopy_atomic
+        d[type] = _deepcopy_atomic
+        d[range] = _deepcopy_atomic
+        d[types.BuiltinFunctionType] = _deepcopy_atomic
+        d[types.FunctionType] = _deepcopy_atomic
+        d[weakref.ref] = _deepcopy_atomic
+        d[property] = _deepcopy_atomic
+        d[types.MethodType] = _deepcopy_method
+        d[dict] = _deepcopy_dict
+        d[tuple] = _deepcopy_tuple
+        d[list] = _deepcopy_list
+
+    def _stage_untyped_storage(
+        self,
+        storage: UntypedStorage,
+        non_blocking: bool = False,
+    ):
+        """
+        Called from the hooked storage_deepcopy function in torch.Tensor.__deepcopy__.
+
+        This method handles the storage optimization logic for the StagingStateDict class.
+        It checks if the storage has already been cached, and if so, reuses it.
+        Otherwise, it creates a new CPU storage and applies memory optimizations.
+
+        Args:
+            storage: The storage to optimize
+
+        Returns:
+            The optimized storage
+        """
+        # Check if we've already cached this storage
+        if storage in self._cached_storage_mapping:
+            cached_storage = self._cached_storage_mapping[storage]
+            if cached_storage.size() != storage.size():
+                raise AssertionError(
+                    "For async checkpointing,  We cache storages in DRAM and reuse them."
+                    "Cached storage size does not match original storage size."
+                    "This should never happen as we track the original storage weakref "
+                    "and clean up the cache storage. Please report this to PyTorch Distributed Checkpointing."
+                )
+            # Reuse cached storage but update with new data
+            cached_storage.copy_(storage, non_blocking=non_blocking)
+            return cached_storage
+
+        # Create new CPU storage
+        if self.share_memory:
+            new_storage = type(storage)._new_shared(storage.size(), device="cpu")
+        else:
+            new_storage = type(storage)(storage.size(), device="cpu")
+
+        # Skip pinning for tensors below the minimum size threshold
+        # Small tensors (e.g., optimizer step counters, scalars) have negligible
+        # transfer time improvement from pinning, but pinning overhead is significant
+        if self.pin_memory and new_storage.nbytes() >= self.pin_memory_min_bytes:
+            pin_memory_utils.pin_memory(new_storage.data_ptr(), new_storage.nbytes())
+            # Set up a weak reference to unpin when cpu storage is garbage collected
+            f = weakref.finalize(
+                new_storage, pin_memory_utils.unpin_memory, new_storage.data_ptr()
+            )
+            # This makes sure that the finalizer is not called after
+            # cuda context is destroyed.
+            f.atexit = False
+
+        new_storage.copy_(storage, non_blocking=non_blocking)
+
+        # Cache the storage - WeakIdKeyDictionary will automatically clean up when storage is garbage collected
+        self._cached_storage_mapping[storage] = new_storage
+        return new_storage
+
+    @torch.no_grad()
+    def stage(
+        self,
+        state_dict: Any,
+        non_blocking: bool = False,
+    ) -> Any:
+        return self.deepcopy_with_tensor_offload(state_dict, None, [], non_blocking)
+
+    def _offload_tensor(self, x, memo, non_blocking=False):
+        """
+        Deep copy a PyTorch tensor with optimized storage handling.
+
+        This method creates a CPU copy of a tensor while applying memory optimizations
+        like sharing and pinning based on the StateDictStager configuration.
+
+        Args:
+            x: The tensor to copy
+            memo: Memo dictionary for tracking already copied objects
+            non_blocking: Whether to perform non-blocking copies where possible
+
+        Returns:
+            A CPU copy of the tensor with optimized storage
+        """
+        # if data_ptr is not 0, we allocate a new storage below. so we can skip
+        # memory allocation by using [] for size.
+        y = x.new_empty([] if x.data_ptr() != 0 else x.size(), device="cpu")
+
+        # Store in memo dict early to handle recursive references
+        d = id(x)
+        memo[d] = y
+
+        if type(x) is torch.Tensor or x.data_ptr() != 0:
+            # Get the untyped storage
+            untyped_storage = x.untyped_storage()
+            storage_id = id(untyped_storage)
+
+            # Check if this storage has already been staged in this deepcopy operation
+            # This handles the case where different tensors share the same storage
+            # (e.g., FSDP state_dict where norm.weight and norm_weight reference same storage)
+            # PyTorch caches untyped_storage() calls, so same storage -> same id
+            if storage_id in memo:
+                copied_storage = memo[storage_id]
+            else:
+                # Storage not seen before in this operation, stage it
+                copied_storage = self._stage_untyped_storage(
+                    untyped_storage, non_blocking=non_blocking
+                )
+                # Add to memo to avoid re-staging if we see this storage again
+                memo[storage_id] = copied_storage
+
+            # Set the tensor data using the staged storage
+            y.set_(copied_storage, x.storage_offset(), x.size(), x.stride())
+
+        # Copy any attributes the tensor might have
+        if hasattr(x, "__dict__"):
+            for attr_name, attr_value in x.__dict__.items():
+                setattr(
+                    y,
+                    attr_name,
+                    self.deepcopy_with_tensor_offload(
+                        attr_value, memo, non_blocking=non_blocking
+                    ),
+                )
+
+        if hasattr(x, "__slots__"):
+            for slot in x.__slots__:
+                if hasattr(x, slot):
+                    setattr(
+                        y,
+                        slot,
+                        self.deepcopy_with_tensor_offload(
+                            getattr(x, slot), memo, non_blocking=non_blocking
+                        ),
+                    )
+
+        return y
+
+    def close(self):
+        """
+        Clean up all cached storages and release associated resources.
+
+        This method clears the internal storage cache, allowing garbage collection
+        of cached CPU storages. Any pinned memory associated with cached storages
+        will be automatically unpinned through weak reference finalizers.
+
+        It also clears the _deepcopy_dispatch dict to break the reference cycle
+        created by closures that capture self. Without this, it may
+        cause memory leaks.
+        """
+        self._cached_storage_mapping.clear()
+        self._deepcopy_dispatch.clear()
+
+    @torch.no_grad()
+    def deepcopy_with_tensor_offload(self, x, memo=None, _nil=[], non_blocking=False):  # noqa: B006
+        """Deep copy operation on arbitrary Python objects with special handling for PyTorch tensors.
+
+        This implementation extends the standard deepcopy functionality to handle PyTorch tensors
+        and their storages in a way that optimizes memory usage and performance, similar to the
+        stage method. It applies memory sharing and pinning optimizations based on the StateDictStager
+        configuration.
+
+        Args:
+            x: The object to deep copy
+            memo: Memo dictionary for tracking already copied objects
+            _nil: Sentinel value for memo dictionary
+            non_blocking: Whether to perform non-blocking copies where possible
+
+        Returns:
+            A deep copy of the input object with optimized tensor storage handling
+        """
+        if memo is None:
+            memo = {}
+
+        d = id(x)
+        y = memo.get(d, _nil)
+        if y is not _nil:
+            return y
+
+        cls = type(x)
+
+        # tensors and subclasses of tensors are handled separately
+        if isinstance(x, torch.Tensor):
+            y = self._offload_tensor(x, memo, non_blocking=non_blocking)
+        else:
+            # Use the dispatch table for standard types
+            copier = self._deepcopy_dispatch.get(cls)
+            if copier is not None:
+                # Check if this is an atomic copier (only accepts x and memo)
+                if copier.__name__ == "_deepcopy_atomic":
+                    y = copier(x, memo)
+                else:
+                    y = copier(x, memo, non_blocking=non_blocking)
+            else:
+                if issubclass(cls, type):
+                    # type copier is also atomic
+                    y = self._deepcopy_dispatch[type](x, memo)
+                else:
+                    copier = getattr(x, "__deepcopy__", None)
+                    if copier is not None:
+                        y = copier(memo)
+                    else:
+                        reductor = dispatch_table.get(cls)
+                        if reductor:
+                            rv = reductor(x)
+                        else:
+                            reductor = getattr(x, "__reduce_ex__", None)
+                            if reductor is not None:
+                                rv = reductor(4)
+                            else:
+                                reductor = getattr(x, "__reduce__", None)
+                                if reductor:
+                                    rv = reductor()
+                                else:
+                                    raise RuntimeError(
+                                        f"un(deep)copyable object of type {cls}"
+                                    )
+                        if isinstance(rv, str):
+                            y = x
+                        else:
+                            # Unpack rv tuple elements (up to 5 from pickle protocol)
+                            # and explicitly pass non_blocking as keyword arg
+                            if len(rv) == 2:
+                                func, args = rv
+                                y = self._reconstruct(
+                                    x, memo, func, args, non_blocking=non_blocking
+                                )
+                            elif len(rv) == 3:
+                                func, args, state = rv
+                                y = self._reconstruct(
+                                    x,
+                                    memo,
+                                    func,
+                                    args,
+                                    state,
+                                    non_blocking=non_blocking,
+                                )
+                            elif len(rv) == 4:
+                                func, args, state, listiter = rv
+                                y = self._reconstruct(
+                                    x,
+                                    memo,
+                                    func,
+                                    args,
+                                    state,
+                                    listiter,
+                                    non_blocking=non_blocking,
+                                )
+                            elif len(rv) == 5:
+                                func, args, state, listiter, dictiter = rv
+                                y = self._reconstruct(
+                                    x,
+                                    memo,
+                                    func,
+                                    args,
+                                    state,
+                                    listiter,
+                                    dictiter,
+                                    non_blocking=non_blocking,
+                                )
+                            else:
+                                raise RuntimeError(
+                                    f"Unexpected pickle protocol return value length: {len(rv)}"
+                                )
+
+        # If is its own copy, don't memoize.
+        if y is not x:
+            memo[d] = y
+            self._keep_alive(x, memo)  # Make sure x lives at least as long as d
+        return y
+
+    def _keep_alive(self, x, memo):
+        """Keeps a reference to the object x in the memo.
+
+        Because we remember objects by their id, we have
+        to assure that possibly temporary objects are kept
+        alive by referencing them.
+        We store a reference at the id of the memo, which should
+        normally not be used unless someone tries to deepcopy
+        the memo itself...
+        """
+        try:
+            memo[id(memo)].append(x)
+        except KeyError:
+            # aha, this is the first one :-)
+            memo[id(memo)] = [x]
+
+    def _reconstruct(
+        self,
+        x,
+        memo,
+        func,
+        args,
+        state=None,
+        listiter=None,
+        dictiter=None,
+        non_blocking=False,
+    ):
+        deep = memo is not None
+        if deep and args:
+            args = tuple(
+                self.deepcopy_with_tensor_offload(arg, memo, non_blocking=non_blocking)
+                for arg in args
+            )
+        y = func(*args)
+        if deep:
+            memo[id(x)] = y
+
+        if state is not None:
+            if deep:
+                state = self.deepcopy_with_tensor_offload(
+                    state, memo, non_blocking=non_blocking
+                )
+            if hasattr(y, "__setstate__"):
+                y.__setstate__(state)
+            else:
+                if isinstance(state, tuple) and len(state) == 2:
+                    state, slotstate = state
+                else:
+                    slotstate = None
+                if state is not None:
+                    y.__dict__.update(state)
+                if slotstate is not None:
+                    for key, value in slotstate.items():
+                        setattr(y, key, value)
+
+        if listiter is not None:
+            if deep:
+                for item in listiter:
+                    item = self.deepcopy_with_tensor_offload(
+                        item, memo, non_blocking=non_blocking
+                    )
+                    y.append(item)
+            else:
+                for item in listiter:
+                    y.append(item)
+        if dictiter is not None:
+            if deep:
+                for key, value in dictiter:
+                    key = self.deepcopy_with_tensor_offload(
+                        key, memo, non_blocking=non_blocking
+                    )
+                    value = self.deepcopy_with_tensor_offload(
+                        value, memo, non_blocking=non_blocking
+                    )
+                    y[key] = value
+            else:
+                for key, value in dictiter:
+                    y[key] = value
+        return y

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_storage_utils.py

@@ -0,0 +1,48 @@
+import os
+
+from .filesystem import FileSystemReader, FileSystemWriter
+from .storage import StorageReader, StorageWriter
+
+
+def _storage_setup(
+    storage: StorageReader | StorageWriter | None,
+    checkpoint_id: str | os.PathLike | None,
+    reader: bool = False,
+) -> StorageReader | StorageWriter | None:
+    if storage:
+        if checkpoint_id is not None:
+            storage.reset(checkpoint_id)
+        return storage
+
+    if not checkpoint_id:
+        raise RuntimeError(
+            "`checkpoint_id` must be specified if "
+            "storage_reader/storage_writer is None."
+        )
+
+    targets: list[type[StorageReader | StorageWriter]] = []
+    if reader:
+        targets = [
+            FileSystemReader,
+        ]
+    else:
+        targets = [
+            FileSystemWriter,
+        ]
+    try:
+        from ._fsspec_filesystem import FsspecReader, FsspecWriter
+
+        targets.append(FsspecReader if reader else FsspecWriter)
+    except Exception:
+        pass
+
+    for target in targets:
+        if target.validate_checkpoint_id(checkpoint_id):
+            storage = target(checkpoint_id)  # type: ignore[call-arg]
+            storage.reset(checkpoint_id)
+            return storage
+
+    raise RuntimeError(
+        "Cannot detect which StorageReader or StorageWriter to use. "
+        "Please specify the storage_reader/storage_writer."
+    )

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_traverse.py

@@ -0,0 +1,208 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+# ruff: noqa: F821
+# flake8: noqa: F821
+from collections.abc import Callable, Collection, Mapping, MutableMapping
+from typing import cast, TypeVar
+from typing_extensions import TypeIs
+
+import torch
+from torch.distributed._shard.sharded_tensor.api import ShardedTensor
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+from torch.distributed.tensor import DTensor
+
+
+PATH_ITEM = str | int
+OBJ_PATH = tuple[PATH_ITEM, ...]
+T = TypeVar("T")
+
+STATE_DICT_ITEM = object
+CONTAINER_TYPE = MutableMapping[PATH_ITEM, STATE_DICT_ITEM]
+
+__all__ = ["traverse_state_dict", "set_element", "get_element", "print_tensor"]
+
+
+def _keep_visiting_tensors(value: STATE_DICT_ITEM) -> TypeIs[torch.Tensor]:
+    return isinstance(value, torch.Tensor)
+
+
+# TODO: update docstring for traverse.py
+def traverse_state_dict(
+    state_dict: STATE_DICT_TYPE,
+    visitor: Callable[[OBJ_PATH, STATE_DICT_ITEM], None],
+    keep_traversing: Callable[[STATE_DICT_ITEM], bool] = _keep_visiting_tensors,
+) -> None:
+    """
+    Invoke ``visitor`` for each value recursively in ``state_dict``.
+    Mapping will be traversed and ``visitor`` will be applied to the leaf elements.
+    ``visitor`` will only be applied to elements in a list or a tuple, if the
+    container contains tensors or mappings.
+    """
+
+    def _is_terminal(value: STATE_DICT_ITEM) -> bool:
+        values: Collection[STATE_DICT_ITEM]
+        if isinstance(value, Mapping):
+            return False
+        elif isinstance(value, list):
+            values = value
+        else:
+            return True
+
+        for entry in values:
+            if isinstance(entry, (Mapping, list)) and not _is_terminal(entry):
+                return False
+            if keep_traversing is not None and keep_traversing(entry):
+                return False
+        return True
+
+    def _traverse_obj(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None:
+        if isinstance(value, Mapping):
+            for k, v in value.items():
+                _traverse_obj(path + (str(k),), v)
+        elif _is_terminal(value):
+            visitor(path, value)
+        elif isinstance(value, (list, tuple)):
+            for i, v in enumerate(value):
+                _traverse_obj(path + (i,), v)
+
+    for key, value in state_dict.items():
+        _traverse_obj((str(key),), value)
+
+    # release reference cycle to prevent memory leaks in async_save
+    del _traverse_obj, _is_terminal
+
+
+def traverse_state_dict_v_2_3(
+    state_dict: STATE_DICT_TYPE,
+    visitor: Callable[[OBJ_PATH, STATE_DICT_ITEM], None],
+    keep_traversing: Callable[[STATE_DICT_ITEM], bool] = _keep_visiting_tensors,
+) -> None:
+    """
+    Traversal is short-circuited when if finds a collection for which ``keep_visiting_tensors`` evaluates
+    to false for all elements.
+    By default, all collections with at least one ``torch.Tensor`` element are traversed.
+    Visitor takes a path argument that is a tuple of the keys used to reach it.
+    """
+
+    # a value is terminal if it has no other containers values inside it
+    def _is_terminal(value: STATE_DICT_ITEM) -> bool:
+        values: Collection[STATE_DICT_ITEM]
+        if isinstance(value, Mapping):
+            values = value.values()
+        elif isinstance(value, list):
+            values = value
+        else:
+            return True
+
+        for entry in values:
+            if isinstance(entry, (Mapping, list)) and not _is_terminal(entry):
+                return False
+            if keep_traversing is not None and keep_traversing(entry):
+                return False
+        return True
+
+    def _traverse_obj(path: OBJ_PATH, value: STATE_DICT_ITEM) -> None:
+        if _is_terminal(value):
+            visitor(path, value)
+        elif isinstance(value, Mapping):
+            for k, v in value.items():
+                _traverse_obj(path + (str(k),), v)
+        elif isinstance(value, list):
+            for i, v in enumerate(value):
+                _traverse_obj(path + (i,), v)
+
+    for key, value in state_dict.items():
+        _traverse_obj((str(key),), value)
+
+    # release reference cycle to prevent memory leaks in async_save
+    del _traverse_obj, _is_terminal
+
+
+def set_element(
+    root_dict: STATE_DICT_TYPE, path: OBJ_PATH, value: STATE_DICT_ITEM
+) -> None:
+    """Set ``value`` in ``root_dict`` along the ``path`` object path."""
+    cur_container = cast(CONTAINER_TYPE, root_dict)
+
+    def extend_list(lst: list[STATE_DICT_ITEM], idx: int) -> None:
+        while len(lst) <= idx:
+            lst.append(None)
+
+    for i in range(1, len(path)):
+        prev_key = path[i - 1]
+        key = path[i]
+        def_val = cast(STATE_DICT_ITEM, {} if type(key) is str else [])
+
+        if isinstance(cur_container, Mapping):
+            cur_container = cast(
+                CONTAINER_TYPE, cur_container.setdefault(prev_key, def_val)
+            )
+        else:
+            # pyrefly: ignore [bad-argument-type]
+            extend_list(cur_container, prev_key)
+            if cur_container[prev_key] is None:
+                cur_container[prev_key] = def_val
+            cur_container = cur_container[prev_key]
+
+    key = path[-1]
+    if type(key) is int:
+        extend_list(cast(list[STATE_DICT_ITEM], cur_container), key)
+
+    cur_container[key] = value
+
+
+def get_element(
+    root_dict: STATE_DICT_TYPE,
+    path: OBJ_PATH,
+    default_value: T | None = None,
+) -> T | None:
+    """Retrieve the value at ``path``from ``root_dict``, returning ``default_value`` if not found."""
+    cur_value = cast(CONTAINER_TYPE, root_dict)
+    for part in path:
+        if type(part) is int:
+            if not isinstance(cur_value, list) or len(cur_value) < part:
+                return default_value
+        elif not isinstance(cur_value, Mapping) or part not in cur_value:
+            return default_value
+
+        cur_value = cast(CONTAINER_TYPE, cur_value[part])
+    return cast(T | None, cur_value)
+
+
+def _print_nested(
+    value: STATE_DICT_ITEM,
+    prefix: str = "",
+    print_fun: Callable[[str], None] = print,
+) -> None:
+    if type(value) is ShardedTensor:
+        print_fun(f"{prefix} ShardedTensor size: {value.size()}")
+        for shard in value.local_shards():
+            _print_nested(
+                shard.tensor,
+                f"{shard.metadata.shard_offsets} ",
+                print_fun=print_fun,
+            )
+    elif type(value) is (DTensor):
+        print_fun(f"{prefix} DistributedTensor size: {value.size()}")
+        # TODO: add local offset for _local_tensor in print_nested.
+        _print_nested(
+            value._local_tensor,
+            print_fun=print_fun,
+        )
+    elif isinstance(value, torch.Tensor):
+        print_fun(f"{prefix} Tensor size: {value.size()}")
+    else:
+        print_fun(f"{prefix} Type: {type(value)}")
+
+
+def print_tensor(
+    path: OBJ_PATH,
+    value: STATE_DICT_ITEM,
+    print_fun: Callable[[str], None] = print,
+) -> None:
+    """
+    Use this callback with traverse_state_dict to print its content.
+
+    By default the content is printed using the builtin ``print`` but this can
+    be change by passing a different ``print_fun` callable.
+    """
+    _print_nested(value, prefix=str(path), print_fun=print_fun)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/_version.py

@@ -0,0 +1,4 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+
+_derived_version: str | None = None

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/api.py

@@ -0,0 +1,50 @@
+import traceback as tb
+from typing import Any
+
+
+WRAPPED_EXCEPTION = tuple[BaseException, tb.StackSummary]
+
+__all__ = ["CheckpointException"]
+
+
+def _wrap_exception(exc: BaseException) -> WRAPPED_EXCEPTION:
+    summary = tb.extract_tb(exc.__traceback__)
+    # Python 3.13+ stores bytecode objects in FrameSummary._code,
+    # which cannot be pickled. Clear them so gather_object succeeds
+    # and the real exception is reported instead of a misleading
+    # "cannot pickle code objects" TypeError.
+    for frame in summary:
+        if hasattr(frame, "_code"):
+            object.__setattr__(frame, "_code", None)
+    return (exc, summary)
+
+
+def _is_wrapped_exception(obj: Any) -> bool:
+    if not isinstance(obj, tuple):
+        return False
+    if len(obj) != 2:
+        return False
+    return isinstance(obj[0], BaseException) and isinstance(obj[1], tb.StackSummary)
+
+
+class CheckpointException(BaseException):
+    """Exception raised if failure was detected as part of a checkpoint load or save."""
+
+    def __init__(self, msg: str, failures: dict[int, WRAPPED_EXCEPTION]):
+        super().__init__(msg, failures)
+        self._failures = failures
+
+    @property
+    def failures(self) -> dict[int, WRAPPED_EXCEPTION]:
+        """Return a dictionary mapping node ranks to their associated exceptions in case of failure."""
+        return self._failures
+
+    def __str__(self) -> str:
+        str = f"CheckpointException ranks:{self._failures.keys()}\n"
+        for rank, exc_pair in self._failures.items():
+            exc, trace = exc_pair
+            str += f"Traceback (most recent call last): (RANK {rank})\n"
+            if trace is not None:
+                str += "".join(tb.format_list(trace))
+            str += "".join(tb.format_exception_only(type(exc), value=exc))
+        return str

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/default_planner.py

@@ -0,0 +1,711 @@
+# mypy: allow-untyped-defs
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import dataclasses
+import io
+import logging
+import math
+import sys
+from bisect import bisect_right, insort
+from collections import ChainMap
+from typing import Any, cast
+
+import torch
+from torch.distributed._shard._utils import narrow_tensor_by_index
+from torch.distributed.checkpoint._dedup_save_plans import dedup_save_plans
+from torch.distributed.checkpoint._nested_dict import (
+    FLATTEN_MAPPING,
+    flatten_state_dict,
+)
+from torch.distributed.checkpoint._sharded_tensor_utils import _flatten_sharded_tensors
+from torch.distributed.checkpoint._traverse import set_element
+from torch.distributed.checkpoint.metadata import (
+    BytesStorageMetadata,
+    ChunkStorageMetadata,
+    Metadata,
+    MetadataIndex,
+    STATE_DICT_TYPE,
+    STORAGE_TYPES,
+    StorageMeta,
+    TensorStorageMetadata,
+)
+from torch.distributed.checkpoint.planner import (
+    LoadPlan,
+    LoadPlanner,
+    ReadItem,
+    SavePlan,
+    SavePlanner,
+    WriteItem,
+    WriteItemType,
+)
+from torch.distributed.checkpoint.planner_helpers import (
+    _compare_save_plans,
+    _contains_usable_plan,
+    _create_default_metadata_only_plan,
+    _create_read_items,
+    _create_write_items,
+    _init_state_dict,
+    _merge_delta_local_plans,
+)
+from torch.distributed.checkpoint.utils import find_state_dict_object
+from torch.distributed.tensor import DTensor
+
+from . import _version
+
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+__all__ = [
+    "DefaultSavePlanner",
+    "DefaultLoadPlanner",
+    "create_default_local_load_plan",
+    "create_default_global_load_plan",
+    "create_default_local_save_plan",
+    "create_default_global_save_plan",
+]
+
+
+# TODO: Update docstrings for default_planner.py
+class DefaultSavePlanner(SavePlanner):
+    mappings: FLATTEN_MAPPING
+
+    def __init__(
+        self,
+        flatten_state_dict: bool = True,
+        flatten_sharded_tensors: bool = True,
+        dedup_replicated_tensors: bool | None = None,
+        dedup_save_to_lowest_rank: bool = False,
+        enable_plan_caching: bool = False,
+    ) -> None:
+        self.flatten_state_dict = flatten_state_dict
+        self.flatten_sharded_tensors = flatten_sharded_tensors
+        self.mappings = {}
+        self.dedup_save_to_lowest_rank = dedup_save_to_lowest_rank
+        if dedup_replicated_tensors is not None:
+            logger.warning(
+                "DefaultSavePlanner's `dedup_replicated_tensors` argument is being "
+                "deprecated, and no longer has any effect. Please remove this argument "
+                "from your call."
+            )
+        self._cached_plans_key: str = self.__class__.__name__
+        self._enable_plan_caching = enable_plan_caching
+
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        storage_meta: StorageMeta | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        if self.flatten_state_dict:
+            state_dict, self.mappings = flatten_state_dict(state_dict)
+        if self.flatten_sharded_tensors:
+            state_dict = _flatten_sharded_tensors(state_dict)
+        self.state_dict = state_dict
+        self.is_coordinator = is_coordinator
+
+    def create_local_plan(self) -> SavePlan:
+        plan = create_default_local_save_plan(self.state_dict, self.is_coordinator)
+        if self.flatten_state_dict:
+            plan = dataclasses.replace(plan, planner_data=self.mappings)
+        self.plan = plan
+
+        if self._enable_plan_caching:
+            # If plans are equal, we can skip sending the plan to the coordinator.
+            if (
+                self._cached_plans_key in SavePlanner._cached_save_plan
+                and _compare_save_plans(
+                    plan, SavePlanner._cached_save_plan[self._cached_plans_key]
+                )
+            ):
+                logger.info(
+                    "No change in the local plan. Skipping sending the plan to the coordinator"
+                )
+                return SavePlan([], usable=False)
+            else:
+                # Store the plan as pending. It will be promoted to the
+                # class-level cache in finish_plan after the global plan
+                # has succeeded. This avoids a stale local cache when
+                # the global plan fails (e.g. validation error) but the
+                # local cache was already populated.
+                self._pending_local_plan = plan
+
+        return self.plan
+
+    def _dedup_save_plans(self, all_plans: list[SavePlan]) -> list[SavePlan]:
+        return dedup_save_plans(all_plans, self.dedup_save_to_lowest_rank)
+
+    def _create_global_plan(
+        self, all_plans: list[SavePlan]
+    ) -> tuple[list[SavePlan], Metadata]:
+        deduped_plans = self._dedup_save_plans(all_plans)
+
+        global_plan, metadata = create_default_global_save_plan(deduped_plans)
+
+        if self.flatten_state_dict:
+            # | does not work for Python 3.8 or older version.
+            # merged_mappings = reduce(
+            #     lambda x, y: x | y, (p.planner_data for p in global_plan)
+            # )
+            planner_data_dict = [p.planner_data for p in global_plan]
+            merged_mappings = dict(ChainMap(*planner_data_dict))
+            metadata = dataclasses.replace(metadata, planner_data=merged_mappings)
+
+        validation_errors = _validate_global_plan(global_plan, metadata)
+        if validation_errors:
+            error_summary = "; ".join(validation_errors)
+            if len(error_summary) > 500:
+                error_summary = error_summary[:500] + "... (truncated)"
+            raise ValueError(f"Failed to validate global plan: {error_summary}")
+
+        return global_plan, metadata
+
+    def _create_global_plan_with_caching(
+        self, all_plans: list[SavePlan]
+    ) -> tuple[list[SavePlan], list[SavePlan], Metadata]:
+        """
+        Create global plan with caching.
+        Returns a tuple of global_plan_delta, global_plan, metadata.
+        """
+        global_plan_delta: list[SavePlan] = []
+
+        if self._cached_plans_key not in SavePlanner._cached_all_plans:
+            # Case 1: If the plans are not cached, the cache will be hydrated with the
+            # all_plans, global_plans (Deduped), and metadata.
+
+            # First create and validate the global plan. Only cache everything
+            # after success to avoid partial cache state
+            global_plan, metadata = self._create_global_plan(all_plans)
+
+            # Cache all plans atomically after successful validation
+            SavePlanner._cached_all_plans[self._cached_plans_key] = all_plans
+            SavePlanner._cached_global_plan[self._cached_plans_key] = global_plan
+            SavePlanner._cached_metadata[self._cached_plans_key] = metadata
+            # If plans are not cached, global_plan delta will be the same as global plan.
+            return global_plan, global_plan, metadata
+
+        # Case 2: Plans are cached
+        if not _contains_usable_plan(all_plans):
+            # Case 2.1: Plans are cached and the local plans have NOT changed (No usable plans).
+            # Global plan delta will be empty plans to avoid the collective overhead.
+            # We can reuse the deduped global plan and metadata from the cache directly.
+            global_plan_delta = [SavePlan([], usable=False)] * len(all_plans)
+            global_plan = SavePlanner._cached_global_plan[self._cached_plans_key]
+            metadata = SavePlanner._cached_metadata[self._cached_plans_key]
+        else:
+            # Case 2.2: Plans are cached but the local plans have changed.
+            # We will merge the changed local plans with the cached local plans.
+            # Updated plans will overwrite the cached plans. New global plan and metadata will be created and cached.
+            # Global plan delta will be created by comparing the new global plan with the cached global plan.
+            # Only the global plan delta (updated ones) will be sent to the coordinator to avoid the collective overhead.
+            merged_plans = _merge_delta_local_plans(
+                SavePlanner._cached_all_plans[self._cached_plans_key], all_plans
+            )
+            # Cache the updated local plans
+            SavePlanner._cached_all_plans[self._cached_plans_key] = merged_plans
+            global_plan, metadata = self._create_global_plan(merged_plans)
+
+            if self._cached_plans_key in self._cached_global_plan:
+                for cached_plan, new_plan in zip(
+                    SavePlanner._cached_global_plan[self._cached_plans_key], global_plan
+                ):
+                    if _compare_save_plans(cached_plan, new_plan):
+                        global_plan_delta.append(SavePlan([], usable=False))
+                    else:
+                        global_plan_delta.append(new_plan)
+
+            # Cache the new global plan and the metadata
+            SavePlanner._cached_global_plan[self._cached_plans_key] = global_plan
+            SavePlanner._cached_metadata[self._cached_plans_key] = metadata
+
+        return global_plan_delta, global_plan, metadata
+
+    def create_global_plan(
+        self, all_plans: list[SavePlan]
+    ) -> tuple[list[SavePlan], Metadata]:
+        global_plan_delta: list[SavePlan] = []
+        if self._enable_plan_caching:
+            # If the plans are cached, we only need to send the global plan delta to be scattered
+            # across ranks. Ranks will use the cached final plans instead.
+            (
+                global_plan_delta,
+                global_plan,
+                metadata,
+            ) = self._create_global_plan_with_caching(all_plans)
+        else:
+            global_plan, metadata = self._create_global_plan(all_plans)
+            # If the caching is not enabled, global delta plan will always be same as the new global plan.
+            global_plan_delta = global_plan
+
+        self.global_plan = global_plan
+        self.metadata = metadata
+
+        return global_plan_delta, self.metadata
+
+    def _finish_plan_with_caching(self, new_plan: SavePlan) -> SavePlan:
+        finished_plan: SavePlan = new_plan
+
+        if not new_plan.usable:
+            finished_plan = SavePlanner._cached_final_save_plan[self._cached_plans_key]
+        else:
+            finished_plan = new_plan
+            SavePlanner._cached_final_save_plan[self._cached_plans_key] = new_plan
+        return finished_plan
+
+    def finish_plan(self, new_plan: SavePlan) -> SavePlan:
+        finished_plan: SavePlan = new_plan
+
+        if self._enable_plan_caching:
+            finished_plan = self._finish_plan_with_caching(new_plan)
+
+            # Promote the pending local plan to the class-level cache now
+            # that the global plan has succeeded and we are finalizing.
+            # This ensures the local cache is only populated after a
+            # successful end-to-end checkpoint plan creation.
+            if hasattr(self, "_pending_local_plan"):
+                SavePlanner._cached_save_plan[self._cached_plans_key] = (
+                    self._pending_local_plan
+                )
+                del self._pending_local_plan
+
+        self.plan = finished_plan
+        return self.plan
+
+    def resolve_data(self, write_item: WriteItem) -> torch.Tensor | io.BytesIO:
+        object = self.lookup_object(write_item.index)
+        return self.transform_object(write_item, object)
+
+    def lookup_object(self, index: MetadataIndex) -> Any:
+        """Extension from the planner interface to make it easy to extend the default planner."""
+        return find_state_dict_object(self.state_dict, index)
+
+    def transform_object(self, write_item: WriteItem, object: Any):
+        """Extension from the planner interface to make it easy to extend the default planner."""
+        if write_item.type == WriteItemType.BYTE_IO:
+            bytes = io.BytesIO()
+            torch.save(object, bytes)
+            object = bytes
+        return object
+
+
+class DefaultLoadPlanner(LoadPlanner):
+    """
+    DefaultLoadPlanner that adds multiple features on top of LoadPlanner.
+
+    In particular it adds the following:
+
+    flatten_state_dict: Handle state_dict with nested dicts
+    flatten_sharded_tensors: For FSDP in 2D parallel mode
+    allow_partial_load: If False, will raise a runtime error if a key is present in state_dict, but not in the checkpoint.
+    """
+
+    original_state_dict: STATE_DICT_TYPE
+    mappings: FLATTEN_MAPPING
+
+    def __init__(
+        self,
+        flatten_state_dict: bool = True,
+        flatten_sharded_tensors: bool = True,
+        allow_partial_load: bool = False,
+    ) -> None:
+        self.flatten_state_dict = flatten_state_dict
+        self.flatten_sharded_tensors = flatten_sharded_tensors
+        self.original_state_dict = {}
+        self.mappings = {}
+        self.allow_partial_load = allow_partial_load
+
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        metadata: Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        _init_state_dict(state_dict)
+        self.original_state_dict = state_dict
+
+        if self.flatten_sharded_tensors:
+            state_dict = _flatten_sharded_tensors(state_dict)
+
+        if self.flatten_state_dict:
+            state_dict, self.mappings = flatten_state_dict(state_dict)
+
+        self.state_dict = state_dict
+        self.metadata = metadata
+        self.is_coordinator = is_coordinator
+
+    def create_local_plan(self) -> LoadPlan:
+        if self.metadata is None:
+            raise AssertionError("self.metadata is not None")
+        if self.flatten_state_dict:
+            # To support checkpoints that are saved before v2.4, we have to
+            # differentiate if the missing keys are due to old checkpoints.
+            # The contracts are:
+            # 1. There are 3 cases when we found a missing key.
+            #    1.1 Actual missing key, but allow_partial_load is False
+            #    1.2 Actual missing key, but allow_partial load is True
+            #    1.3 Old checkpoint, but allow_partial_load is False
+            #    1.4 Old checkpoint, but allow_partial_load is True
+            # 2. If we found a missing key, we first convert the keys back to
+            #    the key format of v2.3
+            # 3. If the previous missing keys are in the v2.3 keys, we assume
+            #    this is a old checkpoint.
+            # 4. Pass the state_dict to `create_default_local_load_plan()`,
+            #    which has the logic to check missing for allow_partial_load.
+            # So for 1.2 and 1.4 cases, we delegate allow_partial_load check to
+            # `create_default_local_load_plan()`. The logic here is to determine
+            # whether the checkpoint belong to 2.3 (or before) or 2.4 (or after).
+            current_keys = set(self.state_dict.keys())
+            load_keys = set(self.metadata.state_dict_metadata.keys())
+            missing_keys = load_keys - current_keys
+            if missing_keys:
+                _version._derived_version = "2_3"
+                old_state_dict, old_mappings = flatten_state_dict(
+                    self.original_state_dict
+                )
+                old_keys = set(old_state_dict.keys())
+                if old_keys & missing_keys:
+                    self.state_dict, self.mappings = old_state_dict, old_mappings
+                # _derived_version is only used by flatten_state_dict now.
+                # Set it back to None so that later we can save to a new version.
+                _version._derived_version = None
+
+        return create_default_local_load_plan(
+            self.state_dict, self.metadata, not self.allow_partial_load
+        )
+
+    def create_global_plan(self, global_plan: list[LoadPlan]) -> list[LoadPlan]:
+        return create_default_global_load_plan(global_plan)
+
+    def finish_plan(self, new_plan: LoadPlan) -> LoadPlan:
+        return new_plan
+
+    def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None:
+        if self.flatten_state_dict:
+            set_element(
+                self.original_state_dict,
+                self.mappings[read_item.dest_index.fqn],
+                torch.load(value, weights_only=False),
+            )
+        else:
+            self.state_dict[read_item.dest_index.fqn] = torch.load(
+                value, weights_only=False
+            )
+
+    def resolve_tensor(self, read_item: ReadItem):
+        tensor = self.lookup_tensor(read_item.dest_index)
+        return self.transform_tensor(read_item, tensor)
+
+    def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None:
+        pass
+
+    def lookup_tensor(self, index: MetadataIndex) -> torch.Tensor:
+        """Extension from the planner interface to make it easy to extend the default planner."""
+        return find_state_dict_object(self.state_dict, index)
+
+    def transform_tensor(self, read_item: ReadItem, tensor: torch.Tensor):
+        """Extension from the planner interface to make it easy to extend the default planner."""
+        return narrow_tensor_by_index(tensor, read_item.dest_offsets, read_item.lengths)
+
+
+class _EmptyStateDictLoadPlanner(DefaultLoadPlanner):
+    """
+    Extension of DefaultLoadPlanner, which rebuilds state_dict from the saved metadata.
+    Useful for loading in state_dict without first initializing a model, such as
+    when converting a DCP checkpoint into a Torch save file.
+
+    . N.B. `state_dict` must be an empty dictionary when used with this LoadPlanner
+
+    .. warning::
+        Because the entire state dict is initialized, It's recommended to only utilize
+        this LoadPlanner on a single rank or process to avoid OOM.
+
+    """
+
+    def __init__(self, keys=None, *args, **kwargs):
+        self.keys = keys
+        super().__init__(*args, **kwargs)
+
+    def _should_include_key(self, key: str, metadata: Metadata) -> bool:
+        if self.keys is None:
+            return True
+
+        if key in self.keys:
+            return True
+
+        unflattened_keys: list[str] = []
+        planner_data = metadata.planner_data.get(key)
+        for unflattened_key in planner_data:
+            if unflattened_keys:
+                unflattened_keys.append(
+                    ".".join([unflattened_keys[-1], str(unflattened_key)])
+                )
+
+            else:
+                unflattened_keys.append(unflattened_key)
+
+        if any(unflattened_key in self.keys for unflattened_key in unflattened_keys):
+            return True
+
+        return False
+
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        metadata: Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        if state_dict:
+            raise AssertionError("not state_dict")
+        if metadata is None:
+            raise AssertionError("metadata is not None")
+
+        # rebuild the state dict from the metadata
+        for k, v in metadata.state_dict_metadata.items():
+            if not self._should_include_key(k, metadata):
+                continue
+
+            if isinstance(v, TensorStorageMetadata):
+                v = torch.empty(v.size, dtype=v.properties.dtype)  # type: ignore[assignment]
+            if metadata.planner_data is not None and k in metadata.planner_data:
+                set_element(state_dict, metadata.planner_data[k], v)
+            else:
+                state_dict[k] = v
+
+        super().set_up_planner(state_dict, metadata, is_coordinator)
+
+
+def create_default_local_load_plan(
+    state_dict: dict[str, Any], metadata: Metadata, strict: bool = True
+) -> LoadPlan:
+    requests = []
+    """
+    Create the ``LoadPlan`` used by DefaultLoadPlanner.
+
+    It produces one read item per value in ``state_dict`` using the metadata in ``metadata``.
+
+    The default behavior is to match key exactly between state_dict and metadata.
+    It handles resharding by issuing multiple read requests against storage in order to match
+    load requirements.
+    """
+
+    for fqn, obj in state_dict.items():
+        # ignore state_dict keys which do not exist in `state_dict` if strict=False
+        if fqn not in metadata.state_dict_metadata:
+            if strict:
+                raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
+            else:
+                continue
+
+        md = metadata.state_dict_metadata[fqn]
+        if (
+            isinstance(md, TensorStorageMetadata)
+            and getattr(obj, "size", None) is not None
+            and md.size != obj.size()
+        ):
+            raise ValueError(
+                f"Size mismatch between saved {md.size} and current: {obj.size()} for {fqn}",
+            )
+        # Since DTensor supports submesh, adding extra check to ensure _create_read_items()
+        # gets called only when the current rank is part of the mesh for the corresponding DTensor.
+        if isinstance(obj, DTensor):
+            if obj.device_mesh.get_coordinate() is not None:
+                requests += _create_read_items(fqn, md, obj)
+        else:
+            requests += _create_read_items(fqn, md, obj)
+
+    return LoadPlan(requests)
+
+
+def create_default_global_load_plan(
+    all_plans: list[LoadPlan],
+) -> list[LoadPlan]:
+    """
+    Create global load plan used by DefaultLoadPlanner.
+
+    The default load behavior involved no global coordination and this function
+    currently doesn't change the local plans.
+    """
+    return all_plans
+
+
+def create_default_local_save_plan(
+    state_dict: dict[str, Any], is_coordinator: bool
+) -> SavePlan:
+    """
+    Create the ``SavePlan`` used by DefaultSavePlanner.
+
+    On non-coordinator ranks, this function ignores tensors and non-tensor objects,
+    only producing writes for ShardedTensor objects.
+
+    On the coordinator rank, produce writes for all values.
+    """
+    requests = []
+    for fqn, obj in state_dict.items():
+        # Since DTensor supports submesh, adding extra check to ensure _create_write_items()
+        # gets called only when the current rank is part of the mesh for the corresponding DTensor.
+        if isinstance(obj, DTensor):
+            if obj.device_mesh.get_coordinate() is not None:
+                requests += _create_write_items(fqn, obj)
+        else:
+            # For the plain tensor and non-tensor values, add the request for all
+            # the ranks. Coordinator will decides whether to deduplicate the
+            # values based on the keys.
+            requests += _create_write_items(fqn, obj)
+
+    return SavePlan(requests)
+
+
+def create_default_global_save_plan(
+    all_plans: list[SavePlan],
+    rewrite_index_hints: bool = True,
+) -> tuple[list[SavePlan], Metadata]:
+    """
+    Create the global plan and metadata used by DefaultSavePlanner.
+
+    Metadata is produced by concatenating the metadata of all ``WriteItem`` from the supplied plans.
+
+    The only global planning change is to update index hints in all ``MetadataIndex`` objects if
+    ``rewrite_index_hints`` is True.
+    """
+    md: dict[str, STORAGE_TYPES] = {}
+    new_plans = []
+    for plan in all_plans:
+        new_items = []
+        for item in plan.items:
+            if item.type != WriteItemType.SHARD:
+                if item.index.fqn in md:
+                    raise AssertionError("item.index.fqn not in md")
+
+            if item.type == WriteItemType.BYTE_IO:
+                md[item.index.fqn] = BytesStorageMetadata()
+                new_items.append(item)
+            else:
+                if item.tensor_data is None:
+                    raise AssertionError("item.tensor_data is not None")
+                tensor_md = cast(
+                    TensorStorageMetadata,
+                    md.setdefault(
+                        item.index.fqn,
+                        TensorStorageMetadata(
+                            properties=item.tensor_data.properties,
+                            size=item.tensor_data.size,
+                            chunks=[],
+                        ),
+                    ),
+                )
+                new_item = item
+                if rewrite_index_hints:
+                    new_index = dataclasses.replace(
+                        item.index, index=len(tensor_md.chunks)
+                    )
+                    new_item = dataclasses.replace(item, index=new_index)
+                new_items.append(new_item)
+
+                if item.tensor_data.chunk is None:
+                    raise AssertionError(f"""
+                    Cannot create MD for tensor without bounds.
+                    FQN: {item.index.fqn}
+                """)
+                tensor_md.chunks.append(item.tensor_data.chunk)
+        new_plans.append(dataclasses.replace(plan, items=new_items))
+    return (new_plans, Metadata(md))
+
+
+def _create_default_local_metadata(state_dict: STATE_DICT_TYPE) -> Metadata:
+    """Return the ``Metadata`` if DefaultSavePlanner was used to checkpoint ``state_dict``."""
+    plan = _create_default_metadata_only_plan(state_dict)
+    _, md = create_default_global_save_plan([plan])
+    return md
+
+
+def _check_box_overlap(box0: ChunkStorageMetadata, box1: ChunkStorageMetadata) -> bool:
+    """Check if two boxes overlap. Tuples are (offset, lengths)."""
+    # For each dim of each shard, check if one shard resides on the other
+    # end of second shard with respect to that dim. As an example for a 2D
+    # shard, we would check if one shard is above or on the left of the
+    # other shard.
+    ndims = len(box0.offsets)
+    for i in range(ndims):
+        if box0.offsets[i] >= box1.offsets[i] + box1.sizes[i]:
+            return False
+        if box1.offsets[i] >= box0.offsets[i] + box0.sizes[i]:
+            return False
+
+    return True
+
+
+def _check_box_bounds(
+    outer_box_size: torch.Size, inner_box: ChunkStorageMetadata
+) -> bool:
+    for i in range(len(outer_box_size)):
+        if inner_box.offsets[i] < 0:
+            return False
+        if inner_box.sizes[i] < 0:
+            return False
+        if inner_box.offsets[i] + inner_box.sizes[i] > outer_box_size[i]:
+            return False
+
+    return True
+
+
+def _validate_global_plan(global_plan: list[SavePlan], metadata: Metadata) -> list[str]:
+    """Validate the global plan and return a list of error messages (empty if valid)."""
+    errors: list[str] = []
+    for key, value in metadata.state_dict_metadata.items():
+        if isinstance(value, BytesStorageMetadata):
+            continue
+        if len(value.size) == 0:
+            continue
+        chunks = value.chunks
+        chunks_volume = 0
+        for chunk in chunks:
+            # Compute the volume
+            if not _check_box_bounds(value.size, chunk):
+                msg = (
+                    f"key:{key} has out of bounds chunk: "
+                    f"tensor-size:{value.size} chunk: {chunk}"
+                )
+                logger.warning(msg)
+                errors.append(msg)
+            chunks_volume += math.prod(chunk.sizes)
+
+        if len(chunks) > 1:
+            dims = len(value.size)
+            sweep_dim = max(range(dims), default=0, key=lambda d: value.size[d])
+            sorted_indices = sorted(
+                range(len(chunks)),
+                key=lambda idx: (
+                    chunks[idx].offsets[sweep_dim],
+                    *(chunks[idx].offsets[d] for d in range(dims)),
+                ),
+            )
+            active: list[tuple[int, int]] = []
+            for idx in sorted_indices:
+                current = chunks[idx]
+                start = current.offsets[sweep_dim]
+                end = start + current.sizes[sweep_dim]
+
+                cutoff = bisect_right(active, (start, sys.maxsize))
+                if cutoff:
+                    del active[:cutoff]
+
+                for _, other_idx in active:
+                    other = chunks[other_idx]
+                    if _check_box_overlap(current, other):
+                        msg = f"key:{key} has overlapping chunks: {current} {other}"
+                        logger.warning(msg)
+                        errors.append(msg)
+
+                insort(active, (end, idx))
+
+        # Check whether combined chunk cover the whole tensor
+        tensor_volume = math.prod(value.size)
+        if len(global_plan) > 1 and chunks_volume != tensor_volume:
+            msg = (
+                f"key:{key} invalid fill tensor-volume: "
+                f"{tensor_volume} chunks-volume: {chunks_volume}"
+            )
+            logger.warning(msg)
+            errors.append(msg)
+
+    return errors

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/format_utils.py

@@ -0,0 +1,292 @@
+# mypy: allow-untyped-defs
+import argparse
+import os
+from enum import Enum
+from typing import cast
+
+import torch
+import torch.distributed as dist
+from torch.distributed._shard._utils import narrow_tensor_by_index
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
+from torch.distributed.checkpoint._nested_dict import flatten_state_dict
+from torch.distributed.checkpoint.default_planner import (
+    _EmptyStateDictLoadPlanner,
+    DefaultLoadPlanner,
+)
+from torch.distributed.checkpoint.metadata import (
+    Metadata,
+    STATE_DICT_TYPE,
+    STORAGE_TYPES,
+    TensorProperties,
+    TensorStorageMetadata,
+)
+from torch.distributed.checkpoint.planner import LoadItemType, LoadPlan, LoadPlanner
+from torch.distributed.checkpoint.planner_helpers import _create_chunk_list
+from torch.distributed.checkpoint.state_dict_loader import _load_state_dict
+from torch.distributed.checkpoint.state_dict_saver import _save_state_dict
+from torch.distributed.checkpoint.storage import StorageReader
+from torch.futures import Future
+
+
+__all__ = [
+    "dcp_to_torch_save",
+    "torch_save_to_dcp",
+    "BroadcastingTorchSaveReader",
+    "DynamicMetaLoadPlanner",
+]
+
+
+class BroadcastingTorchSaveReader(StorageReader):
+    """
+    StorageReader for reading a Torch Save file. This reader will read the entire checkpoint
+    on the coordinator rank, and then broadcast and shard each tensor to all ranks.
+
+    . N.B. Intended to be used with DynamicMetaLoadPlanner
+
+    .. warning::
+        Current implementation only supports loading Tensors.
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> sd = {"mode": model}
+    >>> dcp.load(
+    >>>    sd,
+    >>>    storage_reader=BroadcastingTorchSaveReader(),
+    >>>    planner=DynamicMetaLoadPlanner(),
+    >>>    checkpoint_id="path_to_model.pt"
+    >>> )
+    """
+
+    def __init__(
+        self,
+        checkpoint_id: str | os.PathLike | None = None,
+        coordinator_rank: int = 0,
+    ) -> None:
+        self.checkpoint_id = checkpoint_id
+        self.coordinator_rank = coordinator_rank
+
+    # pyrefly: ignore [bad-override]
+    def read_metadata(self) -> Metadata:
+        """Extends the default StorageReader to support building the metadata file"""
+        # Metadata is built in planner.set_up_planner, since we are not actually reading metadata from
+        # the disk
+        return Metadata(state_dict_metadata={})
+
+    def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
+        """
+        Reads torch save data on the coordinator rank, and broadcast afterwards
+        this incurrs a communication cost, but avoids having to load
+        the entire checkpoint on each rank, hopefully preventing OOM issues
+        """
+        planner = cast(DefaultLoadPlanner, planner)
+
+        # data is read in on the coordinator rank, and broadcast afterwards
+        # this incurs a communication cost, but it avoids having to load
+        # the entire checkpoint on each rank, hopefully preventing OOM issues
+        # TODO: read on each host, instead of only the coordinator
+        if self.is_coordinator:
+            if self.checkpoint_id is None:
+                raise AssertionError("checkpoint_id must be set before reading data")
+            torch_state_dict = torch.load(
+                self.checkpoint_id, map_location="cpu", weights_only=False
+            )
+            if planner.flatten_state_dict:
+                torch_state_dict, _ = flatten_state_dict(torch_state_dict)
+        else:
+            torch_state_dict = None
+
+        for req in plan.items:
+            if req.type == LoadItemType.BYTE_IO:
+                raise RuntimeError(
+                    f"Non-tensor value identified at {req.storage_index.fqn}. "
+                    f"At this time {type(self).__name__} only supports loading Tensors."
+                )
+
+            #  Broadcast the tensor from the coordinator rank
+            if self.is_coordinator:
+                pg_device = dist.distributed_c10d._get_pg_default_device()
+                # pyrefly: ignore [unsupported-operation]
+                tensor = torch_state_dict[req.storage_index.fqn].to(pg_device)
+            else:
+                tensor = torch.empty_like(planner.state_dict[req.storage_index.fqn])
+
+            dist.broadcast(tensor, src=self.coordinator_rank, async_op=False)
+
+            tensor = narrow_tensor_by_index(tensor, req.storage_offsets, req.lengths)
+            target_tensor = planner.resolve_tensor(req).detach()
+            if not target_tensor.size() == tensor.size():
+                raise AssertionError(
+                    f"req {req.storage_index} mismatch sizes, "
+                    f"{target_tensor.size()} vs {tensor.size()}"
+                )
+            target_tensor.copy_(tensor)
+            planner.commit_tensor(req, target_tensor)
+
+        fut: Future = Future()
+        fut.set_result(None)
+        return fut
+
+    # pyrefly: ignore [bad-override]
+    def set_up_storage_reader(self, metadata: Metadata, is_coordinator: bool) -> None:
+        """Implementation of the StorageReader method"""
+        self.is_coordinator = is_coordinator
+        if self.is_coordinator:
+            if not dist.get_rank() == self.coordinator_rank:
+                raise AssertionError(
+                    f"Coordinator rank mismatch: expected {self.coordinator_rank}, "
+                    f"got {dist.get_rank()}"
+                )
+
+        if self.checkpoint_id is None:
+            raise AssertionError(
+                "checkpoint_id must be set before setting up storage reader"
+            )
+
+    def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan:
+        """Implementation of the StorageReader method"""
+        return plan
+
+    def prepare_global_plan(self, global_plan: list[LoadPlan]) -> list[LoadPlan]:
+        """Implementation of the StorageReader method"""
+        return global_plan
+
+    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
+        """Implementation of the StorageReader method"""
+        self.checkpoint_id = checkpoint_id
+
+    @classmethod
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
+        """Implementation of the StorageReader method"""
+        return os.path.isfile(checkpoint_id)
+
+
+class DynamicMetaLoadPlanner(DefaultLoadPlanner):
+    """
+    Extension of DefaultLoadPlanner, which creates a new Metadata object based on the passed in state dict,
+    avoiding the need to read metadata from disk. This is useful when reading formats which don't have a
+    metadata file, like Torch Save files.
+
+    . N.B. Intended to be used with BroadcastingTorchSaveReader
+
+    .. warning::
+        Current implementation only supports loading Tensors.
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> sd = {"mode": model}
+    >>> dcp.load(
+    >>>    sd,
+    >>>    storage_reader=BroadcastingTorchSaveReader(),
+    >>>    planner=DynamicMetaLoadPlanner(),
+    >>>    checkpoint_id="path_to_model.pt"
+    >>> )
+    """
+
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        metadata: Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        """Setups of the planner, extnding default behavior by creating the Metadata object from the state dict"""
+        super().set_up_planner(state_dict, metadata, is_coordinator)
+
+        state_dict_metadata: dict[str, STORAGE_TYPES] = {}
+        for key, tensor in self.state_dict.items():
+            if not torch.is_tensor(tensor):
+                raise RuntimeError(
+                    f"Non-tensor value identified at {key}. "
+                    f"At this time {type(self).__name__} only supports loading Tensors."
+                )
+
+            state_dict_metadata[key] = TensorStorageMetadata(
+                TensorProperties(dtype=tensor.dtype),
+                tensor.size(),
+                _create_chunk_list(tensor),
+            )
+        self.metadata = Metadata(state_dict_metadata=state_dict_metadata)
+
+
+def dcp_to_torch_save(
+    dcp_checkpoint_dir: str | os.PathLike,
+    torch_save_path: str | os.PathLike,
+):
+    """
+    Given a directory containing a DCP checkpoint, this function will convert it into a
+    Torch save file.
+
+    Args:
+        dcp_checkpoint_dir: Directory containing the DCP checkpoint.
+        torch_save_path: Filename to store the converted Torch save file.
+
+    .. warning::
+        To avoid OOM, it's recommended to only run this function on a single rank.
+    """
+    sd: STATE_DICT_TYPE = {}
+    _load_state_dict(
+        sd,
+        storage_reader=FileSystemReader(dcp_checkpoint_dir),
+        planner=_EmptyStateDictLoadPlanner(),
+        no_dist=True,
+    )
+    torch.save(sd, torch_save_path)
+
+
+def torch_save_to_dcp(
+    torch_save_path: str | os.PathLike,
+    dcp_checkpoint_dir: str | os.PathLike,
+):
+    """
+    Given the location of a torch save file, converts it into a DCP checkpoint.
+
+    Args:
+        torch_save_path: Filename of the Torch save file.
+        dcp_checkpoint_dir: Directory to store the DCP checkpoint.
+
+    .. warning::
+        To avoid OOM, it's recommended to only run this function on a single rank.
+    """
+
+    state_dict = torch.load(torch_save_path, weights_only=False)
+    # we don't need stateful behavior here because the expectation is anything loaded by
+    # torch.load would not contain stateful objects.
+    _save_state_dict(
+        state_dict, storage_writer=FileSystemWriter(dcp_checkpoint_dir), no_dist=True
+    )
+
+
+if __name__ == "__main__":
+
+    class FormatMode(Enum):
+        TORCH_TO_DCP = "torch_to_dcp"
+        DCP_TO_TORCH = "dcp_to_torch"
+
+    # Parse command-line arguments
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "mode",
+        type=str,
+        help="Conversion mode",
+        choices=[m.value for m in FormatMode],
+        default=FormatMode.TORCH_TO_DCP,
+    )
+    parser.add_argument("src", type=str, help="Path to the source model")
+    parser.add_argument("dst", type=str, help="Path to the destination model")
+    args = parser.parse_args()
+
+    print(
+        f"Converting checkpoint from {args.src} to {args.dst} using method: '{args.mode}'"
+    )
+    checkpoint_missing_warning = (
+        f"No checkpoint found at {args.src}. Skipping conversion."
+    )
+    if args.mode == FormatMode.TORCH_TO_DCP.value:
+        if os.path.isfile(args.src):
+            torch_save_to_dcp(args.src, args.dst)
+        else:
+            print(checkpoint_missing_warning)
+    elif args.mode == FormatMode.DCP_TO_TORCH.value:
+        if os.path.isdir(args.src):
+            dcp_to_torch_save(args.src, args.dst)
+        else:
+            print(checkpoint_missing_warning)
+    else:
+        raise ValueError(f"Unknown conversion mode: {args.mode}")

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/hf_storage.py

@@ -0,0 +1,391 @@
+# mypy: allow-untyped-defs
+import dataclasses
+import json
+import logging
+import queue
+import threading
+from typing import Any
+
+import torch
+from torch.distributed.checkpoint import FileSystemReader, FileSystemWriter
+from torch.distributed.checkpoint._consolidate_hf_safetensors import (
+    consolidate_safetensors_files,
+)
+from torch.distributed.checkpoint._hf_utils import (
+    _gen_file_name,
+    _HFStorageInfo,
+    _metadata_fn,
+    CUSTOM_METADATA_KEY,
+    SAVED_OFFSETS_KEY,
+    SHARDED_DIR_NAME,
+    SUFFIX,
+)
+from torch.distributed.checkpoint.filesystem import SerializationFormat
+from torch.distributed.checkpoint.metadata import (
+    ChunkStorageMetadata,
+    Metadata,
+    MetadataIndex,
+    StorageMeta,
+    TensorProperties,
+    TensorStorageMetadata,
+)
+from torch.distributed.checkpoint.planner import (
+    LoadPlan,
+    LoadPlanner,
+    ReadItem,
+    SavePlan,
+    SavePlanner,
+    WriteItem,
+)
+from torch.distributed.checkpoint.storage import WriteResult
+from torch.futures import Future
+
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+__all__ = ["HuggingFaceStorageWriter", "HuggingFaceStorageReader"]
+
+
+class HuggingFaceStorageWriter(FileSystemWriter):
+    """
+    A writer that writes to storage in the huggingface safetensors format.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        fqn_to_index_mapping: dict[str, int] | None = None,
+        thread_count: int = 1,
+        save_distributed: bool = False,
+        enable_consolidation: bool = False,
+        thread_count_consolidation: int = 1,
+    ) -> None:
+        """
+        Initialize the huggingface writer pointing to path.
+
+        Args:
+            path: directory where the checkpoint will be read from.
+            fqn_to_index_mapping: A mapping from tensor FQN to the index of the file that the tensor should be written to.
+                              Indices are from 1 to N, where N is the number of files. If not provided,
+                              the tensors will be written to a single file. If none, then all the tensors on the
+                              same rank will be written to the same file.
+            thread_count: Number of threads to use to write distributed checkpoint. Default to 1.
+            save_distributed: If True, save the checkpoint using distributed APIs where every rank saves its own shard.
+                        Default is False which assumes rank-0 checkpointing of the full state_dict.
+            enable_consolidation: If True, consolidate the sharded checkpoint after saving. The sharded tensors will be
+                                saved to path/sharded and the full tensors will be saved to path. Default to False.
+            thread_count_consolidation: Number of threads to use for parallel processing of saving data
+                                to consolidated output files. Default to 1.
+        """
+
+        super().__init__(
+            path=path,
+            serialization_format=SerializationFormat.SAFETENSORS,
+            thread_count=thread_count,
+        )
+        self.fqn_to_index_mapping: dict[str, int] | None = fqn_to_index_mapping
+        self.save_distributed: bool = save_distributed
+        self.enable_consolidation: bool = enable_consolidation
+        self.consolidated_output_path: str | None = None
+        if self.enable_consolidation:
+            self.consolidated_output_path = str(self.path)
+            self.path = self.fs.concat_path(self.path, SHARDED_DIR_NAME)
+        self.thread_count_consolidation = thread_count_consolidation
+
+    def prepare_global_plan(self, plans: list[SavePlan]) -> list[SavePlan]:
+        new_plans = []
+        for i, plan in enumerate(plans, start=1):
+            storage_data: dict[str, Any] = {}
+            if self.fqn_to_index_mapping is not None:
+                storage_data["fqn_to_index_mapping"] = self.fqn_to_index_mapping
+            if self.save_distributed:
+                storage_data["shard_index"] = i
+
+            new_plans.append(dataclasses.replace(plan, storage_data=storage_data))
+
+        return new_plans
+
+    def write_data(
+        self,
+        plan: SavePlan,
+        planner: SavePlanner,
+    ) -> Future[list[WriteResult]]:
+        if len(plan.items) == 0:
+            fut: Future = Future()
+            fut.set_result([])
+            return fut
+
+        # storage_plan is a map from key to file index
+        storage_data: dict[str, Any] = plan.storage_data
+        storage_plan: dict[str, int] | None = None
+        shard_index: int | None = None
+        if "fqn_to_index_mapping" in storage_data:
+            storage_plan = storage_data["fqn_to_index_mapping"]
+        if "shard_index" in storage_data:
+            shard_index = storage_data["shard_index"]
+
+        buckets = self._split_by_storage_plan(storage_plan, plan.items)
+        highest_index = max(storage_plan.values()) if storage_plan is not None else 1
+
+        file_queue: queue.Queue = queue.Queue()
+        for file_index, write_items in buckets.items():
+            file_name = _gen_file_name(file_index, highest_index, shard_index)
+            file_queue.put(
+                (self.fs.concat_path(self.path, file_name), file_name, write_items)
+            )
+
+        return super()._write_data(planner, file_queue)
+
+    def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
+        if self.save_distributed and not self.enable_consolidation:
+            # if we are saving distributed, without consolidating,
+            # then we have no metadata to write because a metadata
+            # file with fqn to file mapping doesn't make sense
+            # in this case, because fqns will be in multiple files
+            logger.info("Not consolidating sharded checkpoint in finish step.")
+            return
+        if self.save_distributed:
+            fqn_to_index_mapping: dict[str, int] = (
+                self.fqn_to_index_mapping
+                if self.fqn_to_index_mapping is not None
+                else dict.fromkeys(metadata.state_dict_metadata.keys(), 1)
+            )
+
+            return consolidate_safetensors_files(
+                input_dir=str(self.path),
+                output_dir=self.consolidated_output_path,  # type: ignore[arg-type]
+                num_threads=self.thread_count_consolidation,
+                fqn_to_index_mapping=fqn_to_index_mapping,
+            )
+
+        # writing a model.index.safetensors.json file with fqn to file mapping
+        # for the rank-0 checkpointing case
+        metadata_to_write = {}
+        storage_md = {}
+        total_size = 0
+        for wr_list in results:
+            storage_md.update(
+                {wr.index.fqn: wr.storage_data.relative_path for wr in wr_list}
+            )
+            total_size += sum([wr.storage_data.length for wr in wr_list])
+        metadata_to_write["metadata"] = {"total_size": total_size}
+        metadata_to_write["weight_map"] = storage_md
+
+        metadata_path = self.fs.concat_path(self.path, f"{_metadata_fn}")
+        with self.fs.create_stream(metadata_path, "w") as metadata_file:
+            json.dump(metadata_to_write, metadata_file, indent=2)
+
+    def _split_by_storage_plan(
+        self, storage_plan: dict[str, int] | None, items: list[WriteItem]
+    ) -> dict[int, list[WriteItem]]:
+        # storage_plan is a map from key to index
+        if storage_plan is None:
+            return {1: items}
+
+        buckets = {}
+        for item in items:
+            key = item.index.fqn
+
+            idx = storage_plan[key]
+            if idx not in buckets:
+                buckets[idx] = [item]
+            else:
+                buckets[idx].append(item)
+
+        return buckets
+
+    @property
+    def metadata_path(self) -> str:
+        return _metadata_fn
+
+
+class HuggingFaceStorageReader(FileSystemReader):
+    """
+    A reader that reads a checkpoint in the huggingface safetensors format.
+    """
+
+    def __init__(self, path: str, thread_count: int = 1) -> None:
+        """
+        Initialize the huggingface reader pointing to path.
+
+        Args:
+            path: directory where the checkpoint will be read from.
+            thread_count: Number of threads to use to read distributed checkpoint. Default to 1.
+        """
+
+        super().__init__(path=path)
+        self.thread_count = thread_count
+
+    def _process_read_request(self, f, req: ReadItem, planner: LoadPlanner) -> None:
+        """Helper function to process a single read request."""
+        # Create slices for each dimension based on offsets and lengths
+        slices = tuple(
+            slice(offset, offset + length)
+            for offset, length in zip(req.storage_offsets, req.lengths)
+        )
+        tensor = f.get_slice(req.storage_index.fqn)[slices]
+        target_tensor = planner.resolve_tensor(req).detach()
+
+        if target_tensor.size() != tensor.size():
+            raise AssertionError(
+                f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}"
+            )
+
+        target_tensor.copy_(tensor)
+        planner.commit_tensor(req, target_tensor)
+
+    def _read_files_from_queue(
+        self,
+        file_queue: queue.Queue,
+        result_queue: queue.Queue,
+        planner: LoadPlanner,
+    ) -> None:
+        from safetensors import safe_open  # type: ignore[import]
+
+        try:
+            while True:
+                file_name, reqs = file_queue.get_nowait()
+                with safe_open(filename=file_name, framework="pt") as f:
+                    for req in reqs:
+                        self._process_read_request(f, req, planner)
+                result_queue.put(True)  # Signal that this file has been processed
+        except queue.Empty:
+            pass
+
+    def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
+        from safetensors import safe_open  # type: ignore[import]
+
+        per_file: dict[str, list[ReadItem]] = {}
+
+        for read_item in plan.items:
+            item_md: _HFStorageInfo = self.storage_data[read_item.storage_index]
+            file_name = item_md.relative_path
+            per_file.setdefault(file_name, []).append(read_item)
+
+        if self.thread_count <= 1 or len(per_file) <= 1:
+            for file_name, reqs in per_file.items():
+                with safe_open(filename=file_name, framework="pt") as f:
+                    for req in reqs:
+                        self._process_read_request(f, req, planner)
+        else:
+            # Use parallel implementation with thread pool
+            file_queue: queue.Queue = queue.Queue()
+            result_queue: queue.Queue = queue.Queue()
+
+            # Fill the queue with files to process
+            for file_name, reqs in per_file.items():
+                file_queue.put((file_name, reqs))
+
+            # Create and start worker threads
+            threads = []
+            num_threads = min(self.thread_count, len(per_file))
+            for _ in range(num_threads):
+                t = threading.Thread(
+                    target=self._read_files_from_queue,
+                    args=(file_queue, result_queue, planner),
+                )
+                t.start()
+                threads.append(t)
+
+            # Wait for all threads to complete
+            for t in threads:
+                t.join()
+
+            # Check if all files were processed
+            processed_count = 0
+            try:
+                while True:
+                    result_queue.get_nowait()
+                    processed_count += 1
+            except queue.Empty:
+                pass
+
+            if processed_count != len(per_file):
+                raise AssertionError(
+                    f"Not all files were processed: {processed_count} out of {len(per_file)}"
+                )
+
+        fut: Future = Future()
+        fut.set_result(None)
+        return fut
+
+    # pyrefly: ignore [bad-override]
+    def read_metadata(self) -> Metadata:
+        from safetensors import safe_open  # type: ignore[import]
+        from safetensors.torch import _getdtype  # type: ignore[import]
+
+        state_dict_metadata: dict[str, TensorStorageMetadata] = {}
+        storage_data: dict[MetadataIndex, _HFStorageInfo] = {}
+
+        safetensors_files = []
+        for file in self.fs.ls(self.path):
+            if file.endswith(SUFFIX):
+                safetensors_files.append(file)
+
+        for safetensor_file in safetensors_files:
+            with safe_open(safetensor_file, framework="pt") as f:
+                keys = f.keys()
+                extra_metadata = f.metadata()
+
+                dcp_sharding_info = None
+                if extra_metadata and extra_metadata.get(CUSTOM_METADATA_KEY):
+                    dcp_sharding_info = json.loads(
+                        extra_metadata.get(CUSTOM_METADATA_KEY)
+                    )
+
+                for key in keys:
+                    shape = f.get_slice(key).get_shape()
+                    dtype = f.get_slice(key).get_dtype()
+                    # construct state_dict_metadata
+                    if dcp_sharding_info is not None:
+                        offset = dcp_sharding_info[key][SAVED_OFFSETS_KEY]
+                    else:
+                        offset = [0] * len(shape)
+
+                    if key not in state_dict_metadata:
+                        state_dict_metadata[key] = TensorStorageMetadata(
+                            properties=TensorProperties(dtype=_getdtype(dtype)),
+                            size=torch.Size(
+                                [saved + offset for saved, offset in zip(shape, offset)]
+                            ),
+                            chunks=[
+                                ChunkStorageMetadata(
+                                    offsets=torch.Size(offset),
+                                    sizes=torch.Size(shape),
+                                )
+                            ],
+                        )
+                    else:
+                        state_dict_metadata[key].chunks.append(
+                            ChunkStorageMetadata(
+                                torch.Size(offset), sizes=torch.Size(shape)
+                            )
+                        )
+                        size = list(state_dict_metadata[key].size)
+                        for i in range(len(size)):
+                            size[i] = max(size[i], shape[i] + offset[i])
+                        state_dict_metadata[key].size = torch.Size(size)
+
+                    # construct storage data
+                    if dcp_sharding_info is not None:
+                        metadata_index = MetadataIndex(
+                            fqn=key, offset=dcp_sharding_info[key][SAVED_OFFSETS_KEY]
+                        )
+                    else:
+                        metadata_index = MetadataIndex(fqn=key, offset=[0] * len(shape))
+                    storage_data[metadata_index] = _HFStorageInfo(
+                        relative_path=safetensor_file,
+                        shape=torch.Size(shape),
+                        dtype=_getdtype(dtype),
+                    )
+
+        metadata = Metadata(
+            state_dict_metadata=state_dict_metadata,  # type: ignore[arg-type]
+            storage_data=storage_data,
+        )
+
+        if getattr(metadata, "storage_meta", None) is None:
+            metadata.storage_meta = StorageMeta()
+        metadata.storage_meta.load_id = self.load_id  # type: ignore[union-attr]
+
+        return metadata

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/logger.py

@@ -0,0 +1,121 @@
+# mypy: allow-untyped-defs
+import functools
+import logging
+import time
+from collections.abc import Callable
+from typing import Any, TypeVar
+from typing_extensions import ParamSpec
+from uuid import uuid4
+
+import torch.distributed.c10d_logger as c10d_logger
+from torch.distributed.checkpoint.logging_handlers import DCP_LOGGER_NAME
+
+
+logger = logging.getLogger()
+
+
+__all__: list[str] = []
+
+# pyrefly: ignore [unknown-name]
+global _dcp_logger
+_dcp_logger = c10d_logger._get_or_create_logger(DCP_LOGGER_NAME)
+
+_T = TypeVar("_T")
+_P = ParamSpec("_P")
+
+
+def _msg_dict_from_dcp_method_args(*args, **kwargs) -> dict[str, Any]:
+    """
+    Extracts log data from dcp method args
+    """
+    msg_dict = {}
+
+    # checkpoint ID can be passed in through the serializer or through the checkpoint id directly
+    storage_writer = kwargs.get("storage_writer")
+    storage_reader = kwargs.get("storage_reader")
+    planner = kwargs.get("planner")
+
+    checkpoint_id = kwargs.get("checkpoint_id")
+    if not checkpoint_id and (serializer := storage_writer or storage_reader):
+        checkpoint_id = getattr(serializer, "checkpoint_id", None)
+
+    msg_dict["checkpoint_id"] = (
+        # pyrefly: ignore [unsupported-operation]
+        str(checkpoint_id) if checkpoint_id is not None else checkpoint_id
+    )
+
+    # Uniquely identify a _dcp_method_logger wrapped function call.
+    msg_dict["uuid"] = str(uuid4().int)
+
+    if storage_writer:
+        msg_dict["storage_writer"] = storage_writer.__class__.__name__
+
+    if storage_reader:
+        msg_dict["storage_reader"] = storage_reader.__class__.__name__
+
+    if planner:
+        msg_dict["planner"] = planner.__class__.__name__
+
+    return msg_dict
+
+
+def _get_msg_dict(func_name, *args, **kwargs) -> dict[str, Any]:
+    msg_dict = _msg_dict_from_dcp_method_args(*args, **kwargs)
+    msg_dict.update(c10d_logger._get_msg_dict(func_name, *args, **kwargs))
+
+    return msg_dict
+
+
+def _dcp_method_logger(
+    log_exceptions: bool = False, **wrapper_kwargs: Any
+) -> Callable[[Callable[_P, _T]], Callable[_P, _T]]:  # pyre-ignore
+    """This method decorator logs the start, end, and exception of wrapped events."""
+
+    def decorator(func: Callable[_P, _T]):
+        @functools.wraps(func)
+        def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> _T:
+            msg_dict = _get_msg_dict(
+                func.__name__, *args, **{**wrapper_kwargs, **kwargs}
+            )
+
+            # log start event
+            msg_dict["event"] = "start"
+            t0 = time.time_ns()
+            msg_dict["time"] = t0
+            msg_dict["log_exceptions"] = log_exceptions
+            _dcp_logger.debug(msg_dict)
+
+            # exceptions
+            try:
+                result = func(*args, **kwargs)
+            except BaseException as error:
+                if log_exceptions:
+                    msg_dict["event"] = "exception"
+                    msg_dict["error"] = f"{error}"
+                    msg_dict["time"] = time.time_ns()
+                    _dcp_logger.error(msg_dict)
+                raise
+
+            # end event
+            msg_dict["event"] = "end"
+            t1 = time.time_ns()
+            msg_dict["time"] = time.time_ns()
+            msg_dict["times_spent"] = t1 - t0
+            _dcp_logger.debug(msg_dict)
+
+            return result
+
+        return wrapper
+
+    return decorator
+
+
+def _init_logger(rank: int):
+    logger.setLevel(logging.INFO)
+    ch = logging.StreamHandler()
+    ch.setLevel(logging.INFO)
+    formatter = logging.Formatter(
+        f"[{rank}] %(asctime)s - %(name)s - %(levelname)s - %(message)s"
+    )
+    ch.setFormatter(formatter)
+    logger.addHandler(ch)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/logging_handlers.py

@@ -0,0 +1,14 @@
+import logging
+
+from torch.distributed.logging_handlers import _log_handlers
+
+
+__all__: list[str] = []
+
+DCP_LOGGER_NAME = "dcp_logger"
+
+_log_handlers.update(
+    {
+        DCP_LOGGER_NAME: logging.NullHandler(),
+    }
+)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/metadata.py

@@ -0,0 +1,185 @@
+# mypy: allow-untyped-defs
+import os
+from collections.abc import Sequence
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any
+
+import torch
+from torch.distributed.checkpoint.stateful import StatefulT
+
+
+__all__ = [
+    "ChunkStorageMetadata",
+    "TensorStorageMetadata",
+    "BytesStorageMetadata",
+    "Metadata",
+    "MetadataIndex",
+    "TensorProperties",
+    "StorageMeta",
+]
+
+
+@dataclass
+class ChunkStorageMetadata:
+    """
+    Each chunk is expected to have the same properties of the TensorStorageMetadata
+    that includes it.
+    """
+
+    offsets: torch.Size
+    sizes: torch.Size
+
+
+class _MEM_FORMAT_ENCODING(Enum):
+    """Describe the memory format of a tensor."""
+
+    TORCH_CONTIGUOUS_FORMAT = 0
+    TORCH_CHANNELS_LAST = 1
+    TORCH_PRESERVE_FORMAT = 2
+
+
+@dataclass
+class TensorProperties:
+    """Properties used to create :class:`Tensor`"""
+
+    # Regular tensor fields
+    dtype: torch.dtype = field(default_factory=torch.get_default_dtype)
+    # This field is deprecated.
+    layout: torch.layout = field(default=torch.strided)
+    # This field is deprecated.
+    requires_grad: bool = False
+    # This field is deprecated.
+    memory_format: torch.memory_format = field(default=torch.contiguous_format)
+    # This field is deprecated.
+    pin_memory: bool = False
+
+    def __getstate__(self):
+        # Since torch.memory_format cannot be pickled!
+        memory_format = self.memory_format
+        if memory_format == torch.contiguous_format:
+            mem_format_encoding = _MEM_FORMAT_ENCODING.TORCH_CONTIGUOUS_FORMAT
+        elif memory_format == torch.channels_last:
+            mem_format_encoding = _MEM_FORMAT_ENCODING.TORCH_CHANNELS_LAST
+        elif memory_format == torch.preserve_format:
+            mem_format_encoding = _MEM_FORMAT_ENCODING.TORCH_PRESERVE_FORMAT
+        else:
+            raise RuntimeError(f"Invalid torch.memory_format: {memory_format}")
+
+        return (
+            self.dtype,
+            self.layout,
+            self.requires_grad,
+            mem_format_encoding,
+            self.pin_memory,
+        )
+
+    def __setstate__(
+        self,
+        state,
+    ):
+        (
+            self.dtype,
+            self.layout,
+            self.requires_grad,
+            mem_format_encoding,
+            self.pin_memory,
+        ) = state
+
+        if mem_format_encoding == _MEM_FORMAT_ENCODING.TORCH_CONTIGUOUS_FORMAT:
+            memory_format = torch.contiguous_format
+        elif mem_format_encoding == _MEM_FORMAT_ENCODING.TORCH_CHANNELS_LAST:
+            memory_format = torch.channels_last
+        elif mem_format_encoding == _MEM_FORMAT_ENCODING.TORCH_PRESERVE_FORMAT:
+            memory_format = torch.preserve_format
+        else:
+            raise RuntimeError(
+                f"Invalid torch.memory_format encoding: {mem_format_encoding}"
+            )
+
+        self.memory_format = memory_format
+
+    @staticmethod
+    def create_from_tensor(tensor: torch.Tensor) -> "TensorProperties":
+        return TensorProperties(
+            dtype=tensor.dtype,
+            layout=tensor.layout,
+            requires_grad=tensor.requires_grad,
+            memory_format=torch.contiguous_format,
+            pin_memory=tensor.is_pinned(),
+        )
+
+
+@dataclass
+class TensorStorageMetadata:
+    properties: TensorProperties
+    size: torch.Size
+    chunks: list[ChunkStorageMetadata]
+
+
+@dataclass
+class BytesStorageMetadata:
+    pass
+
+
+STORAGE_TYPES = TensorStorageMetadata | BytesStorageMetadata
+STATE_DICT_TYPE = dict[str, StatefulT | Any]
+
+
+@dataclass
+class StorageMeta:
+    checkpoint_id: str | os.PathLike | None = None
+    save_id: str | None = None
+    load_id: str | None = None
+    modules: list[str] = field(default_factory=list)
+
+
+@dataclass
+class Metadata:
+    """This class represents the metadata of the checkpoint."""
+
+    # Keys are the same from the `state_dict` used.
+    state_dict_metadata: dict[str, STORAGE_TYPES]
+    # It is the responsibility of the planner and storage plugins to ensure
+    # backward compatibility of the planner_data and storage_data. DCP will
+    # also ensure the backward compatibility of the metadata in this file and
+    # the metadata of the built-in planner and storage plugins.
+    planner_data: Any = None
+    storage_data: Any = None
+    storage_meta: StorageMeta | None = None
+    version: str | None = None
+
+
+@dataclass(frozen=True)
+class MetadataIndex:
+    """This class represents a lookup key for items in a state dict or Metadata."""
+
+    fqn: str
+    """Fully Qualified Name of the object"""
+
+    offset: torch.Size | None = None
+    """If the object is a tensor, offset into the tensor we're looking for"""
+
+    index: int | None = field(hash=False, compare=False, default=None)
+    """
+    Index hint when searching for tensor chunk to speedup lookups (optional)
+
+    A common representation of a sharded tensor is as a list of chunks so to
+    find the index in such a list you need to linear search it.
+
+    When constructing an instance of MetadataIndex that points to that list,
+    one can provide the index as a hint and it will be probed first before
+    the linear search and thus making it significantly faster.
+    """
+
+    def __init__(
+        self,
+        fqn: str,
+        offset: Sequence[int] | None = None,
+        index: int | None = None,
+    ):
+        # We must use object.__setattr__ due to frozen=True
+        object.__setattr__(self, "fqn", fqn)
+        object.__setattr__(self, "index", index)
+        if offset is not None:
+            object.__setattr__(self, "offset", torch.Size(offset))

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/optimizer.py

@@ -0,0 +1,360 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates
+
+import dataclasses
+from collections.abc import Sequence
+from typing import cast
+
+import torch
+import torch.distributed as dist
+from torch._utils import _get_device_module
+from torch.distributed._shard.sharded_tensor.api import ShardedTensor
+from torch.distributed._shard.sharded_tensor.metadata import (
+    TensorProperties as ShardTensorProperties,
+)
+from torch.distributed._shard.sharded_tensor.shard import Shard
+from torch.distributed._shard.sharding_spec.chunk_sharding_spec import ChunkShardingSpec
+from torch.distributed.checkpoint._nested_dict import unflatten_state_dict
+from torch.distributed.checkpoint.default_planner import DefaultLoadPlanner
+from torch.distributed.checkpoint.metadata import (
+    BytesStorageMetadata,
+    ChunkStorageMetadata,
+    Metadata,
+    MetadataIndex,
+    STATE_DICT_TYPE,
+    TensorProperties,
+    TensorStorageMetadata,
+)
+from torch.distributed.checkpoint.planner import LoadPlan, LoadPlanner
+from torch.distributed.checkpoint.planner_helpers import (
+    _create_read_items,
+    create_read_items_for_chunk_list,
+)
+
+# pyrefly: ignore [deprecated]
+from torch.distributed.checkpoint.state_dict_loader import load_state_dict
+from torch.distributed.checkpoint.storage import StorageReader
+from torch.distributed.checkpoint.utils import (
+    _element_wise_add,
+    _element_wise_sub,
+    _normalize_device_info,
+)
+from torch.distributed.distributed_c10d import _get_default_group
+from torch.distributed.fsdp._shard_utils import _create_chunk_sharded_tensor
+from torch.distributed.remote_device import _remote_device
+from torch.distributed.tensor import DTensor
+
+
+STATE_DICT_2D_LAYOUT = dict[str, tuple[Sequence[int] | None, Sequence[int]]]
+
+
+# TODO: Update docstrings for optimizer.py
+__all__ = [
+    "load_sharded_optimizer_state_dict",
+]
+
+
+def _gen_rank_device(global_rank: int, device_type: str = "cuda") -> str:
+    if device_type == "cpu":
+        return "cpu"
+    device_module = _get_device_module(device_type)
+    if device_module.is_available():
+        return _normalize_device_info(
+            device_type, global_rank % device_module.device_count()
+        )
+    return "cpu"
+
+
+def _create_colwise_spec(
+    pg: dist.ProcessGroup | None = None,
+) -> ChunkShardingSpec:
+    pg_device_type = dist.distributed_c10d._get_pg_default_device(pg).type
+    if pg is None:
+        placements = [
+            f"rank:{idx}/{_gen_rank_device(idx, pg_device_type)}"
+            for idx in range(dist.get_world_size())
+        ]
+    else:
+        placements = [
+            f"rank:{idx}/{_gen_rank_device(dist.get_global_rank(pg, idx), pg_device_type)}"
+            for idx in range(pg.size())
+        ]
+    return ChunkShardingSpec(
+        dim=0,
+        placements=cast(list[_remote_device | str], placements),
+    )
+
+
+def _is_nested_tensor(val: torch.Tensor) -> bool:
+    if type(val) is ShardedTensor:
+        if len(val.local_shards()) == 0:
+            return False
+        if type(val.local_shards()[0].tensor) is ShardedTensor:
+            return True
+        if type(val.local_shards()[0].tensor) is DTensor:
+            raise ValueError("Cannot handle DTensor nested inside ShardedTensor")
+    elif type(val) is DTensor and (
+        type(val._local_tensor) is DTensor or type(val._local_tensor) is ShardedTensor
+    ):
+        raise ValueError("Cannot handle nested DTensor")
+    return False
+
+
+def _alloc_tensor(
+    props: TensorProperties, size: Sequence[int], device_type: str = "cuda"
+) -> torch.Tensor:
+    if device_type == "cpu":
+        device = cast(torch.device, _get_device_module(device_type).current_device())
+    else:
+        device = torch.device(
+            device_type, _get_device_module(device_type).current_device()
+        )
+
+    return torch.empty(
+        size=size,
+        dtype=props.dtype,
+        layout=props.layout,
+        requires_grad=props.requires_grad,
+        pin_memory=props.pin_memory,
+        device=device,
+    )
+
+
+def _get_state_dict_2d_layout(
+    state_dict: STATE_DICT_TYPE,
+) -> tuple[STATE_DICT_2D_LAYOUT, dist.ProcessGroup | None]:
+    """
+    Load the right TP slice of the optimizer state.
+
+    This is not easy since the per-tensor slicing can't be inferred from checkpoint metadata.
+    We take advantage of the model state_dict producing a sliced ST to figure out what we need to load.
+    This is pretty fragile and it might be easier for FSDP to compute this info for us.
+    Returns a dictionary where keys are the same of the state_dict and the value is a tuple of
+    (offset, size) for the current rank TP slice.
+    N.B. The state_dict *MUST* come from FSDP.sharded_state_dict.
+    """
+    specs: STATE_DICT_2D_LAYOUT = {}
+    dp_pg: dist.ProcessGroup | None = None
+    for key, value in state_dict.items():
+        specs[key] = (None, value.size())
+        if _is_nested_tensor(value):
+            if not len(value.local_shards()) == 1:
+                raise AssertionError("Cannot handle ST with multiple shards")
+            if not isinstance(value, ShardedTensor):
+                raise AssertionError("Can only handle nested ShardedTensor")
+            shard = value.local_shards()[0]
+            specs[key] = (
+                shard.metadata.shard_offsets,
+                shard.metadata.shard_sizes,
+            )
+            dp_pg = shard.tensor._process_group  # type: ignore[attr-defined]
+
+    return (
+        specs,
+        dp_pg,
+    )
+
+
+class _ReaderWithOffset(DefaultLoadPlanner):
+    translation: dict[MetadataIndex, MetadataIndex]
+    state_dict: STATE_DICT_TYPE
+    # pyrefly: ignore [bad-override]
+    metadata: Metadata
+
+    def __init__(self, fqn_to_offset: dict[str, Sequence[int]]) -> None:
+        super().__init__()
+        self.fqn_to_offset = fqn_to_offset
+        self.metadata = Metadata({})
+        self.state_dict = {}
+        self.translation = {}
+
+    def create_local_plan(self) -> LoadPlan:
+        requests = []
+        self.translation = {}
+        for fqn, obj in self.state_dict.items():
+            md = self.metadata.state_dict_metadata[fqn]
+            if not isinstance(obj, ShardedTensor):
+                requests += _create_read_items(fqn, md, obj)
+                continue
+
+            if fqn not in self.fqn_to_offset:
+                requests += _create_read_items(fqn, md, obj)
+                continue
+
+            offset = self.fqn_to_offset[fqn]
+
+            if not len(obj.local_shards()) == 1:
+                raise AssertionError("Expected exactly one local shard")
+            original_shard = obj.local_shards()[0]
+            local_chunks = [
+                ChunkStorageMetadata(
+                    offsets=torch.Size(
+                        _element_wise_add(original_shard.metadata.shard_offsets, offset)
+                    ),
+                    sizes=torch.Size(original_shard.metadata.shard_sizes),
+                )
+            ]
+
+            reqs = create_read_items_for_chunk_list(
+                fqn, cast(TensorStorageMetadata, md), local_chunks
+            )
+            # TODO: The ReadItems will have a displaced MetadataIndex, fix it.
+            # TODO: we should change _create_sharded_read_items to have more ergonomic API
+            for ri in reqs:
+                if ri.dest_index.offset is None:
+                    raise AssertionError("dest_index.offset must not be None")
+                original_offset = _element_wise_sub(ri.dest_index.offset, offset)
+                original_index = dataclasses.replace(
+                    ri.dest_index, offset=torch.Size(original_offset)
+                )
+                self.translation[ri.dest_index] = original_index
+
+            requests += reqs
+        return LoadPlan(requests)
+
+    def lookup_tensor(self, index: MetadataIndex) -> torch.Tensor:
+        return super().lookup_tensor(self.translation.get(index, index))
+
+
+def load_sharded_optimizer_state_dict(
+    model_state_dict: STATE_DICT_TYPE,
+    optimizer_key: str,
+    storage_reader: StorageReader,
+    planner: LoadPlanner | None = None,
+) -> STATE_DICT_TYPE:
+    """
+    Load a state_dict in conjunction with FSDP sharded optimizer state.
+
+    This is the current recommended way to checkpoint FSDP.
+    >>> # xdoctest: +SKIP
+    >>> import torch.distributed.checkpoint as dist_cp
+    >>> # Save
+    >>> model: torch.nn.Model
+    >>> optim_params = model.parameters()
+    >>> optim = torch.optim.SGD(optim_params, lr=0.01)
+    >>> # Save
+    >>> with FSDP.state_dict_type(model, StateDictType.SHARDED_STATE_DICT):
+    >>>     state_dict = {
+    >>>         "optimizer": FSDP.optim_state_dict(model, optim),
+    >>>         "model": model.state_dict()
+    >>>     }
+    >>>     dist_cp.save_state_dict(
+    >>>         state_dict=optim_state,
+    >>>         storage_writer=dist_cp.FileSystemWriter("checkpoint"),
+    >>>         planner=dist_cp.DefaultSavePlanner(),
+    >>>     )
+    >>>
+    >>> # Load
+    >>> with FSDP.state_dict_type(model_tp, StateDictType.SHARDED_STATE_DICT):
+    >>>     model_state_dict = model_tp.state_dict()
+    >>>     checkpoint = {
+    >>>         "model": model_state_dict
+    >>>     }
+    >>>     dist_cp.load_state_dict(
+    >>>         state_dict=checkpoint,
+    >>>         storage_reader=dist_cp.FileSystemReader(checkpoint_file),
+    >>>         planner=dist_cp.DefaultLoadPlanner(),
+    >>>     )
+    >>>     model.load_state_dict(checkpoint["model_state"])
+    >>>
+    >>>     optim_state = dist_cp.load_sharded_optimizer_state_dict(
+    >>>         model_state_dict,
+    >>>         optimizer_key="optimizer",
+    >>>         storage_reader=dist_cp.FileSystemReader("checkpoint"),
+    >>>     )
+    >>>
+    >>>     flattened_osd = FSDP.optim_state_dict_to_load(
+    >>>        model, optim, optim_state["optimizer"]
+    >>>     )
+    >>>
+    >>>     optim.load_state_dict(flattened_osd)
+    """
+    metadata = storage_reader.read_metadata()
+
+    layout_specs, dp_pg = _get_state_dict_2d_layout(model_state_dict)
+    dp_pg_device_type = dist.distributed_c10d._get_pg_default_device(dp_pg).type
+    device_module = _get_device_module(dp_pg_device_type)
+
+    if dp_pg is None:
+        placements = []
+        for i in range(dist.get_world_size()):
+            device_info = _normalize_device_info(
+                dp_pg_device_type, i % device_module.device_count()
+            )
+            placements.append(f"rank:{i}/{device_info}")
+        sharding_spec = ChunkShardingSpec(dim=0, placements=placements)  # type: ignore[arg-type]
+    else:
+        sharding_spec = _create_colwise_spec(dp_pg)
+
+    # Create a state_dict for optimizer state
+    state_dict: STATE_DICT_TYPE = {}
+
+    fqn_to_offset: dict[str, Sequence[int]] = {}
+    for key, value in metadata.state_dict_metadata.items():
+        key_path = metadata.planner_data[key]
+        if key_path[0] != optimizer_key:
+            continue
+
+        if isinstance(value, BytesStorageMetadata):
+            state_dict[key] = "<bytes_io>"
+            continue
+
+        # value: TensorStorageMetadata
+        if value.size.numel() == 1:
+            state_dict[key] = _alloc_tensor(
+                value.properties, value.size, dp_pg_device_type
+            )
+        elif dp_pg is None:
+            state_dict[key] = _create_chunk_sharded_tensor(
+                _alloc_tensor(value.properties, value.size, dp_pg_device_type),
+                rank=dist.get_rank(),
+                world_size=dist.get_world_size(),
+                num_devices_per_node=device_module.device_count(),
+                pg=_get_default_group(),
+            )
+        else:
+            spec_key = key_path[2]
+            alloc_size = layout_specs.get(spec_key, (None, value.size))[1]
+
+            properties = ShardTensorProperties(
+                dtype=value.properties.dtype,
+                layout=value.properties.layout,
+                requires_grad=value.properties.requires_grad,
+                memory_format=value.properties.memory_format,
+                pin_memory=value.properties.pin_memory,
+            )
+
+            st_md = sharding_spec.build_metadata(torch.Size(alloc_size), properties)
+            local_shards = []
+            current_rank = dist.get_rank(dp_pg)
+            for shard_md in st_md.shards_metadata:
+                if cast(_remote_device, shard_md.placement).rank() != current_rank:
+                    continue
+                local_shards.append(
+                    Shard(
+                        tensor=_alloc_tensor(
+                            value.properties, shard_md.shard_sizes, dp_pg_device_type
+                        ),
+                        metadata=shard_md,
+                    )
+                )
+
+            st = ShardedTensor._init_from_local_shards_and_global_metadata(
+                local_shards, st_md, process_group=dp_pg
+            )
+
+            if spec_key in layout_specs and layout_specs[spec_key][0] is not None:
+                fqn_to_offset[key] = cast(Sequence[int], layout_specs[spec_key][0])
+
+            state_dict[key] = st
+
+    # Whether we unflatten before or after doesn't matter
+    load_state_dict(
+        state_dict=state_dict,
+        storage_reader=storage_reader,
+        # FIXME the type of planner is wrong in load_state_dict
+        planner=_ReaderWithOffset(fqn_to_offset) if dp_pg is not None else planner,
+    )
+
+    state_dict = unflatten_state_dict(state_dict, metadata.planner_data)
+
+    return state_dict

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/planner.py

@@ -0,0 +1,450 @@
+import abc
+import io
+import operator
+from dataclasses import dataclass
+from enum import auto, Enum
+from functools import reduce
+from typing import Any
+
+import torch
+from torch.distributed.checkpoint.metadata import (
+    ChunkStorageMetadata,
+    Metadata,
+    MetadataIndex,
+    STATE_DICT_TYPE,
+    StorageMeta,
+    TensorProperties,
+)
+
+
+__all__ = [
+    "WriteItemType",
+    "LoadItemType",
+    "BytesIOWriteData",
+    "TensorWriteData",
+    "WriteItem",
+    "ReadItem",
+    "SavePlan",
+    "LoadPlan",
+    "SavePlanner",
+    "LoadPlanner",
+]
+
+
+class WriteItemType(Enum):
+    TENSOR = auto()
+    SHARD = auto()
+    BYTE_IO = auto()
+
+
+class LoadItemType(Enum):
+    TENSOR = auto()
+    BYTE_IO = auto()
+
+
+@dataclass(frozen=True)
+class BytesIOWriteData:
+    nbytes: int
+
+
+@dataclass(frozen=True)
+class TensorWriteData:
+    chunk: ChunkStorageMetadata
+    properties: TensorProperties
+    size: torch.Size
+
+
+@dataclass(frozen=True)
+class WriteItem:
+    """Dataclass which holds information about what needs to be written to storage."""
+
+    index: MetadataIndex
+    type: WriteItemType
+
+    # Size of bytesIO data to be written.
+    bytes_io_data: BytesIOWriteData | None = None
+
+    # Value present if it's a tensor write
+    tensor_data: TensorWriteData | None = None
+
+    def tensor_storage_size(self) -> int | None:
+        """
+        Calculates the storage size of the underlying tensor, or None if this is not a tensor write.
+
+        Returns:
+            Optional[int] storage size, in bytes of underlying tensor if any.
+        """
+        if self.tensor_data is None:
+            return None
+
+        numels = reduce(operator.mul, self.tensor_data.size, 1)
+        dtype_size = torch._utils._element_size(self.tensor_data.properties.dtype)
+        return numels * dtype_size
+
+
+@dataclass(frozen=True)
+class ReadItem:
+    # Read Item
+    type: LoadItemType
+
+    # Index into the state_dict
+    dest_index: MetadataIndex
+    # Offsets into destination tensor
+    dest_offsets: torch.Size
+
+    # Index into the checkpoint
+    storage_index: MetadataIndex
+    # Offset into the checkpoint data
+    storage_offsets: torch.Size
+
+    # Size of the hypercube to copy
+    lengths: torch.Size
+
+
+@dataclass(frozen=True)
+class SavePlan:
+    items: list[WriteItem]
+    storage_data: Any = None
+    planner_data: Any = None
+    # This is used to indicate that the ranks should
+    # use the cached plans to write data instead.
+    usable: bool = True
+
+
+@dataclass
+class LoadPlan:
+    items: list[ReadItem]
+    storage_data: Any = None
+    planner_data: Any = None
+
+
+class SavePlanner(abc.ABC):
+    """
+    Abstract class defining the protocol used by save_state_dict to plan the save process.
+
+    SavePlanners are stateful objects that can be used to customize the whole save process.
+
+    SavePlanner acts as an access proxy to the state_dict, so any transformation done to it
+    will be visible to the whole process.
+
+    A planner subclass can expect the following sequence of calls during save_state_dict:
+
+    1) set_up_planner - called on all ranks.
+        Signals the start of a checkpoint save.
+
+    2) create_local_plan - called on all ranks.
+        Process the state_dict and produces a `SavePlan` that will be sent for global planning.
+
+    3) create_global_plan - called on the coordinator rank only.
+        Takes the SavePlan from all ranks and make any global decision.
+
+    4) finish_plan - called on all ranks.
+        This gives each rank a chance to adjust to global planning decisions.
+
+    5) resolve_data - called multiple times on each rank
+        Lookups a value on the `state_dict` for the storage layer to write.
+
+    Users are recommended to extend DefaultSavePlanner instead of this interface directly as
+    most changes can be expressed by changes in a single method.
+
+    There are 3 usual patterns of extension:
+
+    Rewriting state_dict. This is the simplest way to extend the save process as it
+    doesn't requite understanding the intrincacies of how SavePlan works:
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> class RenamePlanner(DefaultSavePlanner):
+    >>>     def set_up_planner(
+    >>>         self,
+    >>>         state_dict: STATE_DICT_TYPE,
+    >>>         storage_meta: Optional[StorageMeta],
+    >>>         is_coordinator: bool,
+    >>>     ) -> None:
+    >>> # prefix all keys with `foo_``
+    >>>         super().set_up_planner({"foo_" + k: v for k, v in state_dict.items()}, storage_meta, is_coordinator)
+
+    Modifying local plan and lookup in tandem. This is useful when fine control of how data is persisted
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> class FP16Planner(DefaultSavePlanner):
+    >>>     def create_local_plan(self):
+    >>>         plan = super().create_local_plan()
+    >>>         for p in plan:
+    >>>             if p.tensor_data is not None:
+    >>>                 p.tensor_data.properties.dtype = torch.float16
+    >>>         return plan
+    >>>
+    >>>     def resolve_data(self, write_item):
+    >>>         item = super().resolve_data(write_item)
+    >>>         return item if write_item.type == WriteItemType.BYTE_IO else item.to(torch.float16)
+
+    Using the global planning step to make central decisions that can't be made individually by each rank
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> from itertools import zip_longest
+    >>> from dataclasses import replace
+    >>> class DDPLoadBalancingPlanner(DefaultSavePlanner):
+    >>> # This uses the default local plan behavior of having all non-sharded writes in rank 0
+    >>> # This sample doesn't handle ShardedTensors
+    >>>     def create_global_plan(self, all_plans):
+    >>>         iters = [iter(all_plans[0].items)] * len(all_plans)
+    >>>         items_per_rank = [
+    >>>             [item for item in items if item is not None]
+    >>>             for items in zip(*zip_longest(*iters), strict=True)
+    >>>         ]
+    >>>         all_plans = [
+    >>>             replace(plan, items=items)
+    >>>             for plan, items in zip(all_plans, items_per_rank, strict=True)
+    >>>         ]
+    >>>         return super().create_global_plan(all_plans)
+
+    Finally, some planners need to save additional metadata in the checkpoint, this is
+    accomplished by having each rank contribute their data items in the local plan and
+    the global planner aggregate them:
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> class SaveExtraDataPlanner(DefaultSavePlanner):
+    >>>     def create_local_plan(self) -> SavePlan:
+    >>>         plan = super().create_local_plan()
+    >>>         return replace(plan, planner_data="per-rank-data")
+    >>>
+    >>>     def create_global_plan(self, all_plans: List[SavePlan]) -> Tuple[List[SavePlan], Metadata]:
+    >>>         global_plan, metadata = super().create_global_plan(all_plans)
+    >>>         merged_data = [p.planner_data for p in global_plan]
+    >>>         metadata = replace(metadata, planner_data=merged_data)
+    >>>         return global_plan, metadata
+    """
+
+    # Save plan for the current rank as computed by `create_local_plan` API
+    # Cached on the local rank.
+    _cached_save_plan: dict[str, SavePlan] = {}
+    # Final save plan for the current rank.
+    # This is created by merging the plan created by `create_local_plan` API
+    # and the result of `create_global_plan` for the given rank.
+    # This is the final plan computed by the `finish_plan` API that gets
+    # sent to the `write_data`.
+    # Cached on the local rank.
+    _cached_final_save_plan: dict[str, SavePlan] = {}
+    # Collection of all the local plans from all the ranks.
+    # This is the input to the `create_global_plan` API.
+    # Cached on the coordinator rank.
+    _cached_all_plans: dict[str, list[SavePlan]] = {}
+    # Global checkpoint plan as computed by `create_global_plan` API.
+    # Cached on the coordinator rank.
+    _cached_global_plan: dict[str, list[SavePlan]] = {}
+    # Metadata for the global checkpoint plan as computed by `create_global_plan` API.
+    # Cached on the coordinator rank.
+    _cached_metadata: dict[str, Metadata] = {}
+
+    @abc.abstractmethod
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        storage_meta: StorageMeta | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        """
+        Initialize this planner to save ``state_dict``.
+
+        Implementations should save those values as they won't be provided lated in the save process.
+
+        This is called on all ranks.
+        """
+
+    @abc.abstractmethod
+    def create_local_plan(self) -> SavePlan:
+        """
+        Compute the save plan for the current rank.
+
+        This will be aggregated and passed to create_global_plan.
+        Planner specific data can be passed through SavePlan::planner_data.
+
+        This is called on all ranks.
+        """
+
+    @abc.abstractmethod
+    def create_global_plan(
+        self, all_plans: list[SavePlan]
+    ) -> tuple[list[SavePlan], Metadata]:
+        """
+        Compute the global checkpoint plan and return the local plan of each rank.
+
+        This is called on the coordinator rank only.
+        """
+
+    @abc.abstractmethod
+    def finish_plan(self, new_plan: SavePlan) -> SavePlan:
+        """
+        Merge the plan created by `create_local_plan` and the result of `create_global_plan`.
+
+        This is called on all ranks.
+        """
+
+    @abc.abstractmethod
+    def resolve_data(self, write_item: WriteItem) -> torch.Tensor | io.BytesIO:
+        """
+        Transform and prepare ``write_item`` from ``state_dict`` for storage, ensuring idempotency and thread-safety.
+
+        Lookup the object associated with ``write_item`` in ``state_dict`` and apply any
+        transformation (such as serialization) prior to the storage layer consuming it.
+
+        Called on each rank multiple times, at least once per WriteItem in the final SavePlan.
+
+        This method should be idempotent and thread-save. StorageWriter implementations
+        are free to call it as frequently as they need.
+
+        Any transformation that allocates memory should be lazily done when his method
+        is called in order to reduce peak memory required by checkpointing.
+
+        When returning tensors, they can be on any device or format, they can be views too.
+        It's the storage layer responsibility to figure out how to save them.
+        """
+
+
+class LoadPlanner:
+    """
+    Abstract class defining the protocol used by load_state_dict to plan the load process.
+
+    LoadPlanner are stateful objects that can be used to customize the whole load process.
+
+    LoadPlanner acts as an access proxy to the state_dict, so any transformation done to it
+    will be visible to the whole process.
+
+    A planner subclass can expect the following sequence of calls during load_state_dict:
+
+    1) set_up_planner - called on all ranks.
+        Signals the start of loading a checkpoint.
+
+    2) create_local_plan - called on all ranks.
+        Process the state_dict and produces a `LoadPlan` that will be sent for global planning.
+
+    3) create_global_plan - called on the coordinator rank only.
+        Takes the LoadPlan from all ranks and make any global decision.
+
+    4) load_bytes - called multiple times on each rank
+        This is called once per non-tensor value in state_dict.
+
+    5) resolve_tensor and commit_tensor - called multiple times on each rank
+        They are called in pair for each Tensor value in state_dict.
+
+    Users are recommended to extend DefaultLoadPlanner instead of this interface directly as
+    most changes can be expressed by changes in a single method.
+
+    There are two usual patterns of extension:
+
+    Rewriting state_dict. This is the simplest way to extend the load process as it
+    doesn't requite understanding the intrincacies of how LoadPlan works. We need
+    to keep a reference to the original state_dict as load happens in place so
+    we need to be able to perform it in place
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> class RenamePlanner(DefaultLoadPlanner):
+    >>>     def set_up_planner(
+    >>>         self,
+    >>>         state_dict: STATE_DICT_TYPE,
+    >>>         metadata: Metadata,
+    >>>         is_coordinator: bool,
+    >>>     ) -> None:
+    >>>         self.original_state_dict = state_dict
+    >>>         state_dict = {"foo_" + k: v for k, v in state_dict.items()}
+    >>>
+    >>>         if self.flatten_sharded_tensors:
+    >>>             state_dict = _flatten_sharded_tensors(state_dict)
+    >>>
+    >>>         if self.flatten_state_dict:
+    >>>             state_dict, self.mappings = flatten_state_dict(state_dict)
+    >>>
+    >>>         self.state_dict = state_dict
+    >>>         self.metadata = metadata
+    >>>         self.is_coordinator = is_coordinator
+    >>>
+    >>>     def load_bytes(self, read_item, value):
+    >>> # Remove the "foo_" prefix
+    >>>         self.original_state_dict[read_item.dest_index.fqn[4:]] = torch.load(value, weights_only=False)
+
+
+    Modifying resolve_tensor and commit_tensor to handle load time transformation.
+
+    >>> # xdoctest: +SKIP("undefined vars")
+    >>> class MetaModelMaterialize(DefaultSavePlanner):
+    >>>     def resolve_tensor(self, read_item):
+    >>>         tensor = super().resolve_tensor(read_item)
+    >>>         return torch.empty_like(tensor, device="cpu")
+    >>>
+    >>>     def commit_tensor(self, read_item, tensor):
+    >>>         self.state_dict[read_item.dest_index.fqn] = tensor
+    """
+
+    @abc.abstractmethod
+    def set_up_planner(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        metadata: Metadata | None = None,
+        is_coordinator: bool = False,
+    ) -> None:
+        """
+        Initialize this instance to load data into ``state_dict``.
+
+        . N.B. This is called on every rank.
+        """
+
+    @abc.abstractmethod
+    def create_local_plan(self) -> LoadPlan:
+        """
+        Create a LoadPlan based on state_dict and metadata provided by set_up_planner.
+
+        . N.B. This is called on every rank.
+        """
+
+    @abc.abstractmethod
+    def create_global_plan(self, global_plan: list[LoadPlan]) -> list[LoadPlan]:
+        """
+        Compute the global load plan and return plans for each rank.
+
+        . N.B. This is called on the coordinator rank only
+        """
+
+    @abc.abstractmethod
+    def finish_plan(self, central_plan: LoadPlan) -> LoadPlan:
+        """Accept the plan from coordinator and return final LoadPlan."""
+
+    @abc.abstractmethod
+    def load_bytes(self, read_item: ReadItem, value: io.BytesIO) -> None:
+        """
+        Load the item described by ``read_item``and ``value``.
+
+        This method is expected to modify in-place the underlying state_dict.
+
+        The contents of ``value`` are defined by the SavePlanner used to produce
+        the checkpoint being loaded.
+        """
+
+    def resolve_bytes(self, read_item: ReadItem) -> io.BytesIO:
+        """
+        Return the BytesIO to be used by the StorageReader to load `read_item`.
+
+        The BytesIO should alias with one on the underlying state_dict as StorageReader will replace its contents.
+        """
+        raise NotImplementedError("LoadPlanner.resolve_bytes is not implemented")
+
+    @abc.abstractmethod
+    def resolve_tensor(self, read_item: ReadItem) -> torch.Tensor:
+        """
+        Return the tensor described by ``read_item`` to be used by the StorageReader to load `read_item`.
+
+        The tensor should alias with one on the underlying state_dict as StorageReader will replace its contents.
+        If, for any reason, that's not possible, the planner can use the ``commit_tensor`` method to copy the data
+        back to the one in state_dict.
+        """
+
+    @abc.abstractmethod
+    def commit_tensor(self, read_item: ReadItem, tensor: torch.Tensor) -> None:
+        """
+        Call once the StorageReader finished loading data into ``tensor``.
+
+        The provided tensor is the same one returned by the call to ``resolve_tensor``.
+        This method is only needed if this LoadPlanner needs to post process ``tensor`` prior to
+        copying it back to the one in the state_dict.
+
+        The contents of tensor will follow its device synchronization model.
+        """

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/planner_helpers.py

@@ -0,0 +1,561 @@
+# mypy: allow-untyped-defs
+import io
+import itertools
+from bisect import bisect_right, insort
+from collections.abc import Callable
+from typing import Any, cast
+
+import torch
+import torch.distributed as dist
+from torch._utils import _get_device_module
+from torch.distributed._shard.metadata import ShardMetadata
+from torch.distributed._shard.sharded_tensor import ShardedTensor
+from torch.distributed.tensor import DTensor
+from torch.distributed.tensor._utils import compute_local_shape_and_global_offset
+
+from .metadata import (
+    BytesStorageMetadata,
+    ChunkStorageMetadata,
+    MetadataIndex,
+    STATE_DICT_TYPE,
+    STORAGE_TYPES,
+    TensorProperties,
+    TensorStorageMetadata,
+)
+from .planner import (
+    LoadItemType,
+    ReadItem,
+    SavePlan,
+    TensorWriteData,
+    WriteItem,
+    WriteItemType,
+)
+from .resharding import (
+    _check_shard_metadata_pair_overlap,
+    _shards_get_overlap_region_wrt_saved_tensor,
+)
+
+
+__all__: list[str] = ["create_read_items_for_chunk_list"]
+
+
+def _compare_save_plans(plan: SavePlan, other_plan: SavePlan) -> bool:
+    """
+    Compare the two Save plans and return True if they are equal.
+
+    Args:
+        plan (SavePlan): First SavePlan to compare.
+        other_plan (SavePlan): Second SavePlan to compare.
+
+    Returns:
+       True if the two plans are equal, False otherwise.
+    """
+    if plan.usable != other_plan.usable:
+        return False
+
+    # Both the plans should have the same number of items
+    if len(plan.items) != len(other_plan.items):
+        return False
+
+    # Both the plans should have the same write items.
+    for plan_item, other_plan_item in zip(plan.items, other_plan.items):
+        # Write item type should be same
+        if plan_item.type != other_plan_item.type:
+            return False
+
+        plan_metadata_index = plan_item.index
+        other_plan_metadata_index = other_plan_item.index
+
+        # Write item metadata_index should be same
+        if (
+            plan_metadata_index.fqn != other_plan_metadata_index.fqn
+            or plan_metadata_index.offset != other_plan_metadata_index.offset
+            or plan_metadata_index.index != other_plan_metadata_index.index
+        ):
+            return False
+
+        # Write item tensor_data should be present in both the write items plans, if it exists in either of them.
+        tensor_data = plan_item.tensor_data
+        other_tensor_data = other_plan_item.tensor_data
+        if (tensor_data and not other_tensor_data) or (
+            not tensor_data and other_tensor_data
+        ):
+            return False
+
+        if tensor_data and other_tensor_data:
+            # Write item tensor_data size should be same
+            if tensor_data.size != other_tensor_data.size:
+                return False
+
+            # Write item tensor_data chunk should be present in both the write items, if it exists in either of them.
+            chunk = tensor_data.chunk
+            other_chunk = other_tensor_data.chunk
+            if (chunk and not other_chunk) or (not chunk and other_chunk):
+                return False
+
+            # Write item tensor_data chunk offsets and sizes should be same
+            if chunk and other_chunk:
+                if (
+                    chunk.offsets != other_chunk.offsets
+                    or chunk.sizes != other_chunk.sizes
+                ):
+                    return False
+
+    return True
+
+
+def _contains_usable_plan(delta_plans: list[SavePlan]) -> bool:
+    """
+    Check if any delta plan is usable, indicating the plan has changed.
+
+    Args:
+        delta_plans (List[SavePlan]): A list of delta plans to check.
+    Returns:
+        True if any delta plan is usable, False otherwise.
+    """
+    return any(delta_plan and delta_plan.usable for delta_plan in delta_plans)
+
+
+def _merge_delta_local_plans(
+    cached_plans: list[SavePlan],
+    delta_plans: list[SavePlan],
+) -> list[SavePlan]:
+    """
+    Merge a list of delta plans into a single plan.
+
+    Args:
+        cached_plans (List[SavePlan]): A list of cached plans.
+        delta_plans (List[SavePlan]): A list of delta plans to merge. It can contain empty plans
+
+    Returns:
+        A single merged plan. If a delta plan is not usable, use the cached plan. Otherwise, use the delta plan.
+    """
+    merged_plans = []
+
+    for cached_plan, delta_plan in zip(cached_plans, delta_plans):
+        if delta_plan and not delta_plan.usable:
+            merged_plans.append(cached_plan)
+        else:
+            merged_plans.append(delta_plan)
+
+    return merged_plans
+
+
+def _create_chunk_from_tensor(tensor: torch.Tensor) -> ChunkStorageMetadata:
+    return ChunkStorageMetadata(
+        offsets=torch.Size([0] * len(tensor.size())), sizes=tensor.size()
+    )
+
+
+def _chunk_for_shard(shard_md: ShardMetadata) -> ChunkStorageMetadata:
+    return ChunkStorageMetadata(
+        offsets=torch.Size(shard_md.shard_offsets),
+        sizes=torch.Size(shard_md.shard_sizes),
+    )
+
+
+def _sharded_tensor_metadata(
+    sharded_tensor: ShardedTensor, shard_md: ShardMetadata
+) -> TensorWriteData:
+    shard_properties = sharded_tensor.metadata().tensor_properties
+
+    properties = TensorProperties(
+        dtype=shard_properties.dtype,
+        layout=shard_properties.layout,
+        requires_grad=shard_properties.requires_grad,
+        memory_format=shard_properties.memory_format,
+        pin_memory=shard_properties.pin_memory,
+    )
+
+    return TensorWriteData(
+        chunk=_chunk_for_shard(shard_md),
+        properties=properties,
+        size=sharded_tensor.metadata().size,
+    )
+
+
+def _create_write_items_for_dtensor(fqn: str, tensor: DTensor) -> WriteItem:
+    sizes, offsets = compute_local_shape_and_global_offset(
+        tensor.shape, tensor.device_mesh, tensor.placements
+    )
+    sizes, offsets = torch.Size(sizes), torch.Size(offsets)
+
+    return WriteItem(
+        index=MetadataIndex(fqn, offsets),
+        type=WriteItemType.SHARD,
+        tensor_data=TensorWriteData(
+            chunk=ChunkStorageMetadata(
+                offsets=offsets,
+                sizes=sizes,
+            ),
+            properties=TensorProperties.create_from_tensor(tensor.to_local()),
+            size=tensor.size(),
+        ),
+    )
+
+
+def _create_write_item_for_shard(
+    fqn: str, sharded_tensor: ShardedTensor, shard_md: ShardMetadata
+) -> WriteItem:
+    offsets = torch.Size(shard_md.shard_offsets)
+    return WriteItem(
+        index=MetadataIndex(fqn, offsets),
+        type=WriteItemType.SHARD,
+        tensor_data=_sharded_tensor_metadata(sharded_tensor, shard_md),
+    )
+
+
+def _create_write_item_for_tensor(fqn: str, tensor: torch.Tensor) -> WriteItem:
+    offsets = torch.Size([0] * len(tensor.size()))
+    return WriteItem(
+        index=MetadataIndex(fqn, offsets),
+        type=WriteItemType.TENSOR,
+        tensor_data=TensorWriteData(
+            chunk=ChunkStorageMetadata(offsets=offsets, sizes=tensor.size()),
+            properties=TensorProperties.create_from_tensor(tensor),
+            size=tensor.size(),
+        ),
+    )
+
+
+def _create_write_item_for_bytesio(fqn: str, bytes: Any):
+    return WriteItem(
+        index=MetadataIndex(fqn),
+        type=WriteItemType.BYTE_IO,
+    )
+
+
+def _create_read_item_for_byteio(
+    dest_index, dest_offset, storage_index, storage_offset, length
+):
+    return ReadItem(
+        type=LoadItemType.BYTE_IO,
+        dest_index=dest_index,
+        dest_offsets=torch.Size((dest_offset,)),
+        storage_index=storage_index,
+        storage_offsets=torch.Size((storage_offset,)),
+        lengths=torch.Size((length,)),
+    )
+
+
+def _create_read_item_for_tensor(
+    dest_index, dest_offsets, storage_index, storage_offsets, lengths
+):
+    return ReadItem(
+        type=LoadItemType.TENSOR,
+        dest_index=dest_index,
+        dest_offsets=torch.Size(dest_offsets),
+        storage_index=storage_index,
+        storage_offsets=torch.Size(storage_offsets),
+        lengths=torch.Size(lengths),
+    )
+
+
+def create_read_items_for_chunk_list(
+    fqn: str,
+    checkpoint_md: TensorStorageMetadata,
+    local_chunks: list[ChunkStorageMetadata],
+) -> list[ReadItem]:
+    """
+    Create a list of ``ReadItem`` based on the checkpoint and local chunks.
+
+    This applies the resharding algorithm and computes the reads needed
+    to satisfy ``local_chunks`` with a checkpoint described by ``checkpoint_md``.
+
+    Args:
+        fqn (str) : The state_dict FQN to pass to ``ReadItem``.
+        checkpoint_md (TensorStorageMetadata): metadata for a given tensor
+            from a checkpoint.
+        local_chunks (List[ChunkStorageMetadata]): Local chunks that needs to be
+            loaded.
+
+    Returns:
+        A list of ``ReadItem`` that will satisfy all input chunks.
+    """
+    read_items: list[ReadItem] = []
+    saved_chunks = checkpoint_md.chunks
+
+    if not local_chunks or not saved_chunks:
+        return read_items
+
+    num_dims = len(local_chunks[0].offsets)
+
+    # Find sweep dimension (dimension with largest extent for better pruning)
+    sweep_dim = 0
+    if num_dims > 1:
+        max_size = 0
+        for dim in range(num_dims):
+            dim_size = max(
+                chunk.offsets[dim] + chunk.sizes[dim]
+                for chunk in itertools.chain(local_chunks, saved_chunks)
+            )
+            if dim_size > max_size:
+                max_size = dim_size
+                sweep_dim = dim
+
+    # Pre-compute bounds: (start, end) for each chunk in sweep dimension
+    # For 0-d tensors, use (0, 1) so all chunks overlap in the sweep line
+    if num_dims == 0:
+        saved_bounds = [(0, 1)] * len(saved_chunks)
+        local_bounds = [(0, 1)] * len(local_chunks)
+    else:
+        saved_bounds = [
+            (c.offsets[sweep_dim], c.offsets[sweep_dim] + c.sizes[sweep_dim])
+            for c in saved_chunks
+        ]
+        local_bounds = [
+            (c.offsets[sweep_dim], c.offsets[sweep_dim] + c.sizes[sweep_dim])
+            for c in local_chunks
+        ]
+
+    saved_sorted_indices = sorted(
+        range(len(saved_chunks)),
+        key=lambda idx: saved_bounds[idx][0],
+    )
+    local_sorted_indices = sorted(
+        range(len(local_chunks)),
+        key=lambda idx: local_bounds[idx][0],
+    )
+
+    active_saved: list[tuple[int, int]] = []
+    saved_ptr = 0
+    num_saved = len(saved_sorted_indices)
+
+    for local_idx in local_sorted_indices:
+        local_chunk = local_chunks[local_idx]
+        local_start, local_end = local_bounds[local_idx]
+
+        cutoff = bisect_right(active_saved, (local_start, -1))
+        if cutoff:
+            del active_saved[:cutoff]
+
+        while saved_ptr < num_saved:
+            storage_idx = saved_sorted_indices[saved_ptr]
+            storage_chunk = saved_chunks[storage_idx]
+            saved_start, saved_end = saved_bounds[storage_idx]
+
+            if saved_start >= local_end:
+                break
+
+            insort(active_saved, (saved_end, storage_idx))
+            saved_ptr += 1
+
+        for _, storage_idx in active_saved:
+            storage_chunk = saved_chunks[storage_idx]
+            if not _check_shard_metadata_pair_overlap(local_chunk, storage_chunk):
+                continue
+
+            storage_offsets = []
+            dest_offsets = []
+            lengths = []
+            for (
+                _dim,
+                offset_for_saved_tensor,
+                offset_for_current_tensor,
+                length,
+            ) in _shards_get_overlap_region_wrt_saved_tensor(
+                saved_shard=storage_chunk, current_shard=local_chunk
+            ):
+                storage_offsets.append(offset_for_saved_tensor)
+                dest_offsets.append(offset_for_current_tensor)
+                lengths.append(length)
+
+            read_items.append(
+                _create_read_item_for_tensor(
+                    dest_index=MetadataIndex(fqn, local_chunk.offsets, local_idx),
+                    dest_offsets=dest_offsets,
+                    storage_index=MetadataIndex(
+                        fqn, storage_chunk.offsets, storage_idx
+                    ),
+                    storage_offsets=storage_offsets,
+                    lengths=lengths,
+                )
+            )
+    return read_items
+
+
+def _create_default_metadata_only_plan(state_dict: STATE_DICT_TYPE) -> SavePlan:
+    requests = []
+    for fqn, obj in state_dict.items():
+        if isinstance(obj, DTensor):
+            requests.append(_create_write_items_for_dtensor(fqn, obj))
+        elif isinstance(obj, ShardedTensor):
+            requests.extend(
+                _create_write_item_for_shard(fqn, obj, shard_md)
+                for shard_md in obj.metadata().shards_metadata
+            )
+        elif isinstance(obj, torch.Tensor):
+            requests.append(_create_write_item_for_tensor(fqn, obj))
+        else:
+            requests.append(_create_write_item_for_bytesio(fqn, obj))
+    return SavePlan(requests)
+
+
+def _create_write_items(fqn: str, object: Any) -> list[WriteItem]:
+    if hasattr(object, "__create_write_items__"):
+        # DTensor implements _Checkpointable
+        return object.__create_write_items__(fqn, object)
+    elif isinstance(object, ShardedTensor):
+        return [
+            _create_write_item_for_shard(fqn, object, shard.metadata)
+            for shard in object.local_shards()
+        ]
+    elif isinstance(object, torch.Tensor):
+        return [_create_write_item_for_tensor(fqn, object)]
+    else:
+        return [_create_write_item_for_bytesio(fqn, object)]
+
+
+def _create_chunk_from_dtensor(tensor: DTensor) -> ChunkStorageMetadata:
+    sizes, offsets = compute_local_shape_and_global_offset(
+        tensor.shape, tensor.device_mesh, tensor.placements
+    )
+    sizes, offsets = torch.Size(sizes), torch.Size(offsets)
+    return ChunkStorageMetadata(
+        offsets=offsets,
+        sizes=sizes,
+    )
+
+
+def _create_chunk_list(tensor: torch.Tensor) -> list[ChunkStorageMetadata]:
+    if hasattr(tensor, "__create_chunk_list__"):
+        # DTensor implements _Checkpointable
+        local_chunks = tensor.__create_chunk_list__()  # type: ignore[attr-defined]
+    elif isinstance(tensor, ShardedTensor):
+        local_chunks = [
+            _chunk_for_shard(shard.metadata) for shard in tensor.local_shards()
+        ]
+    elif isinstance(tensor, torch.Tensor):
+        local_chunks = [_create_chunk_from_tensor(tensor)]
+    else:
+        raise ValueError(
+            "Unsupported Type, expecting one of [Tensor, DTensor, ShardedTensor] "
+            f",but got {type(tensor)}"
+        )
+
+    return local_chunks
+
+
+def _create_read_items(fqn: str, md: STORAGE_TYPES, obj: Any) -> list[ReadItem]:
+    if not isinstance(md, BytesStorageMetadata):
+        try:
+            local_chunks = _create_chunk_list(obj)
+        except ValueError as ex:
+            raise ValueError(
+                f"Invalid checkpoint metadata for {fqn}, "
+                + f"expected BytesStorageMetadata but found {type(md)}",
+            ) from ex
+
+        return create_read_items_for_chunk_list(fqn, md, local_chunks)
+    else:
+        return [
+            _create_read_item_for_byteio(
+                dest_index=MetadataIndex(fqn),
+                dest_offset=0,
+                storage_index=MetadataIndex(fqn),
+                storage_offset=0,
+                length=0,
+            )
+        ]
+
+
+def _init_state_dict(state_dict: dict[str, Any]) -> Any:
+    """
+    Initializes meta tensor if the meta tensor is DTensor or torch.Tensor.
+    """
+
+    def dtensor_func(value: DTensor):
+        device = getattr(value, "device", None)
+        if device == torch.device("meta"):
+            device_type = dist.distributed_c10d._get_pg_default_device().type
+            device = cast(
+                torch.device, _get_device_module(device_type).current_device()
+            )
+            new_local_tensor = torch.empty_like(value.to_local(), device=device)
+            # We need to pass shape and stride explicitly, since DTensor might be
+            # sharded unevenly.
+            dtensor = DTensor.from_local(
+                new_local_tensor,
+                device_mesh=value.device_mesh,
+                placements=value.placements,
+                shape=value.size(),
+                stride=value.stride(),
+            )
+            return dtensor
+        else:
+            return value
+
+    def sharded_tensor_func(value: Any):
+        device = getattr(value, "device", None)
+        if device == torch.device("meta"):
+            raise RuntimeError(
+                f"Found unsupported type {type(value)} for meta device loading."
+            )
+        else:
+            return value
+
+    def tensor_func(value: torch.Tensor):
+        device = getattr(value, "device", None)
+        if device == torch.device("meta"):
+            device_type = dist.distributed_c10d._get_pg_default_device().type
+            device = cast(
+                torch.device, _get_device_module(device_type).current_device()
+            )
+            tensor = torch.empty_like(value, device=device)
+            return tensor
+        else:
+            return value
+
+    _iterate_state_dict(
+        state_dict,
+        dtensor_func,
+        sharded_tensor_func,
+        tensor_func,
+    )
+
+
+def _iterate_state_dict(
+    iter_object: Any,
+    dtensor_func: Callable,
+    sharded_tensor_func: Callable,
+    tensor_func: Callable,
+):
+    """
+    Iterate through the state dict, applying the given functions to each tensor type
+    and update the state dict in place.
+
+    Args:
+        iter_object (Any): the target state_dict.
+        sharded_tensor_func (Callable): the function to apply to ShardedTensor
+        dtensor_func (Callable): the function to apply to DTensor
+        tensor_func (Callable): the function to apply to Tensor
+
+    # TODO: let state_dict_util._iterate_state_dict() to support in place option
+    so we don't need to have two versions of _iterate_state_dict.
+    """
+
+    if isinstance(iter_object, DTensor):
+        return dtensor_func(iter_object)
+    elif isinstance(iter_object, ShardedTensor):
+        return sharded_tensor_func(iter_object)
+    elif isinstance(iter_object, torch.Tensor):
+        return tensor_func(iter_object)
+    elif (
+        isinstance(iter_object, (int, float, str, bytes, io.BytesIO))
+        or iter_object is None
+    ):
+        return iter_object
+    elif isinstance(iter_object, dict):
+        for key, value in iter_object.items():
+            iter_object[key] = _iterate_state_dict(
+                value, dtensor_func, sharded_tensor_func, tensor_func
+            )
+        return iter_object
+    elif isinstance(iter_object, (list, tuple)):
+        ret = [
+            _iterate_state_dict(v, dtensor_func, sharded_tensor_func, tensor_func)
+            for v in iter_object
+        ]
+        if isinstance(iter_object, tuple):
+            ret = tuple(ret)  # type: ignore[assignment]
+        return ret

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/quantized_hf_storage.py

@@ -0,0 +1,506 @@
+# mypy: allow-untyped-defs
+import json
+import logging
+import math
+from pathlib import Path
+from typing import Any
+
+import torch
+from torch.distributed.checkpoint._hf_utils import _metadata_fn
+from torch.distributed.checkpoint.metadata import TensorStorageMetadata
+from torch.distributed.checkpoint.planner import LoadPlanner, ReadItem
+
+from .hf_storage import HuggingFaceStorageReader
+
+
+logger: logging.Logger = logging.getLogger(__name__)
+
+__all__ = ["QuantizedHuggingFaceStorageReader"]
+
+
+class QuantizedHuggingFaceStorageReader(HuggingFaceStorageReader):
+    """
+    Extension of HuggingFaceStorageReader that handles quantized tensors.
+    Checkpoint should have the full tensor in a SafeTensor file. The quantized
+    tensor should not be sharded across multiple files.
+
+    This reader handles the dequantization of tensors during the read process,
+    converting them from quantized blocks to full dequantized tensors before
+    copying to the target tensor.
+    """
+
+    def __init__(
+        self,
+        path: str,
+        thread_count: int = 1,
+        target_dtype: torch.dtype = torch.float32,
+        block_size: int = 128,
+    ):
+        """
+        Initialize the HuggingFace storage reader to load quantized checkpoints
+
+        Args:
+            path: directory where the checkpoint will be read from.
+            thread_count: Number of threads to use to read distributed checkpoint. Defaults to 1.
+            target_dtype: Target dtype for dequantized tensor. Defaults to torch.float32.
+            block_size: Fixed block size for dequantization. Defaults to 128.
+        """
+        super().__init__(path=path, thread_count=thread_count)
+
+        self.target_dtype: torch.dtype = target_dtype
+        self.block_size: int = block_size
+        self._weight_scale_mapping: dict[str, str] = {}
+        # Track which file contains each tensor
+        self._weight_map: dict[str, str] = {}
+        # Cache for full tensor shapes (fqn -> shape)
+        self._tensor_full_shapes: dict[str, torch.Size] = {}
+
+    def read_metadata(self) -> Any:
+        metadata = super().read_metadata()
+
+        # Load quantization metadata first.
+        self._load_quantization_metadata()
+
+        # Build a cache of FQN -> full tensor shape, correcting for quantized tensors.
+        for fqn, tensor_metadata in metadata.state_dict_metadata.items():
+            # Only process TensorStorageMetadata which has size attribute.
+            if isinstance(tensor_metadata, TensorStorageMetadata):
+                # Check if this is a MXFP4 quantized tensor that needs shape correction.
+                if fqn.endswith("_blocks"):
+                    # Save the quantized tensor shapes for lookup when dequantization.
+                    self._tensor_full_shapes[fqn + "_quantized"] = tensor_metadata.size
+                    *prefix_shape, G, B = tensor_metadata.size
+                    dequantized_size = torch.Size([*prefix_shape, G * B * 2])
+
+                    # Update the metadata with the size after dequantization.
+                    # Metadata used by planner to slice state dict.
+                    tensor_metadata.size = dequantized_size
+                    self._tensor_full_shapes[fqn] = dequantized_size
+                else:
+                    self._tensor_full_shapes[fqn] = tensor_metadata.size
+
+        return metadata
+
+    def _load_quantization_metadata(self):
+        """Load quantization metadata from the checkpoint."""
+        checkpoint_path = Path(self.path)
+        # Load weight mapping from index file
+        index_file = checkpoint_path / _metadata_fn
+
+        with open(index_file) as f:
+            index_data = json.load(f)
+            weight_map = index_data.get("weight_map", {})
+            self._build_weight_scale_mapping(weight_map)
+
+    def _build_weight_scale_mapping(self, weight_map: dict[str, str]):
+        """Analyze and build weight-scale tensor pairs from weight mapping."""
+        # Store the complete weight map for file location lookups.
+        self._weight_map = weight_map
+
+        for tensor_name in weight_map:
+            if tensor_name.endswith(".weight_scale_inv"):
+                weight_name = tensor_name.replace(".weight_scale_inv", ".weight")
+                if weight_name in weight_map:
+                    self._weight_scale_mapping[weight_name] = tensor_name
+            # Handle MXFP4 format: _blocks and _scales.
+            elif tensor_name.endswith("_scales"):
+                blocks_name = tensor_name.replace("_scales", "_blocks")
+                if blocks_name in weight_map:
+                    self._weight_scale_mapping[blocks_name] = tensor_name
+
+    def _process_read_request(
+        self, f: Any, req: ReadItem, planner: LoadPlanner
+    ) -> None:
+        """Override the Helper function that processes a single read request."""
+        tensor_fqn = req.storage_index.fqn
+
+        # Check if this is a quantized tensor that needs dequantization
+        if self._is_tensor_quantized(tensor_fqn):
+            tensor = self._read_quantized_tensor_with_block_alignment(req, f)
+        else:
+            # Standard tensor reading
+            slices = tuple(
+                slice(offset, offset + length)
+                for offset, length in zip(req.storage_offsets, req.lengths)
+            )
+            tensor = f.get_slice(tensor_fqn)[slices]
+
+        target_tensor = planner.resolve_tensor(req).detach()
+
+        if target_tensor.size() != tensor.size():
+            raise AssertionError(
+                f"req {req.storage_index} mismatch sizes {target_tensor.size()} vs {tensor.size()}"
+            )
+
+        target_tensor.copy_(tensor)
+        planner.commit_tensor(req, target_tensor)
+
+    def _get_slice_to_block_mapping(
+        self, req: ReadItem
+    ) -> tuple[tuple[int, int], tuple[int, int], slice, slice]:
+        """
+        Calculate which blocks correspond to the requested slice.
+
+        Args:
+            req: Read request containing tensor info and required slices
+
+        Returns:
+            Tuple of (row_block_range, col_block_range, row_slice, col_slice)
+        """
+        # Get the slice information
+        row_slice = slice(
+            req.storage_offsets[0], req.storage_offsets[0] + req.lengths[0]
+        )
+        col_slice = slice(
+            req.storage_offsets[1], req.storage_offsets[1] + req.lengths[1]
+        )
+
+        # Calculate which blocks this slice spans
+        row_start_block = row_slice.start // self.block_size
+        row_end_block = (row_slice.stop - 1) // self.block_size + 1  # Inclusive end
+
+        col_start_block = col_slice.start // self.block_size
+        col_end_block = (col_slice.stop - 1) // self.block_size + 1  # Inclusive end
+
+        return (
+            (row_start_block, row_end_block),
+            (col_start_block, col_end_block),
+            row_slice,
+            col_slice,
+        )
+
+    def _dequantize_tensor_mxfp4(
+        self,
+        blocks: torch.Tensor,
+        scales: torch.Tensor,
+        req: ReadItem,
+        group_start: int,
+        offset_in_first_group: int,
+    ) -> torch.Tensor:
+        """
+        Dequantize a 4D tensor using MXFP4 format.
+        Adapted from openai's implementation:
+        https://github.com/openai/gpt-oss/blob/8890e95919f975a490fc0ba09ffb10890ec7319d/gpt_oss/torch/weights.py#L68
+
+        Args:
+            blocks: Sliced quantized weight tensor of shape [a_slice, b_slice, groups_slice, B] in uint8
+            scales: FULL scale tensor of shape [a, b, c] in uint8 (will be converted to exponents)
+            req: Read request containing slice information
+            group_start: The starting group index in the checkpoint
+            offset_in_first_group: Offset in values within the first group
+
+        Returns:
+            Dequantized tensor matching the requested shape
+        """
+        # FP4 lookup table
+        FP4_VALUES = [
+            +0.0,
+            +0.5,
+            +1.0,
+            +1.5,
+            +2.0,
+            +3.0,
+            +4.0,
+            +6.0,
+            -0.0,
+            -0.5,
+            -1.0,
+            -1.5,
+            -2.0,
+            -3.0,
+            -4.0,
+            -6.0,
+        ]
+
+        # blocks: [a_slice, b_slice, groups_slice, B] uint8.
+        # Read slightly more groups than needed, and slice at the end.
+
+        # Slice the scales to match the blocks dimensions.
+        # [a_full, b_full, c_full] -> [a_slice, b_slice, groups_slice]
+        dim0_start = req.storage_offsets[0]
+        dim0_end = dim0_start + req.lengths[0]
+        dim1_start = req.storage_offsets[1]
+        dim1_end = dim1_start + req.lengths[1]
+        num_groups = blocks.shape[2]
+        scales = scales[
+            dim0_start:dim0_end,
+            dim1_start:dim1_end,
+            group_start : group_start + num_groups,
+        ]
+
+        scales = scales.to(torch.int32) - 127
+
+        if blocks.shape[:-1] != scales.shape:
+            raise AssertionError(f"{blocks.shape=} does not match {scales.shape=}")
+
+        lut = torch.tensor(FP4_VALUES, dtype=self.target_dtype, device=blocks.device)
+
+        *prefix_shape, G, B = blocks.shape
+        rows_total = math.prod(prefix_shape) * G
+
+        blocks = blocks.reshape(rows_total, B)
+        scales = scales.reshape(rows_total, 1)
+
+        out = torch.empty(
+            rows_total, B * 2, dtype=self.target_dtype, device=blocks.device
+        )
+
+        rows_per_chunk = 16384 * 512
+
+        for r0 in range(0, rows_total, rows_per_chunk):
+            r1 = min(r0 + rows_per_chunk, rows_total)
+
+            blk = blocks[r0:r1]
+            exp = scales[r0:r1]
+
+            # nibble indices -> int64
+            idx_lo = (blk & 0x0F).to(torch.long)
+            idx_hi = (blk >> 4).to(torch.long)
+
+            sub = out[r0:r1]
+            sub[:, 0::2] = lut[idx_lo]
+            sub[:, 1::2] = lut[idx_hi]
+
+            torch.ldexp(sub, exp, out=sub)
+
+            del idx_lo, idx_hi, blk, exp
+
+        result = out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2)
+
+        # Slice the last dimension to match the requested range.
+        if offset_in_first_group > 0 or result.shape[-1] > req.lengths[2]:
+            end_offset = offset_in_first_group + req.lengths[2]
+            result = result[..., offset_in_first_group:end_offset]
+
+        return result
+
+    def _dequantize_tensor(
+        self,
+        weight: torch.Tensor,
+        scale_inv: torch.Tensor,
+        full_tensor_shape: torch.Size,
+        slice_info: tuple[tuple[int, int], tuple[int, int], slice, slice],
+    ) -> torch.Tensor:
+        """
+        Dequantize a sliced tensor using the appropriate portion of the scale tensor.
+
+        Args:
+            weight: Sliced quantized weight tensor
+            scale_inv: Full scale inverse tensor for dequantization
+            full_tensor_shape: Shape of the original full tensor
+            slice_info: Block mapping information from _get_slice_to_block_mapping
+
+        Returns:
+            Dequantized tensor
+        """
+        (row_block_range, col_block_range, row_slice, col_slice) = slice_info
+
+        # Convert to float32 for computation
+        # Certain quantized dtypes like Float8_e4m3fn
+        # don't support multiplication on CPU yet in PyTorch.
+        upcasted_weight = weight.to(torch.float32)
+
+        # Create output tensor in target dtype
+        dequantized = weight.detach().to(dtype=self.target_dtype, copy=True)
+
+        # Get the actual slice boundaries
+        row_start_global = row_slice.start
+        row_end_global = row_slice.stop
+        col_start_global = col_slice.start
+        col_end_global = col_slice.stop
+
+        # Apply scaling factors to each block that intersects with our slice
+        for block_i in range(row_block_range[0], row_block_range[1]):
+            for block_j in range(col_block_range[0], col_block_range[1]):
+                # Calculate the block boundaries in global coordinates
+                block_row_start_global = block_i * self.block_size
+                block_row_end_global = min(
+                    block_row_start_global + self.block_size, full_tensor_shape[0]
+                )
+                block_col_start_global = block_j * self.block_size
+                block_col_end_global = min(
+                    block_col_start_global + self.block_size, full_tensor_shape[1]
+                )
+
+                # Find the intersection of the block with our slice
+                intersect_row_start = max(block_row_start_global, row_start_global)
+                intersect_row_end = min(block_row_end_global, row_end_global)
+                intersect_col_start = max(block_col_start_global, col_start_global)
+                intersect_col_end = min(block_col_end_global, col_end_global)
+
+                # Skip if no intersection
+                if (
+                    intersect_row_start >= intersect_row_end
+                    or intersect_col_start >= intersect_col_end
+                ):
+                    continue
+
+                # Convert global coordinates to local coordinates in the sliced tensor
+                local_row_start = intersect_row_start - row_start_global
+                local_row_end = intersect_row_end - row_start_global
+                local_col_start = intersect_col_start - col_start_global
+                local_col_end = intersect_col_end - col_start_global
+
+                # Get the block from the sliced tensor
+                block = upcasted_weight[
+                    local_row_start:local_row_end, local_col_start:local_col_end
+                ]
+
+                # Apply the scale factor
+                scale = scale_inv[block_i, block_j]
+                block = block * scale
+
+                # Convert block to target dtype and store
+                block_converted = block.to(dtype=self.target_dtype)
+                dequantized[
+                    local_row_start:local_row_end, local_col_start:local_col_end
+                ] = block_converted
+
+        return dequantized
+
+    def _is_tensor_quantized(self, tensor_fqn: str) -> bool:
+        """
+        Check if a tensor is a quantized.
+
+        Args:
+            tensor_fqn: Fully qualified name of the tensor
+
+        Returns:
+            True if tensor is quantized and has a corresponding scale tensor,
+            False otherwise
+        """
+        # Skip scale tensors themselves
+        if tensor_fqn.endswith((".weight_scale_inv", "_scales")):
+            return False
+
+        # Check if this weight tensor has a corresponding scale tensor
+        if tensor_fqn not in self._weight_scale_mapping:
+            return False
+
+        return True
+
+    def _read_quantized_tensor_with_block_alignment(
+        self, req: ReadItem, safetensor_file: Any
+    ) -> torch.Tensor:
+        """
+        Read a quantized tensor with block alignment.
+
+        Args:
+            req: Read request containing tensor info and required slices
+            safetensor_file: Open safetensors file handle
+
+        Returns:
+            Dequantized tensor ready for use
+        """
+        tensor_fqn = req.storage_index.fqn
+        scale_fqn = self._weight_scale_mapping[tensor_fqn]
+
+        try:
+            group_start = 0
+            offset_in_first_group = 0
+            if tensor_fqn.endswith("_blocks"):
+                # Full tensor is a 4D MXFP4 quantized tensor: [..., G, B].
+                # Each group G produces B * 2 dequantized values.
+                # Checkpoint [..., G, B] -> dequantized [..., G*B*2].
+
+                # The planner gives 3D requests based on the dequantized shape.
+                # Need to figure out which groups (dimension 2 in checkpoint) to read.
+
+                # Use the quantized checkpoint shape to get the correct B.
+                *prefix_shape, B = self._tensor_full_shapes[tensor_fqn + "_quantized"]
+                values_per_group = B * 2  # Each byte has 2 nibbles (4-bit values).
+
+                # Calculate which groups we need based on the requested range in dim 2.
+                # Ensure the reequest is in 3D.
+                if len(req.storage_offsets) != 3:
+                    raise AssertionError
+
+                # Positions in dequantized space.
+                dim2_start_deq = req.storage_offsets[2]
+                dim2_length_deq = req.lengths[2]
+                dim2_end_deq = dim2_start_deq + dim2_length_deq
+
+                # Convert to group indices.
+                group_start = dim2_start_deq // values_per_group
+                group_end = (dim2_end_deq + values_per_group - 1) // values_per_group
+
+                # Read only the necessary groups from checkpoint.
+                weight_slices_4d = (
+                    slice(
+                        req.storage_offsets[0], req.storage_offsets[0] + req.lengths[0]
+                    ),
+                    slice(
+                        req.storage_offsets[1], req.storage_offsets[1] + req.lengths[1]
+                    ),
+                    slice(group_start, group_end),
+                    slice(None),  # Read all B values for each group.
+                )
+                quantized_tensor = safetensor_file.get_slice(tensor_fqn)[
+                    weight_slices_4d
+                ]
+
+                # Also track the offset within the first group
+                offset_in_first_group = dim2_start_deq - (
+                    group_start * values_per_group
+                )
+            else:
+                # 2D quantized tensor, use 2d block partition.
+                weight_slices = tuple(
+                    slice(offset, offset + length)
+                    for offset, length in zip(req.storage_offsets, req.lengths)
+                )
+                quantized_tensor = safetensor_file.get_slice(tensor_fqn)[weight_slices]
+
+            # Load the corresponding scale inverse tensor (full tensor)
+            scale_file_name = self._weight_map.get(scale_fqn)
+            if scale_file_name is None:
+                raise ValueError(f"Scale tensor {scale_fqn} not found in weight_map")
+
+            # Check if scale tensor is in the same file as the weight tensor
+            weight_file_name = self._weight_map.get(tensor_fqn)
+
+            if scale_file_name == weight_file_name:
+                # Scale tensor is in the same file, use current handle
+                scale_inv = safetensor_file.get_tensor(scale_fqn)
+            else:
+                # Scale tensor is in a different file, need to open it
+                from safetensors import safe_open  # type: ignore[import]
+
+                scale_file_path = Path(self.path) / scale_file_name
+                with safe_open(
+                    scale_file_path, framework="pt", device="cpu"
+                ) as scale_file:
+                    scale_inv = scale_file.get_tensor(scale_fqn)
+
+            # Get the full tensor shape from our O(1) lookup cache
+            full_tensor_shape = self._tensor_full_shapes.get(tensor_fqn)
+            if full_tensor_shape is None:
+                raise ValueError(f"Could not find full tensor shape for {tensor_fqn}")
+
+            # Determine which dequantization function to use.
+            if len(full_tensor_shape) == 2:
+                # 2D block-wise quantization, e.g., used in deepseek v3.1
+                slice_info = self._get_slice_to_block_mapping(req)
+                dequantized_tensor = self._dequantize_tensor(
+                    weight=quantized_tensor,
+                    scale_inv=scale_inv,
+                    full_tensor_shape=full_tensor_shape,
+                    slice_info=slice_info,
+                )
+            elif tensor_fqn.endswith("_blocks"):
+                # 4D with blocks along dimension 2, used in MXFP4, e.g. gpt-oss
+                dequantized_tensor = self._dequantize_tensor_mxfp4(
+                    blocks=quantized_tensor,
+                    scales=scale_inv,
+                    req=req,
+                    group_start=group_start,
+                    offset_in_first_group=offset_in_first_group,
+                )
+            else:
+                raise ValueError("Unsupported quantization types")
+
+            return dequantized_tensor
+
+        except Exception as e:
+            logger.error("Failed to read the quantized tensor!!")
+            raise e

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/resharding.py

@@ -0,0 +1,69 @@
+from torch.distributed.checkpoint.metadata import ChunkStorageMetadata
+
+
+__all__: list[str] = []
+
+
+def _check_shard_metadata_pair_overlap(
+    shard1: ChunkStorageMetadata, shard2: ChunkStorageMetadata
+) -> bool:
+    """Check if two shards overlap."""
+    # For each dim of each shard, check if one shard resides on the other
+    # end of second shard with respect to that dim. As an example for a 2D
+    # shard, we would check if one shard is above or on the left of the
+    # other shard.
+    ndims = len(shard1.offsets)
+    for i in range(ndims):
+        if shard1.offsets[i] >= shard2.offsets[i] + shard2.sizes[i]:
+            return False
+        if shard2.offsets[i] >= shard1.offsets[i] + shard1.sizes[i]:
+            return False
+
+    return True
+
+
+def _shards_get_overlap_region_wrt_saved_tensor(
+    saved_shard: ChunkStorageMetadata, current_shard: ChunkStorageMetadata
+) -> list[tuple[int, int, int, int]]:
+    """
+    Return the overlapping region between saved_shard and current_shard.
+
+    There returned list has the same number of elements as the tensor's dimension.
+    For each element, we produce a tuple with the following contents:
+        (dimension, `saved_shard` offset, `current_shard` offset, length)
+
+    Offsets are relative to each shard.
+    """
+    narrows = []
+    for dim, (
+        saved_shard_offset,
+        current_shard_offset,
+        saved_shard_size,
+        current_shard_size,
+    ) in enumerate(
+        zip(
+            saved_shard.offsets,
+            current_shard.offsets,
+            saved_shard.sizes,
+            current_shard.sizes,
+        )
+    ):
+        min_range_end = min(
+            saved_shard_offset + saved_shard_size,
+            current_shard_offset + current_shard_size,
+        )
+
+        length = min_range_end - max(current_shard_offset, saved_shard_offset)
+
+        if saved_shard_offset > current_shard_offset:
+            offset_for_saved_tensor = 0
+            offset_for_current_tensor = saved_shard_offset - current_shard_offset
+        else:
+            offset_for_saved_tensor = current_shard_offset - saved_shard_offset
+            offset_for_current_tensor = 0
+
+        narrows.append(
+            (dim, offset_for_saved_tensor, offset_for_current_tensor, length)
+        )
+
+    return narrows

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/staging.py

@@ -0,0 +1,477 @@
+import os
+import tempfile
+from concurrent.futures import Future, ThreadPoolExecutor
+from contextlib import nullcontext
+from dataclasses import dataclass
+from datetime import timedelta
+from typing import Any, cast
+from typing_extensions import deprecated, Protocol, runtime_checkable
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+from torch.distributed._state_dict_utils import _copy_state_dict, _create_cpu_state_dict
+from torch.distributed.checkpoint._pg_transport import PGTransport
+from torch.distributed.checkpoint._state_dict_stager import StateDictStager
+from torch.distributed.checkpoint.metadata import STATE_DICT_TYPE
+
+
+__all__ = ["AsyncStager", "BlockingAsyncStager", "DefaultStager", "StagingOptions"]
+
+"""
+Experimental staging module for PyTorch Distributed Checkpointing.
+This module provides advanced staging capabilities for checkpoints including:
+- Asynchronous staging using ThreadPoolExecutor
+- Pinned memory allocation for faster CPU-GPU transfers
+- Shared memory support for multi-process scenarios
+- Non-blocking CUDA operations with stream synchronization
+- Caching of frequently used storages for efficient memory management
+- Automatic resource cleanup and memory management
+Classes:
+    AsyncStager: Protocol defining the staging interface
+    StagingOptions: Configuration dataclass for staging behavior
+    DefaultStager: Default implementation with comprehensive staging features
+    BlockingAsyncStager: Implementation of AsyncStager which stages the state_dict
+    on CPU RAM and blocks until the copy is complete. Please use DefaultStager instead.
+"""
+
+
+@runtime_checkable
+class AsyncStager(Protocol):
+    """
+    This protocol is meant to provide customization and extensibility for dcp.async_save, allowing users
+    to customize how data is staged previous to executing the usual dcp.save path in parallel.
+    The expected order of operations (concretely defined in `torch.distributed.state_dict_saver.async_save`)
+    is the following:
+
+    1. AsyncStager.stage_data(state_dict):
+        This call gives the AsyncStager the opportunity to 'stage'
+        the state_dict. The expectation and purpose of staging in this context is to create a "training-safe"
+        representation of the state dict, meaning that any updates to module data after staging is complete
+        should not be reflected in the state dict returned from this method. For example, in the default
+        case a copy of the entire state dict is created on CPU RAM and returned here, allowing users
+        to continue training without risking changes to data which is being serialized.
+
+    2. dcp.save is called on the state_dict returned from stage in parallel. This call is responsible
+        for serializing the state_dict and writing it to storage.
+
+    3. If AsyncStager.should_synchronize_after_execute is True, this method will be called immediately after
+        the serialization thread starts and before returning from dcp.async_save. If this is set to False,
+        the assumption is the user has defined a custom synchronization point for the purpose of further
+        optimizing save latency in the training loop (for example, by overlapping staging with the
+        forward/backward pass), and it is the respondsibility of the user to call `AsyncStager.synchronize_staging`
+        at the appropriate time.
+
+    """
+
+    # default to True since the common case is to stage synchronously
+    _synchronize_after_execute: bool = True
+
+    @property
+    def should_synchronize_after_execute(self) -> bool:
+        """
+        Whether to synchronize after executing the stage.
+        """
+        return self._synchronize_after_execute
+
+    def stage(
+        self, state_dict: STATE_DICT_TYPE
+    ) -> Future[STATE_DICT_TYPE] | STATE_DICT_TYPE:
+        """
+        Returns a "staged" copy of `state_dict`. The expectation of the staged copy is that it is
+        inoculated from any updates incurred after the stage call is complete.
+        """
+        raise NotImplementedError(
+            f"{self.__class__.__name__} must implement stage method"
+        )
+
+    @deprecated(
+        "`synchronize_staging` is deprecated and will be removed in future versions."
+        "Please use staging_future from AsyncSaveResponse instead.",
+        category=FutureWarning,
+    )
+    def synchronize_staging(self) -> None:
+        """
+        In the case `stage` is async in some way, this method should be called to ensure staging
+        is complete and it is safe to begin modifying the original `state_dict`
+        """
+
+    def close(self) -> None:
+        """
+        Clean up all resources used by the stager.
+        """
+
+
+@dataclass
+class StagingOptions:
+    """
+    Configuration options for checkpoint staging behavior.
+
+    Attributes:
+        use_pinned_memory (bool): Enable pinned memory allocation for faster
+            CPU-GPU transfers. Requires CUDA to be available. Default: True
+        use_shared_memory (bool): Enable shared memory for multi-process
+            scenarios. Useful when multiple processes need access to the
+            same staged data. Default: True
+        use_async_staging (bool): Enable asynchronous staging using a
+            background thread pool. Allows overlapping computation with
+            staging operations. Requires CUDA. Default: True
+        use_non_blocking_copy (bool): Use non-blocking device memory
+            copies with stream synchronization. Improves performance by
+            allowing CPU work to continue during GPU transfers. Default: True
+
+    Note:
+        CUDA-dependent features will raise exception if CUDA is not available.
+    """
+
+    use_pinned_memory: bool = True
+    use_shared_memory: bool = True
+    use_async_staging: bool = True
+    use_non_blocking_copy: bool = True
+
+
+class DefaultStager(AsyncStager):
+    """
+    DefaultStager provides a full-featured staging implementation that combines
+    multiple optimization techniques for efficient checkpoint preparation.
+
+    The staging process works as follows:
+    1. State dictionary is submitted for staging (sync or async)
+    2. Tensors are copied from GPU to optimized CPU storage
+    3. CUDA operations are synchronized if non-blocking copies are used
+    4. Staged state dictionary is returned or made available via Future
+
+    Usage Patterns:
+        # Synchronous staging
+        stager = DefaultStager(StagingOptions(use_async_staging=False))
+        staged_dict = stager.stage(state_dict)
+        stager.close()
+
+        # Asynchronous staging
+        stager = DefaultStager(StagingOptions(use_async_staging=True))
+        future = stager.stage(state_dict)
+        # ... do other work ...
+        staged_dict = future.result()
+        stager.close()
+
+        # Context manager pattern (recommended)
+        stager = DefaultStager(config)
+        with stager:
+        result = stager.stage(state_dict)
+
+    Performance Considerations:
+        - Async staging provides best performance when model computation
+          can overlap with staging operations
+        - Pinned memory improves CPU-GPU transfer speeds but uses more memory
+        - Shared memory allows efficient IPC to checkpoint process
+        - Non-blocking copies reduce GPU idle time during memory transfers
+
+    Thread Safety:
+        DefaultStager is not thread-safe. Each thread should use its own
+        instance, or external synchronization should be provided.
+    """
+
+    def __init__(
+        self,
+        config: StagingOptions = StagingOptions(),
+    ):
+        self._config = config
+        self._state_dict_stager = StateDictStager(
+            pin_memory=config.use_pinned_memory, share_memory=config.use_shared_memory
+        )
+        self._staging_executor = None
+        self._staging_stream = None
+        if self._config.use_async_staging:
+            self._staging_executor = ThreadPoolExecutor(max_workers=1)
+            if torch.accelerator.is_available():
+                # Note: stream needs to be initialized on the main thread after default cuda
+                # stream is setup/used to avoid the risk of accidentally reusing the main
+                # compute stream or in other cases kernels actually launching from the
+                # main thread.
+                self._staging_stream = torch.Stream()
+
+        if self._config.use_non_blocking_copy:
+            if not torch.accelerator.is_available():
+                raise AssertionError(
+                    "Non-blocking copy requires that the current accelerator is available."
+                )
+
+        self._staging_future: Future[STATE_DICT_TYPE] | None = None
+
+    def stage(
+        self,
+        state_dict: STATE_DICT_TYPE,
+        **kwargs: Any,
+    ) -> STATE_DICT_TYPE | Future[STATE_DICT_TYPE]:
+        """
+        This function is responsible for staging staging the state_dict.
+        See class docstring for more details on staging.
+        If use_async_staging is True, it will return a Future object that will be
+        fulfilled when staging is complete.
+        If use_async_staging is False, it will return the fully staged state_dict.
+
+        Args:
+            state_dict (STATE_DICT_TYPE): The state_dict to be staged.
+        """
+        if self._config.use_async_staging:
+            if self._staging_executor is None:
+                raise AssertionError(
+                    "staging_executor should not be None for async staging"
+                )
+            self._staging_future = self._staging_executor.submit(
+                self._stage,
+                state_dict,
+                **kwargs,
+            )
+            return self._staging_future
+        else:
+            return self._stage(state_dict, **kwargs)
+
+    def _stage(self, state_dict: STATE_DICT_TYPE, **kwargs: Any) -> STATE_DICT_TYPE:
+        if self._config.use_non_blocking_copy:
+            if not (self._staging_stream or not self._config.use_async_staging):
+                raise AssertionError(
+                    "Non-blocking copy in a background thread for async staging needs staging_stream to be initialized."
+                )
+            with (
+                self._staging_stream
+                if self._staging_stream is not None
+                else nullcontext()
+            ):
+                state_dict = self._state_dict_stager.stage(
+                    state_dict, non_blocking=self._config.use_non_blocking_copy
+                )
+            # waits for the enqued copy operations to finish.
+            self._staging_stream.synchronize() if self._staging_stream else torch.accelerator.synchronize()
+        else:
+            state_dict = self._state_dict_stager.stage(state_dict, non_blocking=False)
+
+        # release reference cycle to prevent memory leaks in async_save
+        # created by _deepcopy_dispatch that capture self
+        self._state_dict_stager.close()
+
+        return state_dict
+
+    def close(self) -> None:
+        """
+        Clean up all resources used by the DefaultStager. Shuts down the ThreadPoolExecutor
+        used for async staging operations and cleans up the underlying StateDictStager's
+        cached storages. Should be called when the stager is no longer needed to prevent
+        resource leaks, especially in long-running applications. After calling close(),
+        the stager should not be used for further staging operations.
+
+        Example Usage:
+            stager = DefaultStager(StagingOptions(use_async_staging=True))
+            future = stager.stage(state_dict)
+            result = future.result()
+            stager.close()  # Clean up all resources
+        """
+        if self._staging_executor:
+            self._staging_executor.shutdown(wait=True)
+        self._state_dict_stager.close()
+
+    def synchronize_staging(self) -> None:
+        """
+        When use_async_staging is True, this method will wait until staging is complete.
+        If use_async_staging is False, this method is a no-op.
+        """
+        if self._staging_future is not None:
+            self._staging_future.result()
+
+
+class BlockingAsyncStager(AsyncStager):
+    """
+    An implementation of AsyncStager which stages the state_dict on CPU RAM and blocks until the copy is complete.
+    This implementation also provides an option to optimize stage latency using pinned memory.
+
+    N.B. synchronize_staging is a no-op in this case.
+
+
+    """
+
+    # default to True since the common case is to stage synchronously
+    _synchronize_after_execute: bool = False
+
+    def __init__(
+        self,
+        cache_staged_state_dict: bool = False,
+        type_check: bool = False,
+    ):
+        """
+        Initializes the BlockingAsyncStager.
+
+        Args:
+            cache_staged_state_dict: Whether to cache the staged state_dict. This option decreases staging latency
+                at the cost of increases memory usage. Additionally, if this parameter is set to True, it's the expectation
+                that the stager is maintained and reused for multiple dcp.async_save calls. Default to False.
+            type_check: Whether to perform a type check during cpu_offload. Defaults to False.
+
+        """
+        self.cache_staged_state_dict = cache_staged_state_dict
+        self.type_check = type_check
+        self.state_dict_cache: STATE_DICT_TYPE | None = None
+
+    def stage(self, state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE:
+        """
+        Returns a copy of `state_dict` on the CPU.
+        """
+
+        if not self.cache_staged_state_dict:
+            staged_state_dict = _create_cpu_state_dict(state_dict)
+            _copy_state_dict(state_dict, staged_state_dict, type_check=self.type_check)
+            return staged_state_dict
+
+        if self.state_dict_cache is None:
+            self.state_dict_cache = _create_cpu_state_dict(state_dict, pin_memory=True)
+        return _copy_state_dict(state_dict, self.state_dict_cache)
+
+    def synchronize_staging(self) -> None:
+        """
+        No-op function, since staging is blocking.
+        """
+
+    def close(self) -> None:
+        pass
+
+
+class _ReplicationStager(AsyncStager):
+    """
+    An AsyncStager implementation that replicates state_dict across training ranks
+    using PGTransport.
+
+    Args:
+        pg: ProcessGroup for distributed communication
+        timeout: Timeout for communication operations
+        device: Device to use for tensor operations
+        storage_dir: Directory to store persisted state_dicts
+
+    Warning: This is experimental and subject to change.
+    """
+
+    _synchronize_after_execute: bool = False
+
+    def __init__(
+        self,
+        pg: ProcessGroup,
+        timeout: timedelta = timedelta(minutes=30),
+        device: torch.device = torch.device("cpu"),
+        storage_dir: str | None = None,
+    ):
+        self._pg = pg
+        self._timeout = timeout
+        self._device = device
+        self._transport = PGTransport(pg, timeout, device, None)
+
+        # Set up storage directory for persisting exchanged state_dicts
+        if storage_dir is None:
+            self._storage_dir = tempfile.mkdtemp(prefix="replication_stager_")
+        else:
+            self._storage_dir = storage_dir
+        os.makedirs(self._storage_dir, exist_ok=True)
+
+    def stage(
+        self, state_dict: STATE_DICT_TYPE
+    ) -> Future[STATE_DICT_TYPE] | STATE_DICT_TYPE:
+        """
+        Stage the state_dict by replicating it across ranks. Returns a state_dict representing
+        the received replica.
+
+        Perform the actual replication logic. Creates bidirectional pairs where each rank exchanges
+        state_dict with its partner at (rank + world_size//2) % world_size.
+        Uses simple rank-based ordering to prevent deadlocks.
+
+        Assumes world_size is always even.
+        """
+        if not dist.is_initialized():
+            return state_dict
+
+        world_size = dist.get_world_size()
+
+        current_rank = dist.get_rank()
+
+        # Calculate partner rank using half-world offset
+        # creates bidirectional pairs for replication.
+        offset = world_size // 2
+        partner_rank = (current_rank + offset) % world_size
+
+        # Use simple rank-based ordering to prevent deadlocks.
+        # Lower-numbered rank sends first, higher-numbered rank receives first.
+        if current_rank < partner_rank:
+            # Send first, then receive
+            self._transport.send_checkpoint([partner_rank], state_dict)
+            received_state_dict = self._transport.recv_checkpoint(partner_rank)
+        else:
+            # Receive first, then send
+            received_state_dict = self._transport.recv_checkpoint(partner_rank)
+            self._transport.send_checkpoint([partner_rank], state_dict)
+
+        # Persist the received state_dict for future discoverability
+        received_state_dict = cast(STATE_DICT_TYPE, received_state_dict)
+        self._persist_state_dict(received_state_dict, current_rank, partner_rank)
+
+        return received_state_dict
+
+    def _persist_state_dict(
+        self, state_dict: STATE_DICT_TYPE, current_rank: int, partner_rank: int
+    ) -> None:
+        """
+        Persist the received state_dict to disk for future discoverability.
+        Only keeps one replica per rank, overwriting any previous replica.
+        Uses atomic write pattern (temp file + rename).
+
+        Args:
+            state_dict: The state_dict received from partner rank
+            current_rank: Current rank that received the state_dict
+            partner_rank: Rank that sent the state_dict
+        """
+        final_path = self._get_persisted_path(current_rank, partner_rank)
+        temp_path = final_path + ".tmp"
+
+        try:
+            # Ensure parent directory exists and is writable
+            os.makedirs(os.path.dirname(final_path), exist_ok=True)
+
+            # Write to temporary file with explicit flushing
+            with open(temp_path, "wb") as f:
+                torch.save(state_dict, f)
+                # Flush application buffers to OS buffers
+                f.flush()
+                # Force OS buffers to disk for durability
+                os.fsync(f.fileno())
+
+            # Atomic rename to final location
+            os.rename(temp_path, final_path)
+        except Exception as e:
+            # Clean up temp file if it exists
+            try:
+                if os.path.exists(temp_path):
+                    os.remove(temp_path)
+            except Exception:
+                pass  # Ignore cleanup errors
+            # Re-raise the original exception with more context
+            raise RuntimeError(
+                f"Failed to persist state_dict from rank {partner_rank} to rank {current_rank}: {e}"
+            ) from e
+
+    def _get_persisted_path(self, current_rank: int, partner_rank: int) -> str:
+        """
+        Get the file path where a state_dict would be persisted.
+
+        Args:
+            current_rank: Current rank
+
+        Returns:
+            File path for the persisted state_dict
+        """
+        filename = f"rank_{current_rank}_replica_partner_{partner_rank}.pt"
+        return os.path.join(self._storage_dir, filename)
+
+    def synchronize_staging(self) -> None:
+        """
+        No-op function, since staging is blocking.
+        """
+
+    def close(self) -> None:
+        """
+        Clean up resources. Persisted files are intentionally left for future discovery.
+        """

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_loader.py

@@ -0,0 +1,401 @@
+# mypy: allow-untyped-decorators
+# mypy: allow-untyped-defs
+import inspect
+import logging
+import os
+import warnings
+from typing import Any, cast, TYPE_CHECKING
+from typing_extensions import deprecated
+
+import torch
+import torch.distributed as dist
+from torch.distributed.checkpoint.default_planner import _EmptyStateDictLoadPlanner
+from torch.distributed.checkpoint.logger import _dcp_method_logger
+from torch.distributed.checkpoint.stateful import Stateful
+
+from ._storage_utils import _storage_setup
+from .default_planner import DefaultLoadPlanner
+from .planner import LoadPlan, LoadPlanner
+from .storage import StorageReader
+from .utils import _api_bc_check, _DistWrapper, _profile
+
+
+if TYPE_CHECKING:
+    from torch.distributed.checkpoint.metadata import Metadata
+
+__all__ = ["load_state_dict", "load"]
+
+logger = logging.getLogger()
+
+
+@deprecated(
+    "`load_state_dict` is deprecated and will be removed in future versions. "
+    "Please use `load` instead.",
+    category=FutureWarning,
+)
+def load_state_dict(
+    state_dict: dict[str, Any],
+    storage_reader: StorageReader,
+    process_group: dist.ProcessGroup | None = None,
+    coordinator_rank: int = 0,
+    no_dist: bool = False,
+    planner: LoadPlanner | None = None,
+) -> None:
+    """This method is deprecated. Please switch to 'load'."""
+    storage_reader.reset()
+    with _profile():
+        # TODO: test returning `load` here instead.
+        return _load_state_dict(
+            state_dict,
+            storage_reader,
+            process_group,
+            coordinator_rank,
+            no_dist,
+            planner,
+        )
+
+
+@_dcp_method_logger(log_exceptions=True)
+@_api_bc_check
+def load(
+    state_dict: dict[str, Any],
+    *,
+    checkpoint_id: str | os.PathLike | None = None,
+    storage_reader: StorageReader | None = None,
+    planner: LoadPlanner | None = None,
+    process_group: dist.ProcessGroup | None = None,
+    no_dist: bool = False,
+) -> None:
+    """
+    Load a checkpoint into a distributed state dict in SPMD style.
+
+    Each rank must have the same keys in their ``state_dict`` provided to this
+    API. Mismatched keys may result in hangs or errors. If unsure, you can use
+    the ``utils._assert_same_keys`` API to check (but may incur communication
+    costs).
+
+    Each rank will try to read the least amount of data necessary
+    to fulfill the requested `state_dict`. When loading :class:`ShardedTensor`
+    or :class:`DTensor` instances, each rank only reads data for their local shards.
+
+    For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``),
+    load will first call ``state_dict`` before attempting deserialization, followed by
+    ``load_state_dict`` once the deserialization is complete.
+    For each non-``Stateful`` object, load will deserialize the object, and then replace
+    it in the ``state_dict`` with the deserialized object.
+
+    .. warning::
+        All tensors in ``state_dict`` must be allocated on their
+        destination device *prior to* calling this function.
+
+        All non-tensor data is loaded using `torch.load()` and modified in place
+        on state_dict.
+
+    .. warning::
+        Users must call `load_state_dict` on the root module to ensure load
+        pos-processing and non-tensor data properly propagates.
+
+    .. note:
+        If no process group is initialized, this function will assume the intent
+        is to load a checkpoint into the local process. This can be useful in the
+        case of local inference, and when using regular Tensors (as opposed to DTensor
+         or ShardedTensor)
+
+    .. note:
+        Rank 0 is assumed to be the coordinator rank.
+
+    Args:
+        state_dict (Dict[str, Any]): The state_dict to load the checkpoint into.
+        checkpoint_id (Union[str, os.PathLike, None]):
+            The ID of this checkpoint instance. The meaning of the checkpoint_id
+            depends on the storage. It can be a path to a folder or to a file.
+            It can also be a key if the storage is a key-value store.
+            (Default: ``None``)
+        storage_reader (Optional[StorageReader]):
+            Instance of StorageWriter used to perform reads. If this is not
+            specified, DCP will automatically infer the reader based on the
+            checkpoint_id. If checkpoint_id is also None, an exception will
+            be raised. (Default: ``None``)
+        planner (Optional[LoadPlanner]):
+            Instance of LoadPlanner. If this is not specified, the default
+            planner will be used. (Default: ``None``)
+        process_group (Optional[ProcessGroup]):
+            ProcessGroup to be used for cross-rank synchronization.
+            (Default: ``None``)
+        no_dist (bool): If ``True``, this function will assume the intent is to load
+            a checkpoint without using cross-rank synchronization. (Default: ``False``)
+    Returns:
+        None.
+
+    Examples
+        >>> # xdoctest: +SKIP
+        >>> my_model = MyModule()
+        >>> optimizer = Adagrad(my_model.parameters())
+        >>> model_state_dict = my_model.state_dict()
+        >>> fs_storage_reader = torch.distributed.checkpoint.FileSystemReader(
+        ...     "/checkpoint/1"
+        ... )
+
+        >>> torch.distributed.checkpoint.load_state_dict(
+        >>>     state_dict=model_state_dict,
+        >>>     storage_reader=fs_storage_reader,
+        >>> )
+
+        >>> # module.load_state_dict() function might have customized steps
+        >>> # to flush the state_dict, must call it to
+        >>> # ensure correct behavior.
+        >>> my_model.load_state_dict(model_state_dict)
+
+    .. note::
+        load_state_dict uses collectives to coordinate reads across ranks.
+        For NCCL-based process groups, internal tensor representations of
+        objects must be moved to the GPU device before communication takes place.
+        In this case, the device used is given by ``torch.cuda.current_device()``
+        and it is the user's responsibility to ensure that this is set so that each
+        rank has an individual GPU, via ``torch.cuda.set_device()``.
+    """
+
+    no_dist = no_dist or (not dist.is_available()) or (not dist.is_initialized())
+    if no_dist:
+        warnings.warn(
+            "torch.distributed is disabled, unavailable or uninitialized, assuming the intent is to load in a single process.",
+            stacklevel=2,
+        )
+
+    with _profile():
+        storage_reader = cast(
+            StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True)
+        )
+
+        # All ranks must have the same keys in their `state_dict` provided to
+        # this API.  See documentation for more details.
+        # Here we simply sort the keys to ensure that all ranks load values in
+        # the same order.
+        keys = sorted(state_dict.keys())
+
+        stateful_sd = {}
+        for key in keys:
+            if key not in state_dict:
+                continue
+            elem = state_dict[key]
+            stateful_sd[key] = elem.state_dict() if isinstance(elem, Stateful) else elem
+
+        _load_state_dict(
+            state_dict=stateful_sd,
+            storage_reader=storage_reader,
+            process_group=process_group,
+            no_dist=no_dist,
+            planner=planner,
+        )
+        for key in keys:
+            if key not in state_dict:
+                continue
+            elem = state_dict[key]
+            if isinstance(elem, Stateful):
+                # If the state_dict is a Stateful object,
+                # DCP does an in-place load in the original state dict.
+                elem.load_state_dict(stateful_sd[key])
+            else:
+                # Otherwise, replace the state_dict with the loaded state_dict.
+                state_dict[key] = stateful_sd[key]
+
+
+def _load_state_dict(
+    state_dict: dict[str, Any],
+    storage_reader: StorageReader,
+    process_group: dist.ProcessGroup | None = None,
+    coordinator_rank: int = 0,
+    no_dist: bool = False,
+    planner: LoadPlanner | None = None,
+) -> None:
+    torch._C._log_api_usage_once("torch.distributed.checkpoint.load_state_dict")
+
+    distW = _DistWrapper(process_group, not no_dist, coordinator_rank)
+    if planner is None:
+        planner = DefaultLoadPlanner()
+
+    ckpt_kwargs = {}
+    if (ckpt_id := getattr(storage_reader, "checkpoint_id", None)) is not None:
+        ckpt_kwargs["checkpoint_id"] = ckpt_id
+        ckpt_kwargs["process_group"] = distW.group
+
+    use_collectives = True
+    metadata: Metadata | None = None
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def local_step():
+        nonlocal use_collectives
+        nonlocal metadata
+
+        # Use global metadata if available, otherwise fallback to rank local metadata
+        global_metadata_exc: Exception | None = None
+        rank_metadata_exc: Exception | None = None
+        try:
+            metadata = storage_reader.read_metadata()
+        except Exception as e:
+            global_metadata_exc = e
+            logger.warning(
+                "Global metadata is not found. Falling back to rank local metadata.",
+                exc_info=True,
+            )
+
+        if (
+            not metadata
+            and "kwargs" in inspect.signature(storage_reader.read_metadata).parameters
+        ):
+            try:
+                metadata = storage_reader.read_metadata(rank=distW.rank)  # noqa: F841
+                use_collectives = False
+            except Exception as e:
+                rank_metadata_exc = e
+                logger.warning("Rank local metadata is not found.", exc_info=True)
+
+        if planner is None:
+            raise AssertionError("planner is None")
+        if metadata is None:
+            error_parts = ["metadata is None"]
+            if global_metadata_exc is not None:
+                error_parts.append(
+                    f"global metadata read failed: {global_metadata_exc}"
+                )
+            if rank_metadata_exc is not None:
+                error_parts.append(
+                    f"rank local metadata read failed: {rank_metadata_exc}"
+                )
+            raise AssertionError("; ".join(error_parts))
+        planner.set_up_planner(state_dict, metadata, distW.is_coordinator)
+
+        if (
+            "kwargs"
+            in inspect.signature(storage_reader.set_up_storage_reader).parameters
+        ):
+            storage_reader.set_up_storage_reader(
+                metadata,
+                distW.is_coordinator,
+                rank=distW.rank,
+                use_collectives=use_collectives,
+            )
+        else:
+            storage_reader.set_up_storage_reader(metadata, distW.is_coordinator)
+
+        local_plan = planner.create_local_plan()
+        local_plan = storage_reader.prepare_local_plan(local_plan)
+        return local_plan
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def global_step(all_local_plans):
+        if planner is None:
+            raise AssertionError("planner is None")
+        all_local_plans = planner.create_global_plan(all_local_plans)
+        all_local_plans = storage_reader.prepare_global_plan(all_local_plans)
+        return all_local_plans
+
+    central_plan: LoadPlan | None = None
+    if use_collectives:
+        central_plan = distW.reduce_scatter("plan", local_step, global_step)
+    else:
+        local_plan: LoadPlan = local_step()
+        global_plan: list[LoadPlan] = global_step([local_plan])
+        central_plan = global_plan[0]
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def read_data():
+        if planner is None:
+            raise AssertionError("planner is None")
+        if central_plan is None:
+            raise AssertionError("central_plan is None")
+        final_local_plan = planner.finish_plan(central_plan)
+        all_reads = storage_reader.read_data(final_local_plan, planner)
+
+        all_reads.wait()
+        return None
+
+    if use_collectives:
+        _ = distW.all_gather("read", read_data)
+    else:
+        read_data()
+        distW.barrier()
+
+
+def _load_state_dict_from_keys(
+    keys: set[str] | str | None = None,
+    *,
+    checkpoint_id: str | os.PathLike | None = None,
+    storage_reader: StorageReader | None = None,
+    process_group: dist.ProcessGroup | None = None,
+) -> dict[str, Any]:
+    """
+    Load only the specified keys from the checkpoint, if no keys are specified, the entire
+    checkpoint will be loaded. Note, this method completely loads the checkpoint into the
+    current process and is not distributed.
+
+    .. warning::
+
+
+    .. warning::
+
+        All non-tensor data is loaded using `torch.load()`
+
+    .. note:
+        As opposed to the usual pattern, this function does not take a state dict as input
+        and does not load inplace. Instead, a new state dict is directly initialized and read
+        from file.
+
+    .. note:
+        If no process group is initialized, this function will assume the intent
+        is to load a checkpoint into the local process. This can be useful in the
+        case of local inference, and when using regular Tensors (as opposed to DTensor
+         or ShardedTensor)
+
+    .. note:
+        Rank 0 is assumed to be the coordinator rank.
+
+    Args:
+        keys (Optional[Union[set[str], str]]):
+            Loads any key specified in this set. If no keys are specified, the entire checkpoint
+            is loaded.
+        checkpoint_id (Union[str, os.PathLike, None]):
+            The ID of this checkpoint instance. The meaning of the checkpoint_id
+            depends on the storage. It can be a path to a folder or to a file.
+            It can also be a key if the storage is a key-value store.
+            (Default: ``None``)
+        storage_reader (Optional[StorageReader]):
+            Instance of StorageWriter used to perform reads. If this is not
+            specified, DCP will automatically infer the reader based on the
+            checkpoint_id. If checkpoint_id is also None, an exception will
+            be raised. (Default: ``None``)
+        process_group (Optional[ProcessGroup]):
+            ProcessGroup to be used for cross-rank synchronization.
+            (Default: ``None``)
+
+    Returns:
+        State dict from specified keys
+    """
+    torch._C._log_api_usage_once(
+        "torch.distributed.checkpoint._load_state_dict_from_keys"
+    )
+
+    no_dist = not (dist.is_available() and dist.is_initialized())
+    if no_dist:
+        warnings.warn(
+            "torch.distributed is unavailable or uninitialized, assuming the intent is to load in a single process.",
+            stacklevel=2,
+        )
+
+    storage_reader = cast(
+        StorageReader, _storage_setup(storage_reader, checkpoint_id, reader=True)
+    )
+
+    if isinstance(keys, str):
+        keys = {keys}
+
+    sd: dict[str, Any] = {}
+    _load_state_dict(
+        state_dict=sd,
+        storage_reader=storage_reader,
+        process_group=process_group,
+        no_dist=no_dist,
+        planner=_EmptyStateDictLoadPlanner(keys=keys),
+    )
+
+    return sd

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/state_dict_saver.py

@@ -0,0 +1,495 @@
+# mypy: allow-untyped-decorators
+# mypy: allow-untyped-defs
+import inspect
+import os
+import warnings
+from concurrent.futures import Future
+from dataclasses import dataclass
+from enum import Enum
+from typing import cast, TYPE_CHECKING
+from typing_extensions import deprecated
+
+import torch
+import torch.distributed as dist
+from torch.distributed._state_dict_utils import STATE_DICT_TYPE
+from torch.distributed.checkpoint._async_process_executor import (
+    _ProcessBasedAsyncCheckpointExecutor,
+)
+from torch.distributed.checkpoint._async_thread_executor import (
+    _ThreadBasedAsyncCheckpointExecutor,
+)
+from torch.distributed.checkpoint._storage_utils import _storage_setup
+from torch.distributed.checkpoint.default_planner import DefaultSavePlanner
+from torch.distributed.checkpoint.logger import _dcp_method_logger
+from torch.distributed.checkpoint.metadata import Metadata
+from torch.distributed.checkpoint.planner import SavePlan, SavePlanner
+from torch.distributed.checkpoint.staging import (
+    AsyncStager,
+    DefaultStager,
+    StagingOptions,
+)
+from torch.distributed.checkpoint.stateful import Stateful
+from torch.distributed.checkpoint.storage import StorageWriter, WriteResult
+from torch.distributed.distributed_c10d import _get_default_group
+
+from .utils import _api_bc_check, _DistWrapper, _profile
+
+
+if TYPE_CHECKING:
+    from torch.distributed.checkpoint._async_executor import _AsyncCheckpointExecutor
+
+
+__all__ = [
+    "save_state_dict",
+    "save",
+    "async_save",
+    "AsyncCheckpointerType",
+    "AsyncSaveResponse",
+]
+
+
+class AsyncCheckpointerType(Enum):
+    """Enum for async checkpointer type."""
+
+    THREAD = "thread"
+    PROCESS = "process"
+
+
+@deprecated(
+    "`save_state_dict` is deprecated and will be removed in future versions."
+    "Please use `save` instead.",
+    category=FutureWarning,
+)
+def save_state_dict(
+    state_dict: STATE_DICT_TYPE,
+    storage_writer: StorageWriter,
+    process_group: dist.ProcessGroup | None = None,
+    coordinator_rank: int = 0,
+    no_dist: bool = False,
+    planner: SavePlanner | None = None,
+) -> Metadata:
+    """This method is deprecated. Please switch to 'save'."""
+    storage_writer.reset()
+
+    # TODO: test returning `save` here instead.
+    with _profile():
+        return _save_state_dict(
+            state_dict,
+            storage_writer,
+            process_group,
+            coordinator_rank,
+            no_dist,
+            planner,
+        )
+
+
+@_dcp_method_logger(log_exceptions=True)  # type: ignore[arg-type]
+@_api_bc_check
+def save(
+    state_dict: STATE_DICT_TYPE,
+    *,
+    checkpoint_id: str | os.PathLike | None = None,
+    storage_writer: StorageWriter | None = None,
+    planner: SavePlanner | None = None,
+    process_group: dist.ProcessGroup | None = None,
+    no_dist: bool = False,
+    use_collectives: bool = True,
+) -> Metadata:
+    """
+    Save a distributed model in SPMD style.
+
+    This function is different from ``torch.save()`` as it handles
+    ``ShardedTensor`` , and ``DTensor`` by having each rank only save their local shards.
+
+    For each ``Stateful`` object (having both a ``state_dict`` and a ``load_state_dict``),
+    save will call ``state_dict`` before serialization.
+
+    .. warning::
+        There is no guarantees of Backwards Compatibility across PyTorch versions
+        for saved state_dicts.
+
+    .. warning::
+        If using the `process_group` argument, make sure that only its ranks
+        call `save_state_dict` and that all data in state_dict belong to it.
+
+    .. note::
+        When saving checkpoint for FSDP's `ShardingStrategy.HYBRID_SHARD`, only one of
+        the shard_group should be calling `save_state_dict` and the corresponding process
+        group needs to be passed in.
+
+    .. note::
+        If no process group is available, this function assumes the intention is to save the
+         state_dict in the local process.
+
+    .. note:
+        Rank 0 is assumed to be the coordinator rank.
+
+
+    Args:
+        state_dict (Dict[str, Any]): The state_dict to save.
+        checkpoint_id (Union[str, os.PathLike, None]):
+            The ID of this checkpoint instance. The meaning of the checkpoint_id
+            depends on the storage. It can be a path to a folder or to a file.
+            It can also be a key if the storage is a key-value store.
+            (Default: ``None``)
+        storage_writer (Optional[StorageWriter]):
+            Instance of StorageWriter used to perform writes. If this is not
+            specified, DCP will automatically infer the writer based on the
+            checkpoint_id. If checkpoint_id is also None, an exception will
+            be raised. (Default: ``None``)
+        planner (Optional[SavePlanner]):
+            Instance of SavePlanner. If this is not specified, the default
+            planner will be used. (Default: ``None``)
+        process_group (Optional[ProcessGroup]):
+            ProcessGroup to be used for cross-rank synchronization.
+            (Default: ``None``)
+        no_dist (bool):
+            If ``True``, this function will assume the intent is to load
+            a checkpoint on a single rank/process.
+            (Default: ``False``)
+        use_collectives (bool): If ``False``, this function will assume the intent is to save
+            a checkpoint without using cross-rank synchronization.
+            (Default: ``True``)
+            This configuration is experimental and should be used with caution.
+            It will change the format of the saved checkpoint and may not be backward compatible.
+
+    Returns:
+        Metadata: Metadata object for the saved checkpoint.
+
+    Example:
+        >>> # xdoctest: +SKIP
+        >>> my_model = MyModule()
+
+        >>> state_dict = {"model": my_model}
+
+        >>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(
+        ...     "/checkpoint/1"
+        ... )
+        >>> torch.distributed.checkpoint.save(
+        >>>     state_dict=state_dict,
+        >>>     storage_writer=fs_storage_writer,
+        >>> )
+
+    .. note::
+        save_state_dict uses collectives to coordinate writes across ranks.
+        For NCCL-based process groups, internal tensor representations of
+        objects must be moved to the GPU device before communication takes place.
+        In this case, the device used is given by ``torch.cuda.current_device()``
+        and it is the user's responsibility to ensure that this is set so that
+        each rank has an individual GPU, via ``torch.cuda.set_device()``.
+    """
+    torch._C._log_api_usage_once("torch.distributed.checkpoint.save")
+
+    no_dist = no_dist or (not dist.is_available()) or (not dist.is_initialized())
+    if no_dist:
+        warnings.warn(
+            "torch.distributed is disabled, unavailable or uninitialized, assuming the intent is to save in a single process.",
+            stacklevel=2,
+        )
+
+    with _profile():
+        storage_writer = cast(
+            StorageWriter, _storage_setup(storage_writer, checkpoint_id, reader=False)
+        )
+
+        return _save_state_dict(
+            state_dict=_stateful_to_state_dict(state_dict),
+            storage_writer=storage_writer,
+            process_group=process_group,
+            no_dist=no_dist,
+            planner=planner,
+            use_collectives=use_collectives,
+        )
+
+
+@dataclass
+class AsyncSaveResponse:
+    """This class contains futures for staging and upload completion.
+    It is returned by async_save().
+    staging_completion is a future that indicates when local copy
+    of state_dict is complete.
+    upload_completion is a future that indicates when a checkpoint
+    completed saving.
+    """
+
+    staging_completion: Future[None]
+    upload_completion: Future[None]
+
+
+@_dcp_method_logger(log_exceptions=True)
+def async_save(
+    state_dict: STATE_DICT_TYPE,
+    *,
+    checkpoint_id: str | os.PathLike | None = None,
+    storage_writer: StorageWriter | None = None,
+    planner: SavePlanner | None = None,
+    process_group: dist.ProcessGroup | None = None,
+    async_checkpointer_type: AsyncCheckpointerType = AsyncCheckpointerType.THREAD,
+    async_stager: AsyncStager | None = None,
+    no_dist: bool = False,
+    use_collectives: bool = True,
+) -> Future | AsyncSaveResponse:
+    """Asynchronous version of ``save``. This code first de-stages the state_dict on to the
+    staging storage (defaults to CPU memory), and then calls the `save` in a separate thread.
+
+    .. warning::
+        This feature is experimental and subject to change.
+        MUST CALL CLOSE AFTER LAST CHECKPOINT IS SAVED
+
+    Args:
+        state_dict (Dict[str, Any]): The state_dict to save.
+        checkpoint_id (Union[str, os.PathLike, None]):
+            The ID of this checkpoint instance. The meaning of the checkpoint_id
+            depends on the storage. It can be a path to a folder or to a file.
+            It can also be a key if the storage is a key-value store.
+            (Default: ``None``)
+        storage_writer (Optional[StorageWriter]):
+            Instance of StorageWriter used to perform 'stage' and  'save'. If
+            this is not specified, DCP will automatically infer the writer based on the
+            checkpoint_id. If checkpoint_id is also None, an exception will
+            be raised. (Default: ``None``)
+        planner (Optional[SavePlanner]):
+            Instance of SavePlanner. If this is not specified, the default
+            planner will be used. (Default: ``None``)
+        process_group (Optional[ProcessGroup]):
+            ProcessGroup to be used for cross-rank synchronization.
+            (Default: ``None``)
+        async_checkpointer_type (AsyncCheckpointerType):
+            whether to do checkpoint in separate thread or process
+            (Default: ``AsyncCheckpointerType.THREAD``)
+        async_stager (AsyncStager):
+            provides staging implementation. If storage_writer implements AsyncStager
+            and async_stager is provided, async_stager will be used for staging
+        no_dist (bool):
+            If ``True``, this function will assume the intent is to save
+            a checkpoint on a single rank/process.
+            (Default: ``False``)
+        use_collectives: If False, Save the checkpoint without rank coordination. (Default: ``True``)
+            This configuration is experimental and should be used with caution.
+            It will change the format of the saved checkpoint and may not be backward compatible.
+
+    Returns:
+        Future: A future holding the resultant Metadata object from `save`.
+
+    Example:
+        >>> # xdoctest: +SKIP
+        >>> my_model = MyModule()
+
+        >>> state_dict = {"model": my_model}
+
+        >>> fs_storage_writer = torch.distributed.checkpoint.FileSystemWriter(
+        ...     "/checkpoint/1"
+        ... )
+        >>> checkpoint_future = torch.distributed.checkpoint.async_save(
+        >>>     state_dict=state_dict,
+        >>>     storage_writer=fs_storage_writer,
+        >>> )
+        >>>
+        >>> # ... do some work ...
+        >>>
+        >>> checkpoint_future.result()
+
+    """
+    torch._C._log_api_usage_once("torch.distributed.checkpoint.async_save")
+
+    if dist.is_available() and dist.is_initialized():
+        pg = process_group or _get_default_group()
+        if torch.device("cpu") not in pg._device_types:
+            raise AssertionError(
+                "A CPU backend must be enabled for async save; try initializing process group with 'cpu:gloo,cuda:nccl'"
+            )
+
+    if async_stager is None:
+        if storage_writer is not None and isinstance(storage_writer, AsyncStager):
+            # bwc with old storage_writers
+            async_stager = storage_writer
+        else:
+            async_stager = DefaultStager(
+                StagingOptions(
+                    False,
+                    False,
+                    False,
+                    False,
+                )
+            )
+
+    state_dict = _stateful_to_state_dict(state_dict)
+
+    @_dcp_method_logger(log_exceptions=True)
+    def stage_state_dict() -> Future[STATE_DICT_TYPE] | STATE_DICT_TYPE:
+        return async_stager.stage(state_dict)
+
+    staging_future_or_state_dict = stage_state_dict()
+
+    upload_executor: _AsyncCheckpointExecutor = (
+        _ProcessBasedAsyncCheckpointExecutor()
+        if async_checkpointer_type == AsyncCheckpointerType.PROCESS
+        else _ThreadBasedAsyncCheckpointExecutor()
+    )
+
+    upload_future: Future = upload_executor.execute_save(
+        staging_future_or_state_dict,
+        checkpoint_id=checkpoint_id,
+        storage_writer=storage_writer,
+        planner=planner,
+        process_group=process_group,
+        no_dist=no_dist,
+        use_collectives=use_collectives,
+    )
+
+    if isinstance(staging_future_or_state_dict, Future):
+        staging_future = staging_future_or_state_dict
+        return_staging_future: Future[None] = Future()
+
+        def callback(
+            original_staging_future: Future[STATE_DICT_TYPE],
+            return_staging_future: Future[None] = return_staging_future,
+        ):
+            try:
+                original_staging_future.result()
+                return_staging_future.set_result(None)
+            except Exception as e:
+                return_staging_future.set_exception(e)
+
+        if not staging_future.done():
+            staging_future.add_done_callback(callback)
+        else:
+            return_staging_future.set_result(None)
+
+        # return new AsyncSaveResponse for users using new ZOC implementation
+        return AsyncSaveResponse(
+            staging_completion=return_staging_future, upload_completion=upload_future
+        )
+    else:
+
+        @_dcp_method_logger(log_exceptions=True)
+        def maybe_synchronize_staging():
+            if async_stager.should_synchronize_after_execute:
+                async_stager.synchronize_staging()
+
+        maybe_synchronize_staging()
+        return upload_future
+
+
+@_dcp_method_logger(log_exceptions=True)
+def _stateful_to_state_dict(state_dict: STATE_DICT_TYPE) -> STATE_DICT_TYPE:
+    """Creates a shallow copy of `state_dict` where `state_dict` is called for each Stateful object."""
+    stateful_state_dict = {}
+    for key, elem in state_dict.items():
+        # Apply _dcp_method_logger to each state_dict() call
+        def _elem_to_state_dict(elem):
+            return elem.state_dict() if isinstance(elem, Stateful) else elem
+
+        _elem_to_state_dict.__name__ = f"_stateful_to_state_dict.{key}"
+
+        stateful_state_dict[key] = _dcp_method_logger(log_exceptions=True)(
+            _elem_to_state_dict
+        )(elem)
+    return stateful_state_dict
+
+
+def _save_state_dict(
+    state_dict: STATE_DICT_TYPE,
+    storage_writer: StorageWriter,
+    process_group: dist.ProcessGroup | None = None,
+    coordinator_rank: int = 0,
+    no_dist: bool = False,
+    planner: SavePlanner | None = None,
+    use_collectives: bool = True,
+) -> Metadata:
+    torch._C._log_api_usage_once("torch.distributed.checkpoint.save_state_dict")
+
+    distW = _DistWrapper(process_group, not no_dist, coordinator_rank)
+    if planner is None:
+        planner = DefaultSavePlanner()
+    if planner is None:
+        raise AssertionError("planner is None")
+
+    global_metadata = None
+
+    ckpt_kwargs = {}
+    if (ckpt_id := getattr(storage_writer, "checkpoint_id", None)) is not None:
+        ckpt_kwargs["checkpoint_id"] = ckpt_id
+        ckpt_kwargs["process_group"] = distW.group
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def local_step():
+        if planner is None:
+            raise AssertionError("planner is None")
+        storage_meta = storage_writer.storage_meta()
+        if "storage_meta" not in inspect.signature(planner.set_up_planner).parameters:
+            warnings.warn(
+                "The function definition for SavePlanner.set_up_planner has been updated"
+                " to include the storage_meta argument. Please update your implementation"
+                " to include this parameter.",
+                stacklevel=2,
+            )
+            planner.set_up_planner(state_dict, distW.is_coordinator)  # type: ignore[call-arg, arg-type]
+        else:
+            planner.set_up_planner(
+                state_dict=state_dict,
+                storage_meta=storage_meta,
+                is_coordinator=distW.is_coordinator,
+            )
+
+        if (
+            "kwargs"
+            in inspect.signature(storage_writer.set_up_storage_writer).parameters
+        ):
+            storage_writer.set_up_storage_writer(
+                distW.is_coordinator,
+                rank=distW.rank,
+                use_collectives=use_collectives,
+            )
+        else:
+            storage_writer.set_up_storage_writer(distW.is_coordinator)
+
+        local_plan = planner.create_local_plan()
+        local_plan = storage_writer.prepare_local_plan(local_plan)
+        return local_plan
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def global_step(all_local_plans):
+        nonlocal global_metadata
+
+        if planner is None:
+            raise AssertionError("planner is None")
+        all_local_plans, global_metadata = planner.create_global_plan(all_local_plans)
+        all_local_plans = storage_writer.prepare_global_plan(all_local_plans)
+        return all_local_plans
+
+    central_plan: SavePlan | None = None
+    if use_collectives:
+        central_plan = distW.reduce_scatter("plan", local_step, global_step)
+    else:
+        local_plan: SavePlan = local_step()
+        global_plan: list[SavePlan] = global_step([local_plan])
+        central_plan = global_plan[0]
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def write_data():
+        if planner is None:
+            raise AssertionError("planner is None")
+        if central_plan is None:
+            raise AssertionError("central_plan is None")
+        final_local_plan = planner.finish_plan(central_plan)
+        all_writes = storage_writer.write_data(final_local_plan, planner)
+
+        all_writes.wait()
+        return all_writes.value()
+
+    @_dcp_method_logger(**ckpt_kwargs)
+    def finish_checkpoint(all_results):
+        if global_metadata is None:
+            raise AssertionError("global_metadata is None")
+        storage_writer.finish(metadata=global_metadata, results=all_results)
+        return global_metadata
+
+    if use_collectives:
+        metadata = distW.all_reduce("write", write_data, finish_checkpoint)
+    else:
+        write_results: list[WriteResult] = write_data()
+        metadata = finish_checkpoint([write_results])
+        distW.barrier()
+
+    return metadata

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/stateful.py

@@ -0,0 +1,42 @@
+from typing import Any, TypeVar
+from typing_extensions import Protocol, runtime_checkable
+
+
+__all__ = ["Stateful", "StatefulT"]
+
+
+@runtime_checkable
+class Stateful(Protocol):
+    """
+    Stateful protocol for objects that can be checkpointed and restored.
+    """
+
+    def state_dict(self) -> dict[str, Any]:
+        """
+        Objects should return their state_dict representation as a dictionary.
+        The output of this function will be checkpointed, and later restored in
+        `load_state_dict()`.
+
+        .. warning::
+            Because of the inplace nature of restoring a checkpoint, this function
+            is also called during `torch.distributed.checkpoint.load`.
+
+
+        Returns:
+            Dict: The objects state dict
+        """
+
+        ...
+
+    def load_state_dict(self, state_dict: dict[str, Any]) -> None:
+        """
+        Restore the object's state from the provided state_dict.
+
+        Args:
+            state_dict: The state dict to restore from
+        """
+
+        ...
+
+
+StatefulT = TypeVar("StatefulT", bound=Stateful)

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/storage.py

@@ -0,0 +1,288 @@
+import abc
+import os
+from dataclasses import dataclass
+from typing import Any
+
+from torch.distributed.checkpoint.metadata import Metadata, MetadataIndex, StorageMeta
+from torch.distributed.checkpoint.planner import (
+    LoadPlan,
+    LoadPlanner,
+    SavePlan,
+    SavePlanner,
+)
+from torch.futures import Future
+
+
+__all__ = ["WriteResult", "StorageWriter", "StorageReader"]
+
+
+@dataclass(frozen=True)
+class WriteResult:
+    index: MetadataIndex
+
+    size_in_bytes: int
+    storage_data: Any
+
+
+class StorageWriter(abc.ABC):
+    """
+    Interface used by ``save_state_dict`` to write to storage.
+
+    One StorageWriter instance acts as both the coordinator and the follower
+    in a distributed checkpoint. As part of initialization, each instance
+    is told its role.
+
+    A subclass should expect the following sequence of calls.
+
+    0) (all ranks) set checkpoint_id if users pass a valid checkpoint_id.
+    1) (all ranks) set_up_storage_writer()
+    2) (all ranks) prepare_local_plan()
+    3) (coordinator) prepare_global_plan()
+    4) (all ranks) write_data()
+    5) (coordinator) finish()
+    """
+
+    @abc.abstractmethod
+    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
+        """
+        Calls to indicates a brand new checkpoint write is going to happen.
+        A checkpoint_id may be present if users set the checkpoint_id for
+        this checkpoint write. The meaning of the checkpiont_id is
+        storage-dependent. It can be a path to a folder/file or a key for
+        a key-value storage.
+
+        Args:
+            checkpoint_id (Union[str, os.PathLike, None]):
+                The ID of this checkpoint instance. The meaning of the checkpoint_id
+                depends on the storage. It can be a path to a folder or to a file.
+                It can also be a key if the storage is a key-value store.
+                (Default: ``None``)
+        """
+        ...
+
+    @abc.abstractmethod
+    def set_up_storage_writer(
+        self, is_coordinator: bool, *args: Any, **kwargs: Any
+    ) -> None:
+        """
+        Initialize this instance.
+
+        Args:
+            is_coordinator (bool): Whether this instance is responsible for coordinating
+              the checkpoint.
+        """
+
+    @abc.abstractmethod
+    def prepare_local_plan(self, plan: SavePlan) -> SavePlan:
+        """
+        Perform storage-specific local planning.
+
+        While this method can produce a completely different plan, the recommended
+        way is to store storage specific data in SavePlan::storage_data.
+
+        Args:
+            plan (SavePlan): The local plan from the ``SavePlanner`` in use.
+
+        Returns:
+            A transformed ``SavePlan`` after storage local planning
+        """
+
+    @abc.abstractmethod
+    def prepare_global_plan(self, plans: list[SavePlan]) -> list[SavePlan]:
+        """
+        Perform centralized planning of storage.
+
+        This method is only called on the coordinator instance.
+
+        While this method can produce a completely different plan, the preferred
+        way is to store storage specific data in SavePlan::storage_data.
+
+        Args:
+            plans: A list of ``SavePlan`` instances, one for each rank.
+
+        Returns:
+            A list of transformed ``SavePlan`` after storage global planning
+        """
+
+    @abc.abstractmethod
+    def write_data(
+        self, plan: SavePlan, planner: SavePlanner
+    ) -> Future[list[WriteResult]]:
+        """
+        Write all items from ``plan`` using ``planner`` to resolve the data.
+
+        A subclass should call ``SavePlanner::resolve_data`` on each item
+        from the plan to get access to the underlying object to write.
+
+        Subclasses should lazily call `resolve_data` as it can allocate memory.
+        In case of tensors, make following assumptions:
+
+        - They might be on any device, including not matching the one on ``WriteItem::tensor_data``
+        - They might be views or not contiguous. Only the projection needs to be saved.
+
+        Args:
+            plan (SavePlan): The save plan to execute.
+            planner (SavePlanner): Planner object to be used to resolve items to data.
+
+        Returns:
+            A future that completes to a list of WriteResult
+        """
+
+    @abc.abstractmethod
+    def finish(self, metadata: Metadata, results: list[list[WriteResult]]) -> None:
+        """
+        Write the metadata and marks the current checkpoint as successful.
+
+        The actual format/schema used for serializing `metadata` is an
+        implementation detail. The only requirement is that it's recoverable
+        in to the same object graph.
+
+        Args:
+            metadata (Metadata): metadata for the new checkpoint
+            results: A list of WriteResults from all ranks.
+
+        Returns:
+            None
+        """
+
+    @classmethod
+    @abc.abstractmethod
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
+        """
+        Check if the given checkpoint_id is supported by the storage. This allow
+        us to enable automatic storage selection.
+        """
+        ...
+
+    def storage_meta(self) -> StorageMeta | None:
+        """
+        Return the storage-specific metadata. This is used to store additional information
+        in a checkpoint that can be useful for providing request-level observability. StorageMeta
+        is passed to the ``SavePlanner`` during save calls. Returns None by default.
+
+        TODO: provide an example
+        """
+        return None
+
+
+class StorageReader(abc.ABC):
+    """
+    Interface used by ``load_state_dict`` to read from storage.
+
+    One StorageReader instance acts as both the coordinator and the follower
+    in a distributed checkpoint. As part of initialization, each instance
+    is told its role.
+
+    A subclass should expected the following sequence of calls by ``load_state_dict``:
+
+    0) (all ranks) set checkpoint_id if users pass a valid checkpoint_id.
+    1) (all ranks) read_metadata()
+    2) (all ranks) set_up_storage_reader()
+    3) (all ranks) prepare_local_plan()
+    4) (coordinator) prepare_global_plan()
+    5) (all ranks) read_data()
+    """
+
+    @abc.abstractmethod
+    def reset(self, checkpoint_id: str | os.PathLike | None = None) -> None:
+        """
+        Calls to indicates a brand new checkpoint read is going to happen.
+        A checkpoint_id may be present if users set the checkpoint_id for
+        this checkpoint read. The meaning of the checkpiont_id is
+        storage-dependent. It can be a path to a folder/file or a key for
+        a key-value storage.
+
+        Args:
+            checkpoint_id (Union[str, os.PathLike, None]):
+                The ID of this checkpoint instance. The meaning of the checkpoint_id
+                depends on the storage. It can be a path to a folder or to a file.
+                It can also be a key if the storage is more like a key-value store.
+                (Default: ``None``)
+        """
+        ...
+
+    @abc.abstractmethod
+    def read_metadata(self, *args: Any, **kwargs: Any) -> Metadata:
+        """
+        Read the checkpoint metadata.
+
+        Returns:
+            The metadata object associated with the checkpoint being loaded.
+
+        """
+
+    @abc.abstractmethod
+    def set_up_storage_reader(
+        self, metadata: Metadata, is_coordinator: bool, *args: Any, **kwargs: Any
+    ) -> None:
+        """
+        Initialize this instance.
+
+        Args:
+            metadata (Metadata): The metadata schema to use.
+            is_coordinator (bool): Whether this instance is responsible for coordinating
+              the checkpoint.
+        """
+
+    @abc.abstractmethod
+    def prepare_local_plan(self, plan: LoadPlan) -> LoadPlan:
+        """
+        Perform storage-specific local planning.
+
+        While this method can produce a completely different plan, the recommended
+        way is to store storage specific data in LoadPlan::storage_data.
+
+        Args:
+            plan (LoadPlan): The local plan from the ``LoadPlan`` in use.
+
+        Returns:
+            A transformed ``LoadPlan`` after storage local planning
+        """
+
+    @abc.abstractmethod
+    def prepare_global_plan(self, plans: list[LoadPlan]) -> list[LoadPlan]:
+        """
+        Perform centralized planning of storage loading.
+
+        This method is only called on the coordinator instance.
+
+        While this method can produce a completely different plan, the preferred
+        way is to store storage specific data in LoadPlan::storage_data.
+
+        Args:
+            plans: A list of ``LoadPlan`` instances, one for each rank.
+
+        Returns:
+            A list of transformed ``LoadPlan`` after storage global planning
+        """
+
+    @abc.abstractmethod
+    def read_data(self, plan: LoadPlan, planner: LoadPlanner) -> Future[None]:
+        """
+        Read all items from ``plan`` using ``planner`` to resolve the data.
+
+        A subclass should call ``LoadPlanner::load_bytes`` to deserialize a BytesIO
+        object into the right place.
+
+        A subclass should call ``LoadPlanner::resolve_tensor`` to get access to the
+        tensors that in should load data into.
+
+        It's the StorageLayer responsibility to properly schedule any cross device copies
+        required.
+
+        Args:
+            plan (LoadPlan): The local plan to execute on
+            planner (LoadPlanner): The planner object to use to resolve items.
+
+        Returns:
+            A future that completes once all reads are finished.
+        """
+
+    @classmethod
+    @abc.abstractmethod
+    def validate_checkpoint_id(cls, checkpoint_id: str | os.PathLike) -> bool:
+        """
+        Check if the given checkpoint_id is supported by the storage. This allow
+        us to enable automatic storage selection.
+        """
+        ...

kronos-venv/lib/python3.12/site-packages/torch/distributed/checkpoint/utils.py

@@ -0,0 +1,487 @@
+# mypy: allow-untyped-defs
+import cProfile
+import inspect
+import io
+import itertools
+import os
+import warnings
+from collections.abc import Callable, Sequence
+from contextlib import contextmanager
+from functools import wraps
+from pstats import Stats
+from typing import Any, cast, TypeVar
+
+import torch
+import torch.distributed as dist
+from torch.distributed._shard.sharded_tensor import ShardedTensor
+from torch.distributed._shard.sharded_tensor.shard import Shard
+
+from .api import (
+    _is_wrapped_exception,
+    _wrap_exception,
+    CheckpointException,
+    WRAPPED_EXCEPTION,
+)
+from .metadata import MetadataIndex, STATE_DICT_TYPE
+
+
+__all__ = ["find_tensor_shard", "find_state_dict_object"]
+
+T = TypeVar("T")
+R = TypeVar("R")
+
+
+def _get_failure_dict(
+    results: list[T | WRAPPED_EXCEPTION],
+) -> dict[int, WRAPPED_EXCEPTION]:
+    return cast(
+        dict[int, WRAPPED_EXCEPTION],
+        {i: err for i, err in enumerate(results) if _is_wrapped_exception(err)},
+    )
+
+
+def _all_gather_keys(
+    local_dict: dict[str, Any], group: dist.ProcessGroup | None = None
+) -> set[str]:
+    """Gathers all keys, and returns them sorted."""
+    keys = list(local_dict.keys())
+    gathered_keys: list[list[str]] = [None] * dist.get_world_size(group)  # type: ignore[list-item]
+
+    dist.all_gather_object(gathered_keys, keys, group=group)
+    return set(itertools.chain.from_iterable(gathered_keys))
+
+
+def _assert_same_keys(
+    state_dict: dict[str, Any], process_group: dist.ProcessGroup | None = None
+) -> None:
+    """
+    Asserts that all ranks have the same keys in their state dict.
+    This is a collective call which requires all ranks in ``process_group`` to
+    join. It will also induce cross-rank communication and block CPU.
+    """
+
+    if dist.get_world_size(process_group) == 1:
+        return
+
+    all_keys = _all_gather_keys(state_dict, process_group)
+    my_keys = set(state_dict.keys())
+    diff = all_keys - my_keys
+    if len(diff) > 0:
+        raise AssertionError(
+            f"Key(s) present in other ranks but not this one, difference: {diff}"
+        )
+
+
+class _DistWrapper:
+    """
+    This is a wrapper around PG that provides a series of features around object collectives.
+
+    It works without distributed initialized, where most collectives turns into nops.
+
+    All variants that take functions are exception robust, meaning that if one or more
+    ranks raise errors, all ranks will observe those.
+    """
+
+    def __init__(
+        self,
+        group: dist.ProcessGroup | None,
+        use_dist: bool,
+        coordinator_rank: int,
+    ):
+        self.group = group
+        self.use_dist = use_dist
+        self.coordinator_rank = coordinator_rank
+        if self.use_dist:
+            self.global_coordinator_rank = (
+                dist.get_global_rank(group, coordinator_rank)
+                if group is not None
+                else coordinator_rank
+            )
+            self.rank = dist.get_rank(group)
+            self.is_coordinator = self.rank == coordinator_rank
+        else:
+            self.global_coordinator_rank = 0
+            self.rank = 0
+            self.is_coordinator = True
+
+    def get_rank(self) -> int:
+        return self.rank
+
+    def get_world_size(self) -> int:
+        if self.use_dist:
+            return dist.get_world_size(self.group)
+        return 1
+
+    def broadcast_object(self, object: T | None) -> T:
+        """Implement functionality similar to c10d::broadcast_object_list but without distributed enabled."""
+        object_list = [object]
+        if self.use_dist:
+            dist.broadcast_object_list(
+                object_list=object_list,
+                group=self.group,
+                src=self.global_coordinator_rank,
+            )
+        return cast(T, object_list[0])
+
+    def gather_object(self, object: T) -> list[T] | None:
+        """Implement functionality similar to c10d::gather_object but without distributed enabled."""
+        if self.use_dist:
+            gather_objs = (
+                cast(list[T], [None] * dist.get_world_size(self.group))
+                if self.is_coordinator
+                else None
+            )
+
+            dist.gather_object(
+                obj=object,
+                object_gather_list=gather_objs if self.is_coordinator else None,
+                dst=self.global_coordinator_rank,
+                group=self.group,
+            )
+            result = gather_objs
+        else:
+            result = [object]
+        return result
+
+    def all_gather_object(self, object: T) -> list[T]:
+        """Implement functionality similar to c10d::all_gather_object but without distributed enabled."""
+        if self.use_dist:
+            gather_objs = cast(list[T], [None] * dist.get_world_size(self.group))
+
+            dist.all_gather_object(
+                object_list=gather_objs, obj=object, group=self.group
+            )
+        else:
+            gather_objs = [object]
+        return gather_objs
+
+    def scatter_object(self, object_list: list[T] | None) -> T:
+        """Implement functionality similar to c10d::scatter_object but without distributed enabled."""
+        if self.use_dist:
+            gather_result = cast(list[T], [None])
+            dist.scatter_object_list(
+                scatter_object_output_list=gather_result,
+                scatter_object_input_list=object_list if self.is_coordinator else None,
+                src=self.global_coordinator_rank,
+                group=self.group,
+            )
+
+            local_reply = gather_result[0]
+        else:
+            if object_list is None:
+                raise AssertionError("object_list is None")
+            local_reply = object_list[0]
+        return local_reply
+
+    def reduce_scatter(
+        self,
+        step: str,
+        map_fun: Callable[[], T],
+        reduce_fun: Callable[[list[T]], list[R]],
+    ) -> R:
+        """
+        Compute a value on each rank, then do centralized reduce on a single rank, followed by a scatter.
+
+        This method operates in the following way:
+            Run ``map_fun`` on all ranks
+            Gather results on rank 0
+            Call ``reduce_fun`` on all those values
+            Scatter to each rank part of the result.
+        """
+        local_data: WRAPPED_EXCEPTION | T
+        try:
+            local_data = map_fun()
+        except BaseException as e:  # noqa: B036
+            local_data = _wrap_exception(e)
+
+        all_data = self.gather_object(local_data)
+        all_results: list[R | CheckpointException] | None = None
+        if self.is_coordinator:
+            if all_data is None:
+                raise AssertionError("all_data is None")
+            node_failures = _get_failure_dict(all_data)
+
+            if len(node_failures) == 0:
+                try:
+                    # N.B. why can't mypy cast List[R] to List[Union[R, WRAPPED_EXCEPTION]]?
+                    all_results = cast(
+                        list[R | CheckpointException],
+                        reduce_fun(cast(list[T], all_data)),
+                    )
+                except BaseException as e:  # noqa: B036
+                    node_failures[self.rank] = _wrap_exception(e)
+
+            if len(node_failures) > 0:
+                all_results = [
+                    CheckpointException(step, node_failures)
+                ] * self.get_world_size()
+
+        result = self.scatter_object(all_results)
+        if isinstance(result, CheckpointException):
+            raise result
+        return result
+
+    def all_reduce(
+        self,
+        step: str,
+        map_fun: Callable[[], T],
+        reduce_fun: Callable[[list[T]], R],
+    ) -> R:
+        """
+        Compute a value on each rank, then do centralized reduce on a single rank, followed by a broadcast.
+
+        This method operates in the following way:
+            Run ``map_fun`` on all ranks
+            Gather results on rank 0
+            Call ``reduce_fun`` on all those values
+            Broadcast the reduced value to all ranks.
+        """
+        local_data: T | WRAPPED_EXCEPTION
+        try:
+            local_data = map_fun()
+        except BaseException as e:  # noqa: B036
+            local_data = _wrap_exception(e)
+
+        all_data = self.gather_object(local_data)
+        result: R | CheckpointException | None = None
+        if self.is_coordinator:
+            if all_data is None:
+                raise AssertionError("all_data is None")
+            node_failures = _get_failure_dict(all_data)
+            if len(node_failures) == 0:
+                try:
+                    result = reduce_fun(cast(list[T], all_data))
+                except BaseException as e:  # noqa: B036
+                    node_failures[self.rank] = _wrap_exception(e)
+
+            if len(node_failures) > 0:
+                result = CheckpointException(step, node_failures)
+
+        # pyrefly: ignore [bad-argument-type]
+        final_result = self.broadcast_object(result)
+        if isinstance(final_result, CheckpointException):
+            raise final_result
+        # pyrefly: ignore [redundant-cast]
+        return cast(R, final_result)
+
+    def all_gather(
+        self,
+        step: str,
+        map_fun: Callable[[], T],
+    ) -> list[T]:
+        """
+        Compute a value on each rank, then all_gather them.
+
+        This method operates in the following way:
+            Run ``map_cp`` on all ranks
+            all_gather the values to all ranks
+        """
+        result: T | WRAPPED_EXCEPTION
+        try:
+            result = map_fun()
+        except BaseException as e:  # noqa: B036
+            result = _wrap_exception(e)
+
+        all_results = self.all_gather_object(result)
+
+        node_failures = _get_failure_dict(all_results)
+        if len(node_failures) > 0:
+            raise CheckpointException(step, node_failures)
+        return cast(list[T], all_results)
+
+    def broadcast(
+        self,
+        step: str,
+        map_fun: Callable[[], T],
+    ) -> T:
+        """
+        Compute a value on rank 0 and broadcast it.
+
+        This method operates in the following way:
+            Run ``map_cp`` on rank 0
+            broadcast the value
+        """
+        result: T | CheckpointException | None = None
+        if self.is_coordinator:
+            try:
+                result = map_fun()
+            except BaseException as e:  # noqa: B036
+                result = CheckpointException(step, {self.rank: _wrap_exception(e)})
+        # pyrefly: ignore [bad-argument-type]
+        final_result = self.broadcast_object(result)
+        if isinstance(final_result, CheckpointException):
+            raise final_result
+        # pyrefly: ignore [redundant-cast]
+        return cast(T, final_result)
+
+    def barrier(self) -> None:
+        """
+        Add a synchronization point across all processes when using distributed.
+        If torch.distributed is initialized, this function will invoke a barrier across the global process group.
+        If torch.distributed is not initialized, this function is a no-op.
+        """
+        if not self.use_dist:
+            return
+        dist.barrier(group=self.group)
+
+
+def _find_shard(tensor: ShardedTensor, index: MetadataIndex) -> Shard:
+    if index.offset is None:
+        raise ValueError(
+            f"Cannot lookup {index.fqn} since its a ShardedTensor and no offset was provided"
+        )
+
+    shards = tensor.local_shards()
+    # index fast path
+    if index.index is not None:
+        if (
+            len(shards) > index.index
+            and torch.Size(shards[index.index].metadata.shard_offsets) == index.offset
+        ):
+            return shards[index.index]
+
+    for shard in shards:
+        if torch.Size(shard.metadata.shard_offsets) == index.offset:
+            return shard
+    raise ValueError(f"Could not find shard at '{index.offset}' for FQN: '{index.fqn}'")
+
+
+def find_tensor_shard(tensor: torch.Tensor, index: MetadataIndex) -> torch.Tensor:
+    if hasattr(tensor, "__get_tensor_shard__"):
+        # DTensor implements _Checkpointable
+        return tensor.__get_tensor_shard__(index)  # type: ignore[attr-defined]
+    if isinstance(tensor, ShardedTensor):
+        return _find_shard(tensor, index).tensor
+    if index.offset is not None:
+        # special case looking up a tensor by origin
+        if index.offset == torch.Size([0] * len(tensor.size())):
+            return tensor
+        raise ValueError(
+            f"FQN: '{index.fqn}' is not a ShardedTensor, can't find by offset: '{index.offset}'"
+        )
+    return tensor
+
+
+def find_state_dict_object(state_dict: STATE_DICT_TYPE, index: MetadataIndex) -> Any:
+    if index.fqn not in state_dict:
+        raise ValueError(f"Could not find FQN: '{index.fqn}'")
+    obj = state_dict[index.fqn]
+
+    if isinstance(obj, torch.Tensor):
+        return find_tensor_shard(obj, index)
+    elif index.offset is not None:
+        raise ValueError(
+            f"FQN: '{index.fqn}' is not a ShardedTensor, can't find by offset: '{index.offset}'"
+        )
+    return obj
+
+
+def _element_wise_add(a: Sequence[int], b: Sequence[int]) -> list[int]:
+    return [i_a + i_b for i_a, i_b in zip(a, b)]
+
+
+def _element_wise_sub(a: Sequence[int], b: Sequence[int]) -> list[int]:
+    return [i_a - i_b for i_a, i_b in zip(a, b)]
+
+
+class _ReaderView(io.IOBase):
+    def __init__(self, base_stream: io.IOBase, offset: int, len: int):
+        super().__init__()
+        self.offset = offset
+        self.len = len
+        self.base_stream = base_stream
+        self.seek(0)
+
+    def seek(self, offset: int, whence: int = os.SEEK_SET, /) -> int:
+        if whence == os.SEEK_SET:
+            offset = self.offset + offset
+        elif whence == os.SEEK_END:
+            whence = os.SEEK_SET
+            offset = (self.offset + self.len) - offset
+        return self.base_stream.seek(offset, whence)
+
+    def tell(self) -> int:
+        return self.base_stream.tell() - self.offset
+
+    def readable(self) -> bool:
+        return self.base_stream.readable()
+
+    def seekable(self) -> bool:
+        return self.base_stream.seekable()
+
+    def readinto(self, b):
+        max_size = self.len - self.tell()
+        if max_size == 0:
+            return 0
+        if len(b) > max_size:
+            b = memoryview(b)[:max_size]
+        return self.base_stream.readinto(b)  # type: ignore[attr-defined]
+
+    def read(self, size=-1):
+        max_size = self.len - self.tell()
+        if size == -1 or size > max_size:
+            size = max_size
+        return self.base_stream.read(size)
+
+
+def _create_file_view(file: io.IOBase, offset: int, length: int) -> io.IOBase:
+    # FIXME (kumpera) torch.load fails if we wrap with io.BufferedReader
+    return _ReaderView(file, offset, length)
+
+
+def _normalize_device_info(device_type: str, device_id: int) -> str:
+    """Device info normalization."""
+    if device_type == "cpu":
+        return "cpu"
+    return f"{device_type}:{device_id}"
+
+
+# TODO: integrate with distributed logging flag
+ENABLE_PROFILE = False
+
+
+@contextmanager
+def _profile():
+    # Only log the profiling when it is enable and is on rank0  or dist is not
+    # available.
+    if ENABLE_PROFILE and (not dist.is_available() or dist.get_rank() == 0):
+        profiler = cProfile.Profile()
+        profiler.enable()
+        try:
+            yield
+        finally:
+            profiler.disable()
+            stats = Stats(profiler)
+            stats.sort_stats("time").print_stats(10)
+    else:
+        yield
+
+
+def _api_bc_check(func):
+    @wraps(func)
+    def inner_func(*args, **kwargs) -> Any:
+        if len(args) == 2:
+            warnings.warn(
+                f"The argument order of {func.__name__} has been changed. "
+                "Please check the document to avoid future breakages.",
+                stacklevel=2,
+            )
+            sig = inspect.signature(func)
+            kwonlyargs = [
+                p.name for p in sig.parameters.values() if p.kind == p.KEYWORD_ONLY
+            ]
+            if "storage_writer" in kwonlyargs:
+                if "storage_writer" in kwargs:
+                    raise AssertionError(f"storage_writer in kwargs: {(args, kwargs)}")
+                kwargs["storage_writer"] = args[1]
+            elif "storage_reader" in kwonlyargs:
+                if "storage_reader" in kwargs:
+                    raise AssertionError(f"storage_reader in kwargs: {(args, kwargs)}")
+                kwargs["storage_reader"] = args[1]
+            else:
+                raise RuntimeError(f"Unexpected kwonlyargs = {kwonlyargs}")
+            return func(args[0], **kwargs)
+        else:
+            return func(*args, **kwargs)
+
+    return inner_func