Diff Coverage

        use_gloo,
    )

    # After the async save is completed, the user callback will be executed.
    def async_callback() -> None:
        """Run the user callback once the save has been persisted, if one was given."""
        if callback is not None:
            callback()

    Args:
        dests (list): Tensors written into.
        srcs (list): Tensors read from, pairing with ``dests`` by position.
    """
    if dests:
        torch._foreach_copy_(dests, srcs)  # pylint: disable=protected-access


def get_created_group(rank_list: Union[list[int], tuple[int, ...]]) -> Any:
    """

    Returns:
        Any: The cached group, or None when the cache holds none over these ranks.
    """
    return EXISTING_COMM_GROUPS.get(str(tuple(sorted(rank_list))))


def synchronize() -> None:
    """

    Nothing can be queued on a device stream when the process has no device, and there are
    two ways to have none: a CPU-only install carries no ``torch.npu`` at all, and a gloo run
    on an accelerator box never initializes the one it has.
    """
    device = getattr(torch, "npu", None)
    if device is None or not device.is_initialized():
        return
    device.current_stream().synchronize()


# How many broadcasts one rank keeps going at once. Starting the next without waiting on
# the last is what lets a read overlap the send before it, but each one in flight holds

    while len(in_flight) >= _MAX_BROADCASTS_IN_FLIGHT:
        _finish_broadcast(in_flight.popleft())
    handle = dist.broadcast(buffer, source.src_rank, groups[source.group_ranks], async_op=True)
    if handle is None:
        if after is not None:
            after()
        return
    in_flight.append((handle, after))


def _finish_broadcast(pending: tuple) -> None:

        """Send one batch, and arrange for it to be dealt out once it has landed."""
        buffers = self._batches.pop(key, [])
        self._pending_bytes.pop(key, None)
        if not buffers:
            return

        group_ranks, src_rank = key[:2]
        source = BroadcastSource(group_ranks=group_ranks, src_rank=src_rank)
        self.sent += 1

        Any: A tensor view over the shard, writable in place by a collective.
    """
    if isinstance(obj, DTensor):
        if obj.layout is not None and obj.layout.ragged_shard is not None:
            return get_ragged_box_tensor(obj, index).detach()
        return obj.to_local().detach()
    return obj.detach()

    """
    if ckpt_format == "safetensors":
        save_file(tensors=state_dict, filename=file_path)
    else:
        torch.save(obj=state_dict, f=file_path)


def load_checkpoint_file(file_path: str, ckpt_format: str = "safetensors") -> dict:
    """

    Returns:
        dict: What the file holds, keyed by physical name.
    """
    if ckpt_format == "safetensors":
        return load_file(filename=file_path)
    return torch.load(f=file_path)

        return
    for key, value in state_dict.items():
        if isinstance(value, (torch.Tensor, bytes, str)):
            continue
        logger.warning(
            "Unsupported offline checkpoint value type for key %r: %s "
            "(expected torch.Tensor, bytes, or str per project rules).",
            key,
            type(value).__name__,

            "(expected torch.Tensor, bytes, or str per project rules).",
            key,
            type(value).__name__,
        )
        raise TypeError(
            f"Offline checkpoint expects torch.Tensor, bytes, or str per key; "
            f"got {type(value).__name__} for key {key!r}."
        )

    Shard filenames follow Hugging Face conventions (``model.safetensors`` or
    ``model-00001-of-00002.safetensors``).
    """
    get_storage_size = get_tensor_storage_size

    filename_pattern = _HF_SAFE_WEIGHTS_FILENAME_PATTERN

    max_cap: int | str = max_shard_size

    """
    if state_dict is None:
        raise ValueError("state_dict is None.")

    _validate_state_dict(state_dict)

    if os.path.isfile(save_directory):
        raise ValueError(f"The save_directory {save_directory} should be a directory, but a file.")

    filename_to_tensors = state_dict_split.filename_to_tensors.items()

    for shard_file, tensors in filename_to_tensors:
        shard = {tensor: state_dict[tensor] for tensor in tensors}
        save_checkpoint_file(shard, os.path.join(save_directory, shard_file), ckpt_format="safetensors")
    if index is None:
        path_to_weights = os.path.join(save_directory, weights_name)
        logger.info("Model weights saved in %s", path_to_weights)
    else:

    state_dict: dict[str, Any] = {}

    if os.path.isfile(safe_weights_file):
        state_dict = load_checkpoint_file(safe_weights_file, ckpt_format="safetensors")
        logger.info("Loaded safetensors checkpoint from %s", safe_weights_file)

    else:
        safe_index_file = os.path.join(resume_from_checkpoint, _SAFE_WEIGHTS_INDEX_NAME)

        total_size = 0
        for shard_file in shard_files:
            shard_path = os.path.join(resume_from_checkpoint, shard_file)
            shard_state_dict = load_checkpoint_file(shard_path, ckpt_format="safetensors")
            for key, value in shard_state_dict.items():
                if key in state_dict:
                    logger.warning(
                        "Duplicate key %r when merging Hugging Face shards; keeping first occurrence.",

                        key,
                    )
                    continue
                state_dict[key] = value
                if isinstance(value, torch.Tensor):
                    size = get_tensor_storage_size(value)
                    total_size += size
                    logger.debug("Loaded tensor %r, size %s bytes", key, size)
                else:
                    logger.debug("Loaded non-tensor entry %r", key)

                    logger.debug("Loaded tensor %r, size %s bytes", key, size)
                else:
                    logger.debug("Loaded non-tensor entry %r", key)
        logger.info("Merged Hugging Face shards, total tensor bytes (sum per key): %s", total_size)
    _validate_state_dict(state_dict)
    return state_dict


def _full_checkpoint_format_for_path(path: str) -> str:

            raise ValueError(
                "src_platform='torch' requires a single PyTorch checkpoint file path; "
                f"src_ckpt must be an existing file, got {src_ckpt!r}."
            )
        fmt = _full_checkpoint_format_for_path(str(src_ckpt))
        state_dict = load_checkpoint_file(str(src_ckpt), ckpt_format=fmt)
    else:
        raise ValueError(
            f"Unsupported src_platform={src_platform!r}; expected 'huggingface' or 'torch'."
        )

) -> tuple[tuple[tuple[int, ...], tuple[int, ...]], ...]:
    """Decompose a row-major flat interval into ordered N-D boxes."""
    total_numel = prod(shape)
    if not shape:
        raise ValueError("Ragged checkpoint geometry requires a non-scalar global shape")
    if flat_start < 0 or flat_end < flat_start or flat_end > total_numel:
        raise ValueError(
            "Invalid flat interval for Ragged checkpoint geometry, "
            f"got interval=({flat_start}, {flat_end}), shape={shape!r}"
        )
    if flat_start == flat_end:

def compute_ragged_boxes(tensor: DTensor) -> tuple[RaggedCheckpointBox, ...]:
    """Return ordered N-D boxes covering one RaggedShard local flat tensor."""
    layout = tensor.layout
    if layout.ragged_shard is None:
        raise ValueError("compute_ragged_boxes requires a RaggedShard DTensor")

    global_shape = tuple(tensor.shape)
    ragged_slice = _compute_ragged_slice(global_shape, layout)
    raw_boxes = _decompose_flat_interval(

        )
        local_flat_offset += box_numel

    if local_flat_offset != ragged_slice.local_numel:
        raise ValueError(
            "Ragged checkpoint boxes do not cover the local flat tensor, "
            f"covered={local_flat_offset}, expected={ragged_slice.local_numel}"
        )
    return tuple(boxes)

            local_flat = tensor.to_local().reshape((-1,))
            return local_flat[
                box.local_flat_start:box.local_flat_end
            ].reshape(box.sizes)
    raise ValueError(
        "Ragged checkpoint box was not found in the local DTensor, "
        f"fqn={index.fqn!r}, offset={index.offset!r}"
    )

        if item.type == WriteItemType.TENSOR:
            if isinstance(obj, DTensor):
                if obj.layout is not None and obj.layout.ragged_shard is not None:
                    return get_ragged_box_tensor(obj, item.index).detach().to("cpu")
                return obj.to_local().detach().to("cpu")
            if isinstance(obj, Tensor):
                return obj.detach().to("cpu")
            raise TypeError(f"Write item {fqn} expected tensor-like object, got {type(obj)}")
        if item.type == WriteItemType.BYTE_IO:

                and target.layout.ragged_shard is not None
        ):
            local_tensor = get_ragged_box_tensor(target, read_item.dest_index)
        elif isinstance(target, DTensor):
            local_tensor = target.to_local()
        else:
            local_tensor = target

        # Detached because the caller writes the loaded slice in place: an in-place write

    Returns:
        dict: Optimizer state dict with FQN-based keys.
    """
    return _get_optim_state_dict(model, optimizer, options=options)


