Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / huggingface_load_planner.py: 89%
319 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 checkpoints loaded into a finalized model through distributed checkpoint reads.
17:class:`HFLoadPlanner` loads what :meth:`HuggingFaceCheckpointer.load` loads, with the same
18renaming rules, converters and replacement conversions. The difference is in how the values get
19there. Every conversion is run once on :class:`RegionTensor` inputs while the load is planned, which
20turns it into the regions of checkpoint tensors each model tensor is copied from. Each rank then reads
21only the regions its own shards need, and replicated shards are read once and broadcast. A conversion
22that cannot be traced that way reads its checkpoint tensors whole and runs on them for real, on every
23rank, as the legacy loader runs every conversion.
25Rules come from the model's Transformers conversion mapping unless they are passed in, and a subclass
26can override how the mapping is built, how checkpoint keys are named before the rules see them, and
27which converter a matched key goes to. See :class:`HFLoadPlanner`.
28"""
29import logging
30from collections.abc import Callable, Mapping, Sequence
31from copy import deepcopy
32from dataclasses import dataclass, field
33from functools import partial
34from pathlib import Path
35from typing import Any, Optional
37import torch
38import torch.distributed as dist
39from torch import nn
41from hyper_parallel import DTensor
42from hyper_parallel.components.checkpoint.region_tensor import RegionTensor, materialize
43from hyper_parallel.components.checkpoint.weight_conversion import (
44 WeightConverter,
45 WeightRenaming,
46 dot_natural_key,
47 get_model_conversion_mapping,
48)
49from hyper_parallel.core.distributed_checkpoint.api import load
50from hyper_parallel.components.checkpoint.hf_storage import HuggingFaceStorageReader
51from hyper_parallel.core.distributed_checkpoint.metadata import Metadata, TensorStorageMetadata
52from hyper_parallel.components.checkpoint.remap_planner import DeferredRead, RemapBlock, RemapLoadPlanner
53from hyper_parallel.core.distributed_checkpoint.utils import (
54 all_gather_object,
55 create_chunk_list_for_tensor,
56 str_to_dtype,
57)
58from hyper_parallel.components.checkpoint.quantized_hf_storage import build_quantized_hf_reader
59from hyper_parallel.components.checkpoint.load_groups import (
60 LoadGroup,
61 LoadReport,
62 SourceModelView,
63 alias_names_by_target,
64 base_weights_mapping,
65 build_load_groups,
66 build_load_targets,
67 build_replacement_routes,
68 convert_group,
69 copy_into_target,
70 first_tensor,
71 local_target_tensor,
72 checkpoint_logical_shapes,
73 checkpoint_padding_boxes,
74 reject_partial_layout,
75 resolve_checkpoint_location,
76 select_unique_converter,
77 sorted_checkpoint_keys,
78 target_layout,
79 validate_load_result,
80)
82logger = logging.getLogger(__name__)
84Transform = WeightRenaming | WeightConverter
87def _identity(value: Any) -> Any:
88 """``value``, handed to a transform as a loader it calls when it converts."""
89 return value
92@dataclass(eq=False)
93class _Conversion:
94 """
95 One conversion of the load: a transform, and the values it collects in the order it collects them.
97 Attributes:
98 template (Transform): The transform, with nothing collected. Every run converts a copy of it.
99 first_target_name (str): Name the conversion is known by, as its group is.
100 model (Any): The model, or the view of it, handed to the conversion operations.
101 inputs (list[tuple[str, Any]]): ``(source pattern, value)`` of every collected tensor, where a
102 value is a :class:`RegionTensor` or a :class:`_Deferred` output of another conversion.
103 """
105 template: Transform
106 first_target_name: str
107 model: Any
108 inputs: list[tuple[str, Any]]
110 def sources(self) -> frozenset[str]:
111 """Names of the checkpoint tensors the conversion reads, through its inputs."""
112 return frozenset().union(*(value.sources() for _, value in self.inputs))
114 def run(self, values: Sequence[Any]) -> dict[str, Any]:
115 """
116 Convert ``values`` in place of the inputs, on a fresh copy of the transform.
118 Args:
119 values (Sequence[Any]): One tensor per input: regions, meta tensors or real tensors.
121 Returns:
122 dict[str, Any]: What the transform returns, by model tensor name.
123 """
124 transform = deepcopy(self.template)
125 for (pattern, _), value in zip(self.inputs, values):
126 transform.add_tensor(
127 self.first_target_name,
128 self.first_target_name,
129 pattern,
130 partial(_identity, value)
131 )
132 return convert_group(LoadGroup(self.first_target_name, transform), self.model)
135@dataclass(frozen=True, eq=False)
136class _Deferred:
137 """
138 An output of a conversion that could not be traced, computed for real once its sources are read.
140 Attributes:
141 shape (tuple[int, ...]): Shape the conversion gives the output, as run on meta tensors.
142 dtype (torch.dtype): Dtype it gives the output.
143 conversion (_Conversion): The conversion.
144 name (str): Name of the output among those of the conversion.
145 """
147 shape: tuple[int, ...]
148 dtype: torch.dtype
149 conversion: _Conversion
150 name: str
152 def sources(self) -> frozenset[str]:
153 """Names of the checkpoint tensors the output is computed from."""
154 return self.conversion.sources()
157def _meta_value(value: Any) -> torch.Tensor:
158 """A meta tensor of the shape and dtype of a planned value."""
159 return torch.empty(tuple(value.shape), dtype=value.dtype, device="meta")
162def _real_value(value: Any, fetched: Mapping[str, torch.Tensor], cache: dict[int, dict[str, Any]]) -> Any:
163 """The values of a planned value, out of the whole checkpoint tensors it reads."""
164 if isinstance(value, RegionTensor):
165 return materialize(value, fetched)
166 return first_tensor(_run_real(value.conversion, fetched, cache)[value.name])
169def _run_real(
170 conversion: _Conversion, fetched: Mapping[str, torch.Tensor], cache: dict[int, dict[str, Any]]
171) -> dict[str, Any]:
172 """Run ``conversion`` on real tensors, once for however many of its outputs are asked for."""
173 key = id(conversion)
174 if key not in cache:
175 cache[key] = conversion.run([_real_value(value, fetched, cache) for _, value in conversion.inputs])
176 return cache[key]
179def _complete_conversion(
180 conversion: _Conversion,
181 claims: Sequence[tuple[str, str]],
182 targets: Mapping[str, torch.Tensor],
183 logical_shapes: Mapping[str, tuple[int, ...]],
184 fetched: Mapping[str, torch.Tensor],
185) -> None:
186 """
187 Run a conversion that could not be traced, and copy the outputs it was given into their targets.
189 Args:
190 conversion (_Conversion): The conversion.
191 claims (Sequence[tuple[str, str]]): ``(output name, model tensor name)`` of every output to copy.
192 targets (Mapping[str, torch.Tensor]): Model tensors by name.
193 logical_shapes: Model-declared unpadded checkpoint shapes.
194 fetched (Mapping[str, torch.Tensor]): The whole checkpoint tensors the conversion reads.
195 """
196 outputs = _run_real(conversion, fetched, {})
197 for output_name, target_name in claims:
198 copy_into_target(
199 target_name, first_tensor(outputs[output_name]), targets[target_name], logical_shapes.get(target_name),
200 )
203def _dcp_view(target: torch.Tensor) -> Any:
204 """
205 What the distributed checkpoint load fills for a model tensor.
207 A DTensor is filled as it is. A parameter that holds its local shard as a plain tensor beside its
208 layout, which fully sharded parameters do between materialization and their first use, is filled
209 through a DTensor over the same storage, so that its shard is found in the whole tensor.
210 """
211 if isinstance(target, DTensor):
212 return target
213 layout = getattr(target, "_sharding_spec", None)
214 if layout is not None:
215 return DTensor.from_local_with_layout(target.detach(), layout)
216 return target
219@dataclass
220class _Plan:
221 """
222 What planning the load has decided so far.
224 Attributes:
225 table (dict[str, list[RemapBlock]]): Blocks of every model tensor loaded by region reads, by the
226 name the load's state dict gives it.
227 deferred (dict[int, tuple[_Conversion, list[tuple[str, str]]]]): Conversions run for real, by
228 identity, with the outputs each one copies into model tensors.
229 loaded_ids (set[int]): Identity of every model tensor something loads.
230 loaded_keys (set[str]): Every name those tensors are registered under.
231 unexpected (list[str]): Checkpoint keys and converted names that match no model tensor.
232 """
234 table: dict[str, list[RemapBlock]] = field(default_factory=dict)
235 deferred: dict[int, tuple[_Conversion, list[tuple[str, str]]]] = field(default_factory=dict)
236 loaded_ids: set[int] = field(default_factory=set)
237 loaded_keys: set[str] = field(default_factory=set)
238 unexpected: list[str] = field(default_factory=list)
241class HFLoadPlanner(RemapLoadPlanner):
242 """
243 Load planner for a Hugging Face checkpoint read into a finalized model.
245 Tensors are renamed and converted by ``weights_mapping`` exactly as the legacy loader does it, and a
246 model whose replacement modules convert their weights again is loaded in the same two stages. Every
247 rank has to plan the load, and a load that fails to plan on one rank fails on all of them.
249 The rules, in the order they are tried:
251 * ``extra_weights_mapping``, ahead of everything else, for rules a model needs on top of the default.
252 * ``weights_mapping`` when given, which replaces the default. Otherwise :meth:`default_weights_mapping`
253 builds it: the model's Transformers conversion mapping, with ``key_mapping`` renames first.
255 A rule listed twice is used once. A subclass can also override:
257 * :meth:`map_checkpoint_key`, to rename a checkpoint key before any rule sees it, or skip it.
258 * :meth:`select_converter`, to pick the converter of a key several converters match.
259 * :meth:`supports_symbolic`, to have a transform always run on real tensors.
261 Example::
263 planner = HFLoadPlanner(model)
264 load(planner.state_dict_for_load(), storage_reader=HuggingFaceStorageReader(path), planner=planner)
265 report = planner.finish(path)
266 """
268 def __init__(
269 self,
270 model: nn.Module,
271 *,
272 weights_mapping: Optional[Sequence[Transform]] = None,
273 extra_weights_mapping: Sequence[Transform] = (),
274 key_mapping: Optional[Mapping[str, str]] = None,
275 strict: bool = True,
276 broadcast_replicated_tensors: bool = False,
277 external_target_fqns: frozenset[str] = frozenset(),
278 claimed_source_keys: frozenset[str] = frozenset(),
279 ) -> None:
280 """
281 Args:
282 model (nn.Module): The finalized model, its tensors materialized.
283 weights_mapping (Optional[Sequence[Transform]]): Rules replacing the default ones. Default None.
284 extra_weights_mapping (Sequence[Transform]): Rules tried before the others. Default ().
285 key_mapping (Optional[Mapping[str, str]]): Regex renames of checkpoint keys, applied before the
286 other rules. Default None.
287 strict (bool): Fail the load when a model tensor would be left unloaded. Default True.
288 broadcast_replicated_tensors (bool): See :class:`StandardLoadPlanner`. Default False.
289 external_target_fqns: Model tensors already materialized externally.
290 claimed_source_keys: Checkpoint keys already consumed externally.
291 """
292 super().__init__(broadcast_replicated_tensors=broadcast_replicated_tensors)
293 self.model = model
294 self.weights_mapping = None if weights_mapping is None else list(weights_mapping)
295 self.extra_weights_mapping = list(extra_weights_mapping)
296 self.key_mapping = dict(key_mapping or {})
297 self.strict = strict
298 self.report: Optional[LoadReport] = None
299 targets = build_load_targets(model)
300 logical_shapes = checkpoint_logical_shapes(model)
301 self._logical_shapes = {
302 name: logical_shapes[name] for name in targets if name in logical_shapes
303 }
304 self._padded_targets: dict[int, tuple[str, tuple[int, ...]]] = {}
305 if not external_target_fqns <= targets.keys():
306 raise ValueError("External target is absent from the model")
307 self._targets = {key: value for key, value in targets.items() if key not in external_target_fqns}
308 self._claimed_source_keys = claimed_source_keys
309 self._aliases = alias_names_by_target(self._targets)
310 self._state_keys: dict[int, str] = {}
311 for name, target in self._targets.items():
312 self._state_keys.setdefault(id(target), name)
313 self._load_state = {key: _dcp_view(self._targets[key]) for key in self._state_keys.values()}
314 self._loaded_ids: set[int] = set()
315 self._used: tuple[list[Transform], Optional[list[Transform]]] = ([], None)
317 def default_weights_mapping(self, model: nn.Module) -> list[Transform]:
318 """
319 The rules used when none are passed in.
321 Args:
322 model (nn.Module): The model being loaded.
324 Returns:
325 list[Transform]: The model's Transformers conversion mapping, with ``key_mapping`` first.
326 """
327 return get_model_conversion_mapping(model, key_mapping=self.key_mapping or None, hf_quantizer=None)
329 def map_checkpoint_key(self, key: str) -> Optional[str]:
330 """
331 The name the rules see a checkpoint key under.
333 Args:
334 key (str): The key as the checkpoint stores it.
336 Returns:
337 Optional[str]: The key itself by default. None skips the key, which is then neither loaded
338 nor reported as unexpected.
339 """
340 return None if key in self._claimed_source_keys else key
342 def select_converter(
343 self, source_pattern: str, target_name: str, candidates: Sequence[WeightConverter]
344 ) -> Optional[WeightConverter]:
345 """
346 The converter a checkpoint key matched by ``source_pattern`` belongs to.
348 Args:
349 source_pattern (str): The source pattern the key matched.
350 target_name (str): The model tensor the key was renamed to.
351 candidates (Sequence[WeightConverter]): Every converter with that source pattern.
353 Returns:
354 Optional[WeightConverter]: The only converter scoped to a module holding the target, else the
355 only unscoped one, as the legacy loader picks it. None fails the load.
356 """
357 return select_unique_converter(source_pattern, target_name, candidates)
359 def supports_symbolic(self, transform: Transform) -> bool:
360 """
361 Whether a transform may be traced into region reads, rather than always run on real tensors.
363 Args:
364 transform (Transform): The transform.
366 Returns:
367 bool: True by default. A traced transform is still run on real tensors when tracing fails or
368 disagrees with torch on the shapes and dtypes of what it returns.
369 """
370 del transform
371 return True
373 def state_dict_for_load(self) -> dict[str, Any]:
374 """
375 The state dict to pass to :func:`load` with this planner.
377 Returns:
378 dict[str, Any]: Every model tensor once, under the first name it is registered under.
379 """
380 return dict(self._load_state)
382 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs: Any) -> None:
383 """
384 Configure the planner, then plan how every model tensor is loaded out of the checkpoint.
386 Args:
387 state_dict (dict[str, Any]): What :meth:`state_dict_for_load` returned.
388 metadata (Metadata): Metadata of the checkpoint, as :class:`HuggingFaceStorageReader` builds it.
389 **kwargs (Any): As :func:`load` passes them.
391 Raises:
392 RuntimeError: If planning fails on another rank.
393 Exception: Whatever planning raised on this rank.
394 """
395 super().configure_planner(state_dict, metadata, **kwargs)
396 error = None
397 try:
398 self._compile(metadata)
399 except Exception as exc: # pylint: disable=broad-exception-caught
400 error = exc
401 world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
402 message = None if error is None else f"{type(error).__name__}: {error}"
403 failures = all_gather_object(message, world_size, kwargs.get("use_collectives", False))
404 if error is not None:
405 raise error
406 failed = [(rank, failure) for rank, failure in enumerate(failures) if failure is not None]
407 if failed:
408 raise RuntimeError(f"Planning the pretrained load failed on rank {failed[0][0]}: {failed[0][1]}")
410 def finish(self, pretrained_path: str = "") -> LoadReport:
411 """
412 Complete a load this planner planned and :func:`load` executed.
414 Marks the loaded tensors initialized and records the conversions used on the model, where
415 saving converts back through them, as the legacy loader does.
417 Args:
418 pretrained_path (str): Where the checkpoint was loaded from, for the log. Default "".
420 Returns:
421 LoadReport: What was loaded, what is missing and what the checkpoint holds beyond the model.
422 """
423 if self.report is None:
424 raise RuntimeError("HFLoadPlanner.finish() needs a load planned with this planner first")
425 for target in self._targets.values():
426 padded = self._padded_targets.get(id(target))
427 if padded is not None:
428 self._zero_padding(padded[0], target, padded[1])
429 if id(target) in self._loaded_ids:
430 target._is_hf_initialized = True # pylint: disable=protected-access
431 self.apply_used_conversions(self.model)
432 validate_load_result(self.report.missing_keys, self.report.unexpected_keys, strict=False)
433 logger.info("Loaded %d model tensors from %s", len(self.report.loaded_keys), pretrained_path)
434 return self.report
436 @staticmethod
437 def _zero_padding(name: str, target: torch.Tensor, logical_shape: tuple[int, ...]) -> None:
438 """Zero local padding beyond the declared logical checkpoint shape."""
439 view = _dcp_view(target)
440 chunks = create_chunk_list_for_tensor(view)
441 if not chunks:
442 return
443 if len(chunks) != 1:
444 raise ValueError(f"Expected one rectangular local padded parameter shard for {name}")
445 chunk = chunks[0]
446 local = local_target_tensor(target)
447 if tuple(local.shape) != tuple(chunk.sizes):
448 raise ValueError(f"Local padded parameter storage does not match its shard metadata: {name}")
449 with torch.no_grad():
450 for offsets, lengths in checkpoint_padding_boxes(logical_shape, tuple(chunk.offsets), tuple(local.shape)):
451 slices = tuple(slice(start, start + size) for start, size in zip(offsets, lengths))
452 local[slices].zero_()
454 def apply_used_conversions(self, model: nn.Module) -> None:
455 """
456 Record on ``model`` the conversions the load used, where saving in the original format finds them.
458 Args:
459 model (nn.Module): The model loaded.
460 """
461 used, used_replacements = self._used
462 if used_replacements is None:
463 model._weight_conversions = list(used) # pylint: disable=protected-access
464 return
465 model._hp_used_base_weight_conversions = list(used) # pylint: disable=protected-access
466 model._hp_used_replacement_weight_conversions = list(used_replacements) # pylint: disable=protected-access
467 model._weight_conversions = list(used) + list(used_replacements) # pylint: disable=protected-access
469 def _resolved_mapping(self) -> list[Transform]:
470 """The rules in the order they are tried, each once."""
471 if self.weights_mapping is None:
472 base = list(self.default_weights_mapping(self.model))
473 else:
474 renames = [
475 WeightRenaming(source_patterns=source, target_patterns=target)
476 for source, target in self.key_mapping.items()
477 ]
478 base = renames + self.weights_mapping
479 mapping, seen = [], set()
480 for transform in (*self.extra_weights_mapping, *base):
481 if id(transform) not in seen:
482 seen.add(id(transform))
483 mapping.append(transform)
484 return mapping
486 def _checkpoint_keys(self, metadata: Metadata) -> dict[str, str]:
487 """Every checkpoint tensor the rules see, by the name they see it under."""
488 keys: dict[str, str] = {}
489 for key, md in metadata.state_dict_metadata.items():
490 mapped = self.map_checkpoint_key(key) if isinstance(md, TensorStorageMetadata) else None
491 if mapped is None:
492 continue
493 if mapped in keys:
494 raise ValueError(f"Checkpoint keys {keys[mapped]!r} and {key!r} both map to {mapped!r}")
495 keys[mapped] = key
496 return keys
498 def _compile(self, metadata: Metadata) -> None:
499 """Plan the load of every model tensor, and fill the table and deferred reads it is executed from."""
500 keys = self._checkpoint_keys(metadata)
501 make_loader = partial(self._leaf_loader, metadata, keys)
502 mapping = self._resolved_mapping()
503 replacement_mapping = getattr(self.model, "_hp_replacement_weight_conversions", None)
504 source_shapes = getattr(self.model, "_hp_checkpoint_source_shapes", None)
505 plan = _Plan()
506 if replacement_mapping and source_shapes:
507 self._compile_with_replacements(plan, keys, make_loader, mapping, (replacement_mapping, source_shapes))
508 else:
509 groups, unexpected, mapping = build_load_groups(
510 self.model, sorted_checkpoint_keys(keys), self._targets, weights_mapping=mapping,
511 make_loader=make_loader, select_converter=self.select_converter,
512 )
513 plan.unexpected.extend(keys[key] for key in unexpected)
514 for group in groups:
515 for name, value in self._plan_conversion(_conversion_of(group, self.model)).items():
516 self._claim(plan, name, value)
517 self._used = ([transform for transform in mapping if transform.was_used()], None)
518 self._adopt(plan)
520 def _compile_with_replacements(
521 self,
522 plan: _Plan,
523 keys: dict[str, str],
524 make_loader: Callable[[str], Any],
525 mapping: list[Transform],
526 replacements: tuple[list[Transform], dict[str, tuple[int, ...]]],
527 ) -> None:
528 """Plan a model whose replacement modules convert the normalized weights a second time."""
529 replacement_mapping, source_shapes = replacements
530 base_mapping = base_weights_mapping(mapping, replacement_mapping)
531 source_model = SourceModelView(self.model, source_shapes)
532 base_groups, unexpected, _ = build_load_groups(
533 source_model, sorted_checkpoint_keys(keys), source_model.targets, weights_mapping=base_mapping,
534 make_loader=make_loader, select_converter=self.select_converter,
535 )
536 plan.unexpected.extend(keys[key] for key in unexpected)
537 routes = build_replacement_routes(self.model, tuple(source_shapes), self._targets, replacement_mapping)
538 pending: dict[int, list[tuple[str, Any]]] = {}
539 used_replacements = []
540 for group in base_groups:
541 for source_name, value in self._plan_conversion(_conversion_of(group, source_model)).items():
542 route = routes.get(source_name)
543 if route is None:
544 self._claim(plan, source_name, value)
545 continue
546 state, _, source_pattern = route
547 pending.setdefault(id(state), []).append((source_pattern, value))
548 state.received[source_pattern] += 1
549 if state.completed or state.received != state.expected:
550 continue
551 conversion = _Conversion(
552 state.group.transform, state.group.first_target_name, self.model, pending.pop(id(state))
553 )
554 for name, converted in self._plan_conversion(conversion).items():
555 self._claim(plan, name, converted)
556 state.completed = True
557 used_replacements.append(state.group.transform)
558 self._used = ([transform for transform in base_mapping if transform.was_used()], used_replacements)
560 def _leaf_loader(self, metadata: Metadata, keys: dict[str, str], mapped: str) -> Callable[[], RegionTensor]:
561 """A loader handing a transform the whole of the checkpoint tensor the rules see as ``mapped``."""
562 key = keys[mapped]
563 md = metadata.state_dict_metadata[key]
564 return partial(_identity, RegionTensor.leaf(key, md.size, str_to_dtype(md.properties.dtype)))
566 def _plan_conversion(self, conversion: _Conversion) -> dict[str, Any]:
567 """
568 Trace a conversion into regions, or plan it to run on real tensors when it cannot be traced.
570 The conversion is also run on meta tensors, which both checks what tracing returns against what
571 torch computes and gives the shapes of what an untraced conversion will return.
573 Args:
574 conversion (_Conversion): The conversion.
576 Returns:
577 dict[str, Any]: Every output by name: a :class:`RegionTensor` when traced, else a
578 :class:`_Deferred`.
580 Raises:
581 RuntimeError: If the conversion fails on meta tensors as well as when traced.
582 """
583 values = [value for _, value in conversion.inputs]
584 traced = None
585 if self.supports_symbolic(conversion.template) and all(isinstance(v, RegionTensor) for v in values):
586 try:
587 traced = {name: first_tensor(value) for name, value in conversion.run(values).items()}
588 except Exception as exc: # pylint: disable=broad-exception-caught
589 logger.debug("Tracing %s failed, running it on real tensors: %s", conversion.first_target_name, exc)
590 try:
591 expected = {
592 name: first_tensor(value) for name, value in conversion.run([_meta_value(v) for v in values]).items()
593 }
594 except RuntimeError:
595 if traced is not None and all(isinstance(value, RegionTensor) for value in traced.values()):
596 return traced
597 raise
598 if traced is not None and _same_outputs(traced, expected):
599 return traced
600 return {
601 name: _Deferred(tuple(value.shape), value.dtype, conversion, name) for name, value in expected.items()
602 }
604 def _claim(self, plan: _Plan, name: str, value: Any) -> None:
605 """Have ``value`` load the model tensor ``name``, unless something planned earlier already does."""
606 target = self._targets.get(name)
607 if target is None:
608 plan.unexpected.append(name)
609 return
610 target_id = id(target)
611 if target_id not in plan.loaded_ids:
612 key = self._state_keys[target_id]
613 logical_shape = self._logical_shapes.get(name)
614 self._check_target(name, target, self._load_state[key], value, logical_shape)
615 if logical_shape is not None and tuple(value.shape) == logical_shape:
616 if tuple(self._load_state[key].shape) != logical_shape:
617 self._padded_targets[target_id] = (name, logical_shape)
618 if isinstance(value, RegionTensor):
619 plan.table[key] = value.remap_blocks()
620 else:
621 plan.deferred.setdefault(id(value.conversion), (value.conversion, []))[1].append((value.name, name))
622 plan.loaded_ids.add(target_id)
623 plan.loaded_keys.update(self._aliases[target_id])
625 @staticmethod
626 def _check_target(
627 name: str, target: torch.Tensor, view: Any, value: Any,
628 logical_shape: tuple[int, ...] | None = None,
629 ) -> None:
630 """Refuse a model tensor that cannot take ``value``, as the legacy loader refuses it."""
631 if local_target_tensor(target).is_meta:
632 raise ValueError(f"Target must be materialized before loading: {name}")
633 layout = target_layout(target)
634 if layout is not None:
635 reject_partial_layout(name, layout)
636 target_shape = tuple(view.shape)
637 value_shape = tuple(value.shape)
638 if value_shape != target_shape and not (
639 logical_shape is not None
640 and value_shape == logical_shape
641 and len(logical_shape) == len(target_shape)
642 and all(logical <= physical for logical, physical in zip(logical_shape, target_shape))):
643 raise ValueError(
644 f"Shape mismatch for {name}: checkpoint {value_shape} vs target {target_shape}"
645 )
647 def _adopt(self, plan: _Plan) -> None:
648 """Take a finished plan on as the table and deferred reads of the load, and report on it."""
649 missing = tuple(sorted(set(self._targets) - plan.loaded_keys, key=dot_natural_key))
650 unexpected = tuple(sorted(set(plan.unexpected), key=dot_natural_key))
651 validate_load_result(missing, (), self.strict)
652 deferred = []
653 for conversion, claims in plan.deferred.values():
654 complete = partial(_complete_conversion, conversion, tuple(claims), self._targets, self._logical_shapes)
655 sources = tuple(sorted(conversion.sources()))
656 if sources:
657 deferred.append(DeferredRead(sources, complete))
658 else:
659 # Reads nothing, as a conversion of empty tensors does, so there is nothing to wait for.
660 complete({})
661 self.table = {key: tuple(blocks) for key, blocks in plan.table.items()}
662 self.deferred = tuple(deferred)
663 self._loaded_ids = set(plan.loaded_ids)
664 self.report = LoadReport(
665 loaded_keys=tuple(sorted(plan.loaded_keys, key=dot_natural_key)),
666 missing_keys=missing,
667 unexpected_keys=unexpected,
668 )
669 logger.info(
670 "Planned %d model tensors as region reads and %d through %d conversions on whole tensors",
671 len(self.table), sum(len(claims) for _, claims in plan.deferred.values()), len(plan.deferred),
672 )
675def _conversion_of(group: LoadGroup, model: Any) -> _Conversion:
676 """
677 The conversion of a load group, taking over what the group's transform collected.
679 Args:
680 group (LoadGroup): A group :func:`build_load_groups` built with loaders handing out regions.
681 model (Any): The model, or the view of it, handed to the conversion operations.
683 Returns:
684 _Conversion: The conversion, whose template is the group's transform, emptied.
685 """
686 inputs = [
687 (pattern, loader()) for pattern, loaders in group.transform.collected_tensors.items() for loader in loaders
688 ]
689 group.transform.collected_tensors.clear()
690 return _Conversion(group.transform, group.first_target_name, model, inputs)
693def _same_outputs(traced: Mapping[str, Any], expected: Mapping[str, torch.Tensor]) -> bool:
694 """Whether a traced conversion returns what torch computes, by name, shape and dtype."""
695 return traced.keys() == expected.keys() and all(
696 isinstance(value, RegionTensor)
697 and tuple(value.shape) == tuple(expected[name].shape)
698 and value.dtype == expected[name].dtype
699 for name, value in traced.items()
700 )
703class _SingleFileReader(HuggingFaceStorageReader):
704 """A Hugging Face checkpoint that is one safetensors file, under whatever name it has."""
706 def __init__(self, file_path: Path) -> None:
707 """
708 Args:
709 file_path (Path): The safetensors file.
710 """
711 super().__init__(file_path.parent)
712 self._file_name = file_path.name
714 def _weight_files(self) -> list[str]:
715 """The one file."""
716 return [self._file_name]
719def _load_in_process(state_dict: dict[str, Any], reader: HuggingFaceStorageReader, planner: HFLoadPlanner) -> None:
720 """Plan and execute a load in a process with no process group, which has nobody to share reads with."""
721 metadata = reader.load_metadata()
722 planner.configure_planner(state_dict, metadata, is_coordinator=True, rank=0, use_collectives=False)
723 reader.configure_reader(metadata, is_coordinator=True, rank=0)
724 reader.execute_read(planner.build_local_plan(), planner)
727def load_hf_checkpoint(
728 model: nn.Module,
729 pretrained_path: str,
730 *,
731 weights_mapping: Optional[Sequence[Transform]] = None,
732 strict: bool = True,
733 planner: Optional[HFLoadPlanner] = None,
734 **load_kwargs: Any,
735) -> LoadReport:
736 """
737 Load a Hugging Face checkpoint into a finalized model through distributed checkpoint reads.
739 Every rank of the default process group has to call this together. Without a process group the
740 load runs in this process alone.
742 Args:
743 model (nn.Module): The finalized model, its tensors materialized.
744 pretrained_path (str): A safetensors file, a checkpoint directory or a Hub repository id.
745 weights_mapping (Optional[Sequence[Transform]]): Rules replacing the model's default ones.
746 Default None. Ignored when ``planner`` is given.
747 strict (bool): Fail when a model tensor would be left unloaded. Default True. Ignored when
748 ``planner`` is given.
749 planner (Optional[HFLoadPlanner]): The planner to load with, for rules customized beyond
750 ``weights_mapping``. Default None, for ``HFLoadPlanner(model, ...)``.
751 **load_kwargs (Any): Passed on to :func:`load`, such as ``broadcast_batch_bytes``.
753 Returns:
754 LoadReport: What was loaded, what is missing and what the checkpoint holds beyond the model.
755 """
756 location = resolve_checkpoint_location(pretrained_path)
757 reader = build_quantized_hf_reader(location)
758 if reader is None:
759 reader = _SingleFileReader(location) if location.is_file() else HuggingFaceStorageReader(location)
760 if planner is None:
761 planner = HFLoadPlanner(model, weights_mapping=weights_mapping, strict=strict)
762 state_dict = planner.state_dict_for_load()
763 if dist.is_available() and dist.is_initialized():
764 load(state_dict, storage_reader=reader, planner=planner, **load_kwargs)
765 else:
766 _load_in_process(state_dict, reader, planner)
767 return planner.finish(pretrained_path)