Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / load_groups.py: 82%
234 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"""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 make_tensor_loader(index: CheckpointIndex, source_key: str) -> Callable[[], torch.Tensor]:
233 """A loader reading one checkpoint tensor when a conversion asks for it."""
234 return lambda: index.load_tensor(source_key)
237def _scoped_candidates(candidates: Sequence[WeightConverter], target_name: str) -> list[WeightConverter]:
238 """The candidates scoped to a module that holds ``target_name``."""
239 return [
240 converter
241 for converter in candidates
242 if converter.scope_prefix is not None
243 and (
244 target_name == converter.scope_prefix
245 or target_name.startswith(f"{converter.scope_prefix}.")
246 )
247 ]
250def select_unique_converter(
251 source_pattern: str,
252 target_name: str,
253 candidates: Sequence[WeightConverter],
254) -> Optional[WeightConverter]:
255 """
256 Pick the converter a matched checkpoint key belongs to.
258 A converter scoped to a module holding the target wins when it is the only one. Otherwise the only
259 unscoped converter does.
261 Args:
262 source_pattern (str): The source pattern the key matched.
263 target_name (str): The model tensor the key was renamed to.
264 candidates (Sequence[WeightConverter]): Every converter with that source pattern.
266 Returns:
267 Optional[WeightConverter]: The converter, or None when neither rule picks exactly one.
268 """
269 del source_pattern
270 scoped_candidates = _scoped_candidates(candidates, target_name)
271 if len(scoped_candidates) == 1:
272 return scoped_candidates[0]
273 unscoped_candidates = [converter for converter in candidates if converter.scope_prefix is None]
274 return unscoped_candidates[0] if len(unscoped_candidates) == 1 else None
277def _split_transforms(
278 transforms: Sequence[Any],
279) -> tuple[list[WeightRenaming], list[WeightConverter], dict[str, list[WeightConverter]]]:
280 """Renamings, converters, and the converters by each of their source patterns."""
281 renamings = [transform for transform in transforms if isinstance(transform, WeightRenaming)]
282 converters = [transform for transform in transforms if isinstance(transform, WeightConverter)]
283 converters_by_pattern = defaultdict(list)
284 for converter in converters:
285 for pattern in converter.source_patterns:
286 converters_by_pattern[pattern].append(converter)
287 return renamings, converters, converters_by_pattern
290def build_load_groups(
291 model: Any,
292 source_keys: Iterable[str],
293 targets: dict[str, Any],
294 *,
295 weights_mapping: list[WeightRenaming | WeightConverter],
296 make_loader: Callable[[str], Any],
297 select_converter: ConverterSelector = select_unique_converter,
298) -> tuple[
299 tuple[LoadGroup, ...],
300 tuple[str, ...],
301 list[WeightRenaming | WeightConverter],
302]:
303 """
304 Build checkpoint conversion groups and report unmatched checkpoint keys.
306 Args:
307 model (Any): The model, or the view of it the keys are renamed against.
308 source_keys (Iterable[str]): Checkpoint keys, in the order their tensors are collected.
309 targets (dict[str, Any]): Model tensors by name.
310 weights_mapping (list[WeightRenaming | WeightConverter]): Rules renaming and converting keys.
311 make_loader (Callable[[str], Any]): Called with a key, returns what the group collects for it.
312 select_converter (ConverterSelector): Picks the converter of a matched key.
313 Default :func:`select_unique_converter`.
315 Returns:
316 tuple: The groups in order of their first key, the keys renamed to no target, and the mapping.
317 """
318 unsupported = [
319 transform
320 for transform in weights_mapping
321 if not isinstance(transform, (WeightRenaming, WeightConverter))
322 ]
323 if unsupported:
324 names = ", ".join(type(transform).__name__ for transform in unsupported)
325 raise ValueError(f"Unsupported Transformers weight transforms in MVP: {names}")
327 renamings, converters, converters_by_pattern = _split_transforms(weights_mapping)
328 groups: OrderedDict[str, LoadGroup] = OrderedDict()
329 unexpected_keys = []
330 base_model_prefix = getattr(model, "base_model_prefix", None)
332 for source_key in source_keys:
333 target_name, source_pattern = rename_source_key(
334 source_key,
335 renamings,
336 converters,
337 base_model_prefix=base_model_prefix,
338 meta_state_dict=targets,
339 )
340 if target_name not in targets and source_key in targets:
341 target_name, source_pattern = rename_source_key(
342 source_key,
343 [],
344 [],
345 base_model_prefix=base_model_prefix,
346 meta_state_dict=targets,
347 )
348 if target_name not in targets:
349 unexpected_keys.append(source_key)
350 continue
352 if source_pattern is None:
353 source_pattern = source_key
354 transform = WeightRenaming(source_patterns=source_key, target_patterns=target_name)
355 else:
356 converter = select_converter(
357 source_pattern, target_name, converters_by_pattern.get(source_pattern, [])
358 )
359 if converter is None:
360 raise ValueError(
361 "No unique WeightConverter found for matched source pattern "
362 f"{source_pattern!r} and target {target_name!r}"
363 )
364 transform = deepcopy(converter)
366 group = groups.setdefault(
367 target_name,
368 LoadGroup(first_target_name=target_name, transform=transform),
369 )
370 group.transform.add_tensor(target_name, source_key, source_pattern, make_loader(source_key))
372 return tuple(groups.values()), tuple(unexpected_keys), weights_mapping
375def build_replacement_routes(
376 model: nn.Module,
377 source_names: tuple[str, ...],
378 targets: dict[str, torch.Tensor],
379 transforms: list[WeightRenaming | WeightConverter],
380) -> dict[str, tuple[ReplacementLoadGroup, str, str]]:
381 """Route normalized Transformers parameters into replacement converters."""
382 renamings, converters, converters_by_pattern = _split_transforms(deepcopy(transforms))
383 groups: OrderedDict[str, ReplacementLoadGroup] = OrderedDict()
384 routes = {}
385 base_model_prefix = getattr(model, "base_model_prefix", None)
386 for source_name in source_names:
387 target_name, source_pattern = rename_source_key(
388 source_name,
389 renamings,
390 converters,
391 base_model_prefix=base_model_prefix,
392 meta_state_dict=targets,
393 )
394 if source_pattern is None and target_name == source_name:
395 continue
396 if target_name not in targets:
397 continue
399 if source_pattern is None:
400 collected_pattern = source_name
401 transform: WeightRenaming | WeightConverter = WeightRenaming(
402 source_patterns=source_name,
403 target_patterns=target_name,
404 )
405 else:
406 collected_pattern = source_pattern
407 scoped_candidates = _scoped_candidates(converters_by_pattern.get(source_pattern, []), target_name)
408 if len(scoped_candidates) != 1:
409 raise ValueError(
410 "No unique replacement WeightConverter found for source "
411 f"{source_name!r} and target {target_name!r}"
412 )
413 transform = deepcopy(scoped_candidates[0])
415 state = groups.get(target_name)
416 if state is None:
417 state = ReplacementLoadGroup(
418 group=LoadGroup(target_name, transform),
419 expected=Counter(),
420 received=Counter(),
421 )
422 groups[target_name] = state
423 state.expected[collected_pattern] += 1
424 routes[source_name] = (state, target_name, collected_pattern)
425 return routes
428def base_weights_mapping(
429 weights_mapping: list[WeightRenaming | WeightConverter],
430 replacement_mapping: list[WeightRenaming | WeightConverter],
431) -> list[WeightRenaming | WeightConverter]:
432 """The rules of ``weights_mapping`` that are not replacement conversions."""
433 replacement_ids = {id(transform) for transform in replacement_mapping}
434 return [transform for transform in weights_mapping if id(transform) not in replacement_ids]
437def first_tensor(value: Any) -> Any:
438 """The tensor a conversion returned under one name, which a renaming returns in a list."""
439 return value[0] if isinstance(value, list) else value
442def convert_group(group: LoadGroup, model: Any) -> dict[str, Any]:
443 """
444 Convert all checkpoint tensors belonging to one load group.
446 Args:
447 group (LoadGroup): The group, its tensors already collected.
448 model (Any): The model, or the view of it, handed to the conversion operations.
450 Returns:
451 dict[str, Any]: Converted tensors by model tensor name.
453 Raises:
454 RuntimeError: If the conversion fails, naming the group.
455 """
456 try:
457 return group.transform.convert(
458 group.first_target_name,
459 model=model,
460 config=getattr(model, "config", None),
461 hf_quantizer=None,
462 loading_info=None,
463 )
464 except Exception as exc:
465 raise RuntimeError(
466 f"Failed to convert checkpoint tensors for {group.first_target_name}: {exc}"
467 ) from exc
470def local_target_tensor(target: torch.Tensor) -> torch.Tensor:
471 """The part of a model tensor held on this rank."""
472 return target.to_local() if isinstance(target, DTensor) else target
475def target_layout(target: torch.Tensor) -> Any:
476 """The layout a model tensor is sharded with, or None for a whole tensor."""
477 if isinstance(target, DTensor) and target.layout is not None:
478 return target.layout
479 return getattr(target, "_sharding_spec", None)
482def reject_partial_layout(target_name: str, layout: Any) -> None:
483 """Refuse a layout with a Partial placement, which pretrained loads do not support."""
484 if any(isinstance(placement, Partial) for placement in layout.placements):
485 raise ValueError(f"Partial placement is not supported for pretrained loading: {target_name}")
488def shard_for_target(target_name: str, full_tensor: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
489 """Shard a full checkpoint tensor according to its target layout."""
490 layout = target_layout(target)
491 if layout is None:
492 return full_tensor
493 reject_partial_layout(target_name, layout)
495 local_dtensor = distribute_tensor(
496 full_tensor,
497 layout.mesh,
498 layout.alias_placements,
499 src_data_rank=None,
500 )
501 return local_dtensor.to_local()
504def copy_into_target(target_name: str, full_tensor: torch.Tensor, target: torch.Tensor) -> None:
505 """Copy a checkpoint tensor into its materialized local target."""
506 local_tensor = shard_for_target(target_name, full_tensor, target)
507 destination = local_target_tensor(target)
508 if destination.is_meta:
509 raise ValueError(f"Target must be materialized before loading: {target_name}")
510 if tuple(local_tensor.shape) != tuple(destination.shape):
511 raise ValueError(
512 f"Local shape mismatch for {target_name}: checkpoint shard "
513 f"{tuple(local_tensor.shape)} vs target {tuple(destination.shape)}"
514 )
515 local_tensor = local_tensor.to(device=destination.device, dtype=destination.dtype)
516 with torch.no_grad():
517 destination.copy_(local_tensor)
518 target._is_hf_initialized = True # pylint: disable=W0212
521def alias_names_by_target(targets: dict[str, torch.Tensor]) -> dict[int, set[str]]:
522 """Every name each model tensor is registered under, by the identity of the tensor."""
523 aliases = defaultdict(set)
524 for target_name, target in targets.items():
525 aliases[id(target)].add(target_name)
526 return aliases
529def validate_load_result(
530 missing_keys: tuple[str, ...],
531 unexpected_keys: tuple[str, ...],
532 strict: bool,
533) -> None:
534 """Validate missing keys and report ignored checkpoint tensors."""
535 if strict and missing_keys:
536 preview = ", ".join(missing_keys[:10])
537 raise RuntimeError(
538 f"Checkpoint did not load {len(missing_keys)} owned model tensors; "
539 f"first keys: {preview}"
540 )
541 if unexpected_keys:
542 logger.warning(
543 "Ignored %d checkpoint tensors not owned by this model/rank; first keys: %s",
544 len(unexpected_keys),
545 ", ".join(unexpected_keys[:10]),
546 )