Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / components / checkpoint / huggingface_checkpointer.py: 95%
135 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 pretrained-weight backend for :class:`CheckpointerBase`.
17Unlike :class:`~hyper_parallel.components.checkpoint.dcp_checkpointer.DistributedCheckpointer`,
18which round-trips its own format, this backend reads a foreign one: a Transformers checkpoint whose
19tensor names and layouts belong to the released model, not to the finalized distributed one. Keys are
20therefore renamed and tensors converted on the way in, which is what
21``load_groups.py`` exists for.
23The payload this backend fills is ``{"model": <nn.Module>}`` --- a live module rather than a state
24dict, because loading writes into the distributed parameters in place and must not replace them. It
25adds a ``load_report`` entry naming what was loaded, which
26``models/_transformers/checkpoint_loader.py::_finalize_model_loading`` consumes to initialize
27whatever the checkpoint did not cover.
29Only :meth:`HuggingFaceCheckpointer.load` is implemented; writing a Transformers-format checkpoint
30back out is ``models/_transformers/checkpoint_loader.py::CheckpointManager.save_pretrained``.
31"""
33import logging
34import os
35from functools import partial
36from typing import Any, Dict, List, Optional, Union
38import torch # pylint: disable=forbidden-backend-import
39from torch import nn # pylint: disable=forbidden-backend-import
41from hyper_parallel.components.checkpoint.base import CheckpointerBase
42from hyper_parallel.components.checkpoint.huggingface_load_planner import HFLoadPlanner, load_hf_checkpoint
43from hyper_parallel.components.checkpoint.load_groups import (
44 CheckpointIndex,
45 LoadGroup,
46 LoadReport,
47 SourceModelView,
48 alias_names_by_target,
49 base_weights_mapping,
50 build_load_groups,
51 build_load_targets,
52 build_replacement_routes,
53 convert_group,
54 copy_into_target,
55 checkpoint_logical_shapes,
56 make_tensor_loader,
57 resolve_checkpoint_index,
58 validate_load_result,
59)
60from hyper_parallel.components.checkpoint.weight_conversion import (
61 WeightConverter,
62 WeightRenaming,
63 dot_natural_key,
64 get_model_conversion_mapping,
65)
67logger = logging.getLogger(__name__)
69# Picks the pretrained loader when HuggingFaceCheckpointer is not told which one to use.
70HF_LOADER_ENV = "HYPER_PARALLEL_HF_LOADER"
71_HF_LOADERS = ("legacy", "dcp")
73# Keys of the payload :meth:`HuggingFaceCheckpointer.load` reads and fills. ``model`` holds the
74# finalized module to load into; ``load_report`` is written back for the caller that finalizes it.
75_MODEL_KEY = "model"
76_LOAD_REPORT_KEY = "load_report"
78WeightsMapping = List[Union[WeightRenaming, WeightConverter]]
81def resolve_hf_loader(loader: Optional[str] = None) -> str:
82 """
83 Name the pretrained loader to use.
85 Args:
86 loader (str | None): ``"legacy"`` reads whole checkpoint tensors on every rank and shards them in
87 memory. ``"dcp"`` plans the same conversions as distributed checkpoint reads, so that each
88 rank reads only the regions of the checkpoint its shards need. Default None, which reads
89 ``HYPER_PARALLEL_HF_LOADER`` and falls back to ``"dcp"``.
91 Returns:
92 str: ``"legacy"`` or ``"dcp"``.
94 Raises:
95 ValueError: If the loader named is neither.
96 """
97 choice = (loader or os.environ.get(HF_LOADER_ENV) or "dcp").strip().lower()
98 if choice not in _HF_LOADERS:
99 raise ValueError(
100 f"Unknown pretrained loader {choice!r}; expected one of {', '.join(_HF_LOADERS)} "
101 f"(set through the loader argument or {HF_LOADER_ENV})"
102 )
103 return choice
106class HuggingFaceCheckpointer(CheckpointerBase):
107 """Load complete Hugging Face weights into one finalized model.
109 Both loaders selected by ``loader`` reach the same values under the same names --- they differ
110 only in how much of the checkpoint each rank reads:
112 * ``"legacy"`` --- every rank reads whole checkpoint tensors, converts them in host memory and
113 keeps the slice its shards need.
114 * ``"dcp"`` --- the conversions are compiled into checkpoint coordinates first, so each rank
115 reads only the regions its own shards cover.
117 :meth:`save` is not implemented: exporting a Transformers-format checkpoint stays on
118 ``CheckpointManager.save_pretrained`` in ``models/_transformers/checkpoint_loader.py``.
119 """
121 def __init__(
122 self,
123 *,
124 loader: Optional[str] = None,
125 weights_mapping: Optional[WeightsMapping] = None,
126 ) -> None:
127 """Select the loader and the conversion rules :meth:`load` applies.
129 Args:
130 loader: Which loader reads the weights; see :func:`resolve_hf_loader`. Resolved here
131 rather than per load, so an unusable value fails before any checkpoint is opened.
132 weights_mapping: Rules renaming and converting checkpoint tensors. Default None, for the
133 model's Transformers conversion mapping, resolved once the model is known.
134 """
135 self.loader = resolve_hf_loader(loader)
136 self.weights_mapping = weights_mapping
138 # ------------------------------------------------------------------
139 # Save
140 # ------------------------------------------------------------------
142 def save(
143 self,
144 path: str,
145 state: Dict[str, Any],
146 *,
147 global_step: int,
148 save_async: bool = False,
149 ) -> None:
150 """Reject saving; see the class docstring."""
151 raise NotImplementedError(
152 "HuggingFaceCheckpointer only loads pretrained weights. Write a resumable checkpoint "
153 "with DistributedCheckpointer, or export Transformers format through "
154 "CheckpointManager(model).save_pretrained()."
155 )
157 # ------------------------------------------------------------------
158 # Load
159 # ------------------------------------------------------------------
161 def load(
162 self,
163 path: str,
164 state: Dict[str, Any],
165 *,
166 strict_model: bool = True,
167 extra_state_skeleton: Optional[Dict[str, Any]] = None,
168 external_target_fqns: frozenset[str] = frozenset(),
169 claimed_source_keys: frozenset[str] = frozenset(),
170 ) -> Dict[str, Any]:
171 """Load the pretrained weights at ``path`` into ``state["model"]``.
173 Args:
174 path: A safetensors file, a checkpoint directory or a Hub repository id.
175 state: Payload carrying ``model``, the finalized module to load into. Loading writes into
176 its distributed parameters in place, so this is a live module and not a state dict.
177 strict_model: Raise if a model tensor is left unloaded.
178 extra_state_skeleton: Unsupported --- a pretrained checkpoint holds no training progress.
179 external_target_fqns: Model tensors materialized by an external reader.
180 claimed_source_keys: Checkpoint keys consumed by that external reader.
182 Returns:
183 The same ``state`` object, with ``load_report`` set to the :class:`LoadReport` naming what
184 was loaded, what is missing and what the checkpoint holds beyond the model.
186 Raises:
187 ValueError: If ``path`` is empty or ``extra_state_skeleton`` is requested.
188 TypeError: If ``state["model"]`` is not a module.
189 """
190 if extra_state_skeleton is not None:
191 raise ValueError(
192 "A pretrained Hugging Face checkpoint carries no training state; restore progress "
193 "from a resumable checkpoint through DistributedCheckpointer instead."
194 )
195 if not path:
196 raise ValueError("pretrained_path must be provided when load_base_model=True")
197 model = state.get(_MODEL_KEY)
198 if not isinstance(model, nn.Module):
199 raise TypeError(
200 f"HuggingFaceCheckpointer.load expects state[{_MODEL_KEY!r}] to be the finalized "
201 f"nn.Module to load into, got {type(model).__name__}"
202 )
204 state[_LOAD_REPORT_KEY] = self._load_model(
205 model, path, strict=strict_model, external_target_fqns=external_target_fqns,
206 claimed_source_keys=claimed_source_keys,
207 )
208 return state
210 def _load_model(
211 self, model: nn.Module, pretrained_path: str, *, strict: bool,
212 external_target_fqns: frozenset[str] = frozenset(), claimed_source_keys: frozenset[str] = frozenset(),
213 ) -> LoadReport:
214 """Dispatch one model to the configured loader and the conversions it needs."""
215 if self.loader == "dcp":
216 return load_hf_checkpoint(
217 model, pretrained_path, planner=HFLoadPlanner(
218 model, weights_mapping=self.weights_mapping, strict=strict,
219 external_target_fqns=external_target_fqns, claimed_source_keys=claimed_source_keys,
220 ),
221 )
223 weights_mapping = self.weights_mapping
224 if weights_mapping is None:
225 weights_mapping = get_model_conversion_mapping(
226 model,
227 key_mapping=None,
228 hf_quantizer=None,
229 )
231 checkpoint_index = resolve_checkpoint_index(pretrained_path)
232 targets = build_load_targets(model)
233 logical_shapes = checkpoint_logical_shapes(model)
234 if not external_target_fqns <= targets.keys():
235 raise ValueError("External target is absent from the model")
236 targets = {key: value for key, value in targets.items() if key not in external_target_fqns}
237 if claimed_source_keys:
238 checkpoint_index = CheckpointIndex({
239 key: value for key, value in checkpoint_index.files_by_key.items()
240 if key not in claimed_source_keys
241 })
242 replacement_mapping = getattr(
243 model,
244 "_hp_replacement_weight_conversions",
245 None,
246 )
247 source_shapes = getattr(
248 model,
249 "_hp_checkpoint_source_shapes",
250 None,
251 )
252 if replacement_mapping and source_shapes:
253 return self._load_with_replacement_conversions(
254 model,
255 checkpoint_index,
256 targets,
257 weights_mapping,
258 replacement_mapping,
259 source_shapes,
260 pretrained_path,
261 strict,
262 logical_shapes,
263 )
264 return self._load_base_conversions(
265 model,
266 checkpoint_index,
267 targets,
268 weights_mapping,
269 pretrained_path,
270 strict,
271 logical_shapes,
272 )
274 def _load_base_conversions(
275 self,
276 model: nn.Module,
277 checkpoint_index: CheckpointIndex,
278 targets: Dict[str, torch.Tensor],
279 weights_mapping: WeightsMapping,
280 pretrained_path: str,
281 strict: bool,
282 logical_shapes: Dict[str, tuple[int, ...]],
283 ) -> LoadReport:
284 """Load a model whose tensors the Transformers conversion rules reach directly."""
285 groups, unexpected_keys, weight_mapping = build_load_groups(
286 model,
287 checkpoint_index.keys(),
288 targets,
289 weights_mapping=weights_mapping,
290 make_loader=partial(make_tensor_loader, checkpoint_index),
291 )
292 aliases_by_target = alias_names_by_target(targets)
293 loaded_keys = set()
294 loaded_target_ids = set()
296 for group in groups:
297 converted = self._convert_group(group, model)
298 for target_name, tensor in converted.items():
299 target = targets.get(target_name)
300 if target is None:
301 unexpected_keys += (target_name,)
302 continue
303 tensor = tensor[0] if isinstance(tensor, list) else tensor
304 target_id = id(target)
305 if target_id not in loaded_target_ids:
306 copy_into_target(target_name, tensor, target, logical_shapes.get(target_name))
307 loaded_target_ids.add(target_id)
308 loaded_keys.update(aliases_by_target[target_id])
310 missing_keys = tuple(sorted(set(targets) - loaded_keys, key=dot_natural_key))
311 unexpected_keys = tuple(sorted(set(unexpected_keys), key=dot_natural_key))
312 self._validate_load_result(missing_keys, unexpected_keys, strict)
313 used_conversions = [transform for transform in weight_mapping if transform.was_used()]
314 model._weight_conversions = used_conversions # pylint: disable=W0212
315 report = LoadReport(
316 loaded_keys=tuple(sorted(loaded_keys, key=dot_natural_key)),
317 missing_keys=missing_keys,
318 unexpected_keys=unexpected_keys,
319 )
320 logger.info(
321 "Loaded %d model tensors from %s",
322 len(report.loaded_keys),
323 pretrained_path,
324 )
325 return report
327 def _load_with_replacement_conversions(
328 self,
329 model: nn.Module,
330 checkpoint_index: CheckpointIndex,
331 targets: Dict[str, torch.Tensor],
332 weights_mapping: WeightsMapping,
333 replacement_mapping: WeightsMapping,
334 source_shapes: Dict[str, tuple],
335 pretrained_path: str,
336 strict: bool,
337 logical_shapes: Dict[str, tuple[int, ...]],
338 ) -> LoadReport:
339 """Normalize original weights before applying replacement conversions."""
340 base_mapping = base_weights_mapping(weights_mapping, replacement_mapping)
341 source_model = SourceModelView(model, source_shapes)
342 base_groups, unexpected_keys, _ = build_load_groups(
343 source_model,
344 checkpoint_index.keys(),
345 source_model.targets,
346 weights_mapping=base_mapping,
347 make_loader=partial(make_tensor_loader, checkpoint_index),
348 )
349 routes = build_replacement_routes(
350 model,
351 tuple(source_shapes),
352 targets,
353 replacement_mapping,
354 )
355 aliases_by_target = alias_names_by_target(targets)
356 loaded_keys = set()
357 loaded_target_ids = set()
358 used_replacements = []
360 def copy_converted(converted: Dict[str, torch.Tensor]) -> None:
361 """Copy converted tensors into their finalized model targets."""
362 nonlocal unexpected_keys
363 for target_name, tensor in converted.items():
364 target = targets.get(target_name)
365 if target is None:
366 unexpected_keys += (target_name,)
367 continue
368 tensor = tensor[0] if isinstance(tensor, list) else tensor
369 target_id = id(target)
370 if target_id not in loaded_target_ids:
371 copy_into_target(target_name, tensor, target, logical_shapes.get(target_name))
372 loaded_target_ids.add(target_id)
373 loaded_keys.update(aliases_by_target[target_id])
375 for base_group in base_groups:
376 normalized = self._convert_group(base_group, source_model)
377 for source_name, tensor in normalized.items():
378 tensor = tensor[0] if isinstance(tensor, list) else tensor
379 route = routes.get(source_name)
380 if route is None:
381 copy_converted({source_name: tensor})
382 continue
383 state, target_name, source_pattern = route
384 state.group.transform.add_tensor(
385 target_name,
386 source_name,
387 source_pattern,
388 lambda value=tensor: value,
389 )
390 state.received[source_pattern] += 1
391 if not state.completed and state.received == state.expected:
392 copy_converted(self._convert_group(state.group, model))
393 state.completed = True
394 used_replacements.append(state.group.transform)
396 missing_keys = tuple(sorted(set(targets) - loaded_keys, key=dot_natural_key))
397 unexpected_keys = tuple(sorted(set(unexpected_keys), key=dot_natural_key))
398 self._validate_load_result(missing_keys, unexpected_keys, strict)
399 used_base = [transform for transform in base_mapping if transform.was_used()]
400 model._hp_used_base_weight_conversions = used_base # pylint: disable=protected-access
401 model._hp_used_replacement_weight_conversions = ( # pylint: disable=protected-access
402 used_replacements
403 )
404 model._weight_conversions = used_base + used_replacements # pylint: disable=protected-access
405 report = LoadReport(
406 loaded_keys=tuple(sorted(loaded_keys, key=dot_natural_key)),
407 missing_keys=missing_keys,
408 unexpected_keys=unexpected_keys,
409 )
410 logger.info("Loaded %d model tensors from %s", len(report.loaded_keys), pretrained_path)
411 return report
413 @staticmethod
414 def _convert_group(
415 group: LoadGroup,
416 model: Union[nn.Module, SourceModelView],
417 ) -> Dict[str, torch.Tensor]:
418 """Convert all checkpoint tensors belonging to one load group."""
419 return convert_group(group, model)
421 @staticmethod
422 def _validate_load_result(
423 missing_keys: tuple,
424 unexpected_keys: tuple,
425 strict: bool,
426 ) -> None:
427 """Validate missing keys and report ignored checkpoint tensors."""
428 validate_load_result(missing_keys, unexpected_keys, strict)
431__all__ = [
432 "HF_LOADER_ENV",
433 "HuggingFaceCheckpointer",
434 "resolve_hf_loader",
435]