Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / planner.py: 98%
85 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"""Planner interfaces and implementations"""
16import abc
17import dataclasses
18from enum import Enum
19from dataclasses import dataclass, field
20from typing import Any, Optional, Union
22from hyper_parallel.core.distributed_checkpoint.metadata import (
23 Metadata,
24 MetadataIndex,
25 dtype_element_size,
26)
29class WriteItemType(Enum):
30 """Type of write item."""
31 TENSOR = "tensor"
32 BYTE_IO = "byte_io"
35class LoadItemType(Enum):
36 """Type of load item."""
37 TENSOR = "tensor"
38 BYTE_IO = "byte_io"
41@dataclass(frozen=True)
42class WriteItem:
43 """
44 Item to be written to storage.
46 Represents a single logical item (tensor or bytes) to be saved.
48 Attributes:
49 index: Metadata index identifying this item.
50 type: Type of write item (TENSOR or BYTE_IO).
51 tensor_data: Dictionary containing tensor data (for TENSOR type). Default None.
52 bytes_io_data: Bytes data (for BYTE_IO type). Default None.
53 """
54 index: MetadataIndex
55 type: WriteItemType
56 # Keys: 'chunk' (ChunkStorageMetadata), 'properties' (TensorProperties), 'size' (tuple).
57 # Actual tensor data is in planner's tensor cache, not here, to avoid all_gather of tensors.
58 tensor_data: Optional[dict[str, Any]] = None
59 bytes_io_data: Optional[Union[bytes, Any]] = None # Bytes or pickle-serializable object
61 def tensor_storage_size(self) -> Optional[int]:
62 """
63 Best-effort storage size estimation in bytes for tensor items.
65 Returns:
66 Optional[int]: Estimated storage size in bytes for tensor items,
67 or None if estimation cannot be performed (e.g., for non-tensor items).
68 """
69 if self.type != WriteItemType.TENSOR or not self.tensor_data:
70 return None
72 # Try to estimate from metadata
73 chunk = self.tensor_data.get("chunk")
74 properties = self.tensor_data.get("properties")
75 if chunk is None or properties is None:
76 return None
78 # Get size from chunk (local chunk size, not global size)
79 num = 1
80 for dim in chunk.sizes:
81 num *= int(dim)
82 # An item with no dtype recorded counts as one unit per element.
83 dtype_str = getattr(properties, "dtype", None)
84 if dtype_str is None:
85 return int(num)
86 return int(num) * dtype_element_size(dtype_str)
90@dataclass(frozen=True)
91class BroadcastSource:
92 """
93 Where a ReadItem gets its data when more than one rank needs the very same bytes.
95 Only ``src_rank`` reads them from storage; the other ranks of ``group_ranks`` receive
96 them from it. Which rank is the source is decided once, globally, while the load plan is
97 built, so every rank of the group names the same one without having to ask.
99 Attributes:
100 group_ranks: Ascending global ranks that load identical data.
101 src_rank: Member of the group that reads from storage and sends to the rest.
102 """
104 group_ranks: tuple[int, ...]
105 src_rank: int
108@dataclass(frozen=True)
109class ReadItem:
110 """
111 Item to be read from storage.
113 Represents a single logical read operation, mapping from checkpoint storage
114 to destination state_dict location.
116 Attributes:
117 type: Type of load item (TENSOR or BYTE_IO).
118 dest_index: Metadata index identifying the destination in state_dict.
119 dest_offsets: Offsets into the destination tensor (for TENSOR type).
120 storage_index: Metadata index identifying the source in checkpoint.
121 storage_offsets: Offsets into the checkpoint storage data.
122 lengths: Size of the hypercube to copy (dimensions of the data region).
123 broadcastable: Whether this read is worth handing to a collective. False says the
124 destination stayed in host memory, where what a load keeps is little and few --
125 an optimizer's step counter, while the moments beside it follow the parameter
126 onto the accelerator -- so every rank that wants it reads it, like a pickled
127 entry, rather than one reading it and sending it to the rest. Reading it costs
128 less than moving it, and the group it would move through, raised on the
129 accelerator library, holds no backend for host memory in any case. The rank
130 that holds the destination is the only one that can tell, so it says so here
131 and the global plan hears it through the gather. Default True.
132 source: Set when other ranks load the same bytes, naming the one that reads them.
133 None means this rank reads the item itself, which is the only case until a
134 global plan decides otherwise, and the only case at all for an item the line
135 above rules out. Default None.
136 """
137 type: LoadItemType
138 dest_index: MetadataIndex # Index into the state_dict
139 dest_offsets: tuple # Offsets into destination tensor
140 storage_index: MetadataIndex # Index into the checkpoint
141 storage_offsets: tuple # Offset into the checkpoint data
142 lengths: tuple # Size of the hypercube to copy
143 broadcastable: bool = True # Whether sending this is worth what it costs, or possible
144 source: Optional[BroadcastSource] = None # Who reads it, once a global plan has said
147@dataclass(frozen=True)
148class SavePlan:
149 """
150 Plan for saving checkpoint.
152 Contains write items and optional storage/planner-specific data.
154 Frozen so that a plan cannot be edited after it has been handed on: plans are gathered from
155 every rank, kept in the planner's cross-call result cache, and passed to hooks that must treat
156 them as inputs. Build a changed plan with ``dataclasses.replace`` instead. Note this freezes
157 the fields, not the contents -- ``items`` is still a list that could be appended to.
159 Attributes:
160 items: List of WriteItems to be saved. Default [].
161 storage_data: Storage-specific data (optional). Default None.
162 planner_data: Planner-specific data (optional). Default None.
163 """
164 items: list[WriteItem] = field(default_factory=list)
165 storage_data: Optional[dict[MetadataIndex, Any]] = None # Storage-specific data mapping
166 planner_data: Any = None # Planner-specific data (can be any type)
169# Stand-in for the checkpoint location of a read that has been elided from a plan about to
170# cross the wire. One shared instance, so pickling stores it once for the whole plan.
171_STORAGE_ELIDED = MetadataIndex("")
174@dataclass
175class LoadPlan:
176 """
177 Plan for loading checkpoint.
179 Contains read items and optional storage/planner-specific data.
181 Frozen for the same reason as :class:`SavePlan`.
183 Attributes:
184 items: List of ReadItems to be loaded. Default [].
185 storage_data: Storage-specific data (optional). Default None.
186 planner_data: Planner-specific data (optional). Default None.
187 """
188 items: list[ReadItem] = field(default_factory=list)
189 storage_data: Optional[dict[MetadataIndex, Any]] = None # Storage-specific data mapping
190 planner_data: Any = None # Planner-specific data (can be any type)
192 def identity(self) -> "LoadPlan":
193 """
194 A copy carrying only what other ranks need in order to recognize these reads.
196 :meth:`LoadPlanner.build_global_plan` matches reads across ranks by the shard they
197 land in, :attr:`ReadItem.dest_index`, and never looks at where a read comes from in
198 the checkpoint. That part is each rank own business -- it resolves it from metadata
199 it already holds -- so it does not have to travel through the all-gather that feeds
200 the global plan, which is one index and one offset tuple per read saved.
202 The elided fields are replaced by a shared placeholder rather than dropped, so the
203 result is still an ordinary LoadPlan. Only ever pass the result to the gather: the
204 rank that executes a plan keeps its own untouched copy.
206 Everything not named here travels, which is what the global plan matches and decides
207 on: the shard a read lands in, and whether a collective could write into it.
209 Returns:
210 LoadPlan: The same items with their checkpoint location elided.
211 """
212 return dataclasses.replace(
213 self,
214 items=[
215 dataclasses.replace(item, dest_offsets=(), storage_index=_STORAGE_ELIDED, storage_offsets=())
216 for item in self.items
217 ],
218 storage_data=None,
219 planner_data=None,
221 )
224class SavePlanner(abc.ABC):
225 """Abstract base class for save planners."""
227 @abc.abstractmethod
228 def configure_planner(self, state_dict: dict[str, Any], **kwargs: Any) -> None:
229 """
230 Configure the planner with state dict.
232 Args:
233 state_dict (dict[str, Any]): The state_dict to save.
234 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank, remove_redundancy,
235 save_to_minimum_rank).
236 """
238 @abc.abstractmethod
239 def build_local_plan(self) -> SavePlan:
240 """
241 Build local save plan.
243 Creates a plan for saving checkpoint data from the current rank's perspective.
244 This plan contains WriteItems for all tensors and bytes that this rank needs to save.
246 Returns:
247 SavePlan: Local save plan containing WriteItems for this rank.
248 """
250 @abc.abstractmethod
251 def build_global_plan(self, all_plans: list[SavePlan]) -> tuple[SavePlan, Metadata]:
252 """
253 Build this rank's final save plan and the global checkpoint metadata.
255 Combines local plans from all ranks: deduplicates redundant data across ranks and assigns
256 storage indices. Every rank runs this over the same gathered plans and writes only its own
257 shards, so only this rank's plan is returned -- the metadata still describes every rank's.
259 Args:
260 all_plans (list[SavePlan]): Local save plans from all ranks, indexed by rank.
262 Returns:
263 tuple[SavePlan, Metadata]: This rank's save plan with storage indices assigned, and
264 checkpoint metadata containing information about all saved items.
265 """
267 @abc.abstractmethod
268 def finalize_plan(self, plan: SavePlan) -> SavePlan:
269 """
270 Finalize the plan.
272 Performs any final adjustments to the plan before execution, such as updating
273 tensor cache keys or performing planner-specific optimizations.
275 Args:
276 plan (SavePlan): The plan to finalize.
278 Returns:
279 SavePlan: The finalized plan ready for execution.
280 """
282 @abc.abstractmethod
283 def get_data(self, item: WriteItem) -> Any:
284 """
285 Get runtime data for a write item from the current state_dict.
287 Args:
288 item (WriteItem): The write item to get data for.
290 Returns:
291 Any: Runtime object to be written for this item.
292 """
295class LoadPlanner(abc.ABC):
296 """Abstract base class for load planners."""
298 @abc.abstractmethod
299 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs: Any) -> None:
300 """
301 Configure the planner with state dict and metadata.
303 Args:
304 state_dict (dict[str, Any]): The state_dict to load into (modified in-place).
305 metadata (Metadata): Checkpoint metadata.
306 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank).
307 """
309 @abc.abstractmethod
310 def build_local_plan(self) -> LoadPlan:
311 """
312 Build local load plan.
314 Creates a plan for loading checkpoint data from the current rank's perspective.
315 This plan contains ReadItems for all tensors and bytes that this rank needs to load.
317 Returns:
318 LoadPlan: Local load plan containing ReadItems for this rank.
319 """
321 @abc.abstractmethod
322 def build_global_plan(self, all_plans: list[LoadPlan]) -> LoadPlan:
323 """
324 Build this rank's final load plan from all local plans.
326 Every rank runs this over the same gathered plans and reads only its own shards, so only
327 this rank's plan is returned. This method may coordinate across ranks or perform
328 optimizations while doing so.
330 Args:
331 all_plans (list[LoadPlan]): Local load plans from all ranks, indexed by rank.
333 Returns:
334 LoadPlan: This rank's load plan.
335 """
337 @abc.abstractmethod
338 def finalize_plan(self, plan: LoadPlan) -> LoadPlan:
339 """
340 Finalize the plan.
342 Performs any final adjustments to the plan before execution, such as
343 performing planner-specific optimizations or validations.
345 Args:
346 plan (LoadPlan): The plan to finalize.
348 Returns:
349 LoadPlan: The finalized plan ready for execution.
350 """
352 @abc.abstractmethod
353 def acquire_tensor(self, read_item: ReadItem) -> Any:
354 """
355 Acquire tensor for read item.
357 Returns a tensor slice/view where data should be written.
359 The view must be **detached**: the caller writes the loaded slice into it in place,
360 and an in-place write into a view of a leaf parameter that requires grad raises
361 rather than loading. Detaching is left to the implementation because only it knows
362 how the destination was reached.
364 Args:
365 read_item (ReadItem): Read item to acquire tensor for.
367 Returns:
368 Any: Acquired tensor slice/view (tensor-like object), detached from autograd.
369 """
371 @abc.abstractmethod
372 def apply_tensor(self, read_item: ReadItem, tensor: Any) -> None:
373 """
374 Apply tensor after reading.
376 Args:
377 read_item (ReadItem): Read item.
378 tensor (Any): Tensor data to apply (tensor-like object).
379 """
381 @abc.abstractmethod
382 def apply_bytes(self, read_item: ReadItem, value: bytes) -> None:
383 """
384 Apply bytes data.
386 Args:
387 read_item (ReadItem): The read item specifying the destination.
388 value (bytes): The bytes data to deserialize and apply.
389 """