Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / distributed_checkpoint / remap_planner.py: 96%
182 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""Loads of checkpoints whose tensors are named and laid out differently from the state dict they fill.
17A Hugging Face checkpoint stores the query, key and value projections of a layer as three tensors
18where a model may hold one fused, interleaved tensor, and the experts of a block one tensor apiece
19where a model stacks them. :class:`RemapLoadPlanner` loads such a checkpoint without ever putting a
20checkpoint tensor together in memory: it is told, per destination tensor, which box of it comes from
21which region of which checkpoint tensor, as :class:`RemapBlock` entries, and plans each rank's reads
22from the shards that rank holds, exactly as :class:`StandardLoadPlanner` plans them for a checkpoint
23laid out like its state dict. Replicated shards are still read once and broadcast.
25What a block cannot describe is handed to :class:`DeferredRead`: its checkpoint tensors are read
26whole and given to a function that puts the result in place itself.
27"""
28from collections.abc import Callable, Mapping, Sequence
29from dataclasses import dataclass
30from typing import Any, Optional
32from hyper_parallel.core.distributed_checkpoint.metadata import (
33 MetadataIndex,
34 TensorStorageMetadata,
35 dtype_element_size,
36)
37from hyper_parallel.core.distributed_checkpoint.planner import LoadItemType, LoadPlan, ReadItem
38from hyper_parallel.core.distributed_checkpoint.ragged import get_ragged_box_tensor
39from hyper_parallel.core.distributed_checkpoint.standard_planner import StandardLoadPlanner, _is_on_host
40from hyper_parallel.core.distributed_checkpoint.utils import (
41 chunk_to_area,
42 create_chunk_list_for_tensor,
43 dcp_timer_decorator,
44 infer_intersection,
45 narrow_tensor_by_index,
46)
47from hyper_parallel.core.dtensor.dtensor import DTensor
49# Deferred reads land in no state dict entry, so their items name a shard of their own. The prefix
50# cannot start a state dict key, which is a dotted module path.
51_DEFERRED_FQN_PREFIX = "<deferred>/"
54@dataclass(frozen=True)
55class RemapBlock:
56 """
57 A box of a destination tensor, and where in one checkpoint tensor each element of it is read from.
59 The element at index ``x`` of the destination tensor, for ``offsets <= x < offsets + lengths``, is
60 read from index ``base + coeff @ (x - offsets)`` of the checkpoint tensor ``source``: ``coeff[s][a]``
61 is how far along checkpoint dimension ``s`` one step along destination dimension ``a`` moves.
63 A block whose every varying destination dimension walks a checkpoint dimension of its own one
64 element at a time is a plain region of the checkpoint tensor in another order, and the read is
65 copied straight into the destination. Anything else, and any block with ``post`` operations, is
66 read into a buffer of its own and put in place from there.
68 Attributes:
69 offsets (tuple[int, ...]): Where the box starts in the whole destination tensor.
70 lengths (tuple[int, ...]): Size of the box, positive in every dimension.
71 source (str): Name of the checkpoint tensor the box is read from.
72 base (tuple[int, ...]): Index of the checkpoint tensor the first element of the box comes from.
73 coeff (tuple[tuple[int, ...], ...]): Non-negative steps, one row per checkpoint dimension and one
74 column per destination dimension.
75 post (tuple[Callable[[Any], Any], ...]): Elementwise operations applied in order to the values
76 read, in the dtype the checkpoint stores them in, before they are copied into place.
77 Default ().
78 """
80 offsets: tuple[int, ...]
81 lengths: tuple[int, ...]
82 source: str
83 base: tuple[int, ...]
84 coeff: tuple[tuple[int, ...], ...]
85 post: tuple[Callable[[Any], Any], ...] = ()
87 def __post_init__(self) -> None:
88 """Reject a block whose parts disagree on the number of dimensions or step backwards."""
89 ndim = len(self.offsets)
90 if len(self.lengths) != ndim or any(length <= 0 for length in self.lengths):
91 raise ValueError(
92 f"RemapBlock of {self.source!r} needs one positive length per offset, "
93 f"got offsets={self.offsets} lengths={self.lengths}"
94 )
95 if len(self.coeff) != len(self.base) or any(
96 len(row) != ndim or any(step < 0 for step in row) for row in self.coeff
97 ):
98 raise ValueError(
99 f"RemapBlock of {self.source!r} needs one row of {ndim} non-negative steps per checkpoint "
100 f"dimension, got base={self.base} coeff={self.coeff}"
101 )
104@dataclass(frozen=True)
105class DeferredRead:
106 """
107 Checkpoint tensors read whole, and what to do with them once all of them are.
109 For destination tensors no block can describe. Every rank that plans the read reads the tensors
110 into host memory itself, never through a broadcast, and ``complete`` is called once with all of
111 them to put whatever it computes out of them in place.
113 Attributes:
114 sources (tuple[str, ...]): Names of the checkpoint tensors to read, each once.
115 complete (Callable[[dict[str, Any]], None]): Called with ``{name: tensor}`` once every source
116 has been read.
117 """
119 sources: tuple[str, ...]
120 complete: Callable[[dict[str, Any]], None]
122 def __post_init__(self) -> None:
123 """Reject an empty read and a source named twice, which would never be counted complete."""
124 if not self.sources or len(set(self.sources)) != len(self.sources):
125 raise ValueError(f"DeferredRead needs distinct sources, got {self.sources}")
128@dataclass(frozen=True)
129class _Recipe:
130 """
131 How the region one read item stands for is put in place.
133 Attributes:
134 block (RemapBlock): The block the item was cut from.
135 sub_lengths (tuple[int, ...]): Size of the destination box the item fills.
136 axes (Optional[tuple[tuple[int, int], ...]]): ``(checkpoint dimension, destination dimension)``
137 pairs of a read copied straight into place, or None for one gathered from a buffer.
138 """
140 block: RemapBlock
141 sub_lengths: tuple[int, ...]
142 axes: Optional[tuple[tuple[int, int], ...]]
145class _Capture:
146 """
147 Stands where a read would be copied to, and keeps what was read for the planner to place itself.
149 The storage reader checks the size of what it acquires against the size of what it read, then
150 copies with ``copy_``. A capture answers both, holding on to the tensor instead of copying it.
151 """
153 __slots__ = ("_size", "value")
155 def __init__(self, size: Sequence[int]) -> None:
156 """
157 Args:
158 size (Sequence[int]): Size of the region the read returns.
159 """
160 self._size = tuple(size)
161 self.value: Any = None
163 def size(self) -> tuple[int, ...]:
164 """Size of the region the read returns."""
165 return self._size
167 def copy_(self, tensor: Any) -> "_Capture":
168 """Keep ``tensor``, the region as read."""
169 self.value = tensor
170 return self
173def _recipe_key(item: ReadItem) -> tuple:
174 """What tells the reads of one plan apart, before and after a global plan has marked them."""
175 return item.dest_index, item.dest_offsets, item.storage_index, item.storage_offsets
178def _source_start(block: RemapBlock, sub_offsets: Sequence[int]) -> tuple[int, ...]:
179 """Checkpoint index the first element of the destination box starting at ``sub_offsets`` comes from."""
180 steps = [offset - origin for offset, origin in zip(sub_offsets, block.offsets)]
181 return tuple(start + sum(c * step for c, step in zip(row, steps)) for start, row in zip(block.base, block.coeff))
184def _source_lengths(block: RemapBlock, sub_lengths: Sequence[int]) -> tuple[int, ...]:
185 """Size of the checkpoint region a destination box of ``sub_lengths`` reads from, gaps included."""
186 return tuple(1 + sum(c * (length - 1) for c, length in zip(row, sub_lengths)) for row in block.coeff)
189def _direct_axes(block: RemapBlock, sub_lengths: Sequence[int]) -> Optional[tuple[tuple[int, int], ...]]:
190 """
191 Pair the dimensions of a read that can be copied straight into its destination box.
193 Args:
194 block (RemapBlock): The block the box belongs to.
195 sub_lengths (Sequence[int]): Size of the box.
197 Returns:
198 Optional[tuple[tuple[int, int], ...]]: ``(checkpoint dimension, destination dimension)`` for
199 every destination dimension the box is longer than one along, when each of them walks a
200 checkpoint dimension of its own one element at a time and nothing is left to apply. None
201 otherwise.
202 """
203 if block.post:
204 return None
205 pairs: dict[int, int] = {}
206 for axis, length in enumerate(sub_lengths):
207 if length == 1:
208 continue
209 walked = [dim for dim, row in enumerate(block.coeff) if row[axis]]
210 if len(walked) != 1 or block.coeff[walked[0]][axis] != 1 or walked[0] in pairs:
211 return None
212 pairs[walked[0]] = axis
213 return tuple(sorted(pairs.items()))
216def _direct_view(region: Any, axes: tuple[tuple[int, int], ...], source_lengths: tuple[int, ...]) -> Any:
217 """
218 View a destination box in the order and shape of the checkpoint region it is copied from.
220 Args:
221 region (Any): The destination box.
222 axes (tuple[tuple[int, int], ...]): Pairs from :func:`_direct_axes`.
223 source_lengths (tuple[int, ...]): Size of the checkpoint region.
225 Returns:
226 Any: A view of ``region``. Only dimensions of size one are added or dropped, so it still
227 writes through to the destination.
228 """
229 walked = [axis for _, axis in axes]
230 order = walked + [axis for axis in range(len(region.shape)) if axis not in walked]
231 if order != sorted(order):
232 region = region.permute(order)
233 return region.reshape(source_lengths)
236def _gather(block: RemapBlock, sub_lengths: tuple[int, ...], read: Any) -> Any:
237 """
238 Pick the elements of a destination box out of the checkpoint region read for it.
240 Args:
241 block (RemapBlock): The block the box belongs to.
242 sub_lengths (tuple[int, ...]): Size of the box.
243 read (Any): The checkpoint region, starting at the element the box starts with.
245 Returns:
246 Any: The values of the box in destination order, with the block's post operations applied.
247 """
248 strides = read.stride()
249 stride = tuple(
250 sum(row[axis] * step for row, step in zip(block.coeff, strides)) for axis in range(len(sub_lengths))
251 )
252 values = read.as_strided(sub_lengths, stride, read.storage_offset())
253 for operation in block.post:
254 values = operation(values)
255 return values
258class RemapLoadPlanner(StandardLoadPlanner):
259 """
260 Load planner for a checkpoint whose tensors are named and laid out differently from the state dict.
262 Every state dict entry to fill is described by ``table``: blocks that tile it, each saying which
263 region of which checkpoint tensor it is read from. Entries without blocks are left alone, and
264 ``deferred`` reads whole checkpoint tensors for whatever the blocks leave out.
266 Each rank intersects the shards it holds with the blocks and reads only those regions. A read that
267 is a plain region of a checkpoint tensor is copied straight into place; one that has to be gathered
268 or changed is read into a buffer first. Only whole checkpoint tensors are read, one chunk each, as
269 :class:`HuggingFaceStorageReader` describes them.
271 Example::
273 planner = RemapLoadPlanner({"linear.weight": [block_from_q, block_from_k]})
274 load(state_dict, storage_reader=HuggingFaceStorageReader(path), planner=planner)
275 """
277 def __init__(
278 self,
279 table: Optional[Mapping[str, Sequence[RemapBlock]]] = None,
280 deferred: Sequence[DeferredRead] = (),
281 broadcast_replicated_tensors: bool = False,
282 ) -> None:
283 """
284 Args:
285 table (Optional[Mapping[str, Sequence[RemapBlock]]]): Blocks of every state dict entry to fill,
286 by its flattened name. Default None, for a subclass that fills ``table`` while configuring.
287 deferred (Sequence[DeferredRead]): Whole-tensor reads for what no block describes. Default ().
288 broadcast_replicated_tensors (bool): See :class:`StandardLoadPlanner`; ``load`` overrides it.
289 Default False.
290 """
291 super().__init__(broadcast_replicated_tensors=broadcast_replicated_tensors)
292 self.table: dict[str, tuple[RemapBlock, ...]] = {
293 name: tuple(blocks) for name, blocks in (table or {}).items()
294 }
295 self.deferred: tuple[DeferredRead, ...] = tuple(deferred)
296 self._recipes: dict[tuple, _Recipe] = {}
297 self._deferred_slots: dict[tuple, tuple[int, str]] = {}
298 self._captured: dict[int, dict[str, Any]] = {}
300 @dcp_timer_decorator
301 def build_local_plan(self) -> LoadPlan:
302 """
303 Plan this rank's reads: the part of every block inside a shard it holds, and the deferred reads.
305 Returns:
306 LoadPlan: One read per (shard, block) overlap, then one per deferred checkpoint tensor.
308 Raises:
309 RuntimeError: If the planner is not configured.
310 ValueError: If the table names an entry the state dict does not have, or a checkpoint
311 tensor that is missing, stored in shards, or smaller than a block reads.
312 """
313 if self.state_dict is None or self.metadata is None:
314 raise RuntimeError("Planner not configured")
315 self._recipes.clear()
316 self._deferred_slots.clear()
317 self._captured.clear()
319 items: list[ReadItem] = []
320 for fqn, blocks in self.table.items():
321 if fqn not in self.state_dict:
322 raise ValueError(f"Remap table names {fqn!r}, which the state dict does not have")
323 items.extend(self._plan_entry(fqn, self.state_dict[fqn], blocks))
324 for group_index, group in enumerate(self.deferred):
325 items.extend(self._plan_deferred(group_index, source) for source in group.sources)
326 return LoadPlan(items=items)
328 def acquire_tensor(self, read_item: ReadItem) -> Any:
329 """
330 Where one read is copied to.
332 Args:
333 read_item (ReadItem): A read of this planner's plan.
335 Returns:
336 Any: A view of the destination for a read copied straight into place, otherwise a
337 capture that keeps the read for :meth:`apply_tensor`.
338 """
339 key = _recipe_key(read_item)
340 if key in self._deferred_slots:
341 return _Capture(read_item.lengths)
342 recipe = self._recipe(key)
343 if recipe.axes is None:
344 return _Capture(read_item.lengths)
345 return _direct_view(self._destination_box(read_item, recipe), recipe.axes, read_item.lengths)
347 def apply_tensor(self, read_item: ReadItem, tensor: Any) -> None:
348 """
349 Put a captured read in place: gather a box out of it, or keep it for its deferred read.
351 Args:
352 read_item (ReadItem): A read of this planner's plan.
353 tensor (Any): What :meth:`acquire_tensor` returned, the read already copied into it.
354 """
355 key = _recipe_key(read_item)
356 slot = self._deferred_slots.get(key)
357 if slot is not None:
358 self._keep_deferred(slot, tensor.value)
359 return
360 recipe = self._recipe(key)
361 if recipe.axes is not None:
362 return
363 self._destination_box(read_item, recipe).copy_(_gather(recipe.block, recipe.sub_lengths, tensor.value))
365 def _shard_element_size(self, index: MetadataIndex) -> int:
366 """
367 Bytes one element of a destination shard takes, which is what a broadcast of it moves.
369 Args:
370 index (MetadataIndex): Names the shard.
372 Returns:
373 int: Element size of the state dict entry, or the default size for a deferred read.
374 """
375 dtype = getattr(self.state_dict.get(index.fqn), "dtype", None) if self.state_dict else None
376 return dtype_element_size(None if dtype is None else str(dtype))
378 def _plan_entry(self, fqn: str, obj: Any, blocks: Sequence[RemapBlock]) -> list[ReadItem]:
379 """Plan the reads of one state dict entry, for the shards of it this rank holds."""
380 if not self._rank_owns_dtensor_shard(obj):
381 return []
382 broadcastable = not _is_on_host(obj)
383 items = []
384 for local_index, chunk in enumerate(create_chunk_list_for_tensor(obj)):
385 dest_index = MetadataIndex(fqn=fqn, offset=chunk.offsets, index=local_index)
386 chunk_area = chunk_to_area(chunk)
387 for block in blocks:
388 block_area = tuple((start, start + length) for start, length in zip(block.offsets, block.lengths))
389 overlap = infer_intersection(chunk_area, block_area)
390 if overlap is not None:
391 items.append(self._plan_block_read(dest_index, chunk.offsets, block, overlap, broadcastable))
392 return items
394 def _plan_block_read(
395 self,
396 dest_index: MetadataIndex,
397 chunk_offsets: tuple[int, ...],
398 block: RemapBlock,
399 overlap: tuple[tuple[int, int], ...],
400 broadcastable: bool,
401 ) -> ReadItem:
402 """Plan the read of the part of ``block`` that lands in one shard, and remember how to place it."""
403 sub_offsets = tuple(start for start, _ in overlap)
404 sub_lengths = tuple(stop - start for start, stop in overlap)
405 metadata = self._whole_tensor(block.source)
406 storage_offsets = _source_start(block, sub_offsets)
407 lengths = _source_lengths(block, sub_lengths)
408 if len(lengths) != len(metadata.size) or any(
409 start + length > size for start, length, size in zip(storage_offsets, lengths, metadata.size)
410 ):
411 raise ValueError(
412 f"A block reads {lengths} at {storage_offsets} of checkpoint tensor {block.source!r}, "
413 f"which has size {tuple(metadata.size)}"
414 )
415 item = ReadItem(
416 type=LoadItemType.TENSOR,
417 dest_index=dest_index,
418 dest_offsets=tuple(offset - origin for offset, origin in zip(sub_offsets, chunk_offsets)),
419 storage_index=MetadataIndex(fqn=block.source, offset=metadata.chunks[0].offsets),
420 storage_offsets=storage_offsets,
421 lengths=lengths,
422 broadcastable=broadcastable,
423 )
424 self._recipes[_recipe_key(item)] = _Recipe(block, sub_lengths, _direct_axes(block, sub_lengths))
425 return item
427 def _plan_deferred(self, group_index: int, source: str) -> ReadItem:
428 """Plan the whole-tensor read of one source of a deferred read."""
429 metadata = self._whole_tensor(source)
430 size = tuple(metadata.size)
431 item = ReadItem(
432 type=LoadItemType.TENSOR,
433 dest_index=MetadataIndex(fqn=f"{_DEFERRED_FQN_PREFIX}{group_index}/{source}"),
434 dest_offsets=(),
435 storage_index=MetadataIndex(fqn=source, offset=metadata.chunks[0].offsets),
436 storage_offsets=(0,) * len(size),
437 lengths=size,
438 # Kept in host memory for the conversion that follows, which no collective can write into.
439 broadcastable=False,
440 )
441 self._deferred_slots[_recipe_key(item)] = (group_index, source)
442 return item
444 def _whole_tensor(self, source: str) -> TensorStorageMetadata:
445 """
446 Metadata of a checkpoint tensor stored whole, in one chunk.
448 Raises:
449 ValueError: If the checkpoint has no such tensor, or stores it in shards.
450 """
451 metadata = self.metadata.state_dict_metadata.get(source)
452 if not isinstance(metadata, TensorStorageMetadata):
453 raise ValueError(f"Checkpoint has no tensor {source!r}")
454 chunks = metadata.chunks
455 if len(chunks) != 1 or any(chunks[0].offsets) or tuple(chunks[0].sizes) != tuple(metadata.size):
456 raise ValueError(
457 f"Checkpoint tensor {source!r} is stored in {len(chunks)} shards; "
458 "RemapLoadPlanner reads whole tensors only"
459 )
460 return metadata
462 def _recipe(self, key: tuple) -> _Recipe:
463 """The recipe of a read of the current plan."""
464 recipe = self._recipes.get(key)
465 if recipe is None:
466 raise RuntimeError(f"Read of {key[2].fqn!r} into {key[0].fqn!r} is not part of the current plan")
467 return recipe
469 def _destination_box(self, read_item: ReadItem, recipe: _Recipe) -> Any:
470 """The box of the destination shard one read fills, detached so it is written in place."""
471 target = self.state_dict[read_item.dest_index.fqn]
472 if isinstance(target, DTensor) and target.layout is not None and target.layout.ragged_shard is not None:
473 local = get_ragged_box_tensor(target, read_item.dest_index)
474 elif isinstance(target, DTensor):
475 local = target.to_local()
476 else:
477 local = target
478 return narrow_tensor_by_index(local.detach(), read_item.dest_offsets, recipe.sub_lengths)
480 def _keep_deferred(self, slot: tuple[int, str], value: Any) -> None:
481 """Keep one source of a deferred read, and complete the read once it has all of them."""
482 group_index, source = slot
483 captured = self._captured.setdefault(group_index, {})
484 captured[source] = value
485 group = self.deferred[group_index]
486 if len(captured) == len(group.sources):
487 del self._captured[group_index]
488 group.complete(captured)