def set_optim_state_dict(
    model: Any,

        optim_state_dict: The optimizer state dict to load.
        options: Optional configuration (full_state_dict, cpu_offload,
            strict, broadcast_from_rank0, etc.).
    """
    _set_optim_state_dict(model, optimizer, optim_state_dict, options=options)

    Returns:
        int: Element count times element size, which is what it occupies on disk.
    """
    if not isinstance(tensor, torch.Tensor):
        raise TypeError(f"get_tensor_storage_size expects torch.Tensor, got {type(tensor)!r}")
    return int(tensor.numel()) * int(tensor.element_size())


def str_to_dtype(dtype_str: str) -> torch.dtype:
    """

        torch.dtype: The dtype it names.
    """
    parts = dtype_str.split(".", 1)
    if len(parts) != 2:
        raise ValueError(f"Expected dtype string like 'torch.float32', got {dtype_str!r}.")
    prefix, name = parts
    if prefix != "torch":
        raise ValueError(f"Expected PyTorch dtype string with prefix 'torch', got {dtype_str!r}.")
    dtype = getattr(torch, name, None)
    if isinstance(dtype, torch.dtype):
        return dtype
    raise ValueError(f"{dtype_str!r} does not resolve to a torch.dtype.")


def dcp_timer_decorator(func: Callable) -> Callable:
    """

    Returns:
        Any: The narrowed tensor slice (tensor-like object).
    """
    if not offsets or not lengths:
        return tensor
    slices = tuple(
        slice(int(off), int(off) + int(ln))
        for off, ln in zip(offsets, lengths)
    )

    # Validate input formats
    def is_valid_axis_list(axis_list: Any) -> None:
        """Reject an area that is not a sequence of two-element ranges."""
        if not isinstance(axis_list, (tuple, list)):
            raise TypeError("Area must be a tuple of ranges")
        for axis_range in axis_list:
            if (not isinstance(axis_range, (tuple, list)) \
                or len(axis_range) != 2):
                raise TypeError("Each axis range must be a 2-element tuple")

    is_valid_axis_list(area_a)
    is_valid_axis_list(area_b)

    """
    if isinstance(obj, DTensor):
        layout = obj.layout
        if layout is None:
            shape = obj.shape if hasattr(obj, "shape") else obj.to_local().shape
            return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=tuple(shape))]
        if layout.ragged_shard is not None:
            return [
                ChunkStorageMetadata(offsets=box.offsets, sizes=box.sizes)
                for box in compute_ragged_boxes(obj)

        tensor_map = getattr(layout, "tensor_map", None) or getattr(layout, "_tensor_map", None)
        rank_list = getattr(layout, "rank_list", None) or getattr(layout, "_rank_list", None)

        if mesh_shape is None or tensor_map is None or rank_list is None:
            shape = obj.shape if hasattr(obj, "shape") else obj.to_local().shape
            return [ChunkStorageMetadata(offsets=(0,) * len(shape), sizes=tuple(shape))]

        current_rank = dist.get_rank()
        if current_rank not in rank_list:
            return []

        inner_rank_id = rank_list.index(current_rank)
        full_shape = obj.shape
        slice_area = infer_slice_area_by_layout(

        holder_count = (len(occurrences) - 1) // 2
        if holder_count > 1:
            duplicates.append((item_size, occurrences))
            continue
        masks[occurrences[1]][occurrences[2]] = 1
        storage_sizes[occurrences[1]] += item_size

    # Largest first, so the shards with room to unbalance the plans are placed while every plan is
    # still a candidate. Python's sort is stable, including with reverse=True, so equally sized
    # duplicates keep their gather order and every rank still agrees on the owner. Sorting is

    for item_size, occurrences in duplicates:
        # Occurrences were appended in ascending plan order, so slot 1 is the lowest plan index
        # and the storage-size search breaks ties towards it.
        if save_to_minimum_rank:
            owner_slot = 1
        else:
            owner_slot = _least_loaded_owner_slot(occurrences, storage_sizes)
        owner_idx = occurrences[owner_slot]
        masks[owner_idx][occurrences[owner_slot + 1]] = 1

    best_size = storage_sizes[occurrences[1]]
    for slot in range(3, len(occurrences), 2):
        size = storage_sizes[occurrences[slot]]
        if size < best_size:
            best_slot, best_size = slot, size
    return best_slot


def traverse_state_dict(

        for entry in values:
            if isinstance(entry, (Mapping, list, tuple)) and not _is_terminal(entry):
                return False
            if isinstance(entry, Tensor):
                return False
        return True

    def _traverse_obj(path: tuple[Any, ...], value: Any) -> None:
        if isinstance(value, Mapping):

            cur_container = cur_container[prev_key]

    last_key = path[-1]
    if isinstance(last_key, int):
        extend_list(cur_container, last_key)

    cur_container[last_key] = value