Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / loader.py: 0%
151 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"""Import a generated modeling file and resolve its model class.
17The bundle at ``<yaml_dir>/generated/`` is not on ``sys.path`` and
18is not a normal package, so it is loaded by file location. The loaded module
19is registered in ``sys.modules`` under a stable synthetic name, which is what
20lets ``dataclasses``, ``pickle``, and (importantly) ``torch.load`` resolve
21classes defined inside it.
22"""
23from __future__ import annotations
25import hashlib
26import importlib.util
27import logging
28import os
29import sys
30from types import ModuleType
31from typing import Any, Optional
33logger = logging.getLogger(__name__)
35# Synthetic package prefix for generated modules. Keyed by artifact dir so two
36# YAMLs generating for the same model do not collide in ``sys.modules``.
37_GENERATED_MODULE_PREFIX = "hyper_parallel_generated"
40def _module_name_for(artifact_dir: str, stem: str) -> str:
41 """Build a stable, collision-free ``sys.modules`` name for one bundle.
43 The basename alone is not unique — two YAMLs in different directories both
44 produce ``generated/train/``, which would map to the same synthetic module
45 and make the second import silently return the first bundle's classes. The
46 absolute path is hashed instead, so the name is unique per bundle and still
47 stable across runs (a prerequisite for ``torch.load``/``pickle`` to resolve
48 classes defined in the generated file).
49 """
50 abs_dir = os.path.abspath(os.path.normpath(artifact_dir))
51 digest = hashlib.sha256(abs_dir.encode("utf-8")).hexdigest()[:12]
52 marker = os.path.basename(abs_dir)
53 safe = "".join(c if c.isalnum() or c == "_" else "_" for c in marker)
54 return f"{_GENERATED_MODULE_PREFIX}.{safe}_{digest}.{stem}"
57def find_generated_modeling_file(artifact_dir: str) -> str:
58 """Locate the single ``modeling_*_gen_npu.py`` in the bundle.
60 Raises when absent or ambiguous rather than picking one — a bundle with two
61 modeling files means generation left something stale behind, and silently
62 choosing would produce a model that does not match the meta.
64 Args:
65 artifact_dir (str): Directory holding the generated bundle.
67 Returns:
68 str: Absolute path of the generated modeling file.
69 """
70 if not os.path.isdir(artifact_dir):
71 raise FileNotFoundError(
72 f"codegen: artifact dir {artifact_dir} does not exist; "
73"run generation before building from the generated artifact"
74 )
75 candidates = sorted(
76 name for name in os.listdir(artifact_dir)
77 if name.startswith("modeling_") and name.endswith("_gen_npu.py")
78 )
79 if not candidates:
80 raise FileNotFoundError(
81 f"codegen: no modeling_*_gen_npu.py in {artifact_dir}; "
82 "the bundle is incomplete"
83 )
84 if len(candidates) > 1:
85 raise RuntimeError(
86 f"codegen: {len(candidates)} generated modeling files in "
87 f"{artifact_dir} ({candidates}); remove the stale bundle and "
88 "regenerate"
89 )
90 return os.path.join(artifact_dir, candidates[0])
93def import_generated_module(artifact_dir: str) -> ModuleType:
94 """Import the generated modeling module from ``artifact_dir``.
96 Repeat calls for the same bundle return the cached module.
98 Args:
99 artifact_dir (str): Directory holding the generated bundle.
101 Returns:
102 ModuleType: The imported (and cached) generated modeling module.
103 """
104 modeling_path = find_generated_modeling_file(artifact_dir)
105 stem = os.path.splitext(os.path.basename(modeling_path))[0]
106 module_name = _module_name_for(artifact_dir, stem)
108 cached = sys.modules.get(module_name)
109 if cached is not None:
110 return cached
112 # A ``<remote>`` source uses
113 # single-dot relative imports (``from .configuration_deepseek import ...``)
114 # that resolve against its own package. spec_from_file_location compiles
115 # the generated file as the synthetic top-level module
116 # ``hyper_parallel_generated.<digest>.<stem>`` — setting its ``__package__``
117 # to the synthetic *package* (hyper_parallel_generated.<digest>) and
118 # registering that package + its sibling modules before exec makes those
119 # relative imports resolve inside the bundle. Both are keyed on the same
120 # path-derived name as the module, so they stay collision-free per bundle
121 # and the registered modules are reused on the next import.
122 _register_synthetic_package(module_name, artifact_dir)
124 spec = importlib.util.spec_from_file_location(module_name, modeling_path)
125 if spec is None or spec.loader is None:
126 raise ImportError(f"codegen: cannot build an import spec for {modeling_path}")
127 module = importlib.util.module_from_spec(spec)
128 module.__package__ = _package_name_for(module_name)
129 # Registered before exec so a self-referential import inside the generated
130 # file resolves to this partially-initialized module rather than looping.
131 sys.modules[module_name] = module
132 try:
133 spec.loader.exec_module(module)
134 except Exception:
135 sys.modules.pop(module_name, None)
136 raise
137 logger.info("codegen: imported generated modeling file %s", modeling_path)
138 return module
141def _package_name_for(module_name: str) -> str:
142 """The synthetic package a generated module lives in (minus the stem).
144 ``hyper_parallel_generated.<safe>_<digest>.<stem>`` -> ``...<safe>_<digest>``.
145 This is the package the generated file's relative imports resolve against,
146 and it is stable per bundle (the digest is the artifact dir's path hash).
147 """
148 return module_name.rsplit(".", 1)[0]
151def _register_synthetic_package(module_name: str, artifact_dir: str) -> None:
152 """Register the synthetic package + its sibling modules in ``sys.modules``.
154 A plain ``ModuleType`` package with ``__path__`` set to the bundle dir lets
155 ``from .name import x`` resolve against the bundle's own files without any
156 ``sys.path`` mutation. Only the package root is created; sibling modules
157 are loaded on demand by the import machinery (exec'd here) and registered
158 under the synthetic package name so ``torch.load``/``pickle`` resolve
159 classes defined in them the same way they resolve the generated module.
160 """
161 package_name = _package_name_for(module_name)
162 package = sys.modules.get(package_name)
163 if package is None:
164 package = ModuleType(package_name)
165 package.__package__ = package_name
166 package.__path__ = [os.path.abspath(artifact_dir)]
167 package.__file__ = "<codegen-synthetic-package>"
168 sys.modules[package_name] = package
169 elif not hasattr(package, "__path__"):
170 package.__path__ = [os.path.abspath(artifact_dir)]
172 # Pre-register sibling .py files (the ``<remote>`` relative-import
173 # closure, copied into the bundle by ``emit_bundle``) so their own
174 # cross-references resolve inside the same package. The generated
175 # modeling file is the one being imported — skip it.
176 prefix = package_name + "."
177 modeling_name = find_generated_modeling_file(artifact_dir)
178 for name in sorted(os.listdir(artifact_dir)):
179 if not name.endswith(".py") or name == "__init__.py":
180 continue
181 if name == os.path.basename(modeling_name):
182 continue
183 sibling_name = prefix + os.path.splitext(name)[0]
184 if sibling_name in sys.modules:
185 continue
186 spec = importlib.util.spec_from_file_location(
187 sibling_name, os.path.join(artifact_dir, name),
188 )
189 if spec is None or spec.loader is None:
190 continue
191 sibling = importlib.util.module_from_spec(spec)
192 sibling.__package__ = package_name
193 sys.modules[sibling_name] = sibling
194 try:
195 spec.loader.exec_module(sibling)
196 except Exception:
197 sys.modules.pop(sibling_name, None)
198 raise
201def resolve_generated_model_class(module: ModuleType, hf_config: Any) -> type:
202 """Find the model class in ``module`` matching the config's architecture.
204 The generated file keeps the original class names, so the architecture name
205 from the HF config is the lookup key.
207 Args:
208 module (ModuleType): The imported generated modeling module.
209 hf_config (Any): HF config whose ``architectures`` names the class.
211 Returns:
212 type: The model class to instantiate.
213 """
214 architectures = getattr(hf_config, "architectures", None) or []
215 if not architectures:
216 raise ValueError(
217 "codegen: hf_config has no 'architectures'; cannot pick a model "
218 f"class from {module.__name__}"
219 )
220 arch_name = architectures[0]
221 model_cls = getattr(module, arch_name, None)
222 if model_cls is None:
223 available = sorted(
224 name for name, value in vars(module).items() if isinstance(value, type)
225 )
226 raise AttributeError(
227 f"codegen: generated module {module.__name__} has no class "
228 f"{arch_name!r}; defined classes: {available}"
229 )
230 if not isinstance(model_cls, type):
231 raise TypeError(
232 f"codegen: {module.__name__}.{arch_name} is not a class"
233 )
234 return model_cls
237def _register_generated_model_conversions(model: Any) -> None:
238 """Let the generated model share the library's native weight conversions.
240 The generated modeling file is loaded as custom code (its module name does
241 not start with ``transformers.``), so ``PreTrainedModel.is_custom_code()``
242 is True. ``get_model_conversion_mapping`` skips custom-code submodules
243 unless their class name or ``model_type`` is in
244 ``transformers.conversion_mapping.USER_REGISTERED_MAPPINGS``. The generated
245 model only patches the parallel-strategy layer — its weight layout is the
246 native one — so it *should* use the same checkpoint conversions as the
247 installed model (e.g. the qwen2_moe / qwen3_moe fusion of per-expert
248 ``gate_proj``/``up_proj``/``down_proj`` into ``experts.gate_up_proj``/
249 ``experts.down_proj``). Registering the identifiers here marks the model
250 as user-registered so those library conversions apply instead of being
251 silently skipped, which would leave every fused expert target unmatched
252 during checkpoint loading.
254 Only marks identifiers already covered by the library mapping; an unknown
255 ``model_type`` gets no entry and registration is a no-op. The lookup is by
256 ``model_type`` (class-name lookup falls back to it), so registering the
257 class name is not required for the root model.
258 """
259 from transformers.conversion_mapping import ( # pylint: disable=import-outside-toplevel
260 USER_REGISTERED_MAPPINGS,
261 )
263 model_type = getattr(getattr(model, "config", None), "model_type", None)
264 if model_type:
265 # Guard against API drift across transformers versions: if the set is
266 # renamed or removed the whole registration is skipped rather than
267 # breaking the generated-model build.
268 if getattr(USER_REGISTERED_MAPPINGS, "add", None):
269 USER_REGISTERED_MAPPINGS.add(model_type)
270 logger.info(
271 "codegen: registered generated model_type %r for checkpoint "
272 "conversion mapping", model_type,
273 )
276def _load_generated_checkpoint(
277 model_cls: type,
278 pretrained_path: str,
279 hf_config: Any,
280 model_args: tuple,
281 torch_dtype: Any,
282 kwargs: dict,
283) -> Any:
284 """Construct a generated model and load its weights with layout conversions.
286 The generated modeling file embeds the final (post-replacement) modules
287 directly -- notably ``GroupedExperts`` (expert weights stored transposed
288 relative to the checkpoint) and the fused ``GQAAttention`` (``linear_qkv``
289 built from ``q_proj``/``k_proj``/``v_proj``). The native path records these
290 conversions through ``apply_module_replacements``; no replacement runs for
291 the generated artifact (the fused modules are literals in the generated
292 code), so synthesise the same wiring here: collect every module's
293 ``make_transforms()`` (scoped to its FQN) and publish the *pre-replacement*
294 parameter shapes, then let ``CheckpointManager`` apply the conversions
295 through the replacement-conversion route.
296 """
297 from hyper_parallel.models._transformers.checkpoint_loader import ( # pylint: disable=import-outside-toplevel
298 CheckpointManager,
299 )
301 model = model_cls._from_config(hf_config, *model_args, **kwargs)
302 _register_generated_model_conversions(model)
304 transforms: list[Any] = []
305 for module_fqn, module in model.named_modules():
306 make_transforms = getattr(module, "make_transforms", None)
307 if make_transforms is None:
308 continue
309 for transform in make_transforms():
310 transform.scope_prefix = module_fqn
311 transforms.append(transform)
312 if transforms:
313 model._hp_checkpoint_source_shapes = _source_model_shapes(hf_config) # pylint: disable=protected-access
314 model._hp_replacement_weight_conversions = transforms # pylint: disable=protected-access
316 # from_pretrained materializes meta tensors internally; _from_config
317 # does not. Materialize the meta model before loading so checkpoint
318 # copy targets are real tensors (init_device: meta in the training
319 # config). Downstream infra moves it to the accelerator afterwards.
320 if any(p.is_meta for p in model.parameters()):
321 model.to_empty(device="cpu")
323 CheckpointManager(model).load_checkpoint(pretrained_path, strict=False)
324 if torch_dtype not in (None, "auto"):
325 _cast_parameters_to(model, torch_dtype)
326 return model
329def _cast_parameters_to(model: Any, torch_dtype: Any) -> None:
330 """Cast the model's parameters to ``torch_dtype`` the way native loading does.
332 ``model.to(dtype)`` also down-casts floating-point *buffers*, but the native
333 ``from_pretrained(dtype=...)`` path only casts parameters: buffers such as
334 the RoPE ``inv_freq`` table keep the precision they were computed with.
335 Casting them here rounds the table to bf16, which changes every cos/sin
336 value and makes the generated model diverge numerically from the native one
337 even though both consume the same checkpoint. Modules listed in
338 ``_keep_in_fp32_modules`` stay in fp32 on both paths, so they are skipped.
340 Args:
341 model: Model whose parameters are cast in place.
342 torch_dtype: Configured dtype, either ``"bfloat16"`` or
343 ``"torch.bfloat16"``.
344 """
345 import torch # pylint: disable=import-outside-toplevel
347 dtype = getattr(torch, str(torch_dtype).rsplit(".", 1)[-1])
348 keep_fp32 = tuple(getattr(model, "_keep_in_fp32_modules", None) or ())
349 for name, parameter in model.named_parameters():
350 if keep_fp32 and any(module_name in name for module_name in keep_fp32):
351 continue
352 if parameter.is_floating_point() and parameter.dtype != dtype:
353 parameter.data = parameter.data.to(dtype)
356def _source_model_shapes(hf_config: Any) -> dict[str, tuple[int, ...]]:
357 """Capture the pre-replacement parameter shapes for the conversion route.
359 The native path records ``_named_tensor_shapes(model)`` *before* swapping in
360 the fused modules, so its converters are routed by the original parameter
361 names and layouts. The generated artifact has the fused modules baked in
362 and no longer owns the converter source parameters (``q_proj``/``k_proj``/
363 ``v_proj`` exist only as converter source patterns), so rebuild the original
364 parameter tree from the config on the meta device -- no storage is
365 allocated -- and publish its shapes instead.
366 """
367 import transformers # pylint: disable=import-outside-toplevel
368 from hyper_parallel import init_empty_weights # pylint: disable=import-outside-toplevel
369 from hyper_parallel.models.replacement import ( # pylint: disable=import-outside-toplevel
370 _named_tensor_shapes,
371 )
373 with init_empty_weights():
374 source = transformers.AutoModelForCausalLM.from_config(hf_config)
375 return _named_tensor_shapes(source)
378def init_generated_model(
379 artifact_dir: str,
380 pretrained_model_name_or_path: Optional[str],
381 hf_config: Any,
382 *model_args: Any,
383 torch_dtype: Any = "auto",
384 **kwargs: Any,
385) -> Any:
386 """Build the model from the generated modeling file.
388 Mirrors the custom-model path in ``model_init``: ``from_pretrained`` when a
389 checkpoint path is given, ``from_config`` otherwise. After the model is
390 built, its identifiers are registered for the library weight conversions
391 (``_register_generated_model_conversions``) and the frozen param plan is
392 verified against the *live* parameter tree (``verify_param_plan``) — this is
393 the check that catches a stale bundle: generation-time failures that degrade
394 to an empty plan, or a model whose parameter tree drifted from what the plan
395 records, fail here instead of training with a stale sharding contract.
397 Args:
398 artifact_dir (str): Directory holding the generated bundle.
399 pretrained_model_name_or_path (Optional[str]): Checkpoint to load, or
400 None to build from the config only.
401 hf_config (Any): HF config of the model being generated.
402 *model_args (Any): Extra positional args for the constructor.
403 torch_dtype (Any): Target dtype, an HF-style string, or "auto".
404 **kwargs (Any): Extra keyword args for the constructor.
406 Returns:
407 Any: The instantiated model with converted weights loaded.
408 """
409 module = import_generated_module(artifact_dir)
410 model_cls = resolve_generated_model_class(module, hf_config)
412 if pretrained_model_name_or_path is not None:
413 model = _load_generated_checkpoint(
414 model_cls,
415 pretrained_model_name_or_path,
416 hf_config,
417 model_args,
418 torch_dtype,
419 kwargs,
420 )
421 else:
422 # transformers >= 5.14 removed ``PreTrainedModel.from_config`` in
423 # favour of ``_from_config`` (a classmethod on the concrete model
424 # class). The generated model class only inherits ``PreTrainedModel``,
425 # so ``from_config`` no longer exists on it. ``_from_config`` is the
426 # canonical entry in both 5.14.x and the pinned 5.15.0; the AutoModel
427 # family's ``from_config`` (which the native path borrows) is just a
428 # resolver that ultimately calls ``_from_config`` anyway.
429 model = model_cls._from_config(hf_config, *model_args, **kwargs)
430 _register_generated_model_conversions(model)
432 from hyper_parallel.codegen.check.preflight import ( # pylint: disable=import-outside-toplevel
433 verify_param_plan,
434 )
435 from hyper_parallel.codegen.meta import ( # pylint: disable=import-outside-toplevel
436 load_codegen_meta,
437 )
439 meta = load_codegen_meta(os.path.join(artifact_dir, "codegen_meta.json"))
440 if meta is None:
441 raise FileNotFoundError(
442 f"codegen: missing meta at {os.path.join(artifact_dir, 'codegen_meta.json')}; "
443 "the bundle is incomplete — regenerate the artifact"
444 )
445 verify_param_plan(meta, model)
446 return model
449__all__ = [
450 "find_generated_modeling_file",
451 "import_generated_module",
452 "init_generated_model",
453 "resolve_generated_model_class",
454]