Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / huggingface_load_planner.py: 94%
288 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 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.core.distributed_checkpoint.hf_storage import HuggingFaceStorageReader
51from hyper_parallel.core.distributed_checkpoint.metadata import Metadata, TensorStorageMetadata
52from hyper_parallel.core.distributed_checkpoint.remap_planner import DeferredRead, RemapBlock, RemapLoadPlanner
53from hyper_parallel.core.distributed_checkpoint.utils import all_gather_object, str_to_dtype
54from hyper_parallel.components.checkpoint.load_groups import (
55 LoadGroup,
56 LoadReport,
57 SourceModelView,
58 alias_names_by_target,
59 base_weights_mapping,
60 build_load_groups,
61 build_load_targets,
62 build_replacement_routes,
63 convert_group,
64 copy_into_target,
65 first_tensor,
66 local_target_tensor,
67 reject_partial_layout,
68 resolve_checkpoint_location,
69 select_unique_converter,
70 sorted_checkpoint_keys,
71 target_layout,
72 validate_load_result,
73)
75logger = logging.getLogger(__name__)
77Transform = WeightRenaming | WeightConverter
80def _identity(value: Any) -> Any:
81 """``value``, handed to a transform as a loader it calls when it converts."""
82 return value
85@dataclass(eq=False)
86class _Conversion:
87 """
88 One conversion of the load: a transform, and the values it collects in the order it collects them.
90 Attributes:
91 template (Transform): The transform, with nothing collected. Every run converts a copy of it.
92 first_target_name (str): Name the conversion is known by, as its group is.
93 model (Any): The model, or the view of it, handed to the conversion operations.
94 inputs (list[tuple[str, Any]]): ``(source pattern, value)`` of every collected tensor, where a
95 value is a :class:`RegionTensor` or a :class:`_Deferred` output of another conversion.
96 """
98 template: Transform
99 first_target_name: str
100 model: Any
101 inputs: list[tuple[str, Any]]
103 def sources(self) -> frozenset[str]:
104 """Names of the checkpoint tensors the conversion reads, through its inputs."""
105 return frozenset().union(*(value.sources() for _, value in self.inputs))
107 def run(self, values: Sequence[Any]) -> dict[str, Any]:
108 """
109 Convert ``values`` in place of the inputs, on a fresh copy of the transform.
111 Args:
112 values (Sequence[Any]): One tensor per input: regions, meta tensors or real tensors.
114 Returns:
115 dict[str, Any]: What the transform returns, by model tensor name.
116 """
117 transform = deepcopy(self.template)
118 for (pattern, _), value in zip(self.inputs, values):
119 transform.add_tensor(
120 self.first_target_name,
121 self.first_target_name,
122 pattern,
123 partial(_identity, value)
124 )
125 return convert_group(LoadGroup(self.first_target_name, transform), self.model)
128@dataclass(frozen=True, eq=False)
129class _Deferred:
130 """
131 An output of a conversion that could not be traced, computed for real once its sources are read.
133 Attributes:
134 shape (tuple[int, ...]): Shape the conversion gives the output, as run on meta tensors.
135 dtype (torch.dtype): Dtype it gives the output.
136 conversion (_Conversion): The conversion.
137 name (str): Name of the output among those of the conversion.
138 """
140 shape: tuple[int, ...]
141 dtype: torch.dtype
142 conversion: _Conversion
143 name: str
145 def sources(self) -> frozenset[str]:
146 """Names of the checkpoint tensors the output is computed from."""
147 return self.conversion.sources()
150def _meta_value(value: Any) -> torch.Tensor:
151 """A meta tensor of the shape and dtype of a planned value."""
152 return torch.empty(tuple(value.shape), dtype=value.dtype, device="meta")
155def _real_value(value: Any, fetched: Mapping[str, torch.Tensor], cache: dict[int, dict[str, Any]]) -> Any:
156 """The values of a planned value, out of the whole checkpoint tensors it reads."""
157 if isinstance(value, RegionTensor):
158 return materialize(value, fetched)
159 return first_tensor(_run_real(value.conversion, fetched, cache)[value.name])
162def _run_real(
163 conversion: _Conversion, fetched: Mapping[str, torch.Tensor], cache: dict[int, dict[str, Any]]
164) -> dict[str, Any]:
165 """Run ``conversion`` on real tensors, once for however many of its outputs are asked for."""
166 key = id(conversion)
167 if key not in cache:
168 cache[key] = conversion.run([_real_value(value, fetched, cache) for _, value in conversion.inputs])
169 return cache[key]
172def _complete_conversion(
173 conversion: _Conversion,
174 claims: Sequence[tuple[str, str]],
175 targets: Mapping[str, torch.Tensor],
176 fetched: Mapping[str, torch.Tensor],
177) -> None:
178 """
179 Run a conversion that could not be traced, and copy the outputs it was given into their targets.
181 Args:
182 conversion (_Conversion): The conversion.
183 claims (Sequence[tuple[str, str]]): ``(output name, model tensor name)`` of every output to copy.
184 targets (Mapping[str, torch.Tensor]): Model tensors by name.
185 fetched (Mapping[str, torch.Tensor]): The whole checkpoint tensors the conversion reads.
186 """
187 outputs = _run_real(conversion, fetched, {})
188 for output_name, target_name in claims:
189 copy_into_target(target_name, first_tensor(outputs[output_name]), targets[target_name])
192def _dcp_view(target: torch.Tensor) -> Any:
193 """
194 What the distributed checkpoint load fills for a model tensor.
196 A DTensor is filled as it is. A parameter that holds its local shard as a plain tensor beside its
197 layout, which fully sharded parameters do between materialization and their first use, is filled
198 through a DTensor over the same storage, so that its shard is found in the whole tensor.
199 """
200 if isinstance(target, DTensor):
201 return target
202 layout = getattr(target, "_sharding_spec", None)
203 if layout is not None:
204 return DTensor.from_local_with_layout(target.detach(), layout)
205 return target
208@dataclass
209class _Plan:
210 """
211 What planning the load has decided so far.
213 Attributes:
214 table (dict[str, list[RemapBlock]]): Blocks of every model tensor loaded by region reads, by the
215 name the load's state dict gives it.
216 deferred (dict[int, tuple[_Conversion, list[tuple[str, str]]]]): Conversions run for real, by
217 identity, with the outputs each one copies into model tensors.
218 loaded_ids (set[int]): Identity of every model tensor something loads.
219 loaded_keys (set[str]): Every name those tensors are registered under.
220 unexpected (list[str]): Checkpoint keys and converted names that match no model tensor.
221 """
223 table: dict[str, list[RemapBlock]] = field(default_factory=dict)
224 deferred: dict[int, tuple[_Conversion, list[tuple[str, str]]]] = field(default_factory=dict)
225 loaded_ids: set[int] = field(default_factory=set)
226 loaded_keys: set[str] = field(default_factory=set)
227 unexpected: list[str] = field(default_factory=list)
230class HFLoadPlanner(RemapLoadPlanner):
231 """
232 Load planner for a Hugging Face checkpoint read into a finalized model.
234 Tensors are renamed and converted by ``weights_mapping`` exactly as the legacy loader does it, and a
235 model whose replacement modules convert their weights again is loaded in the same two stages. Every
236 rank has to plan the load, and a load that fails to plan on one rank fails on all of them.
238 The rules, in the order they are tried:
240 * ``extra_weights_mapping``, ahead of everything else, for rules a model needs on top of the default.
241 * ``weights_mapping`` when given, which replaces the default. Otherwise :meth:`default_weights_mapping`
242 builds it: the model's Transformers conversion mapping, with ``key_mapping`` renames first.
244 A rule listed twice is used once. A subclass can also override:
246 * :meth:`map_checkpoint_key`, to rename a checkpoint key before any rule sees it, or skip it.
247 * :meth:`select_converter`, to pick the converter of a key several converters match.
248 * :meth:`supports_symbolic`, to have a transform always run on real tensors.
250 Example::
252 planner = HFLoadPlanner(model)
253 load(planner.state_dict_for_load(), storage_reader=HuggingFaceStorageReader(path), planner=planner)
254 report = planner.finish(path)
255 """
257 def __init__(
258 self,
259 model: nn.Module,
260 *,
261 weights_mapping: Optional[Sequence[Transform]] = None,
262 extra_weights_mapping: Sequence[Transform] = (),
263 key_mapping: Optional[Mapping[str, str]] = None,
264 strict: bool = True,
265 broadcast_replicated_tensors: bool = False,
266 external_target_fqns: frozenset[str] = frozenset(),
267 claimed_source_keys: frozenset[str] = frozenset(),
268 ) -> None:
269 """
270 Args:
271 model (nn.Module): The finalized model, its tensors materialized.
272 weights_mapping (Optional[Sequence[Transform]]): Rules replacing the default ones. Default None.
273 extra_weights_mapping (Sequence[Transform]): Rules tried before the others. Default ().
274 key_mapping (Optional[Mapping[str, str]]): Regex renames of checkpoint keys, applied before the
275 other rules. Default None.
276 strict (bool): Fail the load when a model tensor would be left unloaded. Default True.
277 broadcast_replicated_tensors (bool): See :class:`StandardLoadPlanner`. Default False.
278 external_target_fqns: Model tensors already materialized externally.
279 claimed_source_keys: Checkpoint keys already consumed externally.
280 """
281 super().__init__(broadcast_replicated_tensors=broadcast_replicated_tensors)
282 self.model = model
283 self.weights_mapping = None if weights_mapping is None else list(weights_mapping)
284 self.extra_weights_mapping = list(extra_weights_mapping)
285 self.key_mapping = dict(key_mapping or {})
286 self.strict = strict
287 self.report: Optional[LoadReport] = None
288 targets = build_load_targets(model)
289 if not external_target_fqns <= targets.keys():
290 raise ValueError("External target is absent from the model")
291 self._targets = {key: value for key, value in targets.items() if key not in external_target_fqns}
292 self._claimed_source_keys = claimed_source_keys
293 self._aliases = alias_names_by_target(self._targets)
294 self._state_keys: dict[int, str] = {}
295 for name, target in self._targets.items():
296 self._state_keys.setdefault(id(target), name)
297 self._load_state = {key: _dcp_view(self._targets[key]) for key in self._state_keys.values()}
298 self._loaded_ids: set[int] = set()
299 self._used: tuple[list[Transform], Optional[list[Transform]]] = ([], None)
301 def default_weights_mapping(self, model: nn.Module) -> list[Transform]:
302 """
303 The rules used when none are passed in.
305 Args:
306 model (nn.Module): The model being loaded.
308 Returns:
309 list[Transform]: The model's Transformers conversion mapping, with ``key_mapping`` first.
310 """
311 return get_model_conversion_mapping(model, key_mapping=self.key_mapping or None, hf_quantizer=None)
313 def map_checkpoint_key(self, key: str) -> Optional[str]:
314 """
315 The name the rules see a checkpoint key under.
317 Args:
318 key (str): The key as the checkpoint stores it.
320 Returns:
321 Optional[str]: The key itself by default. None skips the key, which is then neither loaded
322 nor reported as unexpected.
323 """
324 return None if key in self._claimed_source_keys else key
326 def select_converter(
327 self, source_pattern: str, target_name: str, candidates: Sequence[WeightConverter]
328 ) -> Optional[WeightConverter]:
329 """
330 The converter a checkpoint key matched by ``source_pattern`` belongs to.
332 Args:
333 source_pattern (str): The source pattern the key matched.
334 target_name (str): The model tensor the key was renamed to.
335 candidates (Sequence[WeightConverter]): Every converter with that source pattern.
337 Returns:
338 Optional[WeightConverter]: The only converter scoped to a module holding the target, else the
339 only unscoped one, as the legacy loader picks it. None fails the load.
340 """
341 return select_unique_converter(source_pattern, target_name, candidates)
343 def supports_symbolic(self, transform: Transform) -> bool:
344 """
345 Whether a transform may be traced into region reads, rather than always run on real tensors.
347 Args:
348 transform (Transform): The transform.
350 Returns:
351 bool: True by default. A traced transform is still run on real tensors when tracing fails or
352 disagrees with torch on the shapes and dtypes of what it returns.
353 """
354 del transform
355 return True
357 def state_dict_for_load(self) -> dict[str, Any]:
358 """
359 The state dict to pass to :func:`load` with this planner.
361 Returns:
362 dict[str, Any]: Every model tensor once, under the first name it is registered under.
363 """
364 return dict(self._load_state)
366 def configure_planner(self, state_dict: dict[str, Any], metadata: Metadata, **kwargs: Any) -> None:
367 """
368 Configure the planner, then plan how every model tensor is loaded out of the checkpoint.
370 Args:
371 state_dict (dict[str, Any]): What :meth:`state_dict_for_load` returned.
372 metadata (Metadata): Metadata of the checkpoint, as :class:`HuggingFaceStorageReader` builds it.
373 **kwargs (Any): As :func:`load` passes them.
375 Raises:
376 RuntimeError: If planning fails on another rank.
377 Exception: Whatever planning raised on this rank.
378 """
379 super().configure_planner(state_dict, metadata, **kwargs)
380 error = None
381 try:
382 self._compile(metadata)
383 except Exception as exc: # pylint: disable=broad-exception-caught
384 error = exc
385 world_size = dist.get_world_size() if dist.is_available() and dist.is_initialized() else 1
386 message = None if error is None else f"{type(error).__name__}: {error}"
387 failures = all_gather_object(message, world_size, kwargs.get("use_collectives", False))
388 if error is not None:
389 raise error
390 failed = [(rank, failure) for rank, failure in enumerate(failures) if failure is not None]
391 if failed:
392 raise RuntimeError(f"Planning the pretrained load failed on rank {failed[0][0]}: {failed[0][1]}")
394 def finish(self, pretrained_path: str = "") -> LoadReport:
395 """
396 Complete a load this planner planned and :func:`load` executed.
398 Marks the loaded tensors initialized and records the conversions used on the model, where
399 saving converts back through them, as the legacy loader does.
401 Args:
402 pretrained_path (str): Where the checkpoint was loaded from, for the log. Default "".
404 Returns:
405 LoadReport: What was loaded, what is missing and what the checkpoint holds beyond the model.
406 """
407 if self.report is None:
408 raise RuntimeError("HFLoadPlanner.finish() needs a load planned with this planner first")
409 for target in self._targets.values():
410 if id(target) in self._loaded_ids:
411 target._is_hf_initialized = True # pylint: disable=protected-access
412 self.apply_used_conversions(self.model)
413 validate_load_result(self.report.missing_keys, self.report.unexpected_keys, strict=False)
414 logger.info("Loaded %d model tensors from %s", len(self.report.loaded_keys), pretrained_path)
415 return self.report
417 def apply_used_conversions(self, model: nn.Module) -> None:
418 """
419 Record on ``model`` the conversions the load used, where saving in the original format finds them.
421 Args:
422 model (nn.Module): The model loaded.
423 """
424 used, used_replacements = self._used
425 if used_replacements is None:
426 model._weight_conversions = list(used) # pylint: disable=protected-access
427 return
428 model._hp_used_base_weight_conversions = list(used) # pylint: disable=protected-access
429 model._hp_used_replacement_weight_conversions = list(used_replacements) # pylint: disable=protected-access
430 model._weight_conversions = list(used) + list(used_replacements) # pylint: disable=protected-access
432 def _resolved_mapping(self) -> list[Transform]:
433 """The rules in the order they are tried, each once."""
434 if self.weights_mapping is None:
435 base = list(self.default_weights_mapping(self.model))
436 else:
437 renames = [
438 WeightRenaming(source_patterns=source, target_patterns=target)
439 for source, target in self.key_mapping.items()
440 ]
441 base = renames + self.weights_mapping
442 mapping, seen = [], set()
443 for transform in (*self.extra_weights_mapping, *base):
444 if id(transform) not in seen:
445 seen.add(id(transform))
446 mapping.append(transform)
447 return mapping
449 def _checkpoint_keys(self, metadata: Metadata) -> dict[str, str]:
450 """Every checkpoint tensor the rules see, by the name they see it under."""
451 keys: dict[str, str] = {}
452 for key, md in metadata.state_dict_metadata.items():
453 mapped = self.map_checkpoint_key(key) if isinstance(md, TensorStorageMetadata) else None
454 if mapped is None:
455 continue
456 if mapped in keys:
457 raise ValueError(f"Checkpoint keys {keys[mapped]!r} and {key!r} both map to {mapped!r}")
458 keys[mapped] = key
459 return keys
461 def _compile(self, metadata: Metadata) -> None:
462 """Plan the load of every model tensor, and fill the table and deferred reads it is executed from."""
463 keys = self._checkpoint_keys(metadata)
464 make_loader = partial(self._leaf_loader, metadata, keys)
465 mapping = self._resolved_mapping()
466 replacement_mapping = getattr(self.model, "_hp_replacement_weight_conversions", None)
467 source_shapes = getattr(self.model, "_hp_checkpoint_source_shapes", None)
468 plan = _Plan()
469 if replacement_mapping and source_shapes:
470 self._compile_with_replacements(plan, keys, make_loader, mapping, (replacement_mapping, source_shapes))
471 else:
472 groups, unexpected, mapping = build_load_groups(
473 self.model, sorted_checkpoint_keys(keys), self._targets, weights_mapping=mapping,
474 make_loader=make_loader, select_converter=self.select_converter,
475 )
476 plan.unexpected.extend(keys[key] for key in unexpected)
477 for group in groups:
478 for name, value in self._plan_conversion(_conversion_of(group, self.model)).items():
479 self._claim(plan, name, value)
480 self._used = ([transform for transform in mapping if transform.was_used()], None)
481 self._adopt(plan)
483 def _compile_with_replacements(
484 self,
485 plan: _Plan,
486 keys: dict[str, str],
487 make_loader: Callable[[str], Any],
488 mapping: list[Transform],
489 replacements: tuple[list[Transform], dict[str, tuple[int, ...]]],
490 ) -> None:
491 """Plan a model whose replacement modules convert the normalized weights a second time."""
492 replacement_mapping, source_shapes = replacements
493 base_mapping = base_weights_mapping(mapping, replacement_mapping)
494 source_model = SourceModelView(self.model, source_shapes)
495 base_groups, unexpected, _ = build_load_groups(
496 source_model, sorted_checkpoint_keys(keys), source_model.targets, weights_mapping=base_mapping,
497 make_loader=make_loader, select_converter=self.select_converter,
498 )
499 plan.unexpected.extend(keys[key] for key in unexpected)
500 routes = build_replacement_routes(self.model, tuple(source_shapes), self._targets, replacement_mapping)
501 pending: dict[int, list[tuple[str, Any]]] = {}
502 used_replacements = []
503 for group in base_groups:
504 for source_name, value in self._plan_conversion(_conversion_of(group, source_model)).items():
505 route = routes.get(source_name)
506 if route is None:
507 self._claim(plan, source_name, value)
508 continue
509 state, _, source_pattern = route
510 pending.setdefault(id(state), []).append((source_pattern, value))
511 state.received[source_pattern] += 1
512 if state.completed or state.received != state.expected:
513 continue
514 conversion = _Conversion(
515 state.group.transform, state.group.first_target_name, self.model, pending.pop(id(state))
516 )
517 for name, converted in self._plan_conversion(conversion).items():
518 self._claim(plan, name, converted)
519 state.completed = True
520 used_replacements.append(state.group.transform)
521 self._used = ([transform for transform in base_mapping if transform.was_used()], used_replacements)
523 def _leaf_loader(self, metadata: Metadata, keys: dict[str, str], mapped: str) -> Callable[[], RegionTensor]:
524 """A loader handing a transform the whole of the checkpoint tensor the rules see as ``mapped``."""
525 key = keys[mapped]
526 md = metadata.state_dict_metadata[key]
527 return partial(_identity, RegionTensor.leaf(key, md.size, str_to_dtype(md.properties.dtype)))
529 def _plan_conversion(self, conversion: _Conversion) -> dict[str, Any]:
530 """
531 Trace a conversion into regions, or plan it to run on real tensors when it cannot be traced.
533 The conversion is also run on meta tensors, which both checks what tracing returns against what
534 torch computes and gives the shapes of what an untraced conversion will return.
536 Args:
537 conversion (_Conversion): The conversion.
539 Returns:
540 dict[str, Any]: Every output by name: a :class:`RegionTensor` when traced, else a
541 :class:`_Deferred`.
543 Raises:
544 RuntimeError: If the conversion fails on meta tensors as well as when traced.
545 """
546 values = [value for _, value in conversion.inputs]
547 traced = None
548 if self.supports_symbolic(conversion.template) and all(isinstance(v, RegionTensor) for v in values):
549 try:
550 traced = {name: first_tensor(value) for name, value in conversion.run(values).items()}
551 except Exception as exc: # pylint: disable=broad-exception-caught
552 logger.debug("Tracing %s failed, running it on real tensors: %s", conversion.first_target_name, exc)
553 try:
554 expected = {
555 name: first_tensor(value) for name, value in conversion.run([_meta_value(v) for v in values]).items()
556 }
557 except RuntimeError:
558 if traced is not None and all(isinstance(value, RegionTensor) for value in traced.values()):
559 return traced
560 raise
561 if traced is not None and _same_outputs(traced, expected):
562 return traced
563 return {
564 name: _Deferred(tuple(value.shape), value.dtype, conversion, name) for name, value in expected.items()
565 }
567 def _claim(self, plan: _Plan, name: str, value: Any) -> None:
568 """Have ``value`` load the model tensor ``name``, unless something planned earlier already does."""
569 target = self._targets.get(name)
570 if target is None:
571 plan.unexpected.append(name)
572 return
573 target_id = id(target)
574 if target_id not in plan.loaded_ids:
575 key = self._state_keys[target_id]
576 self._check_target(name, target, self._load_state[key], value)
577 if isinstance(value, RegionTensor):
578 plan.table[key] = value.remap_blocks()
579 else:
580 plan.deferred.setdefault(id(value.conversion), (value.conversion, []))[1].append((value.name, name))
581 plan.loaded_ids.add(target_id)
582 plan.loaded_keys.update(self._aliases[target_id])
584 @staticmethod
585 def _check_target(name: str, target: torch.Tensor, view: Any, value: Any) -> None:
586 """Refuse a model tensor that cannot take ``value``, as the legacy loader refuses it."""
587 if local_target_tensor(target).is_meta:
588 raise ValueError(f"Target must be materialized before loading: {name}")
589 layout = target_layout(target)
590 if layout is not None:
591 reject_partial_layout(name, layout)
592 if tuple(value.shape) != tuple(view.shape):
593 raise ValueError(
594 f"Shape mismatch for {name}: checkpoint {tuple(value.shape)} vs target {tuple(view.shape)}"
595 )
597 def _adopt(self, plan: _Plan) -> None:
598 """Take a finished plan on as the table and deferred reads of the load, and report on it."""
599 missing = tuple(sorted(set(self._targets) - plan.loaded_keys, key=dot_natural_key))
600 unexpected = tuple(sorted(set(plan.unexpected), key=dot_natural_key))
601 validate_load_result(missing, (), self.strict)
602 deferred = []
603 for conversion, claims in plan.deferred.values():
604 complete = partial(_complete_conversion, conversion, tuple(claims), self._targets)
605 sources = tuple(sorted(conversion.sources()))
606 if sources:
607 deferred.append(DeferredRead(sources, complete))
608 else:
609 # Reads nothing, as a conversion of empty tensors does, so there is nothing to wait for.
610 complete({})
611 self.table = {key: tuple(blocks) for key, blocks in plan.table.items()}
612 self.deferred = tuple(deferred)
613 self._loaded_ids = set(plan.loaded_ids)
614 self.report = LoadReport(
615 loaded_keys=tuple(sorted(plan.loaded_keys, key=dot_natural_key)),
616 missing_keys=missing,
617 unexpected_keys=unexpected,
618 )
619 logger.info(
620 "Planned %d model tensors as region reads and %d through %d conversions on whole tensors",
621 len(self.table), sum(len(claims) for _, claims in plan.deferred.values()), len(plan.deferred),
622 )
625def _conversion_of(group: LoadGroup, model: Any) -> _Conversion:
626 """
627 The conversion of a load group, taking over what the group's transform collected.
629 Args:
630 group (LoadGroup): A group :func:`build_load_groups` built with loaders handing out regions.
631 model (Any): The model, or the view of it, handed to the conversion operations.
633 Returns:
634 _Conversion: The conversion, whose template is the group's transform, emptied.
635 """
636 inputs = [
637 (pattern, loader()) for pattern, loaders in group.transform.collected_tensors.items() for loader in loaders
638 ]
639 group.transform.collected_tensors.clear()
640 return _Conversion(group.transform, group.first_target_name, model, inputs)
643def _same_outputs(traced: Mapping[str, Any], expected: Mapping[str, torch.Tensor]) -> bool:
644 """Whether a traced conversion returns what torch computes, by name, shape and dtype."""
645 return traced.keys() == expected.keys() and all(
646 isinstance(value, RegionTensor)
647 and tuple(value.shape) == tuple(expected[name].shape)
648 and value.dtype == expected[name].dtype
649 for name, value in traced.items()
650 )
653class _SingleFileReader(HuggingFaceStorageReader):
654 """A Hugging Face checkpoint that is one safetensors file, under whatever name it has."""
656 def __init__(self, file_path: Path) -> None:
657 """
658 Args:
659 file_path (Path): The safetensors file.
660 """
661 super().__init__(file_path.parent)
662 self._file_name = file_path.name
664 def _weight_files(self) -> list[str]:
665 """The one file."""
666 return [self._file_name]
669def _load_in_process(state_dict: dict[str, Any], reader: HuggingFaceStorageReader, planner: HFLoadPlanner) -> None:
670 """Plan and execute a load in a process with no process group, which has nobody to share reads with."""
671 metadata = reader.load_metadata()
672 planner.configure_planner(state_dict, metadata, is_coordinator=True, rank=0, use_collectives=False)
673 reader.configure_reader(metadata, is_coordinator=True, rank=0)
674 reader.execute_read(planner.build_local_plan(), planner)
677def load_hf_checkpoint(
678 model: nn.Module,
679 pretrained_path: str,
680 *,
681 weights_mapping: Optional[Sequence[Transform]] = None,
682 strict: bool = True,
683 planner: Optional[HFLoadPlanner] = None,
684 **load_kwargs: Any,
685) -> LoadReport:
686 """
687 Load a Hugging Face checkpoint into a finalized model through distributed checkpoint reads.
689 Every rank of the default process group has to call this together. Without a process group the
690 load runs in this process alone.
692 Args:
693 model (nn.Module): The finalized model, its tensors materialized.
694 pretrained_path (str): A safetensors file, a checkpoint directory or a Hub repository id.
695 weights_mapping (Optional[Sequence[Transform]]): Rules replacing the model's default ones.
696 Default None. Ignored when ``planner`` is given.
697 strict (bool): Fail when a model tensor would be left unloaded. Default True. Ignored when
698 ``planner`` is given.
699 planner (Optional[HFLoadPlanner]): The planner to load with, for rules customized beyond
700 ``weights_mapping``. Default None, for ``HFLoadPlanner(model, ...)``.
701 **load_kwargs (Any): Passed on to :func:`load`, such as ``broadcast_batch_bytes``.
703 Returns:
704 LoadReport: What was loaded, what is missing and what the checkpoint holds beyond the model.
705 """
706 location = resolve_checkpoint_location(pretrained_path)
707 reader = _SingleFileReader(location) if location.is_file() else HuggingFaceStorageReader(location)
708 if planner is None:
709 planner = HFLoadPlanner(model, weights_mapping=weights_mapping, strict=strict)
710 state_dict = planner.state_dict_for_load()
711 if dist.is_available() and dist.is_initialized():
712 load(state_dict, storage_reader=reader, planner=planner, **load_kwargs)
713 else:
714 _load_in_process(state_dict, reader, planner)
715 return planner.finish(pretrained_path)