Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / load_groups.py: 83%
271 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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"""Hugging Face checkpoint names and layouts, mapped onto the tensors of a finalized model.
17Both pretrained loaders are built on what is here. :class:`HuggingFaceCheckpointer` reads whole checkpoint
18tensors and converts them in memory, while :class:`HFLoadPlanner` plans the same conversions as
19distributed checkpoint reads. Keys are renamed, converters are chosen and replacement conversions are
20routed the same way in both, so the two load the same values under the same names.
21"""
22import json
23import logging
24from collections import Counter, OrderedDict, defaultdict
25from collections.abc import Callable, Iterable, Sequence
26from copy import deepcopy
27from dataclasses import dataclass
28from pathlib import Path
29from typing import Any, Optional
31import torch
32from huggingface_hub import snapshot_download
33from safetensors import safe_open
34from torch import nn
36from hyper_parallel import DTensor, Partial, distribute_tensor
37from hyper_parallel.components.checkpoint.weight_conversion import (
38 WeightConverter,
39 WeightRenaming,
40 dot_natural_key,
41 rename_source_key,
42)
44logger = logging.getLogger(__name__)
46SAFE_WEIGHTS_NAME = "model.safetensors"
47SAFE_WEIGHTS_INDEX_NAME = "model.safetensors.index.json"
48_SNAPSHOT_PATTERNS = ("*.safetensors", "*.safetensors.index.json", "engram_host_rows.json")
50# Picks the converter a checkpoint key matched by ``source pattern`` belongs to, for the model tensor
51# ``target name``, out of every converter with that source pattern. None when no single one fits.
52ConverterSelector = Callable[[str, str, Sequence[WeightConverter]], Optional[WeightConverter]]
55@dataclass(frozen=True)
56class LoadReport:
57 """Summary of one pretrained-weight load."""
59 loaded_keys: tuple[str, ...]
60 missing_keys: tuple[str, ...]
61 unexpected_keys: tuple[str, ...]
64@dataclass(frozen=True)
65class CheckpointIndex:
66 """Map checkpoint tensor names to their safetensors shard files."""
68 files_by_key: dict[str, Path]
70 def keys(self) -> tuple[str, ...]:
71 """Return checkpoint keys in deterministic natural order."""
72 return sorted_checkpoint_keys(self.files_by_key)
74 def load_tensor(self, key: str) -> torch.Tensor:
75 """Materialize one checkpoint tensor on CPU."""
76 file_path = self.files_by_key.get(key)
77 if file_path is None:
78 raise ValueError(f"Checkpoint key is not indexed: {key}")
79 with safe_open(str(file_path), framework="pt", device="cpu") as checkpoint:
80 return checkpoint.get_tensor(key)
83@dataclass
84class LoadGroup:
85 """Checkpoint tensors converted together, and the transform converting them."""
87 first_target_name: str
88 transform: WeightRenaming | WeightConverter
91@dataclass(frozen=True)
92class _TensorShape:
93 """Shape of a model tensor, as shape-aware conversion operations ask for it."""
95 shape: torch.Size
98class SourceModelView:
99 """Expose pre-replacement parameter shapes to Transformers conversion ops."""
101 def __init__(self, model: nn.Module, shapes: dict[str, tuple[int, ...]]) -> None:
102 """Build a lightweight model view from captured tensor shapes."""
103 self.config = getattr(model, "config", None)
104 self.base_model_prefix = getattr(model, "base_model_prefix", None)
105 self._targets = {
106 name: _TensorShape(torch.Size(shape)) for name, shape in shapes.items()
107 }
109 @property
110 def targets(self) -> dict[str, _TensorShape]:
111 """Shape of every pre-replacement tensor, by name."""
112 return self._targets
114 def get_parameter(self, name: str) -> _TensorShape:
115 """Return source parameter metadata used by shape-aware converters."""
116 try:
117 return self._targets[name]
118 except KeyError as exc:
119 raise AttributeError(f"source model has no parameter {name!r}") from exc
122@dataclass
123class ReplacementLoadGroup:
124 """A replacement conversion waiting for the normalized tensors it is routed."""
126 group: LoadGroup
127 expected: Counter[str]
128 received: Counter[str]
129 completed: bool = False
132def sorted_checkpoint_keys(keys: Iterable[str]) -> tuple[str, ...]:
133 """Checkpoint keys in natural order, experts 2 before 10, which is the order loads collect them in."""
134 return tuple(sorted(keys, key=dot_natural_key))
137def _index_single_file(file_path: Path) -> CheckpointIndex:
138 """Index every tensor of one safetensors file."""
139 with safe_open(str(file_path), framework="pt", device="cpu") as checkpoint:
140 files_by_key = {key: file_path for key in checkpoint.keys()}
141 return CheckpointIndex(files_by_key)
144def _index_sharded_checkpoint(directory: Path, index_path: Path) -> CheckpointIndex:
145 """Index tensors described by a sharded safetensors index file."""
146 try:
147 index_data = json.loads(index_path.read_text(encoding="utf-8"))
148 except (OSError, json.JSONDecodeError) as exc:
149 raise ValueError(f"Failed to read safetensors index {index_path}: {exc}") from exc
150 weight_map = index_data.get("weight_map")
151 if not isinstance(weight_map, dict) or not weight_map:
152 raise ValueError(f"Safetensors index has no non-empty weight_map: {index_path}")
154 files_by_key = {}
155 for key, relative_path in weight_map.items():
156 file_path = directory / relative_path
157 if not file_path.is_file():
158 raise ValueError(f"Safetensors shard for {key} does not exist: {file_path}")
159 files_by_key[key] = file_path
160 return CheckpointIndex(files_by_key)
163def resolve_checkpoint_location(pretrained_path: str) -> Path:
164 """
165 Resolve a pretrained path into a local safetensors file or checkpoint directory.
167 Args:
168 pretrained_path (str): A ``.safetensors`` file, a checkpoint directory, or a Hub repository id,
169 which is downloaded first.
171 Returns:
172 Path: The file or the directory.
174 Raises:
175 ValueError: If the path is a file of another kind.
176 """
177 path = Path(pretrained_path).expanduser()
178 if path.is_file():
179 if path.suffix != ".safetensors":
180 raise ValueError(f"MVP only supports safetensors checkpoints, got: {path}")
181 return path
182 if path.is_dir():
183 return path
184 return Path(
185 snapshot_download(
186 repo_id=pretrained_path,
187 allow_patterns=list(_SNAPSHOT_PATTERNS),
188 )
189 )
192def resolve_checkpoint_index(pretrained_path: str) -> CheckpointIndex:
193 """Resolve a local or Hub checkpoint into a tensor-to-file index."""
194 location = resolve_checkpoint_location(pretrained_path)
195 if location.is_file():
196 return _index_single_file(location)
198 index_path = location / SAFE_WEIGHTS_INDEX_NAME
199 if index_path.is_file():
200 return _index_sharded_checkpoint(location, index_path)
202 single_file = location / SAFE_WEIGHTS_NAME
203 if single_file.is_file():
204 return _index_single_file(single_file)
205 raise ValueError(
206 "MVP requires model.safetensors or model.safetensors.index.json under "
207 f"{location}"
208 )
211def join_fqn(module_name: str, tensor_name: str) -> str:
212 """Name of a tensor registered on the module named ``module_name``."""
213 return f"{module_name}.{tensor_name}" if module_name else tensor_name
216def build_load_targets(model: nn.Module) -> dict[str, torch.Tensor]:
217 """Collect persistent parameters and buffers owned by the model."""
218 targets = {}
219 # Direct module registries preserve tied aliases and let us exclude
220 # non-persistent buffers without materializing an FSDP state_dict.
221 for module_name, module in model.named_modules(remove_duplicate=False):
222 for tensor_name, parameter in module._parameters.items(): # pylint: disable=W0212
223 if parameter is not None:
224 targets[join_fqn(module_name, tensor_name)] = parameter
225 non_persistent = module._non_persistent_buffers_set # pylint: disable=W0212
226 for tensor_name, buffer in module._buffers.items(): # pylint: disable=W0212
227 if buffer is not None and tensor_name not in non_persistent:
228 targets[join_fqn(module_name, tensor_name)] = buffer
229 return targets
232def checkpoint_logical_shapes(model: nn.Module) -> dict[str, tuple[int, ...]]:
233 """Collect model-declared logical parameter shapes for trailing padding on any dimensions.
235 Modules may expose ``checkpoint_logical_shapes()`` returning logical shapes
236 keyed by parameter names relative to that module. Undeclared parameters keep
237 the normal strict checkpoint shape checks.
239 Args:
240 model: Model whose modules declare logical checkpoint shapes.
242 Returns:
243 Logical shapes keyed by fully qualified model parameter names.
244 """
245 shapes = {}
246 for module_name, module in model.named_modules():
247 provider = getattr(module, "checkpoint_logical_shapes", None)
248 if not callable(provider):
249 continue
250 declared = provider()
251 if not isinstance(declared, dict):
252 raise ValueError(f"Invalid checkpoint padding declaration on {module_name}")
253 for parameter_name, shape in declared.items():
254 parameter = module.get_parameter(parameter_name)
255 if (not isinstance(shape, tuple) or len(shape) != parameter.ndim
256 or any(not isinstance(size, int) or isinstance(size, bool) or size < 0 for size in shape)):
257 raise ValueError(f"Invalid logical checkpoint shape on {module_name}: {shape}")
258 name = join_fqn(module_name, parameter_name)
259 if name in shapes and shapes[name] != shape:
260 raise ValueError(f"Conflicting logical checkpoint shapes for {name}")
261 shapes[name] = shape
262 return shapes
265def checkpoint_padding_boxes(
266 logical_shape: tuple[int, ...], offsets: tuple[int, ...], sizes: tuple[int, ...],
267) -> list[tuple[tuple[int, ...], tuple[int, ...]]]:
268 """Partition a local shard's trailing padding into disjoint local boxes.
270 Args:
271 logical_shape: Unpadded global checkpoint dimensions.
272 offsets: Global origin of this rank's rectangular shard.
273 sizes: Dimensions of this rank's local shard.
275 Returns:
276 Local offsets and dimensions of non-overlapping padding boxes.
277 """
278 if len(logical_shape) != len(offsets) or len(offsets) != len(sizes):
279 raise ValueError("Logical shape and local shard must have the same number of dimensions")
280 valid = tuple(max(0, min(size, logical - offset))
281 for logical, offset, size in zip(logical_shape, offsets, sizes))
282 boxes = []
283 for axis, (length, size) in enumerate(zip(valid, sizes)):
284 lengths = valid[:axis] + (size - length,) + sizes[axis + 1:]
285 if all(lengths):
286 starts = (0,) * axis + (length,) + (0,) * (len(sizes) - axis - 1)
287 boxes.append((starts, lengths))
288 return boxes
291def make_tensor_loader(index: CheckpointIndex, source_key: str) -> Callable[[], torch.Tensor]:
292 """A loader reading one checkpoint tensor when a conversion asks for it."""
293 return lambda: index.load_tensor(source_key)
296def _scoped_candidates(candidates: Sequence[WeightConverter], target_name: str) -> list[WeightConverter]:
297 """The candidates scoped to a module that holds ``target_name``."""
298 return [
299 converter
300 for converter in candidates
301 if converter.scope_prefix is not None
302 and (
303 target_name == converter.scope_prefix
304 or target_name.startswith(f"{converter.scope_prefix}.")
305 )
306 ]
309def select_unique_converter(
310 source_pattern: str,
311 target_name: str,
312 candidates: Sequence[WeightConverter],
313) -> Optional[WeightConverter]:
314 """
315 Pick the converter a matched checkpoint key belongs to.
317 A converter scoped to a module holding the target wins when it is the only one. Otherwise the only
318 unscoped converter does.
320 Args:
321 source_pattern (str): The source pattern the key matched.
322 target_name (str): The model tensor the key was renamed to.
323 candidates (Sequence[WeightConverter]): Every converter with that source pattern.
325 Returns:
326 Optional[WeightConverter]: The converter, or None when neither rule picks exactly one.
327 """
328 del source_pattern
329 scoped_candidates = _scoped_candidates(candidates, target_name)
330 if len(scoped_candidates) == 1:
331 return scoped_candidates[0]
332 unscoped_candidates = [converter for converter in candidates if converter.scope_prefix is None]
333 return unscoped_candidates[0] if len(unscoped_candidates) == 1 else None
336def _split_transforms(
337 transforms: Sequence[Any],
338) -> tuple[list[WeightRenaming], list[WeightConverter], dict[str, list[WeightConverter]]]:
339 """Renamings, converters, and the converters by each of their source patterns."""
340 renamings = [transform for transform in transforms if isinstance(transform, WeightRenaming)]
341 converters = [transform for transform in transforms if isinstance(transform, WeightConverter)]
342 converters_by_pattern = defaultdict(list)
343 for converter in converters:
344 for pattern in converter.source_patterns:
345 converters_by_pattern[pattern].append(converter)
346 return renamings, converters, converters_by_pattern
349def build_load_groups(
350 model: Any,
351 source_keys: Iterable[str],
352 targets: dict[str, Any],
353 *,
354 weights_mapping: list[WeightRenaming | WeightConverter],
355 make_loader: Callable[[str], Any],
356 select_converter: ConverterSelector = select_unique_converter,
357) -> tuple[
358 tuple[LoadGroup, ...],
359 tuple[str, ...],
360 list[WeightRenaming | WeightConverter],
361]:
362 """
363 Build checkpoint conversion groups and report unmatched checkpoint keys.
365 Args:
366 model (Any): The model, or the view of it the keys are renamed against.
367 source_keys (Iterable[str]): Checkpoint keys, in the order their tensors are collected.
368 targets (dict[str, Any]): Model tensors by name.
369 weights_mapping (list[WeightRenaming | WeightConverter]): Rules renaming and converting keys.
370 make_loader (Callable[[str], Any]): Called with a key, returns what the group collects for it.
371 select_converter (ConverterSelector): Picks the converter of a matched key.
372 Default :func:`select_unique_converter`.
374 Returns:
375 tuple: The groups in order of their first key, the keys renamed to no target, and the mapping.
376 """
377 unsupported = [
378 transform
379 for transform in weights_mapping
380 if not isinstance(transform, (WeightRenaming, WeightConverter))
381 ]
382 if unsupported:
383 names = ", ".join(type(transform).__name__ for transform in unsupported)
384 raise ValueError(f"Unsupported Transformers weight transforms in MVP: {names}")
386 renamings, converters, converters_by_pattern = _split_transforms(weights_mapping)
387 groups: OrderedDict[str, LoadGroup] = OrderedDict()
388 unexpected_keys = []
389 base_model_prefix = getattr(model, "base_model_prefix", None)
391 for source_key in source_keys:
392 target_name, source_pattern = rename_source_key(
393 source_key,
394 renamings,
395 converters,
396 base_model_prefix=base_model_prefix,
397 meta_state_dict=targets,
398 )
399 if target_name not in targets and source_key in targets:
400 target_name, source_pattern = rename_source_key(
401 source_key,
402 [],
403 [],
404 base_model_prefix=base_model_prefix,
405 meta_state_dict=targets,
406 )
407 if target_name not in targets:
408 unexpected_keys.append(source_key)
409 continue
411 if source_pattern is None:
412 source_pattern = source_key
413 transform = WeightRenaming(source_patterns=source_key, target_patterns=target_name)
414 else:
415 converter = select_converter(
416 source_pattern, target_name, converters_by_pattern.get(source_pattern, [])
417 )
418 if converter is None:
419 raise ValueError(
420 "No unique WeightConverter found for matched source pattern "
421 f"{source_pattern!r} and target {target_name!r}"
422 )
423 transform = deepcopy(converter)
425 group = groups.setdefault(
426 target_name,
427 LoadGroup(first_target_name=target_name, transform=transform),
428 )
429 group.transform.add_tensor(target_name, source_key, source_pattern, make_loader(source_key))
431 return tuple(groups.values()), tuple(unexpected_keys), weights_mapping
434def build_replacement_routes(
435 model: nn.Module,
436 source_names: tuple[str, ...],
437 targets: dict[str, torch.Tensor],
438 transforms: list[WeightRenaming | WeightConverter],
439) -> dict[str, tuple[ReplacementLoadGroup, str, str]]:
440 """Route normalized Transformers parameters into replacement converters."""
441 renamings, converters, converters_by_pattern = _split_transforms(deepcopy(transforms))
442 groups: OrderedDict[str, ReplacementLoadGroup] = OrderedDict()
443 routes = {}
444 base_model_prefix = getattr(model, "base_model_prefix", None)
445 for source_name in source_names:
446 target_name, source_pattern = rename_source_key(
447 source_name,
448 renamings,
449 converters,
450 base_model_prefix=base_model_prefix,
451 meta_state_dict=targets,
452 )
453 if source_pattern is None and target_name == source_name:
454 continue
455 if target_name not in targets:
456 continue
458 if source_pattern is None:
459 collected_pattern = source_name
460 transform: WeightRenaming | WeightConverter = WeightRenaming(
461 source_patterns=source_name,
462 target_patterns=target_name,
463 )
464 else:
465 collected_pattern = source_pattern
466 scoped_candidates = _scoped_candidates(converters_by_pattern.get(source_pattern, []), target_name)
467 if len(scoped_candidates) != 1:
468 raise ValueError(
469 "No unique replacement WeightConverter found for source "
470 f"{source_name!r} and target {target_name!r}"
471 )
472 transform = deepcopy(scoped_candidates[0])
474 state = groups.get(target_name)
475 if state is None:
476 state = ReplacementLoadGroup(
477 group=LoadGroup(target_name, transform),
478 expected=Counter(),
479 received=Counter(),
480 )
481 groups[target_name] = state
482 state.expected[collected_pattern] += 1
483 routes[source_name] = (state, target_name, collected_pattern)
484 return routes
487def base_weights_mapping(
488 weights_mapping: list[WeightRenaming | WeightConverter],
489 replacement_mapping: list[WeightRenaming | WeightConverter],
490) -> list[WeightRenaming | WeightConverter]:
491 """The rules of ``weights_mapping`` that are not replacement conversions."""
492 replacement_ids = {id(transform) for transform in replacement_mapping}
493 return [transform for transform in weights_mapping if id(transform) not in replacement_ids]
496def first_tensor(value: Any) -> Any:
497 """The tensor a conversion returned under one name, which a renaming returns in a list."""
498 return value[0] if isinstance(value, list) else value
501def convert_group(group: LoadGroup, model: Any) -> dict[str, Any]:
502 """
503 Convert all checkpoint tensors belonging to one load group.
505 Args:
506 group (LoadGroup): The group, its tensors already collected.
507 model (Any): The model, or the view of it, handed to the conversion operations.
509 Returns:
510 dict[str, Any]: Converted tensors by model tensor name.
512 Raises:
513 RuntimeError: If the conversion fails, naming the group.
514 """
515 try:
516 return group.transform.convert(
517 group.first_target_name,
518 model=model,
519 config=getattr(model, "config", None),
520 hf_quantizer=None,
521 loading_info=None,
522 )
523 except Exception as exc:
524 raise RuntimeError(
525 f"Failed to convert checkpoint tensors for {group.first_target_name}: {exc}"
526 ) from exc
529def local_target_tensor(target: torch.Tensor) -> torch.Tensor:
530 """The part of a model tensor held on this rank."""
531 return target.to_local() if isinstance(target, DTensor) else target
534def target_layout(target: torch.Tensor) -> Any:
535 """The layout a model tensor is sharded with, or None for a whole tensor."""
536 if isinstance(target, DTensor) and target.layout is not None:
537 return target.layout
538 return getattr(target, "_sharding_spec", None)
541def reject_partial_layout(target_name: str, layout: Any) -> None:
542 """Refuse a layout with a Partial placement, which pretrained loads do not support."""
543 if any(isinstance(placement, Partial) for placement in layout.placements):
544 raise ValueError(f"Partial placement is not supported for pretrained loading: {target_name}")
547def shard_for_target(target_name: str, full_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
548 """Shard a full checkpoint tensor according to its target layout."""
549 layout = target_layout(target)
550 if layout is None:
551 return full_tensor
552 reject_partial_layout(target_name, layout)
554 local_dtensor = distribute_tensor(
555 full_tensor,
556 layout.mesh,
557 layout.alias_placements,
558 src_data_rank=None,
559 )
560 return local_dtensor.to_local()
563def copy_into_target(
564 target_name: str, full_tensor: torch.Tensor, target: torch.Tensor,
565 logical_shape: tuple[int, ...] | None = None,
566) -> None:
567 """Copy a checkpoint tensor into its materialized local target."""
568 layout = target_layout(target)
569 target_shape = tuple(getattr(layout, "tensor_shape", None) or target.shape)
570 if logical_shape is not None and tuple(full_tensor.shape) == logical_shape and target_shape != logical_shape:
571 if (len(logical_shape) != len(target_shape)
572 or any(logical > physical for logical, physical in zip(logical_shape, target_shape))):
573 raise ValueError(f"Invalid checkpoint padding shape for {target_name}")
574 padded = full_tensor.new_zeros(target_shape)
575 padded[tuple(slice(0, size) for size in logical_shape)].copy_(full_tensor)
576 full_tensor = padded
577 local_tensor = shard_for_target(target_name, full_tensor, target)
578 destination = local_target_tensor(target)
579 if destination.is_meta:
580 raise ValueError(f"Target must be materialized before loading: {target_name}")
581 if tuple(local_tensor.shape) != tuple(destination.shape):
582 raise ValueError(
583 f"Local shape mismatch for {target_name}: checkpoint shard "
584 f"{tuple(local_tensor.shape)} vs target {tuple(destination.shape)}"
585 )
586 local_tensor = local_tensor.to(device=destination.device, dtype=destination.dtype)
587 with torch.no_grad():
588 destination.copy_(local_tensor)
589 target._is_hf_initialized = True # pylint: disable=W0212
592def alias_names_by_target(targets: dict[str, torch.Tensor]) -> dict[int, set[str]]:
593 """Every name each model tensor is registered under, by the identity of the tensor."""
594 aliases = defaultdict(set)
595 for target_name, target in targets.items():
596 aliases[id(target)].add(target_name)
597 return aliases
600def validate_load_result(
601 missing_keys: tuple[str, ...],
602 unexpected_keys: tuple[str, ...],
603 strict: bool,
604) -> None:
605 """Validate missing keys and report ignored checkpoint tensors."""
606 if strict and missing_keys:
607 preview = ", ".join(missing_keys[:10])
608 raise RuntimeError(
609 f"Checkpoint did not load {len(missing_keys)} owned model tensors; "
610 f"first keys: {preview}"
611 )
612 if unexpected_keys:
613 logger.warning(
614 "Ignored %d checkpoint tensors not owned by this model/rank; first keys: %s",
615 len(unexpected_keys),
616 ", ".join(unexpected_keys[:10]),
617 )