Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / standard_planner.py: 77%
273 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
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"""Standard planner implementations for checkpoint save and load."""
16from dataclasses import dataclass
17import dataclasses
18import pickle
19from typing import Any, Optional, Union
21from hyper_parallel.core.distributed_checkpoint.metadata import (
22 CHUNK_INFO,
23 Metadata,
24 MetadataIndex,
25 ChunkStorageMetadata,
26 ChunkInfo,
27 TensorStorageMetadata,
28 TensorProperties,
29 BytesStorageMetadata
30)
31from hyper_parallel.core.distributed_checkpoint.planner import (
32 SavePlan,
33 SavePlanner,
34 LoadPlan,
35 LoadPlanner,
36 WriteItem,
37 WriteItemType,
38 ReadItem,
39 LoadItemType
40)
41from hyper_parallel.core.distributed_checkpoint.reshard import infer_slice_area_by_rank, infer_intersection
42from hyper_parallel.core.distributed_checkpoint.util import (
43 narrow_tensor_by_index,
44 chunk_to_area,
45 create_chunk_list_for_tensor,
46 remove_redundant_plans,
47 flatten_state_dict,
48 set_element,
49)
50from hyper_parallel.core.dtensor.dtensor import DTensor
51from hyper_parallel.core.dtensor.layout import Layout
52from hyper_parallel.platform import get_platform
54platform = get_platform()
55Tensor = platform.Tensor
58@dataclass(frozen=True)
59class CachedSaveResult:
60 """Cached finalized save result keyed by planner cache namespace."""
62 final_plan: SavePlan
63 metadata: Metadata
66class StandardSavePlanner(SavePlanner):
67 """Standard implementation of SavePlanner for distributed checkpoint saving."""
69 _cached_save_result: dict[str, CachedSaveResult] = {}
71 def __init__(
72 self,
73 enable_plan_caching: bool = True,
74 remove_redundancy: bool = True,
75 save_to_minimum_rank: bool = False,
76 ):
77 self.state_dict: Optional[dict[str, Any]] = None
78 self.is_coordinator: bool = False
79 self.rank: int = 0
80 self.remove_redundancy: bool = remove_redundancy
81 self.save_to_minimum_rank: bool = save_to_minimum_rank
82 self.flatten_state_dict: bool = True
83 self._enable_plan_caching: bool = enable_plan_caching
84 self._cached_plans_key: str = self.__class__.__name__
86 def configure_planner(self, state_dict: dict[str, Any], **kwargs) -> None:
87 """
88 Configure planner.
90 Args:
91 state_dict (dict[str, Any]): The state_dict to save.
92 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank, remove_redundancy,
93 save_to_minimum_rank).
94 """
95 self.is_coordinator = kwargs.get("is_coordinator", False)
96 self.rank = kwargs.get("rank", 0)
97 self.remove_redundancy = kwargs.get("remove_redundancy", self.remove_redundancy)
98 self.save_to_minimum_rank = kwargs.get("save_to_minimum_rank", self.save_to_minimum_rank)
99 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
101 use_collectives = bool(kwargs.get("use_collectives", True))
102 if not use_collectives:
103 self.remove_redundancy = False
104 self._enable_plan_caching = False
105 elif "enable_plan_caching" in kwargs:
106 self._enable_plan_caching = bool(kwargs["enable_plan_caching"])
108 if self.flatten_state_dict:
109 state_dict, self.name_mapping = flatten_state_dict(state_dict)
110 self.state_dict = state_dict
111 self._cached_plans_key = self._build_cache_key(state_dict)
113 def _build_cache_key(self, state_dict: dict[str, Any]) -> str:
114 """Build a stable cache namespace from sorted state_dict keys."""
115 return f"{self.__class__.__name__}:{'||'.join(state_dict.keys())}"
117 def build_local_plan(self) -> SavePlan:
118 """
119 Create local save plan.
121 Returns:
122 SavePlan: Local save plan containing WriteItems for this rank.
123 """
124 if self.state_dict is None:
125 raise RuntimeError("Planner not set up")
127 def compute_global_offsets(global_shape: tuple[int, ...], dtensor_layout: Layout) -> tuple[int, ...]:
128 """
129 Compute the offsets of local tensor in global tensor based on layout.
131 Args:
132 global_shape (tuple[int, ...]): Global shape of the tensor.
133 dtensor_layout (Layout): Layout of the DTensor.
135 Returns:
136 tuple[int, ...]: Tuple of offsets for each dimension.
137 """
138 if dtensor_layout is None:
139 # If layout is None, return all zeros (no sharding)
140 return tuple(0 for _ in global_shape)
142 # Validate layout attributes
143 if not hasattr(dtensor_layout, 'mesh_shape') or dtensor_layout.mesh_shape is None:
144 raise ValueError("Layout must have mesh_shape attribute")
145 if not hasattr(dtensor_layout, 'tensor_map') or dtensor_layout.tensor_map is None:
146 raise ValueError("Layout must have tensor_map attribute")
147 if not hasattr(dtensor_layout, 'rank_list') or dtensor_layout.rank_list is None:
148 raise ValueError("Layout must have rank_list attribute")
150 current_rank = self.rank
151 if current_rank not in dtensor_layout.rank_list:
152 raise ValueError(
153 f"Current rank {current_rank} not found in layout's rank_list {dtensor_layout.rank_list}")
155 inner_rank_id = dtensor_layout.rank_list.index(current_rank)
156 # Calculate slice area using infer_slice_area_by_rank
157 slice_area = infer_slice_area_by_rank(
158 mesh_shape=dtensor_layout.mesh_shape,
159 tensor_map=dtensor_layout.tensor_map,
160 rank_id=inner_rank_id,
161 full_shape=global_shape
162 )
163 # Extract offsets (start values) from slice_area
164 return tuple(start for start, _ in slice_area)
166 items = []
167 for fqn, obj in self.state_dict.items():
168 # Check if it's a DTensor
169 if isinstance(obj, DTensor):
170 # Create write item for DTensor
171 local_tensor = obj.to_local()
172 layout = obj.layout
174 # Get chunk metadata with offsets
175 if layout:
176 offsets = compute_global_offsets(obj.shape, layout)
177 else:
178 offsets = (0,) * len(local_tensor.shape)
180 sizes = local_tensor.shape
181 chunk = ChunkStorageMetadata(offsets=offsets, sizes=sizes)
182 # Get tensor properties
183 dtype_str = str(local_tensor.dtype) if hasattr(local_tensor, 'dtype') else 'unknown'
184 properties = TensorProperties(dtype=dtype_str)
185 # Create write item for this tensor
186 index = MetadataIndex(fqn=fqn, offset=offsets, index=None)
187 write_item = WriteItem(
188 index=index,
189 type=WriteItemType.TENSOR,
190 tensor_data={
191 'chunk': chunk,
192 'properties': properties,
193 'size': obj.shape,
194 }
195 )
196 items.append(write_item)
197 elif isinstance(obj, Tensor):
198 # Create write item for platform.Tensor: build single chunk with tensor's own size
199 dtype_str = str(obj.dtype) if hasattr(obj, 'dtype') else 'unknown'
200 properties = TensorProperties(dtype=dtype_str)
201 # handle Tensor with shard information
202 if hasattr(obj, CHUNK_INFO):
203 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo):
204 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance")
205 chunk = getattr(obj, CHUNK_INFO).chunk
206 # Single chunk covering the whole tensor (offsets=0, sizes=shape)
207 else:
208 chunk = ChunkStorageMetadata(
209 offsets=(0,) * len(obj.shape),
210 sizes=obj.shape,
211 )
212 index = MetadataIndex(fqn=fqn, offset=chunk.offsets, index=None)
213 write_item = WriteItem(
214 index=index,
215 type=WriteItemType.TENSOR,
216 tensor_data={
217 'chunk': chunk,
218 'properties': properties,
219 'size': getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) else obj.shape,
220 }
221 )
222 items.append(write_item)
223 else:
224 # Handle non-tensor types (bytes, etc.)
225 index = MetadataIndex(fqn=fqn)
226 write_item = WriteItem(
227 index=index,
228 type=WriteItemType.BYTE_IO,
229 bytes_io_data=None
230 )
231 items.append(write_item)
233 plan = SavePlan(items=items)
234 if self.flatten_state_dict:
235 plan.planner_data = self.name_mapping
236 return plan
238 def build_global_plan(self, all_plans: list[SavePlan]) -> tuple[list[SavePlan], Metadata]:
239 """
240 Build global plan from all local plans.
242 Collects chunks from all ranks, validates consistency, and creates metadata for the checkpoint.
244 Args:
245 all_plans (list[SavePlan]): List of local plans from all ranks.
247 Returns:
248 tuple[list[SavePlan], Metadata]: Updated plans and checkpoint metadata.
249 """
250 # Deduplicate plans if redundancy removal is enabled
251 if self.remove_redundancy and len(all_plans) > 1:
252 all_plans = remove_redundant_plans(all_plans, save_to_minimum_rank=self.save_to_minimum_rank)
254 # Collect all write items by FQN
255 fqn_to_chunks: dict[str, list[ChunkStorageMetadata]] = {}
256 fqn_to_properties: dict[str, TensorProperties] = {}
257 fqn_to_size: dict[str, tuple] = {}
258 state_dict_metadata: dict[str, Union[TensorStorageMetadata, BytesStorageMetadata]] = {}
260 final_global_plans: list[SavePlan] = []
261 for plan in all_plans:
262 with_index_items = []
263 for item in plan.items:
264 if item.type == WriteItemType.TENSOR and item.tensor_data:
265 fqn = item.index.fqn
266 chunk = item.tensor_data['chunk']
267 properties = item.tensor_data['properties']
268 size = item.tensor_data['size']
270 # Validate consistency across ranks
271 if fqn in fqn_to_chunks and (fqn_to_properties[fqn] != properties or fqn_to_size[fqn] != size):
272 raise ValueError(f"The {fqn} in different rank has different properties and size.")
274 # Initialize FQN entry if not exists
275 if fqn not in fqn_to_chunks:
276 fqn_to_properties[fqn] = properties
277 fqn_to_size[fqn] = size
278 fqn_to_chunks[fqn] = []
280 # Append chunk and set index (platform.Tensor has exactly one chunk)
281 new_index = dataclasses.replace(item.index, index=len(fqn_to_chunks[fqn]))
282 with_index_item = dataclasses.replace(item, index=new_index)
283 with_index_items.append(with_index_item)
284 fqn_to_chunks[fqn].append(chunk)
286 elif item.type == WriteItemType.BYTE_IO:
287 with_index_items.append(item)
288 state_dict_metadata[item.index.fqn] = BytesStorageMetadata()
289 else:
290 raise ValueError(f"Unsupported write item type: {item.type}")
292 final_global_plans.append(dataclasses.replace(plan, items=with_index_items))
294 # Create metadata for all tensors
295 for fqn, chunks in fqn_to_chunks.items():
296 state_dict_metadata[fqn] = TensorStorageMetadata(
297 properties=fqn_to_properties[fqn],
298 size=fqn_to_size[fqn],
299 chunks=chunks
300 )
302 metadata = Metadata(state_dict_metadata=state_dict_metadata)
303 if self.flatten_state_dict:
304 merged_mapping = {}
305 for p in all_plans:
306 merged_mapping.update(p.planner_data)
307 metadata.planner_data = merged_mapping
308 return final_global_plans, metadata
310 def finalize_plan(self, plan: SavePlan) -> SavePlan:
311 """
312 Finalize the plan.
314 Args:
315 plan (SavePlan): Plan to finalize.
317 Returns:
318 SavePlan: Finalized plan.
319 """
320 return plan
322 def get_cached(self) -> Optional[CachedSaveResult]:
323 """Return cached finalized plan and metadata when plan caching is enabled."""
324 if (
325 not self._enable_plan_caching
326 or self._cached_plans_key not in StandardSavePlanner._cached_save_result
327 ):
328 return None
329 return StandardSavePlanner._cached_save_result[self._cached_plans_key]
331 def cache_result(self, final_plan: SavePlan, metadata: Metadata) -> None:
332 """Store finalized plan and metadata in the class-level planner cache."""
333 if not self._enable_plan_caching:
334 return
335 StandardSavePlanner._cached_save_result[self._cached_plans_key] = CachedSaveResult(
336 final_plan=final_plan,
337 metadata=metadata,
338 )
340 def get_data(self, item: WriteItem) -> Any:
341 """
342 Get current runtime data from state_dict for a write item.
344 Args:
345 item (WriteItem): Write item describing what to write.
347 Returns:
348 Any: Runtime object to be written.
349 """
350 if self.state_dict is None:
351 raise RuntimeError("Planner not set up")
352 fqn = item.index.fqn
353 if fqn not in self.state_dict:
354 raise KeyError(f"Key {fqn} not found in state_dict")
355 obj = self.state_dict[fqn]
356 if item.type == WriteItemType.TENSOR:
357 if isinstance(obj, DTensor):
358 return obj.to_local().detach().cpu()
359 if isinstance(obj, Tensor):
360 return obj.detach().cpu()
361 raise TypeError(f"Write item {fqn} expected tensor-like object, got {type(obj)}")
362 if item.type == WriteItemType.BYTE_IO:
363 return obj
364 raise TypeError(f"Unsupported write item type: {item.type}")
367def create_read_items_for_chunk_list(
368 fqn: str,
369 checkpoint_md: TensorStorageMetadata,
370 local_chunks: list[ChunkStorageMetadata],
371) -> list[ReadItem]:
372 """
373 Create ReadItems by matching local chunks (what this rank needs) with
374 saved chunks (checkpoint_md.chunks), including resharding overlaps.
376 Mirrors torch create_read_items_for_chunk_list behavior.
378 Args:
379 fqn (str): Fully qualified name of the tensor.
380 checkpoint_md (TensorStorageMetadata): Tensor storage metadata from checkpoint.
381 local_chunks (list[ChunkStorageMetadata]): List of local chunks needed by this rank.
383 Returns:
384 list[ReadItem]: List of ReadItems for loading the required data.
385 """
386 read_items: list[ReadItem] = []
387 saved_chunks = checkpoint_md.chunks
388 if not local_chunks or not saved_chunks:
389 return read_items
391 for local_idx, local_chunk in enumerate(local_chunks):
392 local_area = chunk_to_area(local_chunk)
393 for storage_idx, storage_chunk in enumerate(saved_chunks):
394 saved_area = chunk_to_area(storage_chunk)
395 overlap = infer_intersection(local_area, saved_area)
396 if overlap is None:
397 continue
399 dest_offsets = tuple(overlap[i][0] - local_chunk.offsets[i] for i in range(len(overlap)))
400 storage_offsets = tuple(overlap[i][0] - storage_chunk.offsets[i] for i in range(len(overlap)))
401 lengths = tuple(overlap[i][1] - overlap[i][0] for i in range(len(overlap)))
403 read_items.append(
404 ReadItem(
405 type=LoadItemType.TENSOR,
406 dest_index=MetadataIndex(fqn=fqn, offset=local_chunk.offsets, index=local_idx),
407 dest_offsets=dest_offsets,
408 storage_index=MetadataIndex(fqn=fqn, offset=storage_chunk.offsets, index=storage_idx),
409 storage_offsets=storage_offsets,
410 lengths=lengths,
411 )
412 )
413 return read_items
416class StandardLoadPlanner(LoadPlanner):
417 """
418 Standard implementation of LoadPlanner.
420 Iterate state_dict and creates load plans via chunk list for resharding support.
421 """
423 def __init__(self, allow_partial_load: bool = False):
424 """
425 Args:
426 allow_partial_load (bool): If True, allow loading when checkpoint has fewer keys than state_dict.
427 Default False.
428 """
429 self.state_dict: Optional[dict[str, Any]] = None
430 self.metadata: Optional[Metadata] = None
431 self.is_coordinator: bool = False
432 self.rank: int = 0
433 self.allow_partial_load = allow_partial_load
434 self.flatten_state_dict: bool = True
436 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
437 """
438 Configure planner with state dict and metadata.
440 Args:
441 state_dict (dict[str, Any]): The state_dict to load into (modified in-place).
442 metadata (Metadata): Checkpoint metadata.
443 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank).
444 """
445 self.state_dict = state_dict
446 self.metadata = metadata
447 self.is_coordinator = kwargs.get("is_coordinator", False)
448 self.rank = kwargs.get("rank", 0)
449 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
450 self.original_state_dict = state_dict
451 if self.flatten_state_dict:
452 state_dict, self.name_mapping = flatten_state_dict(state_dict)
453 self.state_dict = state_dict
455 def build_local_plan(self) -> LoadPlan:
456 """
457 Build local load plan.
459 Iterate state_dict and creates load plans via chunk list for resharding support.
461 Returns:
462 LoadPlan: Local load plan containing ReadItems for this rank.
463 """
464 if self.state_dict is None or self.metadata is None:
465 raise RuntimeError("Planner not configured")
467 requests: list[ReadItem] = []
468 strict = not self.allow_partial_load
469 for fqn, obj in self.state_dict.items():
470 if fqn not in self.metadata.state_dict_metadata:
471 if fqn.endswith(('matched_adamw_rms', 'step')):
472 continue
473 if strict:
474 raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
475 continue
476 md = self.metadata.state_dict_metadata[fqn]
477 if isinstance(md, TensorStorageMetadata):
478 obj_size = getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) \
479 else getattr(obj, "shape", None)
480 if obj_size is None or md.size != tuple(obj_size):
481 raise ValueError(
482 f"Size mismatch between saved {md.size} and current: {obj_size} for {fqn}",
483 )
484 if isinstance(obj, DTensor):
485 layout = getattr(obj, "layout", None)
486 rank_list = getattr(layout, "rank_list", None) if layout else None
487 if rank_list is None and layout is not None:
488 rank_list = getattr(layout, "_rank_list", None)
489 if layout is not None and rank_list is not None:
490 if get_platform().get_rank() not in rank_list:
491 continue
492 # Both DTensor and platform.Tensor: create local chunks and read items
493 local_chunks = create_chunk_list_for_tensor(obj)
494 requests += create_read_items_for_chunk_list(fqn, md, local_chunks)
495 else:
496 requests.append(
497 ReadItem(
498 type=LoadItemType.BYTE_IO,
499 dest_index=MetadataIndex(fqn=fqn),
500 dest_offsets=(0,),
501 storage_index=MetadataIndex(fqn=fqn),
502 storage_offsets=(0,),
503 lengths=(0,),
504 )
505 )
506 return LoadPlan(items=requests)
508 def build_global_plan(self, all_plans: list[LoadPlan]) -> list[LoadPlan]:
509 """
510 Build global plan from all local plans.
512 For now, returns plans as-is. In a more sophisticated implementation, you might need to coordinate across ranks.
514 Args:
515 all_plans (list[LoadPlan]): List of local plans from all ranks.
517 Returns:
518 list[LoadPlan]: Global plans (currently returns plans as-is).
519 """
520 return all_plans
522 def finalize_plan(self, plan: LoadPlan) -> LoadPlan:
523 """
524 Finalize the plan (no-op for default implementation).
526 Args:
527 plan (LoadPlan): Plan to finalize.
529 Returns:
530 LoadPlan: Finalized plan.
531 """
532 return plan
534 def acquire_tensor(self, read_item: ReadItem) -> Any:
535 """
536 Acquire the destination slice (narrow view) for this read_item.
538 StorageReader uses this to copy loaded data into the correct region.
539 Torch-aligned behavior.
541 Args:
542 read_item (ReadItem): The read item specifying what to load.
544 Returns:
545 Any: The destination tensor slice where data should be written
546 (tensor-like object).
547 """
548 if self.state_dict is None:
549 raise RuntimeError("Planner not configured")
551 fqn = read_item.dest_index.fqn
552 if fqn not in self.state_dict:
553 raise KeyError(f"Key {fqn} not found in state_dict")
555 target = self.state_dict[fqn]
556 local_tensor = target.to_local().detach() if isinstance(target, DTensor) else target.detach()
557 return narrow_tensor_by_index(
558 local_tensor,
559 read_item.dest_offsets,
560 read_item.lengths,
561 )
563 def apply_tensor(self, read_item: ReadItem, tensor: Any) -> None:
564 """
565 Apply tensor after reading.
567 After read_data copies into the slice, this is no-op when tensor is the
568 same slice. When the backend has no copy_ (e.g. mindspore), read_data
569 passes the loaded slice here; we copy it into the destination slice.
571 Args:
572 read_item (ReadItem): The read item that was processed.
573 tensor (Any): The tensor data to apply (tensor-like object).
574 """
575 if tensor is None:
576 return
577 dest_slice = self.acquire_tensor(read_item)
578 if dest_slice is tensor:
579 return
580 if hasattr(dest_slice, "copy_"):
581 dest_slice.copy_(tensor)
582 else:
583 # Fallback: assign into state_dict if supported
584 dest_slice[...] = tensor
586 def apply_bytes(self, read_item: ReadItem, value: bytes) -> None:
587 """
588 Load bytes data into state_dict.
590 Args:
591 read_item (ReadItem): The read item specifying the destination.
592 value (bytes): The bytes data to deserialize and load.
593 """
594 if self.state_dict is None:
595 raise RuntimeError("Planner not set up")
597 fqn = read_item.dest_index.fqn
598 # Deserialize bytes
599 obj = pickle.loads(value)
600 self.state_dict[fqn] = obj
601 if self.flatten_state_dict:
602 set_element(self.original_state_dict, self.name_mapping[fqn], obj)
606class _DcpMergeLoadPlanner(StandardLoadPlanner):
607 """Load planner that builds distributed checkpoint from dcp into fully ``state_dict`` (in-place)."""
609 def __init__(self) -> None:
610 super().__init__()
612 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs) -> None:
613 if len(state_dict) > 0:
614 raise ValueError(
615 "state_dict must be empty for _DcpMergeLoadPlanner; "
616 "it is populated in-place from checkpoint metadata."
617 )
619 if metadata is None:
620 raise ValueError("metadata must not be None for _DcpMergeLoadPlanner.")
622 self.is_coordinator = kwargs.get("is_coordinator", False)
623 for k, v in metadata.state_dict_metadata.items():
624 if isinstance(v, TensorStorageMetadata):
625 v = platform.empty(
626 platform.list_to_size(v.size),
627 dtype=platform.str_to_dtype(v.properties.dtype),
628 )
630 state_dict[k] = v
631 if metadata.planner_data is not None and k in metadata.planner_data:
632 set_element(state_dict, metadata.planner_data[k], v)
634 super().configure_planner(
635 state_dict,
636 metadata,
637 is_coordinator=self.is_coordinator,
638 flatten_state_dict=True,
639 )