Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / standard_planner.py: 83%
321 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# ============================================================================
16"""Standard planner implementations for checkpoint save and load."""
18from dataclasses import dataclass
19from collections import defaultdict
20import dataclasses
21import math
22from itertools import compress
23import pickle
24from typing import Any, Optional, Union
26import torch
27from torch import Tensor
29from hyper_parallel.core.distributed_checkpoint.metadata import (
30 CHUNK_INFO,
31 Metadata,
32 MetadataIndex,
33 dtype_element_size,
34 ChunkStorageMetadata,
35 ChunkInfo,
36 TensorStorageMetadata,
37 TensorProperties,
38 BytesStorageMetadata
39)
40from hyper_parallel.core.distributed_checkpoint.planner import (
41 SavePlan,
42 SavePlanner,
43 LoadPlan,
44 LoadPlanner,
45 WriteItem,
46 WriteItemType,
47 ReadItem,
48 LoadItemType,
49 BroadcastSource,
50)
51from hyper_parallel.core.distributed_checkpoint.ragged import (
52 create_ragged_write_items,
53 get_ragged_box_tensor,
54)
55from hyper_parallel.core.distributed_checkpoint.utils import (
56 narrow_tensor_by_index,
57 chunk_to_area,
58 create_chunk_list_for_tensor,
59 infer_intersection,
60 plan_ownership_masks,
61 flatten_state_dict,
62 set_element,
63 dcp_timer_decorator,
64 logger,
65 str_to_dtype,
66)
67from hyper_parallel.core.dtensor.dtensor import DTensor
68from hyper_parallel.core.dtensor.layout import Layout, infer_slice_area_by_layout
71def _own_plan_index(all_plans: Union[list[SavePlan], list[LoadPlan]], rank: int) -> int:
72 """
73 Locate this rank's plan in the gathered list.
75 Both gather paths return one entry per rank in rank order; without collectives the
76 "gather" is just this rank's own plan in a one element list.
78 Args:
79 all_plans (Union[list[SavePlan], list[LoadPlan]]): Local plans from all ranks.
80 rank (int): Rank looking for its own plan.
82 Returns:
83 int: Index of this rank's plan.
85 Raises:
86 ValueError: If the gathered list holds no plan for this rank.
87 """
88 own_index = rank if len(all_plans) > 1 else 0
89 if not 0 <= own_index < len(all_plans):
90 raise ValueError(
91 f"Rank {rank} has no plan of its own among the {len(all_plans)} gathered "
92 "plans; the gathered list must hold one plan per rank, in rank order."
93 )
94 return own_index
97def _compute_global_offsets(
98 global_shape: tuple[int, ...],
99 dtensor_layout: Layout,
100 current_rank: int,
101) -> tuple[int, ...]:
102 """
103 Compute the offsets of local tensor in global tensor based on layout.
105 Args:
106 global_shape (tuple[int, ...]): Global shape of the tensor.
107 dtensor_layout (Layout): Layout of the DTensor.
108 current_rank (int): Rank whose shard offsets are wanted.
110 Returns:
111 tuple[int, ...]: Tuple of offsets for each dimension.
112 """
113 if dtensor_layout is None:
114 # If layout is None, return all zeros (no sharding)
115 return tuple(0 for _ in global_shape)
117 # Validate layout attributes
118 if not hasattr(dtensor_layout, 'mesh_shape') or dtensor_layout.mesh_shape is None:
119 raise ValueError("Layout must have mesh_shape attribute")
120 if not hasattr(dtensor_layout, 'tensor_map') or dtensor_layout.tensor_map is None:
121 raise ValueError("Layout must have tensor_map attribute")
122 if not hasattr(dtensor_layout, 'rank_list') or dtensor_layout.rank_list is None:
123 raise ValueError("Layout must have rank_list attribute")
125 if current_rank not in dtensor_layout.rank_list:
126 raise ValueError(
127 f"Current rank {current_rank} not found in layout's rank_list {dtensor_layout.rank_list}")
129 inner_rank_id = dtensor_layout.rank_list.index(current_rank)
130 # Calculate slice area using infer_slice_area_by_rank
131 slice_area = infer_slice_area_by_layout(
132 dtensor_layout,
133 inner_rank_id,
134 global_shape,
135 )
136 # Extract offsets (start values) from slice_area
137 return tuple(start for start, _ in slice_area)
140@dataclass(frozen=True)
141class CachedSaveResult:
142 """Cached finalized save result keyed by planner cache namespace."""
144 final_plan: SavePlan
145 metadata: Metadata
148class StandardSavePlanner(SavePlanner):
149 """Standard implementation of SavePlanner for distributed checkpoint saving."""
151 cached_save_result: dict[str, CachedSaveResult] = {}
153 def __init__(
154 self,
155 enable_plan_caching: bool = True,
156 remove_redundancy: bool = True,
157 save_to_minimum_rank: bool = False,
158 ) -> None:
159 """
160 Args:
161 enable_plan_caching (bool): Reuse the plan and metadata of an earlier save when the
162 state dict has not changed shape. Default True.
163 remove_redundancy (bool): Let only one rank write each replicated item. Default True.
164 save_to_minimum_rank (bool): Give a redundant item to the lowest-numbered plan that
165 holds it rather than to the one with the least planned storage. Default False.
166 """
167 self.state_dict: Optional[dict[str, Any]] = None
168 self.is_coordinator: bool = False
169 self.rank: int = 0
170 self.remove_redundancy: bool = remove_redundancy
171 self.save_to_minimum_rank: bool = save_to_minimum_rank
172 self.flatten_state_dict: bool = True
173 self._enable_plan_caching: bool = enable_plan_caching
174 self._cached_plans_key: str = self.__class__.__name__
176 def configure_planner(self, state_dict: dict[str, Any], **kwargs: Any) -> None:
177 """
178 Configure planner.
180 Args:
181 state_dict (dict[str, Any]): The state_dict to save.
182 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank, remove_redundancy,
183 save_to_minimum_rank).
184 """
185 self.is_coordinator = kwargs.get("is_coordinator", False)
186 self.rank = kwargs.get("rank", 0)
187 self.remove_redundancy = kwargs.get("remove_redundancy", self.remove_redundancy)
188 self.save_to_minimum_rank = kwargs.get("save_to_minimum_rank", self.save_to_minimum_rank)
189 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
191 use_collectives = bool(kwargs.get("use_collectives", True))
192 if not use_collectives:
193 self.remove_redundancy = False
194 self._enable_plan_caching = False
195 elif "enable_plan_caching" in kwargs:
196 self._enable_plan_caching = bool(kwargs["enable_plan_caching"])
198 if self.flatten_state_dict:
199 state_dict, self.name_mapping = flatten_state_dict(state_dict)
200 self.state_dict = state_dict
201 self._cached_plans_key = self._build_cache_key(state_dict)
203 def _build_cache_key(self, state_dict: dict[str, Any]) -> str:
204 """Build a stable cache namespace from state_dict keys."""
205 return f"{self.__class__.__name__}:{'||'.join(state_dict.keys())}"
207 @dcp_timer_decorator
208 def build_local_plan(self) -> SavePlan:
209 """
210 Create local save plan.
212 Returns:
213 SavePlan: Local save plan containing WriteItems for this rank.
214 """
215 if self.state_dict is None:
216 raise RuntimeError("Planner not set up")
218 items = []
219 for fqn, obj in self.state_dict.items():
220 # Check if it's a DTensor
221 if isinstance(obj, DTensor):
222 if obj.layout is not None and obj.layout.ragged_shard is not None:
223 items.extend(create_ragged_write_items(fqn, obj))
224 continue
225 # Create write item for DTensor
226 local_tensor = obj.to_local()
227 layout = obj.layout
229 # Get chunk metadata with offsets
230 if layout:
231 offsets = _compute_global_offsets(obj.shape, layout, self.rank)
232 else:
233 offsets = (0,) * len(local_tensor.shape)
235 sizes = local_tensor.shape
236 chunk = ChunkStorageMetadata(offsets=offsets, sizes=sizes)
237 # Get tensor properties
238 dtype_str = str(local_tensor.dtype) if hasattr(local_tensor, 'dtype') else 'unknown'
239 properties = TensorProperties(dtype=dtype_str)
240 # Create write item for this tensor
241 index = MetadataIndex(fqn=fqn, offset=offsets, index=None)
242 write_item = WriteItem(
243 index=index,
244 type=WriteItemType.TENSOR,
245 tensor_data={
246 'chunk': chunk,
247 'properties': properties,
248 'size': obj.shape,
249 }
250 )
251 items.append(write_item)
252 elif isinstance(obj, Tensor):
253 # Create write item for torch.Tensor: build single chunk with tensor's own size
254 dtype_str = str(obj.dtype) if hasattr(obj, 'dtype') else 'unknown'
255 properties = TensorProperties(dtype=dtype_str)
256 # handle Tensor with shard information
257 if hasattr(obj, CHUNK_INFO):
258 if not isinstance(getattr(obj, CHUNK_INFO), ChunkInfo):
259 raise ValueError("The attr CHUNK_INFO should be a ChunkInfo instance")
260 chunk = getattr(obj, CHUNK_INFO).chunk
261 # Single chunk covering the whole tensor (offsets=0, sizes=shape)
262 else:
263 chunk = ChunkStorageMetadata(
264 offsets=(0,) * len(obj.shape),
265 sizes=obj.shape,
266 )
267 index = MetadataIndex(fqn=fqn, offset=chunk.offsets, index=None)
268 write_item = WriteItem(
269 index=index,
270 type=WriteItemType.TENSOR,
271 tensor_data={
272 'chunk': chunk,
273 'properties': properties,
274 'size': getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) else obj.shape,
275 }
276 )
277 items.append(write_item)
278 else:
279 # Handle non-tensor types (bytes, etc.)
280 index = MetadataIndex(fqn=fqn)
281 write_item = WriteItem(
282 index=index,
283 type=WriteItemType.BYTE_IO,
284 bytes_io_data=None
285 )
286 items.append(write_item)
288 return SavePlan(
289 items=items,
290 planner_data=self.name_mapping if self.flatten_state_dict else None,
291 )
293 @dcp_timer_decorator
294 def build_global_plan(self, all_plans: list[SavePlan]) -> tuple[SavePlan, Metadata]:
295 """
296 Build this rank's final save plan and the global checkpoint metadata.
298 Every rank receives all local plans from the gather and runs this itself, and each one
299 goes on to write only its own shards, so only this rank's plan is rebuilt here. The other
300 ranks' items are still walked -- the metadata is global, and an item's chunk index is its
301 position in that global chunk list -- but no plan object is built for them.
303 Args:
304 all_plans (list[SavePlan]): Local plans from all ranks, indexed by rank.
306 Returns:
307 tuple[SavePlan, Metadata]: This rank's plan with chunk indices assigned, and the
308 checkpoint metadata describing every rank's chunks.
310 Raises:
311 ValueError: If an item has an unsupported type.
312 """
313 own_index = _own_plan_index(all_plans, self.rank)
315 # Redundant items are skipped through a per-plan mask rather than by materialising
316 # deduplicated plans: the loop below walks plan.items anyway.
317 masks = (
318 plan_ownership_masks(all_plans, save_to_minimum_rank=self.save_to_minimum_rank)
319 if self.remove_redundancy and len(all_plans) > 1
320 else None
321 )
323 # FQN -> (properties, size, chunks), collecting every rank's chunks in gather order.
324 fqn_info: dict[str, tuple[TensorProperties, tuple, list[ChunkStorageMetadata]]] = {}
325 state_dict_metadata: dict[str, Union[TensorStorageMetadata, BytesStorageMetadata]] = {}
326 own_items: list[WriteItem] = []
328 for plan_index, plan in enumerate(all_plans):
329 is_own_plan = plan_index == own_index
330 items = plan.items if masks is None else compress(plan.items, masks[plan_index])
331 for item in items:
332 if item.type == WriteItemType.TENSOR and item.tensor_data:
333 tensor_data = item.tensor_data
334 chunks = self._global_chunks_for(
335 item.index.fqn, tensor_data['properties'], tensor_data['size'], fqn_info)
336 if is_own_plan:
337 # Set index (torch.Tensor has exactly one chunk). replace() copies the
338 # remaining fields by name, so adding a field to either dataclass cannot
339 # silently drop it here; only this rank's items are rebuilt, so the cost
340 # is per-rank, not per-cluster.
341 new_index = dataclasses.replace(item.index, index=len(chunks))
342 own_items.append(dataclasses.replace(item, index=new_index))
343 chunks.append(tensor_data['chunk'])
345 elif item.type == WriteItemType.BYTE_IO:
346 if is_own_plan:
347 own_items.append(item)
348 state_dict_metadata[item.index.fqn] = BytesStorageMetadata()
349 else:
350 raise ValueError(f"Unsupported write item type: {item.type}")
352 # Create metadata for all tensors
353 for fqn, (properties, size, chunks) in fqn_info.items():
354 state_dict_metadata[fqn] = TensorStorageMetadata(
355 properties=properties,
356 size=size,
357 chunks=chunks
358 )
360 metadata = Metadata(state_dict_metadata=state_dict_metadata)
361 if self.flatten_state_dict:
362 merged_mapping = {}
363 for p in all_plans:
364 merged_mapping.update(p.planner_data)
365 metadata.planner_data = merged_mapping
366 return dataclasses.replace(all_plans[own_index], items=own_items), metadata
368 @staticmethod
369 def _global_chunks_for(fqn: str, properties: TensorProperties, size: tuple, fqn_info: dict) -> list:
370 """
371 Return the global chunk list for ``fqn``, registering it on first sight.
373 Args:
374 fqn (str): Fully qualified name of the tensor.
375 properties (TensorProperties): Properties this rank reported for it.
376 size (tuple): Global shape this rank reported for it.
377 fqn_info (dict): FQN -> (properties, size, chunks), grown in place.
379 Returns:
380 list: The chunk list every rank's chunks for this FQN are appended to.
382 Raises:
383 ValueError: If another rank described this FQN with different properties or size.
384 """
385 info = fqn_info.get(fqn)
386 if info is None:
387 chunks = []
388 fqn_info[fqn] = (properties, size, chunks)
389 return chunks
390 known_properties, known_size, chunks = info
391 if known_properties != properties or known_size != size:
392 raise ValueError(f"The {fqn} in different rank has different properties and size, "
393 f"properties: {known_properties} != {properties}, "
394 f"size: or {known_size} != {size}.")
395 return chunks
397 def finalize_plan(self, plan: SavePlan) -> SavePlan:
398 """
399 Finalize the plan.
401 Args:
402 plan (SavePlan): Plan to finalize.
404 Returns:
405 SavePlan: Finalized plan.
406 """
407 return plan
409 def get_cached(self) -> Optional[CachedSaveResult]:
410 """Return cached finalized plan and metadata when plan caching is enabled."""
411 if (
412 not self._enable_plan_caching
413 or self._cached_plans_key not in StandardSavePlanner.cached_save_result
414 ):
415 return None
416 return StandardSavePlanner.cached_save_result[self._cached_plans_key]
418 def cache_result(self, final_plan: SavePlan, metadata: Metadata) -> None:
419 """Store finalized plan and metadata in the class-level planner cache."""
420 if not self._enable_plan_caching:
421 return
422 StandardSavePlanner.cached_save_result[self._cached_plans_key] = CachedSaveResult(
423 final_plan=final_plan,
424 metadata=metadata,
425 )
427 def get_data(self, item: WriteItem) -> Any:
428 """
429 Get current runtime data from state_dict for a write item.
431 Args:
432 item (WriteItem): Write item describing what to write.
434 Returns:
435 Any: Runtime object to be written.
436 """
437 if self.state_dict is None:
438 raise RuntimeError("Planner not set up")
439 fqn = item.index.fqn
440 if fqn not in self.state_dict:
441 raise KeyError(f"Key {fqn} not found in state_dict")
442 obj = self.state_dict[fqn]
443 if item.type == WriteItemType.TENSOR:
444 if isinstance(obj, DTensor):
445 if obj.layout is not None and obj.layout.ragged_shard is not None:
446 return get_ragged_box_tensor(obj, item.index).detach().to("cpu")
447 return obj.to_local().detach().to("cpu")
448 if isinstance(obj, Tensor):
449 return obj.detach().to("cpu")
450 raise TypeError(f"Write item {fqn} expected tensor-like object, got {type(obj)}")
451 if item.type == WriteItemType.BYTE_IO:
452 return obj
453 raise TypeError(f"Unsupported write item type: {item.type}")
456def create_read_items_for_chunk_list(
457 fqn: str,
458 checkpoint_md: TensorStorageMetadata,
459 local_chunks: list[ChunkStorageMetadata],
460 broadcastable: bool = True,
461) -> list[ReadItem]:
462 """
463 Create ReadItems by matching local chunks (what this rank needs) with
464 saved chunks (checkpoint_md.chunks), including resharding overlaps.
466 Mirrors torch create_read_items_for_chunk_list behavior.
468 Args:
469 fqn (str): Fully qualified name of the tensor.
470 checkpoint_md (TensorStorageMetadata): Tensor storage metadata from checkpoint.
471 local_chunks (list[ChunkStorageMetadata]): List of local chunks needed by this rank.
472 broadcastable (bool): Whether a collective could write into where these reads land.
473 Default True; see :attr:`ReadItem.broadcastable`.
475 Returns:
476 list[ReadItem]: List of ReadItems for loading the required data.
477 """
478 read_items: list[ReadItem] = []
479 saved_chunks = checkpoint_md.chunks
480 if not local_chunks or not saved_chunks:
481 return read_items
483 for local_idx, local_chunk in enumerate(local_chunks):
484 local_area = chunk_to_area(local_chunk)
485 for storage_idx, storage_chunk in enumerate(saved_chunks):
486 saved_area = chunk_to_area(storage_chunk)
487 overlap = infer_intersection(local_area, saved_area)
488 if overlap is None:
489 continue
491 dest_offsets = tuple(overlap[i][0] - local_chunk.offsets[i] for i in range(len(overlap)))
492 storage_offsets = tuple(overlap[i][0] - storage_chunk.offsets[i] for i in range(len(overlap)))
493 lengths = tuple(overlap[i][1] - overlap[i][0] for i in range(len(overlap)))
495 read_items.append(
496 ReadItem(
497 type=LoadItemType.TENSOR,
498 dest_index=MetadataIndex(fqn=fqn, offset=local_chunk.offsets, index=local_idx),
499 dest_offsets=dest_offsets,
500 storage_index=MetadataIndex(fqn=fqn, offset=storage_chunk.offsets, index=storage_idx),
501 storage_offsets=storage_offsets,
502 lengths=lengths,
503 broadcastable=broadcastable,
504 )
505 )
506 return read_items
509def _is_on_host(obj: Any) -> bool:
510 """
511 Whether a state dict entry's local buffer sits in host memory rather than on the device.
513 An entry that stayed on the host is read by every rank that wants it rather than read
514 once and sent, because sending it is not worth what it costs and often is not possible
515 at all. What stays on the host is little and few: an optimizer keeps its step counter
516 there while the moments beside it follow the parameter onto the device, one fp32 scalar
517 per parameter tensor, so a model of nine thousand of them leaves 36 KiB a rank -- read
518 out of a file that is already open, beside the other shards of that file. The broadcast
519 that would replace it costs a collective and two passes through a staging buffer, and
520 on the host network of a large job that collective is a tree over TCP whose latency
521 alone outweighs the read.
523 And a load that broadcasts at all does so through groups raised on the accelerator
524 library, which holds no backend for host memory and refuses a tensor kept there.
526 What a broadcast carries is the local buffer, the same one :func:`_shard_buffer` hands
527 the collective, so a DTensor is asked about its shard rather than about the whole tensor
528 it is part of. An entry that does not say where it lives is taken to be where the rest
529 of the load is.
531 Args:
532 obj (Any): State dict entry.
534 Returns:
535 bool: True only when the entry says it is in host memory.
536 """
537 local = obj.to_local() if isinstance(obj, DTensor) else obj
538 return getattr(getattr(local, "device", None), "type", None) == "cpu"
541class StandardLoadPlanner(LoadPlanner):
542 """
543 Standard implementation of LoadPlanner.
545 Iterate state_dict and creates load plans via chunk list for resharding support.
546 """
548 def __init__(self, allow_partial_load: bool = False, broadcast_replicated_tensors: bool = False) -> None:
549 """
550 Args:
551 allow_partial_load (bool): If True, allow loading when checkpoint has fewer keys than state_dict.
552 Default False.
553 broadcast_replicated_tensors (bool): If True, a tensor that several ranks load
554 identically is read by one of them and sent to the rest. Off by default;
555 ``configure_planner`` overrides this from :func:`load`.
556 """
557 self.state_dict: Optional[dict[str, Any]] = None
558 self.metadata: Optional[Metadata] = None
559 self.is_coordinator: bool = False
560 self.rank: int = 0
561 self.allow_partial_load = allow_partial_load
562 self.broadcast_replicated_tensors: bool = broadcast_replicated_tensors
563 self.flatten_state_dict: bool = True
565 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs: Any) -> None:
566 """
567 Configure planner with state dict and metadata.
569 Args:
570 state_dict (dict[str, Any]): The state_dict to load into (modified in-place).
571 metadata (Metadata): Checkpoint metadata.
572 **kwargs: Additional keyword arguments (e.g., is_coordinator, rank).
573 """
574 self.state_dict = state_dict
575 self.metadata = metadata
576 self.is_coordinator = kwargs.get("is_coordinator", False)
577 self.rank = kwargs.get("rank", 0)
578 self.broadcast_replicated_tensors = kwargs.get(
579 "broadcast_replicated_tensors", self.broadcast_replicated_tensors
580 )
581 self.flatten_state_dict = kwargs.get("flatten_state_dict", True)
582 self.original_state_dict = state_dict
583 if self.flatten_state_dict:
584 state_dict, self.name_mapping = flatten_state_dict(state_dict)
585 self.state_dict = state_dict
587 def _rank_owns_dtensor_shard(self, obj: Any) -> bool:
588 """
589 Check whether the current rank appears in the rank list of a DTensor layout.
591 Objects that are not DTensors, and DTensors whose layout does not carry a rank list,
592 are always owned by the current rank.
594 Args:
595 obj (Any): State dict entry to check.
597 Returns:
598 bool: False only when the layout has a rank list the current rank is absent from.
599 """
600 if not isinstance(obj, DTensor):
601 return True
602 layout = getattr(obj, "layout", None)
603 if layout is None:
604 return True
605 rank_list = getattr(layout, "rank_list", None) if layout else None
606 if rank_list is None:
607 rank_list = getattr(layout, "_rank_list", None)
608 if rank_list is None:
609 return True
610 return self.rank in rank_list
612 @dcp_timer_decorator
613 def build_local_plan(self) -> LoadPlan:
614 """
615 Build local load plan.
617 Iterate state_dict and creates load plans via chunk list for resharding support.
619 Returns:
620 LoadPlan: Local load plan containing ReadItems for this rank.
621 """
622 if self.state_dict is None or self.metadata is None:
623 raise RuntimeError("Planner not configured")
625 requests: list[ReadItem] = []
626 strict = not self.allow_partial_load
627 for fqn, obj in self.state_dict.items():
628 if fqn not in self.metadata.state_dict_metadata:
629 if fqn.endswith(('matched_adamw_rms', 'step')):
630 continue
631 if strict:
632 raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
633 continue
634 md = self.metadata.state_dict_metadata[fqn]
635 if isinstance(md, TensorStorageMetadata):
636 obj_size = getattr(obj, CHUNK_INFO).global_shape if hasattr(obj, CHUNK_INFO) \
637 else getattr(obj, "shape", None)
638 if obj_size is None or md.size != tuple(obj_size):
639 raise ValueError(
640 f"Size mismatch between saved {md.size} and current: {obj_size} for {fqn}",
641 )
642 if not self._rank_owns_dtensor_shard(obj):
643 continue
644 # Both DTensor and torch.Tensor: create local chunks and read items
645 local_chunks = create_chunk_list_for_tensor(obj)
646 requests += create_read_items_for_chunk_list(
647 fqn, md, local_chunks, broadcastable=not _is_on_host(obj),
648 )
649 else:
650 requests.append(
651 ReadItem(
652 type=LoadItemType.BYTE_IO,
653 dest_index=MetadataIndex(fqn=fqn),
654 dest_offsets=(0,),
655 storage_index=MetadataIndex(fqn=fqn),
656 storage_offsets=(0,),
657 lengths=(0,),
658 )
659 )
660 return LoadPlan(items=requests)
662 @dcp_timer_decorator
663 def build_global_plan(self, all_plans: list[LoadPlan]) -> LoadPlan:
664 """
665 Decide, across all ranks, who reads which shard and who is sent it instead.
667 A shard held by several ranks needs only one of them to reach storage. This finds
668 those groups in the gathered plans, picks the reader of each, and marks the items of
669 this rank's plan with that decision. Shards are handed out one at a time and are not
670 tied to one another: the two shards of a tensor go to two groups and two broadcasts,
671 even when the same ranks hold both. Nothing is dropped from the plan either, since a
672 rank that receives a shard still has to take part in the broadcast carrying it.
674 Every rank runs this over the same gathered plans and reads only its own shards, so
675 only this rank's plan is built. The assignment is a pure function of the gathered
676 plans, so the ranks of a group name the same reader without one more round of
677 communication.
679 Args:
680 all_plans (list[LoadPlan]): Local plan of every rank, indexed by global rank.
682 Returns:
683 LoadPlan: This rank's plan, with the items of its replicated shards marked.
684 """
685 own_plan = all_plans[_own_plan_index(all_plans, self.rank)]
687 # A world of one has no one to broadcast with.
688 if not self.broadcast_replicated_tensors or len(all_plans) < 2:
689 return own_plan
691 # Which ranks hold each shard, and how many elements reading it moves. dest_index
692 # names the shard -- which local chunk of which tensor -- and ranks naming the same
693 # one hold the same buffer, since a chunk and the checkpoint layout together settle
694 # what has to be read for it. BYTE_IO items are left out: each rank rebuilds its
695 # own pickled object, not storage a collective could write into.
696 #
697 # A shard some rank reported as not worth sending is left out for much the same
698 # reason: what sits in host memory is little enough that reading it costs less than
699 # moving it, and the group a broadcast would go through may hold no backend for it
700 # anyway. One rank saying so is enough, and every rank reads the same report out of
701 # the gather, so they all drop the same shards and go on naming the same readers
702 # for the rest.
703 replicas: dict[MetadataIndex, set[int]] = defaultdict(set)
704 group_elements: dict[MetadataIndex, int] = defaultdict(int)
705 unsendable: set[MetadataIndex] = set()
706 for rank, plan in enumerate(all_plans):
707 for item in plan.items:
708 if item.type is LoadItemType.TENSOR:
709 replicas[item.dest_index].add(rank)
710 group_elements[item.dest_index] += math.prod(item.lengths)
711 if not item.broadcastable:
712 unsendable.add(item.dest_index)
714 # Every member of a group counted its own read into group_elements, and they all read
715 # the same, so dividing by the group size gives back what the one reader will move.
716 shards = [
717 (self._shard_element_size(index) * group_elements[index] // len(ranks), index, tuple(sorted(ranks)))
718 for index, ranks in replicas.items() if len(ranks) > 1 and index not in unsendable
719 ]
720 if not shards:
721 return own_plan
723 # Heaviest shard first, each to the rank of its group that has been given the fewest
724 # bytes so far: the costly reads are placed while the ranks are still even, and the
725 # small ones fill the gaps afterwards. Shards are weighed rather than counted so
726 # that one rank does not end up with all the large ones. Equal ones keep the order
727 # replicas was filled in, which follows the gathered plans and so is the same on
728 # every rank -- keep it that way, since ranks that disagreed here would name
729 # different readers and hang on one another broadcasts.
730 shards.sort(key=lambda shard: -shard[0])
731 load_bytes = [0] * len(all_plans)
732 sources: dict[MetadataIndex, BroadcastSource] = {}
733 for nbytes, index, group_ranks in shards:
734 src_rank = min(group_ranks, key=lambda rank: (load_bytes[rank], rank))
735 load_bytes[src_rank] += nbytes
736 sources[index] = BroadcastSource(group_ranks=group_ranks, src_rank=src_rank)
738 logger.info(
739 "[rank=%d] >>> %d replicated shards: busiest rank reads %d of %d bytes",
740 self.rank, len(shards), max(load_bytes), sum(load_bytes),
741 )
743 # Mark this rank's items whether it reads or receives: the item is what tells a rank
744 # which side of the broadcast it is on. A shard is held by one group, so an item
745 # naming it is enough -- which rank the plan belongs to does not come into it.
746 return dataclasses.replace(own_plan, items=[
747 dataclasses.replace(item, source=sources[item.dest_index])
748 if item.dest_index in sources else item
749 for item in own_plan.items
750 ])
752 def _shard_element_size(self, index: MetadataIndex) -> int:
753 """
754 Bytes one element of the shard ``index`` names takes, which weighs it when readers are picked.
756 Args:
757 index (MetadataIndex): Names the shard, as a read item destination does.
759 Returns:
760 int: Element size of the checkpoint tensor the shard is read from.
761 """
762 md = self.metadata.state_dict_metadata.get(index.fqn)
763 return dtype_element_size(getattr(getattr(md, "properties", None), "dtype", None))
765 def finalize_plan(self, plan: LoadPlan) -> LoadPlan:
766 """
767 Finalize the plan (no-op for default implementation).
769 Args:
770 plan (LoadPlan): Plan to finalize.
772 Returns:
773 LoadPlan: Finalized plan.
774 """
775 return plan
777 def acquire_tensor(self, read_item: ReadItem) -> Any:
778 """
779 Acquire the destination slice (narrow view) for this read_item.
781 StorageReader uses this to copy loaded data into the correct region.
782 Torch-aligned behavior.
784 Args:
785 read_item (ReadItem): The read item specifying what to load.
787 Returns:
788 Any: The destination tensor slice where data should be written
789 (tensor-like object).
790 """
791 if self.state_dict is None:
792 raise RuntimeError("Planner not configured")
794 fqn = read_item.dest_index.fqn
795 if fqn not in self.state_dict:
796 raise KeyError(f"Key {fqn} not found in state_dict")
798 target = self.state_dict[fqn]
799 # Each branch only picks out this rank's piece of the entry; detaching and narrowing
800 # it is the same for all three, so it is done once below rather than per branch.
801 if (
802 isinstance(target, DTensor)
803 and target.layout is not None
804 and target.layout.ragged_shard is not None
805 ):
806 local_tensor = get_ragged_box_tensor(target, read_item.dest_index)
807 elif isinstance(target, DTensor):
808 local_tensor = target.to_local()
809 else:
810 local_tensor = target
812 # Detached because the caller writes the loaded slice in place: an in-place write
813 # into a view of a leaf parameter that requires grad raises rather than loading.
814 return narrow_tensor_by_index(
815 local_tensor.detach(),
816 read_item.dest_offsets,
817 read_item.lengths,
818 )
820 def apply_tensor(self, read_item: ReadItem, tensor: Any) -> None:
821 """
822 Apply tensor after reading.
824 Nothing to do: ``_copy_tensor_to_target`` already wrote the loaded slice into the
825 destination, through ``copy_`` when the backend has one and through ``[...]``
826 otherwise, so by the time this runs the state dict entry already holds the data.
827 Kept as the hook other planners override.
829 Args:
830 read_item (ReadItem): The read item that was processed.
831 tensor (Any): The tensor data to apply (tensor-like object).
832 """
834 def apply_bytes(self, read_item: ReadItem, value: bytes) -> None:
835 """
836 Load bytes data into state_dict.
838 Args:
839 read_item (ReadItem): The read item specifying the destination.
840 value (bytes): The bytes data to deserialize and load.
841 """
842 if self.state_dict is None:
843 raise RuntimeError("Planner not set up")
845 fqn = read_item.dest_index.fqn
846 # Deserialize bytes
847 obj = pickle.loads(value)
848 self.state_dict[fqn] = obj
849 if self.flatten_state_dict:
850 set_element(self.original_state_dict, self.name_mapping[fqn], obj)
854class _DcpMergeLoadPlanner(StandardLoadPlanner):
855 """Load planner that builds distributed checkpoint from dcp into fully ``state_dict`` (in-place)."""
857 def __init__(self) -> None:
858 """Build a planner that merges every shard of a checkpoint into one state dict."""
859 super().__init__()
861 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs: Any) -> None:
862 """
863 Populate the empty ``state_dict`` from the checkpoint metadata, then configure.
865 Args:
866 state_dict (dict[str, Any]): Must be empty; filled in place with a full-size tensor
867 per entry of the metadata.
868 metadata (Metadata): Metadata of the checkpoint being merged.
869 **kwargs (Any): Additional keyword arguments (e.g. is_coordinator).
870 """
871 if len(state_dict) > 0:
872 raise ValueError(
873 "state_dict must be empty for _DcpMergeLoadPlanner; "
874 "it is populated in-place from checkpoint metadata."
875 )
877 if metadata is None:
878 raise ValueError("metadata must not be None for _DcpMergeLoadPlanner.")
880 self.is_coordinator = kwargs.get("is_coordinator", False)
881 for k, v in metadata.state_dict_metadata.items():
882 if isinstance(v, TensorStorageMetadata):
883 v = torch.empty(
884 torch.Size(v.size),
885 dtype=str_to_dtype(v.properties.dtype),
886 )
888 state_dict[k] = v
889 if metadata.planner_data is not None and k in metadata.planner_data:
890 set_element(state_dict, metadata.planner_data[k], v)
892 super().configure_planner(
893 state_dict,
894 metadata,
895 is_coordinator=self.is_coordinator,
896 flatten_state_dict=True,
897 )