Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / metadata.py: 100%
53 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd. All rights reserved.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Checkpoint metadata structures for distributed checkpoint save and load."""
16from dataclasses import dataclass, field
17from typing import Any, Optional, Union
20CHUNK_INFO = "chunk_info"
23@dataclass(frozen=True)
24class MetadataIndex:
25 """
26 Index to identify a specific piece of data in the checkpoint.
28 ``index`` is a positional hint into the chunk list, not part of the identity: a chunk is
29 identified by its ``fqn`` and ``offset``, and the same chunk is described with ``index=None``
30 while plans are being deduplicated and with its ordinal afterwards. Keeping it out of
31 ``__eq__``/``__hash__`` (as torch's MetadataIndex does) makes those two spellings compare
32 equal, and drops the field from every hash computed over the millions of indices a global
33 plan builds.
35 Attributes:
36 fqn: Fully qualified name of the tensor/object.
37 offset: Offset in the tensor (for sharded tensors). Default ().
38 index: Index for sharded tensors (None for non-sharded). Not compared or hashed.
39 Default None.
40 """
41 fqn: str
42 offset: tuple = field(default_factory=tuple)
43 index: Optional[int] = field(default=None, compare=False, hash=False)
46@dataclass(frozen=True)
47class ChunkStorageMetadata:
48 """
49 Metadata for a chunk of storage.
51 Represents a portion of a distributed tensor stored in the checkpoint.
53 Attributes:
54 offsets: Offsets in the global tensor for each dimension.
55 sizes: Sizes of the chunk for each dimension.
56 """
57 offsets: tuple
58 sizes: tuple
61@dataclass(frozen=True)
62class ChunkInfo:
63 """
64 Info for a tensor chunk.
66 Represents a portion of a distributed tensor stored in the checkpoint.
68 Attributes:
69 chunk: Offsets in the global tensor for each dimension.
70 global_shape: Sizes of the chunk for each dimension.
71 replica_rank_list: Have the same sharded tensor ranks list.
72 """
73 chunk: ChunkStorageMetadata
74 global_shape: tuple[int]
75 replica_rank_list: Optional[tuple[int]] = None
78# Ordered longest-match-first: "bfloat16" has to be tried before "float16", which is a
79# substring of it. A dict would work today but hides that the order is load-bearing.
80_DTYPE_ELEMENT_SIZES = (
81 ("bfloat16", 2),
82 ("float16", 2),
83 ("float32", 4),
84 ("float64", 8),
85 ("int8", 1),
86 ("int16", 2),
87 ("int32", 4),
88 ("int64", 8),
89 ("bool", 1),
90)
91DEFAULT_DTYPE_ELEMENT_SIZE = 4
93# A checkpoint names only a handful of distinct dtypes, while this is asked once per plan
94# item - tens of thousands of times on a large model - so each name is scanned once and
95# answered from here afterwards.
96_dtype_element_size_cache: dict[str, int] = {}
99def dtype_element_size(dtype: Optional[str]) -> int:
100 """
101 Bytes one element of ``dtype`` takes.
103 The dtype reaches the checkpoint as a framework-specific string such as
104 ``"torch.bfloat16"`` or ``"Float32"``, so it is matched by substring rather than by
105 equality. Anything unrecognized falls back to :data:`DEFAULT_DTYPE_ELEMENT_SIZE`, which
106 keeps the sizes usable as relative weights even for a dtype this table does not name.
108 Args:
109 dtype (Optional[str]): Dtype name from :class:`TensorProperties`, or None.
111 Returns:
112 int: Size of one element in bytes.
113 """
114 if dtype is None:
115 return DEFAULT_DTYPE_ELEMENT_SIZE
116 key = str(dtype)
117 cached = _dtype_element_size_cache.get(key)
118 if cached is not None:
119 return cached
120 lowered = key.lower()
121 element_size = DEFAULT_DTYPE_ELEMENT_SIZE
122 for name, size in _DTYPE_ELEMENT_SIZES:
123 if name in lowered:
124 element_size = size
125 break
126 _dtype_element_size_cache[key] = element_size
127 return element_size
130@dataclass(frozen=True)
131class TensorProperties:
132 """
133 Properties of a tensor.
135 Attributes:
136 dtype: Data type of the tensor (as string).
137 requires_grad: Whether the tensor requires gradients. Default False.
138 memory_format: Memory format (optional). Default None.
139 """
140 dtype: str
141 requires_grad: bool = False
142 memory_format: Optional[str] = None
145@dataclass
146class BytesStorageMetadata:
147 """Metadata for bytes data stored in checkpoint."""
150@dataclass(frozen=True)
151class TensorStorageMetadata:
152 """
153 Metadata for a distributed tensor.
155 Contains properties, global size, and list of chunks stored across ranks.
157 Attributes:
158 properties: Tensor properties (dtype, etc.).
159 size: Global size of the tensor.
160 chunks: List of chunks stored in the checkpoint. Default [].
161 """
162 properties: TensorProperties
163 size: tuple
164 chunks: list[ChunkStorageMetadata] = field(default_factory=list)
167@dataclass
168class Metadata:
169 """
170 Global metadata for a checkpoint.
172 Contains metadata for all items in the state_dict, along with planner and storage-specific data.
174 Attributes:
175 state_dict_metadata: Mapping from FQN to storage metadata.
176 planner_data: Planner-specific data (optional). Default None.
177 storage_data: Storage-specific data (optional). Default None.
178 version: Checkpoint format version. Default "1.0".
179 """
180 state_dict_metadata: dict[str, Union[TensorStorageMetadata, BytesStorageMetadata]]
181 planner_data: Any = None # Planner-specific data (can be any type)
182 storage_data: Optional[dict[MetadataIndex, Any]] = None # Storage mapping: MetadataIndex -> StorageInfo
183 version: str = "1.0"