Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / manager.py: 14%
355 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"""Codegen manager: decide, generate (rank0), reuse, and preflight.
17The manager is the only entry point the trainer / CLI talks to. It does NOT
18do AST patching or source emitting itself — that is delegated to the
19``emit`` layer — it decides whether an artifact is stale, who generates it,
20and whether it passes preflight before training starts.
21"""
23from __future__ import annotations
25import hashlib
26from fnmatch import fnmatchcase
27import logging
28import os
29from typing import Any, Optional
31from hyper_parallel.codegen.artifact import (
32 ArtifactLayout,
33 clean_temp_artifacts,
34 ensure_artifact_dir,
35 resolve_default_artifact_layout,
36 resolve_artifact_layout,
37 write_bundle_atomic,
38)
39from hyper_parallel.codegen.emit import emit_bundle
40from hyper_parallel.codegen.hash import canonical_json, sha256_file, signature_from_spec
41from hyper_parallel.codegen.inline.framework_spec import (
42 declaration_summary,
43 external_state_classes_for,
44)
45from hyper_parallel.codegen.meta import (
46 CodegenMeta,
47 load_codegen_meta,
48)
50logger = logging.getLogger(__name__)
52CODEGEN_VERSION = "0.1.1"
55def ensure_codegen_artifact(
56 config: Any,
57 yaml_path: str | None = None,
58 *,
59 artifact_dir: str | None = None,
60 hf_config: Any = None,
61) -> Optional[ArtifactLayout]:
62 """Ensure an up-to-date artifact bundle.
64 Returns the layout when codegen is active, ``None`` otherwise. If the
65 signature hits an existing bundle it is reused; otherwise ``emit_bundle``
66 is called to produce the bundle.
68 ``hf_config`` is the resolved HuggingFace config of the model; when given,
69 it is folded into the signature and meta ``source`` (so a transformers
70 modeling-file change regenerates). When omitted, the source identity is
71 resolved from the model id alone.
73 Rank rendezvous: every path out of this function — reuse, wait, and
74 generate — must cross exactly one ``barrier_after_artifact`` call. The
75 reuse decision is a per-rank disk read, so ranks can disagree while rank0
76 swaps in a fresh bundle (a late-starting rank reads the new meta and takes
77 the reuse path while its peers hold the stale one and take the wait path);
78 a rank that skips the barrier on reuse deadlocks the whole run.
79 """
80 if not getattr(config, "codegen", False):
81 return None
83 if hf_config is None:
84 hf_config = _resolve_hf_config_for_layout(config)
85 layout = _resolve_layout(config, yaml_path, artifact_dir, hf_config=hf_config)
86 ensure_artifact_dir(layout)
87 is_rank0 = _is_rank0()
88 if is_rank0:
89 clean_temp_artifacts(layout)
91 signature = _compute_signature(config, layout, hf_config=hf_config)
92 existing = load_codegen_meta(layout.meta_path)
94 if existing is not None and not should_regenerate(existing, signature):
95 if is_rank0:
96 logger.info(
97 "codegen: reuse artifact %s (signature %s)",
98 layout.artifact_dir,
99 signature,
100 )
101 # Ranks can disagree on reuse while rank0 swaps the bundle in; this
102 # path must cross the barrier too, or the peers deadlock inside theirs.
103 barrier_after_artifact()
104 return layout
106 if not is_rank0:
107 wait_for_rank0_artifact(layout, signature)
108 barrier_after_artifact()
109 return layout
111 # rank0 generates.
112 logger.info(
113 "codegen: generating artifact %s (signature %s)", layout.artifact_dir, signature
114 )
115 meta = _build_meta(config, layout, signature, hf_config=hf_config)
116 _fill_plan_fields(meta, config, layout, hf_config=hf_config)
117 _record_inline_declarations(meta)
118 files = emit_bundle(layout, meta)
119 write_bundle_atomic(layout, files, meta)
120 barrier_after_artifact()
121 return layout
124def preflight_integrity_check(
125 config: Any,
126 yaml_path: str | None = None,
127 *,
128 artifact_dir: str | None = None,
129 hf_config: Any = None,
130) -> None:
131 """Validate the existing bundle before training starts; fail fast on drift.
133 Composes the check-layer per-zone verifiers: signature, output hashes,
134 importability, meta schema, boundary-form structure, and a warning for
135 skipped overrides. The frozen ``param_plan`` is verified for internal
136 consistency (an empty plan fails — generation always freezes a derived
137 plan); full model-side coverage against the live parameter tree runs
138 later in :func:`init_generated_model`, where a model handle exists.
139 """
140 if not getattr(config, "codegen", False):
141 return
142 from hyper_parallel.codegen.check.preflight import (
143 verify_boundary_forms,
144 verify_generated_import,
145 verify_meta_required_fields,
146 verify_output_hashes,
147 verify_param_plan,
148 verify_signature,
149 warn_skipped_overrides,
150 )
152 if hf_config is None:
153 hf_config = _resolve_hf_config_for_layout(config)
154 layout = _resolve_layout(config, yaml_path, artifact_dir, hf_config=hf_config)
155 meta = load_codegen_meta(layout.meta_path)
156 if meta is None:
157 raise FileNotFoundError(
158 f"codegen preflight: missing meta at {layout.meta_path}; run generate first"
159 )
160 verify_meta_required_fields(meta)
161 signature = _compute_signature(config, layout, hf_config=hf_config)
162 verify_signature(meta, signature)
163 verify_output_hashes(meta, layout)
164 verify_generated_import(layout)
165 verify_boundary_forms(meta, layout)
166 verify_param_plan(meta, model=None)
167 warn_skipped_overrides(meta)
170def should_regenerate(meta: CodegenMeta, signature: str) -> bool:
171 """Whether the stored meta no longer matches the current signature."""
172 return meta.signature != signature
175def wait_for_rank0_artifact(
176 layout: ArtifactLayout,
177 current_signature: str,
178 *,
179 timeout_s: float = 600.0,
180) -> None:
181 """Non-rank0: poll until rank0 has produced a fully-written artifact.
183 ``current_signature`` is this rank's own computed signature. The wait is
184 signature-aware: a parseable meta is only accepted when its signature
185 matches, so a stale bundle left on disk from a previous regeneration is
186 never mistaken for "rank0 is done" — the poll continues until the new
187 bundle's meta (written last inside the staging dir, then swapped in) lands.
188 The meta file is written last, so a matching signature also means the
189 modeling/diff/init files are already in place.
190 """
191 import time
193 deadline = time.monotonic() + timeout_s
194 while time.monotonic() < deadline:
195 meta = load_codegen_meta(layout.meta_path)
196 if meta is not None and meta.signature == current_signature:
197 return
198 time.sleep(0.5)
199 raise TimeoutError(
200 f"codegen: timed out waiting for rank0 artifact at {layout.artifact_dir}"
201 )
204def barrier_after_artifact() -> None:
205 """Rejoin all ranks after the artifact is ready.
207 Every rank must call it exactly once on the way out of
208 ``ensure_codegen_artifact``, on every path: rank0 after writing the
209 bundle, each other rank right after ``wait_for_rank0_artifact`` returns,
210 and every rank that takes the reuse path (the reuse decision is a per-rank
211 disk read, so ranks can disagree while rank0 swaps in a fresh bundle — a
212 rank that skips the barrier deadlocks its peers inside theirs). It is the
213 collective that lets all ranks leave the codegen entry together before
214 training starts. The half-written-bundle race it was written to guard
215 against is already closed by ``wait_for_rank0_artifact`` (signature aware,
216 meta written last); the barrier is purely a rendezvous point, so it must
217 be reached by every rank.
219 Uses ``torch.distributed`` directly (the process group was initialized in
220 ``trainer._setup`` before the codegen manager runs); a no-op without a
221 live process group.
222 """
223 try:
224 import torch.distributed as dist
226 if dist.is_initialized():
227 dist.barrier()
228 except Exception: # pragma: no cover - best-effort without a process group
229 pass
232# ---------------------------------------------------------------------------
233# helpers
234# ---------------------------------------------------------------------------
237def _is_rank0() -> bool:
238 """True when this process is global rank 0 (or is a single process).
240 Training config is parsed before the process group is initialized, so a
241 torchrun worker must use its ``RANK`` environment value. Other callers use
242 the initialized process group when available and otherwise act as rank 0.
243 """
244 rank = os.environ.get("RANK")
245 if rank is not None:
246 try:
247 return int(rank) == 0
248 except ValueError as exc:
249 raise ValueError(f"codegen: invalid RANK value {rank!r}") from exc
250 try:
251 from hyper_parallel.trainer.runtime.distributed import (
252 get_global_rank_safe,
253 )
255 return get_global_rank_safe() == 0
256 except Exception:
257 return True
260def _model_name(config: Any, hf_config: Any = None) -> Optional[str]:
261 """Resolve the model name from an HF config or the YAML model target."""
262 model_type = getattr(hf_config, "model_type", None)
263 if model_type:
264 return _model_type_to_modeling_token(str(model_type))
265 target = getattr(config, "model", None)
266 path = getattr(target, "_target_", None)
267 if isinstance(path, str):
268 name = path.rsplit(".", 1)[-1]
269 if name.startswith(("AutoModel", "Pretrained", "from_")):
270 return None
271 return name
272 return getattr(config, "model_name", None)
275def _model_type_to_modeling_token(model_type: str) -> str:
276 """Convert HF ``model_type`` into the generated modeling file token."""
277 parts = [part for part in model_type.replace("-", "_").split("_") if part]
278 if not parts:
279 return "model"
280 return "_".join(part[:1].upper() + part[1:] for part in parts)
283def _model_class(config: Any, hf_config: Any = None) -> Optional[str]:
284 """The entry model class name the override rewrite targets in ``__init__``.
286 The HF config's ``architectures`` list is the lookup key the loader uses to
287 pick the class from the generated file, so it is the architecture's class
288 name (e.g. ``AnthropicV3ForCausalLM``). Falls back to the loaded config's
289 own attribute when ``hf_config`` was not passed in.
291 ``ensure_codegen_artifact`` is reached from config preparation and the CLI,
292 neither of which has a resolved ``hf_config`` in hand (the model resolves
293 it only inside ``from_pretrained``, after generation). So when
294 neither ``hf_config`` nor ``config`` exposes ``architectures``, resolve the
295 model path with the same rule ``spec/project`` uses and load the config —
296 lazily, to keep this module importable without ``transformers``.
297 """
298 for candidate in (
299 getattr(hf_config, "architectures", None),
300 getattr(config, "architectures", None),
301 ):
302 if isinstance(candidate, (list, tuple)) and candidate:
303 return str(candidate[0])
305 model = getattr(config, "model", None)
306 path = getattr(model, "pretrained_model_name_or_path", None)
307 if path is None:
308 path = getattr(model, "model_name_or_path", None)
309 if path is None:
310 return None
311 try:
312 from hyper_parallel.models._transformers.config_resolver import get_hf_config
313 except ImportError:
314 return None
315 try:
316 arch = getattr(get_hf_config(path), "architectures", None)
317 except Exception:
318 return None
319 if isinstance(arch, (list, tuple)) and arch:
320 return str(arch[0])
321 return None
324def _resolve_layout(
325 config: Any,
326 yaml_path: str | None,
327 artifact_dir: str | None,
328 *,
329 hf_config: Any = None,
330) -> ArtifactLayout:
331 """Resolve either YAML-anchored or model-construction artifact layout."""
332 model_name = _model_name(config, hf_config)
333 if artifact_dir is not None:
334 return resolve_default_artifact_layout(artifact_dir, model_name=model_name)
335 model = getattr(config, "model", None)
336 configured_dir = getattr(model, "codegen_artifact_dir", None)
337 if configured_dir:
338 return resolve_default_artifact_layout(str(configured_dir), model_name=model_name)
339 if yaml_path is not None:
340 return resolve_artifact_layout(str(yaml_path), model_name=model_name)
341 configured_yaml_path = getattr(config, "_yaml_path", None)
342 if configured_yaml_path:
343 return resolve_artifact_layout(str(configured_yaml_path), model_name=model_name)
344 return resolve_default_artifact_layout(model_name=model_name)
347def _resolve_hf_config_for_layout(config: Any) -> Any:
348 """Best-effort HF config lookup used only for artifact naming/layout."""
349 model = getattr(config, "model", None)
350 if model is None:
351 return None
352 path = getattr(model, "pretrained_model_name_or_path", None)
353 if path is None:
354 path = getattr(model, "model_name_or_path", None)
355 if path is None:
356 return None
357 kwargs = _hf_config_kwargs(model)
358 try:
359 from hyper_parallel.models._transformers.config_resolver import get_hf_config
360 except ImportError:
361 return None
362 try:
363 return get_hf_config(
364 path,
365 getattr(model, "attn_implementation", "sdpa"),
366 getattr(model, "torch_dtype", "auto"),
367 **kwargs,
368 )
369 except Exception:
370 return None
373def _hf_config_kwargs(model: Any) -> dict[str, Any]:
374 """Collect model target kwargs that affect ``AutoConfig.from_pretrained``."""
375 names = (
376 "cache_dir",
377 "force_download",
378 "local_files_only",
379 "proxies",
380 "resume_download",
381 "revision",
382 "subfolder",
383 "token",
384 "trust_remote_code",
385 "use_auth_token",
386 )
387 kwargs = {
388 name: getattr(model, name)
389 for name in names
390 if getattr(model, name, None) is not None
391 }
392 overrides = getattr(model, "config_overrides", None)
393 if isinstance(overrides, dict):
394 kwargs.update(overrides)
395 return kwargs
398def _compute_signature(
399 config: Any,
400 layout: ArtifactLayout,
401 *,
402 hf_config: Any = None,
403) -> str:
404 """Canonical signature over the projected spec, parallel dims, overrides,
405 codegen switches, and the resolved modeling source.
407 The full ``project_codegen_spec`` projection ensures the digest
408 covers the same fields the plan/emit layers consume. The YAML file bytes
409 are deliberately excluded (see ``_project_spec``).
410 """
411 spec = _project_spec(config, layout, hf_config=hf_config)
412 return signature_from_spec(spec)
415def _project_spec(
416 config: Any,
417 layout: ArtifactLayout,
418 *,
419 hf_config: Any = None,
420) -> dict:
421 """CodegenSpec canonical dict used for the signature.
423 The payload covers the projected CodegenSpec plus the resolved modeling
424 source (installed transformers module or remote code), ``sha256`` included
425 — a change to the HF modeling file changes the signature and regenerates
426 the bundle. The YAML file bytes are intentionally NOT hashed: the spec
427 projection already captures every codegen-affecting field, so an unrelated
428 YAML edit (comment, logging, optimizer tweak) must not invalidate a
429 bundle.
430 """
431 from hyper_parallel.codegen.spec.project import project_codegen_spec
432 from hyper_parallel.codegen.source.resolver import resolve_model_source
434 payload = project_codegen_spec(config, layout=layout).to_dict()
435 model_id = payload.get("source", {}).get("model_name_or_path") or ""
436 is_gen = _gen_backend(config)
438 source = resolve_model_source(config, model_id, hf_config=hf_config)
439 if source is not None:
440 # Replace the SourceSpec projection with the resolved file identity;
441 # the sha256 here is what makes a modeling-file change regenerate.
442 # ``transformers_version`` is a SourceSpec field the pure projection
443 # leaves empty; record the actual env version whenever we resolved a
444 # real file (an environment upgrade that changes the modeling file's
445 # behavior must regenerate, not silently reuse).
446 resolved = source.to_dict()
447 # ``SourceSpec.to_dict`` keeps only architecture / model_type / sha256 —
448 # neither the checkpoint identity nor its structural config. Without
449 # them, two checkpoints that share a modeling file but differ in shape
450 # (e.g. the same architecture at different layer counts) would produce
451 # the same signature and silently reuse a frozen per-layer plan written
452 # for the other shape. Fold the model id and a structural-config digest
453 # back in so a shape change invalidates the bundle.
454 resolved["model_id"] = model_id
455 if hf_config is not None:
456 resolved["structure"] = _structure_summary(hf_config)
457 payload["source"] = resolved
458 payload["source"]["transformers_version"] = _transformers_version()
459 elif is_gen:
460 # codegen=True resolves to the ``gen`` backend, whose whole point is
461 # producing an artifact FROM the original modeling file. A bundle
462 # generated from a guessed/unresolved source would be silently wrong,
463 # so fail fast instead of degrading — the user must fix the model id,
464 # install transformers, or make the checkpoint reachable.
465 raise RuntimeError(
466 f"codegen: cannot resolve modeling source for {model_id!r} — no "
467 "AutoConfig, no installed transformers module, no remote code. "
468 "The gen backend generates the artifact from the original "
469 "modeling file; fix the model id, install transformers, or make "
470 "the checkpoint reachable"
471 )
473 payload["codegen_implementation"] = {
474 "version": CODEGEN_VERSION,
475 "sha256": _codegen_implementation_digest(),
476 }
477 identity = payload.get("source", {}).get("architecture")
478 # The declarations a target resolves to are derived from structure (the
479 # family's runtime factories + the framework archetype table), which live
480 # outside the codegen implementation tree; summarizing them here folds
481 # their content into the signature so a change to a factory's component or
482 # a strategy's body still regenerates an artifact.
483 payload["inline_declarations"] = {
484 "specs": declaration_summary(identity, payload.get("overrides")),
485 }
487 return payload
490_STRUCTURE_FIELDS = (
491 # Architectural (shape-defining) config fields. A change in any of these
492 # — most importantly the layer/experts counts a per-layer sharding plan is
493 # expanded against — must invalidate a generated bundle even when the
494 # underlying modeling file is unchanged.
495 "num_hidden_layers",
496 "num_layers",
497 "hidden_size",
498 "intermediate_size",
499 "num_attention_heads",
500 "num_key_value_heads",
501 "num_experts",
502 "num_local_experts",
503 "n_routed_experts",
504 "first_k_dense_replace",
505 "moe_intermediate_size",
506 "n_group",
507 "topk_group",
508 "num_experts_per_tok",
509 "top_k",
510)
513def _structure_summary(hf_config: Any) -> dict[str, Any]:
514 """A stable digest of the config fields that define the model's structure."""
515 return {
516 name: getattr(hf_config, name)
517 for name in _STRUCTURE_FIELDS
518 if getattr(hf_config, name, None) is not None
519 }
522def _codegen_implementation_digest(package_dir: Optional[str] = None) -> str:
523 """Hash the Python implementation that can affect generated bundles.
525 The relative path is part of the digest so moving content between modules
526 is treated as an implementation change. Non-Python files and bytecode
527 caches are excluded because they do not define codegen semantics.
528 """
529 root = os.path.abspath(package_dir or os.path.dirname(__file__))
530 digest = hashlib.sha256()
531 for dirpath, dirnames, filenames in os.walk(root):
532 dirnames[:] = sorted(name for name in dirnames if name != "__pycache__")
533 for filename in sorted(filenames):
534 if not filename.endswith(".py"):
535 continue
536 path = os.path.join(dirpath, filename)
537 relative = os.path.relpath(path, root).replace(os.sep, "/")
538 digest.update(relative.encode("utf-8"))
539 digest.update(b"\0")
540 with open(path, "r", encoding="utf-8", newline=None) as handle:
541 digest.update(handle.read().encode("utf-8"))
542 digest.update(b"\0")
543 return digest.hexdigest()
546def _build_meta(
547 config: Any,
548 layout: ArtifactLayout,
549 signature: str,
550 *,
551 hf_config: Any = None,
552) -> CodegenMeta:
553 """Build the meta record from the same projection the signature used.
555 ``meta.source`` is taken verbatim from the signature payload's ``source``
556 (resolved ``ResolvedSource`` dict, or the ``SourceSpec`` projection when
557 nothing resolved) so the recorded identity always matches what the
558 signature covers.
559 """
560 spec = _project_spec(config, layout, hf_config=hf_config)
561 from hyper_parallel.codegen.meta import (
562 COVERED_DEFAULTS,
563 NOT_COVERED_DEFAULT,
564 )
566 return CodegenMeta(
567 codegen_version=CODEGEN_VERSION,
568 signature=signature,
569 yaml_path=layout.yaml_path,
570 yaml_sha256=_config_sha256(spec, layout),
571 source=dict(spec.get("source") or {}),
572 parallel_dims=spec.get("parallel_dims", {}),
573 model_name=_model_name(config, hf_config),
574 model_class=_model_class(config, hf_config),
575 # COVERED_DEFAULTS is the starting copy; _fill_plan_fields sets
576 # covered.sharding_plan True once the plan is frozen, so a real
577 # artifact short-circuits the planner and the runtime shards it. The
578 # entrypoint name is the contract the runtime dispatches on; an inline
579 # artifact that defines no such function is parallelized from meta by
580 # ``runtime._parallelize_inline_from_meta``.
581 covered=dict(COVERED_DEFAULTS),
582 not_covered=list(NOT_COVERED_DEFAULT),
583 entrypoints={"parallelize": "hp_parallelize"},
584 )
587def _fill_plan_fields(
588 meta: CodegenMeta,
589 config: Any,
590 layout: ArtifactLayout,
591 *,
592 hf_config: Any = None,
593) -> None:
594 """Derive + freeze the sharding plan into ``meta``'s plan slots.
596 Generation-time pipeline: build the meta-device model, run the real
597 ``ShardingPlanner`` over the offline mesh, and freeze the result into the
598 ``param_plan`` / ``frozen_sharded_params`` / ``injections`` /
599 ``tied_pairs`` / ``mesh_dim_names`` fields. Derivation failures are NOT
600 degraded: a None plan or an empty frozen plan would produce a plan-less
601 artifact that still passes preflight (nothing verifies an empty
602 ``param_plan``), so this stage raises and generation fails instead of
603 silently continuing.
604 """
605 from hyper_parallel.codegen.plan.derive import (
606 build_meta_model,
607 build_plan_overrides,
608 derive_sharding_plan,
609 )
610 from hyper_parallel.codegen.plan.freeze import freeze_plan
611 from hyper_parallel.codegen.spec.project import project_codegen_spec
613 spec = project_codegen_spec(config, layout=layout)
614 model = build_meta_model(spec, config_overrides=_config_overrides(config), hf_config=hf_config)
615 if model is None:
616 raise RuntimeError(
617 "codegen: cannot derive a sharding plan — the spec carries no "
618 "model identity to build the meta model from; fix the model "
619 "source in the YAML"
620 )
621 replace_meta = _apply_generation_replacements(meta, config, model, spec)
622 plan_overrides = build_plan_overrides(config, spec)
623 # S3: a plan override whose matched MoE boundary carries no explicit
624 # ``local_compute_fn`` has its EP compute factory inferred from structure
625 # (never by model name) and written into the spec so it lands in
626 # ``meta.injections``. An explicit ``_target_`` stays the opt-out escape
627 # hatch and is never overwritten.
628 inferred_ep_targets = _fill_inferred_ep_targets(model, plan_overrides)
629 plan = derive_sharding_plan(
630 model,
631 spec,
632 plan_overrides=plan_overrides,
633 )
634 _drop_inferred_ep_injections_on_non_moe(plan, model, inferred_ep_targets)
635 frozen = freeze_plan(plan, model)
636 if not frozen.param_plan:
637 raise RuntimeError(
638 "codegen: plan derivation produced an empty param plan for "
639 f"{spec.source.model_name_or_path!r} — the artifact would carry "
640 "no sharding contract and pass preflight; check the generation "
641 "logs for the planner failure and regenerate"
642 )
643 meta.param_plan = frozen.param_plan
644 meta.frozen_sharded_params = frozen.frozen_sharded_params
645 meta.injections = frozen.injections
646 meta.tied_pairs = frozen.tied_pairs
647 meta.special_handlers = frozen.special_handlers
648 meta.boundary_classes = frozen.boundary_classes
649 meta.mesh_dim_names = list(frozen.mesh_dim_names) or None
651 meta.module_overrides = replace_meta
652 # The inline pipeline lowers the matched replacements into the generated
653 # file, so the artifact owns ``replace_module``. Declare coverage so the
654 # trainer's HF replacement step is skipped on the gen path — but only when
655 # something was actually matched. ``covered`` is always a
656 # COVERED_DEFAULTS copy, so the key is present either way.
657 meta.covered["module_overrides"] = bool(replace_meta)
659 # A non-empty frozen plan is the artifact's full parallel logic; the
660 # runtime shards from ``codegen_meta.json``. ``covered`` is always a
661 # COVERED_DEFAULTS copy, so the key is present either way.
662 meta.covered["sharding_plan"] = True
665def _record_inline_declarations(meta: CodegenMeta) -> None:
666 """Record the facts the *runtime* also needs, derived from structure.
668 The class names whose ``forward`` an inline strategy body replaces (the EP
669 MoE shell selected by the compute factories this plan injects) are proven
670 by the injected targets' own resolved specs. Recording them in ``meta``
671 (the file the runtime loads at training time) is what keeps the runtime
672 from carrying its own copy that drifts when a family adds a model.
673 """
674 identity = (meta.source or {}).get("architecture")
675 meta.external_state_classes = list(
676 external_state_classes_for(identity, getattr(meta, "injections", None))
677 )
680def _apply_generation_replacements(
681 meta: CodegenMeta,
682 config: Any,
683 model: Any,
684 spec: Any,
685) -> list[dict[str, Any]]:
686 """Apply YAML ``replace_module`` actions to the generation meta model.
688 Generated model ``__init__`` applies module replacements before runtime
689 sharding. Generation must mirror that order: compile the raw YAML entries
690 against the original meta model to record serializable FQNs, apply the same
691 replacements to the meta model, then derive the sharding plan from the
692 replaced parameter tree. Otherwise the frozen plan may reference
693 pre-replacement parameters (for example ``q_proj``/``k_proj``/``v_proj``)
694 that no longer exist after a generated attention module installs fused
695 ``qkv_proj`` weights.
696 """
697 raw_entries = getattr(config, "plan_overrides", None) or []
698 if not raw_entries:
699 return []
701 from hyper_parallel.models.replacement import (
702 _apply_module_replacement_actions,
703 )
704 from hyper_parallel.codegen.emit.replacement import (
705 compile_overrides_for_meta,
706 )
707 from hyper_parallel.trainer.config import (
708 entries_to_module_replacements,
709 )
711 specs = entries_to_module_replacements(raw_entries)
712 factory_paths = []
713 target_configs = []
714 for entry in raw_entries:
715 replace_module = getattr(entry, "replace_module", None)
716 if replace_module is None:
717 continue
718 factory_paths.append(_target_path(replace_module))
719 # Capture the YAML Target's static args so the runtime can rebuild the
720 # same pre-bound factory (see ``_target_replacement_factory``) instead
721 # of importing the bare ``@module_replacement`` fn — otherwise a
722 # configured arg would fall back to its default or fail to build.
723 target_configs.append(
724 {
725 name: value
726 for name, value in replace_module.to_dict().items()
727 if name != "_target_"
728 }
729 )
730 records = compile_overrides_for_meta(
731 model,
732 specs,
733 factory_paths=factory_paths,
734 configs_by_spec={
735 id(spec): config for spec, config in zip(specs, target_configs)
736 },
737 )
738 # ``compile_overrides_for_meta`` already raised on any no-match rule (a
739 # configuration error); every returned record here matched a live module.
740 replace_meta = list(records)
741 if replace_meta:
742 model, _ = _apply_module_replacement_actions(
743 model,
744 specs,
745 context={
746 "low_precision": None,
747 "tp": spec.parallel_dims.tp_size > 1,
748 "cp": spec.parallel_dims.cp_size > 1,
749 "ep": spec.parallel_dims.ep_size > 1,
750 "pp": spec.parallel_dims.pp_size > 1,
751 },
752 capture_checkpoint_metadata=False,
753 )
754 return replace_meta
757_STRUCTURAL_CONFIG_KEYS = ("num_hidden_layers", "num_layers", "use_cache")
760def _config_overrides(config: Any) -> dict[str, Any]:
761 """Return config-only kwargs passed to AutoConfig for generated models.
763 ``from_pretrained`` configs carry these explicitly under
764 ``model.config_overrides``. A ``from_config`` builder instead declares its
765 structure as model-target kwargs (e.g. ``num_hidden_layers``), so when
766 ``config_overrides`` is absent the recognized structural keys are folded in
767 — otherwise the frozen plan would be derived at the source checkpoint's
768 shape and disagree with the (cropped) model the builder constructs.
769 """
770 model = getattr(config, "model", None)
771 overrides = getattr(model, "config_overrides", None)
772 if overrides is None:
773 overrides = {
774 key: getattr(model, key)
775 for key in _STRUCTURAL_CONFIG_KEYS
776 if getattr(model, key, None) is not None
777 }
778 if not isinstance(overrides, dict):
779 raise TypeError(
780 "codegen: model.config_overrides must be a mapping of "
781 f"AutoConfig keyword arguments, got {type(overrides).__name__}"
782 )
783 return dict(overrides)
786def _target_path(target: Any) -> Optional[str]:
787 """Return the serialized target path for a YAML Target-like object."""
788 if target is None or not hasattr(target, "to_dict"):
789 return None
790 data = target.to_dict()
791 return data.get("_target_") if isinstance(data, dict) else None
794#: Structure fingerprint -> model-family MoE EP compute factory.
795#:
796#: S3 maps a recognized ``MoeStructure`` (router / shared-expert branch /
797#: expert-storage) to the factory that would have been wired through a YAML
798#: ``local_compute_fn._target_``. The mapping is not duplicated here: it is the
799#: same framework structure->archetype table the model adapters project from,
800#: so the fingerprint is the *only* discriminator and a model with the same
801#: structure is covered without new registry work. Built lazily (the archetype
802#: table lives in the distributed package) and cached for the module lifetime.
803_EP_FACTORY_BY_STRUCTURE: dict[tuple[str, str, str], str] | None = None
806def _build_ep_factory_table() -> dict[tuple[str, str, str], str]:
807 """Project the framework archetype table into a fingerprint->path map."""
808 from hyper_parallel.distributed.expert_parallel.archetypes import ( # pylint: disable=C0415
809 moe_archetypes,
810 )
812 return {
813 fingerprint: archetype.compute_factory
814 for fingerprint, archetype in moe_archetypes().items()
815 }
818def _ep_factory_table() -> dict[tuple[str, str, str], str]:
819 """Return the cached fingerprint->factory ``_target_`` map."""
820 global _EP_FACTORY_BY_STRUCTURE # pylint: disable=global-statement
821 if _EP_FACTORY_BY_STRUCTURE is None:
822 _EP_FACTORY_BY_STRUCTURE = _build_ep_factory_table()
823 return _EP_FACTORY_BY_STRUCTURE
826def _fill_inferred_ep_targets(model: Any, plan_overrides: dict[str, Any]) -> set[str]:
827 """S3: fill a missing ``local_compute_fn`` on MoE EP overrides.
829 Every desugared plan override whose ``match`` resolves to a MoE boundary
830 but which declares no explicit ``local_compute_fn`` gets its EP compute
831 factory inferred from structure (``detect_moe_structure``) instead of a
832 YAML ``_target_``. Entries that already carry a ``local_compute_fn`` are
833 left untouched — the explicit ``_target_`` remains the escape hatch for
834 structures detection cannot cover. Ordinary (non-MoE) overrides are
835 untouched. This is a no-op on a struct construct; the ``plan_overrides``
836 dict's spec values are mutated in place so ``freeze_injections`` records
837 the inferred factory in ``meta.injections``.
839 Returns:
840 The factory paths that were inferred. They identify the injections the
841 framework chose itself, which the plan build must keep off boundaries
842 that cannot host experts (see
843 ``_drop_inferred_ep_injections_on_non_moe``).
844 """
845 inferred: set[str] = set()
846 for match, spec in plan_overrides.items():
847 if getattr(spec, "local_compute_fn", None) is not None:
848 continue
849 target = _infer_ep_compute_for_match(model, match)
850 if target is not None:
851 spec.local_compute_fn = target
852 path = _target_path(target)
853 if path:
854 inferred.add(path)
855 return inferred
858def _infer_ep_compute_for_match(
859 model: Any, match: str
860) -> Optional["Any"]:
861 """Return the EP compute ``_target_`` covering every MoE block ``match`` hits.
863 Resolves ``match`` (an FQN or FQN glob, ``fnmatchcase``) against the meta
864 model's module tree. Returns a ``Target`` for the structure-inferred EP
865 factory when at least one matched module is a MoE boundary; ``None`` when
866 ``match`` touches no MoE boundary (an ordinary sharding override needs no
867 compute injection). An unrecognized MoE structure raises
868 ``UnsupportedModuleStructure`` so the artifact never silently picks a wrong
869 strategy — the escape-hatch text points the user at writing an explicit
870 ``_target_``.
871 """
872 candidate = None
873 candidate_path = None
874 for fqn, module in model.named_modules():
875 if not fnmatchcase(fqn, match):
876 continue
877 if not _looks_like_moe_boundary(module):
878 continue
879 target = _import_ep_target_for(module, match)
880 # Compare the serialized factory path, never the Target objects: every
881 # ``_build_ep_target`` call wraps a fresh closure and ``Target`` defines no
882 # equality, so identity would make any model with two MoE boundaries look
883 # ambiguous even when they share one factory.
884 path = _target_path(target)
885 if candidate is None:
886 candidate, candidate_path = target, path
887 elif path != candidate_path:
888 from hyper_parallel.distributed.expert_parallel.structure import ( # pylint: disable=C0415
889 UnsupportedModuleStructure,
890 )
892 raise UnsupportedModuleStructure(
893 f"{match!r} matched multiple MoE boundaries with different "
894 "EP compute factories; declare an explicit "
895 "local_compute_fn._target_ in the YAML"
896 )
897 return candidate
900def _drop_inferred_ep_injections_on_non_moe(
901 plan: Any, model: Any, inferred_targets: set[str]
902) -> None:
903 """Keep an *inferred* EP injection off boundaries that cannot host experts.
905 The S3 inference fills the factory once per ``match``; the planner's glob
906 merge then hands that spec to every matched boundary, including boundaries
907 with no experts to shard (DeepSeek-V3's dense ``first_k_dense_replace``
908 MLPs). Such a boundary cannot be an expert-parallel region — there is nothing
909 to all-to-all — and the artifact would render a local-region forward whose
910 compute the runtime never binds (the applier resolves an EP factory against a
911 non-MoE module and fails its interface assertion). Dropping the injection
912 here, while the plan is still codegen's own, keeps the artifact coherent.
914 Only targets the framework inferred are affected: a user-declared
915 ``local_compute_fn`` stays the escape hatch and is never dropped.
916 """
917 if not inferred_targets:
918 return
919 for fqn, spec in (getattr(plan, "modules", {}) or {}).items():
920 injection = getattr(spec, "local_compute_fn", None)
921 if injection is None or _target_path(injection) not in inferred_targets:
922 continue
923 module = model.get_submodule(fqn) if model is not None else None
924 if module is None or _looks_like_moe_boundary(module):
925 continue
926 logger.info(
927 "codegen: inferred EP injection dropped on %s — %s carries no "
928 "experts/router, so it cannot be an expert-parallel region",
929 fqn, type(module).__name__)
930 spec.local_compute_fn = None
931 spec.region_dispatch = None
934def _looks_like_moe_boundary(module: Any) -> bool:
935 """Cheap gate: does the module carry an expert plus a router gate?"""
936 return hasattr(module, "experts") and (
937 hasattr(module, "gate") or hasattr(module, "router")
938 )
941def _import_ep_target_for(module: Any, match: str) -> "Any":
942 """Import the delayed target bound to the structure-inferred factory."""
943 from hyper_parallel.distributed.expert_parallel.structure import ( # pylint: disable=C0415
944 UnsupportedModuleStructure,
945 detect_moe_structure,
946 )
948 structure = detect_moe_structure(module)
949 fingerprint = (
950 structure.router,
951 structure.shared_experts,
952 structure.expert_storage,
953 )
954 factory_path = _ep_factory_table().get(fingerprint)
955 if factory_path is None:
956 raise UnsupportedModuleStructure(
957 f"{match!r} matched {type(module).__name__} with structure "
958 f"{fingerprint} that has no EP compute factory. Declare an "
959 "explicit local_compute_fn._target_ to opt out of structure "
960 "detection."
961 )
962 return _build_ep_target(factory_path)
965def _build_ep_target(factory_path: str) -> "Any":
966 """Build a delayed ``Target`` for the EP factory at ``factory_path``.
968 ``local_compute_fn`` is a *factory* Target — the framework calls
969 ``build()`` (which resolves the import) only at apply time, never during
970 generation. So the callable here resolves lazily; generation serializes
971 only the ``_target_`` path into meta, and the runtime re-resolves it from
972 that path. Deferring the import keeps codegen free of a hard dependency on
973 any concrete model package.
974 """
975 from hyper_parallel.trainer.config.target import Target # pylint: disable=C0415
977 def _factory(**kwargs: Any) -> Any:
978 import importlib # pylint: disable=C0415
980 module_path, _, attr = factory_path.rpartition(".")
981 fn = getattr(importlib.import_module(module_path), attr)
982 return fn(**kwargs)
984 return Target(_factory, target_path=factory_path)
987def _gen_backend(config: Any) -> bool:
988 """Whether the codegen pipeline is active for this config.
990 ``codegen=True`` drives artifact generation regardless of the modeling
991 backend (a ``force_hf`` run still generates), so this is simply the
992 codegen switch — and those are exactly the runs that must fail fast
993 when the modeling source cannot be resolved.
994 """
995 return bool(getattr(config, "codegen", False))
998def _transformers_version() -> str:
999 """Installed transformers version, or ``""`` when not importable.
1001 Folded into the signature/meta so an environment upgrade that changes the
1002 modeling file's behavior regenerates instead of silently reusing.
1003 """
1004 try:
1005 import transformers
1007 return getattr(transformers, "__version__", "")
1008 except ImportError:
1009 return ""
1012def _config_sha256(spec: dict, layout: ArtifactLayout) -> str:
1013 """Return YAML file hash when available, otherwise hash the projected spec."""
1014 if layout.yaml_path:
1015 return sha256_file(layout.yaml_path)
1016 return hashlib.sha256(canonical_json(spec).encode("utf-8")).hexdigest()
1019__all__ = [
1020 "CODEGEN_VERSION",
1021 "barrier_after_artifact",
1022 "ensure_codegen_artifact",
1023 "preflight_integrity_check",
1024 "should_regenerate",
1025 "wait_for_rank0_artifact",
1026]