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