Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / runtime.py: 0%
512 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 runtime: literal-plan execution helpers for generated modules.
17The generated modeling file must not re-derive the sharding plan — it only
18executes the frozen literal plan (``_HYPER_PARAM_PLAN`` etc.) that generation
19recorded in ``codegen_meta.json``. Execution is **convergent with the native
20applier**: parameter sharding (Phase A ``_shard_planned_parameters``), special
21handlers (Phase B) and tied weights (Phase D) run against a live
22``ShardingPlan`` rebuilt from the meta (:func:`_rebuild_live_plan_from_meta`),
23so codegen no longer keeps a second literal sharding implementation. Only the
24forward assembly (the codegen Phase C helpers ``hp_install_boundaries`` /
25``hp_bind_compute`` / ``hp_apply_inner_wrapper``) stays literal, because
26the generated file pre-rewrites forward and the native Phase C would
27double-wrap that redistribution.
29Helpers fail fast: any unknown placement form, missing entrypoint, or absent
30meta errors instead of silently degrading to a ``Replicate()`` no-op.
31"""
32from __future__ import annotations
34import functools
35import inspect
36import logging
37import os
38from typing import Any, Optional
40from hyper_parallel.codegen.loader import import_generated_module
41from hyper_parallel.codegen.meta import (
42 CodegenMeta,
43 load_codegen_meta as _load_codegen_meta_from_path,
44 meta_to_dict,
45 validate_meta_schema,
46)
47# Shared boundary-form classifier (module-level imports are stdlib-only, so
48# this stays importable without torch). The EP-dependence criterion must be
49# the same one generation used to pick the boundary's forward form: an entry
50# whose ``ep`` placement actually *changes* routes to the expert mesh; an
51# identity ``R -> R`` key (the frozen plan's "not EP-dependent" marker) is not
52# a dependency, so a degenerate topology (tp=cp=ep=1) prunes such boundaries
53# to the identity form while the install path is skipped entirely.
54from hyper_parallel.codegen.plan.boundary_forms import (
55 _entry_has_ep_placement,
56 declared_out_names,
57 is_boundary_entry,
58)
59# Pure placement parsing (freeze.py imports no torch), needed by both the
60# meta->plan rebuild and its per-field deserializer below.
61from hyper_parallel.codegen.plan.freeze import parse_named_placement
63# ``sharding/apply`` imports torch at module level, so it is deliberately not
64# imported here: runtime.py must stay importable without torch (generation-time
65# hosts and the local smoke tests import this module for its pure helpers). The
66# apply helpers are imported lazily inside the one function that needs each.
68logger = logging.getLogger(__name__)
71def load_codegen_meta(artifact_dir: str) -> Optional[CodegenMeta]:
72 """Read ``codegen_meta.json`` from an artifact directory.
74 Thin wrapper over ``codegen/meta.load_codegen_meta``: the meta module
75 takes a full path, the runtime contract is the artifact dir. ``None``
76 when absent or malformed.
77 """
78 return _load_codegen_meta_from_path(os.path.join(artifact_dir, "codegen_meta.json"))
81def verify_codegen_signature(meta: CodegenMeta, hf_config: Any = None) -> None:
82 """Validate the generated bundle metadata before model construction.
84 The full recompute (canonical signature over YAML + env) is the trainer
85 preflight's job — it runs before ``_build_model`` with the trainer config.
86 The gate cannot see ``TrainerConfig`` (infrastructure only holds
87 ``distributed_setup`` / ``hf_config``), so this check covers what it can
88 and fails fast on anything inconsistent:
90 - meta passes the schema (``validate_meta_schema``);
91 - ``covered`` / ``covered.sharding_plan`` / ``entrypoints["parallelize"]``
92 exist — the gate decision would be meaningless without them;
93 - the meta's recorded architecture matches ``hf_config.architectures[0]``,
94 when the hf_config is available (catches a bundle generated for a
95 different model class).
96 """
97 validate_meta_schema(meta_to_dict(meta))
98 covered = meta.covered or {}
99 if not isinstance(covered.get("sharding_plan"), bool):
100 raise TypeError(
101 "codegen gate: meta.covered.sharding_plan must be a bool; "
102 f"got {covered.get('sharding_plan')!r}"
103 )
104 if not meta.entrypoints.get("parallelize"):
105 raise KeyError(
106 "codegen gate: meta.entrypoints['parallelize'] is missing — the "
107 "generated module cannot be parallelized; regenerate the artifact"
108 )
109 if hf_config is not None:
110 architectures = getattr(hf_config, "architectures", None) or []
111 arch_name = architectures[0] if architectures else None
112 source_arch = (meta.source or {}).get("architecture")
113 if arch_name and source_arch and arch_name != source_arch:
114 raise ValueError(
115 f"codegen gate: meta records architecture {source_arch!r} but "
116 f"hf_config resolves {arch_name!r}; the bundle does not match "
117 "the requested model"
118 )
121def publish_codegen_mesh_context(artifact_dir: str, mesh_context: Any) -> None:
122 """Publish ``mesh_context`` into the generated module namespace.
124 Lowered boundary forwards reference ``mesh_context`` as a module global.
125 A pure data-parallel/FSDP run does not need generated parameter sharding,
126 but it still executes those lowered forwards; publishing the context keeps
127 identity redistributions well-defined without calling ``hp_parallelize``.
128 """
129 module = import_generated_module(artifact_dir)
130 module.__dict__["mesh_context"] = mesh_context
133# ---------------------------------------------------------------------------
134# Mesh resolution (mirrors apply_sharding_plan's source-mesh selection)
135# ---------------------------------------------------------------------------
137def _active_mesh(mesh: Any, mesh_dim_names):
138 """Slice a full mesh down to the plan's active axes (``_get_active_mesh``).
140 Delegates to the native ``_get_active_mesh_with_names`` — the slicing and
141 reordering contract lives in one place; this wrapper keeps no mirrored
142 implementation to stay in sync.
144 ``apply_sharding_plan`` aligns its sharding mesh to ``plan.mesh_dim_names``
145 (the axes the placements were resolved against) by slicing the device
146 mesh. The runtime receives the same
147 full meshes (a raw ``DeviceMesh`` carries ``dp``/``cp`` axes the plan
148 never shards on; the trainer's ``fsdp_non_moe_mesh`` carries
149 ``fsdp_replicate``/``fsdp_shard``) — slicing by the same names keeps the
150 dense mesh identical to the applier's active mesh, so a placement never
151 lands on a wrong axis.
153 A mesh without the given axes (e.g. an ``OfflineMesh`` built without
154 ``ep``) returns itself, matching ``_get_active_mesh``'s fallback. A mesh
155 that HAS the axes but cannot be sliced raises — never silently returns the
156 unsliced mesh (that would shard along a wrong axis).
158 The frozen plan records its active axes in *offline-mesh* order
159 (``build_offline_mesh`` appends ``tp`` then ``cp``), which is NOT
160 guaranteed to be the live ``DeviceMesh``'s declared order (a trainer mesh
161 is ``("dp", "cp", "tp")``). ``DeviceMesh.__getitem__`` only accepts a
162 slice whose axis indices are ascending (device_mesh.py), so the names are
163 reordered into the mesh's own declared order before slicing. The reordered
164 tuple is returned alongside the mesh
165 (``(sub_mesh, active_dim_names)``): the caller must resolve placements
166 against the *returned* names, because the sliced sub-mesh carries exactly
167 that axis order and a placement tuple must line up with it.
168 """
169 from hyper_parallel.distributed._builder.applier import ( # pylint: disable=C0415
170 _get_active_mesh_with_names,
171 )
172 return _get_active_mesh_with_names(mesh, mesh_dim_names)
175def _resolve_runtime_meshes(mesh: Any, mesh_dim_names=None):
176 """Resolve the active (dense) mesh and the FSDP source meshes.
178 Mirrors ``apply_sharding_plan`` / ``_resolve_parameter_source_meshes``.
179 ``dense_mesh`` is the mesh parameters shard across and boundaries
180 redistribute against: the **device** mesh sliced down to the caller's
181 ``mesh_dim_names`` (the plan's active axes), exactly like
182 ``_get_active_mesh(mesh_context.device_mesh, plan.mesh_dim_names)`` in the
183 applier. ``dense_source_mesh`` is a different thing — the FSDP2 source
184 metadata mesh: the non-MoE topology's ``tp`` child, with the device mesh
185 ``tp`` child as the no-FSDP-context fallback. ``expert_source_mesh`` is
186 the MoE topology's ``ep`` child. The source meshes stay un-sliced (they
187 are the true ``tp`` / ``ep`` child meshes).
189 Returns ``(dense_mesh, dense_source_mesh, expert_source_mesh,
190 active_dim_names)``. ``active_dim_names`` is the mesh-major reordering of
191 the caller's ``mesh_dim_names`` (the order ``dense_mesh`` actually carries
192 after slicing; ``_active_mesh`` reorders the frozen plan's offline-mesh
193 order into the live mesh's declared order). Callers that resolve
194 placements against ``dense_mesh`` must use this returned order, never the
195 raw ``mesh_dim_names`` argument.
196 """
197 device_mesh = getattr(mesh, "device_mesh", mesh)
198 dense_mesh, active_dim_names = _active_mesh(device_mesh, mesh_dim_names)
200 fsdp_non_moe_mesh = getattr(mesh, "fsdp_non_moe_mesh", None)
201 if fsdp_non_moe_mesh is not None:
202 dense_source_mesh = (
203 fsdp_non_moe_mesh["tp"]
204 if "tp" in tuple(getattr(fsdp_non_moe_mesh, "mesh_dim_names", ()) or ())
205 else None
206 )
207 else:
208 dense_source_mesh = (
209 device_mesh["tp"]
210 if "tp" in tuple(getattr(device_mesh, "mesh_dim_names", ()) or ())
211 else None
212 )
213 fsdp_moe_mesh = getattr(mesh, "fsdp_moe_mesh", None)
214 expert_mesh = (
215 fsdp_moe_mesh["ep"]
216 if fsdp_moe_mesh is not None
217 and "ep" in tuple(getattr(fsdp_moe_mesh, "mesh_dim_names", ()) or ())
218 else None
219 )
220 if expert_mesh is not None:
221 # The expert mesh is also sliced to the plan's active axes, so it gets
222 # the same mesh-major reordering (and must arrive with it).
223 expert_mesh, _ = _active_mesh(expert_mesh, mesh_dim_names)
224 return dense_mesh, dense_source_mesh, expert_mesh, active_dim_names
227# ---------------------------------------------------------------------------
228# TP gradient metadata (delegated to the native builder:
229# parameter_sharding._build_runtime_source_shard_info ->
230# source_shard.build_source_shard_info)
231# ---------------------------------------------------------------------------
233def hp_build_tp_grad_info(
234 model,
235 plan,
236 dense_source_mesh,
237 expert_source_mesh,
238) -> Optional[dict[str, tuple[Any, Any]]]:
239 """Build FSDP2 source metadata through the native builder.
241 Delegates to the SAME entry the applier runs —
242 ``parameter_sharding._build_runtime_source_shard_info``, which unwraps the
243 production parameters once (``_local_params_context``) and calls
244 ``source_shard.build_source_shard_info``. Codegen holds no second copy of
245 that logic, so the two paths cannot drift:
247 - every planned parameter maps to its own declared placement, resolved over
248 the source mesh's non-FSDP axes, recorded as a tuple (FSDP2 iterates the
249 placements, so a bare ``Shard(0)`` would crash with "'Shard' object is
250 not iterable");
251 - ``experts.*`` parameters of an EP spec (``spec._ep_size > 0``) map to the
252 expert EP source mesh;
253 - ``build_source_shard_info``'s FSDP-owned-axis rejection applies here too;
254 - tied pairs normalize to one placement (``Shard`` takes precedence).
256 ``plan`` is the live ``ShardingPlan`` rebuilt from the frozen meta
257 (:func:`_rebuild_live_plan_from_meta`); the two source meshes come from
258 ``parameter_sharding._resolve_parameter_source_meshes``, exactly as
259 ``apply_sharding_plan`` resolves them (``build_source_shard_info`` strips
260 the FSDP-owned axes itself).
262 The one-shot ``_local_params_context`` unwrap happens inside the native
263 builder — after this call the model's DTensor parameters are permanently
264 plain locals, so ``hp_build_tp_grad_info`` must be the LAST sharding
265 helper the generated ``hp_parallelize`` runs (the applier's own ordering
266 invariant).
267 """
268 from hyper_parallel.distributed._builder.parameter_sharding import (
269 _build_runtime_source_shard_info,
270 )
272 return _build_runtime_source_shard_info(
273 [model], plan, dense_source_mesh, expert_source_mesh, False
274 )
277# ---------------------------------------------------------------------------
278# Boundary redistribution through PrecompiledBoundary
279# ---------------------------------------------------------------------------
281def _load_target_path(path: str):
282 """Import the callable at a dotted ``_target_`` path.
284 The frozen injection serializes a ``trainer.config.Target`` to
285 ``{"_target_": <path>}``; the runtime re-imports the callable and rebuilds
286 the ``Target`` so the applier's resolution helpers (``_is_delayed_target``,
287 ``Target.build``) accept it unchanged. ``_target_`` paths may carry
288 ``(obj, qualname)`` for classes, but injection targets are module-level
289 functions, so the dotted path splits at the last dot into
290 (module, attribute).
291 """
292 from importlib import import_module
294 if not isinstance(path, str) or not path:
295 raise ValueError(f"frozen injection target path must be a non-empty string, got {path!r}")
296 module_name, _, attr = path.rpartition(".")
297 if not module_name or not attr:
298 raise ValueError(f"frozen injection target path {path!r} is not a dotted path")
299 fn = getattr(import_module(module_name), attr, None)
300 if fn is None:
301 raise ValueError(
302 f"frozen injection target path {path!r} does not name a callable "
303 f"({module_name!r} has no attribute {attr!r})"
304 )
305 if not callable(fn):
306 raise TypeError(f"frozen injection target path {path!r} resolves to {type(fn).__name__}, not callable")
307 return fn
310def _injection_target(value: Any) -> Any:
311 """Rebuild an injection factory/wrapper from its frozen form.
313 ``freeze_injections`` serializes a ``Target`` via ``to_dict()``
314 (``{"_target_": path, **kwargs}``) and a plain callable by ``__name__``.
315 A dict is reconstructed into a real ``trainer.config.Target`` so the
316 applier's ``_is_delayed_target`` / ``Target.build`` path applies; a string
317 is left as the registry name the applier resolves the same way it would
318 have at plan time.
319 """
320 if isinstance(value, dict) and "_target_" in value:
321 from hyper_parallel.trainer.config import Target
323 path = value["_target_"]
324 return Target(
325 _load_target_path(path),
326 target_path=path,
327 **{k: v for k, v in value.items() if k != "_target_"},
328 )
329 return value
332class _FrozenBoundarySpec:
333 """Minimal spec view over one frozen ``param_plan`` boundary entry.
335 ``PrecompiledBoundary`` reads ``in_src`` / ``in_dst`` / ``out_src`` /
336 ``out_dst`` / ``out_names`` off the spec object; the frozen entry stores
337 them as JSON strings, so the adapter re-parses them into real placements
338 once, up front.
340 The injection helpers (``hp_bind_compute`` and
341 ``hp_apply_inner_wrapper``) also hand this to the applier's resolution
342 chain (``_resolve_local_compute_fn`` / ``_resolve_inner_wrapper``), so it
343 additionally mirrors the injection fields those read. ``injection`` (a
344 frozen ``freeze_injections`` entry) supplies ``inner_wrapper`` /
345 ``inner_target`` / ``inner_out_src`` / ``local_compute_fn`` (each
346 re-imported from the frozen form); ``entry`` (the frozen ``param_plan``
347 entry) supplies ``region_dispatch`` / ``out_src`` / ``out_names``.
348 """
350 def __init__(self, entry: dict[str, Any], injection: Optional[dict[str, Any]] = None):
351 self.params = parse_named_placement(entry.get("params") or {}) or {}
352 self.in_src = parse_named_placement(entry.get("in_src") or {}) or {}
353 self.in_dst = parse_named_placement(entry.get("in_dst") or {}) or {}
354 self.out_src = (
355 parse_named_placement(entry["out_src"]) if entry.get("out_src") else None
356 )
357 self.out_dst = (
358 parse_named_placement(entry["out_dst"]) if entry.get("out_dst") else None
359 )
360 self.out_names = list(entry["out_names"]) if entry.get("out_names") else None
361 self.region_dispatch = entry.get("region_dispatch")
362 # D-22: the native suppression/restore pair reads this off the spec
363 # object it is handed, so the frozen view mirrors it.
364 self._deferred_bias_params = tuple(entry.get("deferred_bias_params") or ())
365 for key in ("inner_wrapper", "inner_target", "inner_out_src", "local_compute_fn"):
366 if injection is not None and key in injection:
367 setattr(self, key, _injection_target(injection[key]))
370def _build_rewrap_plan(
371 plan_entry: dict[str, Any],
372 dim_names: tuple[str, ...],
373) -> list[tuple[int, str, tuple]]:
374 """Pre-parse one frozen entry's ``out_src`` into executable rewrap ops.
376 Returns ``[(output_index, out_name, placements), ...]``. ``InstalledBoundary``
377 resolves once at install time so the per-forward path never re-parses.
378 """
379 from hyper_parallel.distributed.recipe_spec import resolve_placements
381 declared = plan_entry.get("out_src") or {}
382 out_names = declared_out_names(plan_entry)
383 name_to_idx = {name: index for index, name in enumerate(out_names)}
384 plan: list[tuple[int, str, tuple]] = []
385 for out_name, named_placement in declared.items():
386 index = name_to_idx.get(out_name)
387 if index is None:
388 raise ValueError(
389 f"InstalledBoundary.rewrap_outputs: out_src declares output {out_name!r}, "
390 f"but out_names={list(out_names)!r} does not contain it"
391 )
392 # Out_src is a frozen JSON contract: its placement leaves are
393 # canonical strings ("S(1)"/"R"/"P" ...). resolve_placements would
394 # hand those strings straight to DTensor.from_local, and a string
395 # sequence trips _build_layout's alias branch — which demands the
396 # placement count equal the tensor rank. The extended MoE identity
397 # contract declares 3 axes ({cp,ep,tp}) but the dense mesh carries only
398 # 2, so resolving leaves 2 strings against a 3-D (B,S,H) output and
399 # fails. Parse the leaves back into real Placement objects first (the
400 # same step _FrozenBoundarySpec applies for the boundary compile / the
401 # native _rewrap_local_outputs runs on already-object contracts); the
402 # non-alias path then maps the 2-axis MeshPlacement onto the 3-D
403 # tensor, padding the unpredicted dim with Replicate().
404 named_placement = parse_named_placement(named_placement)
405 plan.append(
406 (index, out_name, tuple(resolve_placements(named_placement, dim_names)))
407 )
408 return plan
411def hp_to_local_if_dtensor(output):
412 """Convert a local-region boundary output back to a local tensor after exit.
414 Named by the emitted local-region forward template (``emit.parallel``), so
415 it stays until that template goes away — see the boundary-ownership item in
416 the restructuring plan.
417 """
418 to_local = getattr(output, "to_local", None)
419 if callable(to_local):
420 return to_local()
421 return output
424# ---------------------------------------------------------------------------
425# Install-time boundary compilation (one-shot; generated forwards bind these)
426# ---------------------------------------------------------------------------
428def _compile_boundary(entry: dict[str, Any], mesh: Any, mesh_dim_names):
429 """Compile one frozen entry into a ``PrecompiledBoundary`` (uncached).
431 The generated path compiles per module at install time: ``_bind_input_indices``
432 mutates ``op.arg_index`` against one module's forward signature, and two classes
433 with identical contracts but different signatures must not share a plan.
434 """
435 from hyper_parallel.distributed._builder.precompiled_boundary import (
436 PrecompiledBoundary,
437 )
438 from hyper_parallel.distributed._builder.tp_collective_lowering import (
439 create_tp_collective_lowerer,
440 )
442 dim_names = (
443 mesh_dim_names
444 or tuple(getattr(mesh, "mesh_dim_names", ()) or ())
445 )
446 return PrecompiledBoundary(
447 _FrozenBoundarySpec(entry),
448 mesh,
449 dim_names,
450 op_lowerer=create_tp_collective_lowerer(mesh, dim_names),
451 )
454class InstalledBoundary:
455 """One boundary's compiled redistribute plan, bound to a module instance.
457 ``hp_install_boundaries`` resolves everything once, at install time
458 rather than on every forward call: the dense/expert mesh routing (the
459 empty-active-axes no-op, the ep-entry expert-mesh compile, the ep-key drop
460 when the dense mesh is active), the active-axis slicing, the
461 rank-order/backend lowerer validation, and the input-index binding.
462 Generated forwards then call
463 ``self._hp_boundary.redistribute_inputs/outputs`` with zero per-call
464 resolution; the execution semantics — identity passthrough, DTensor
465 unwrap on identity ops, local collectives, ``DTensor.redistribute``
466 fallback — are the compiled ``PrecompiledBoundary``'s own.
467 """
469 def __init__(
470 self,
471 entry: dict[str, Any],
472 mesh_context: Any,
473 mesh_dim_names: Optional[tuple[str, ...]] = None,
474 *,
475 module: Any = None,
476 ) -> None:
477 dense_mesh, _, expert_mesh, active_dim_names = (
478 _resolve_runtime_meshes(mesh_context, mesh_dim_names)
479 )
480 if dense_mesh is None:
481 raise ValueError(
482 "hp_install_boundaries requires a DeviceMesh / MeshContext "
483 "to compile boundaries against"
484 )
485 # Routing facts the tp_collective re-validation reads: the dense mesh
486 # and the active axis names this boundary actually compiled against
487 # (empty = the no-op / expert-mesh branches below).
488 self.dense_mesh = dense_mesh
489 self.active_dim_names = tuple(active_dim_names or ())
490 self.noop = False
491 self._boundary = None
492 if not active_dim_names:
493 if expert_mesh is not None and _entry_has_ep_placement(entry):
494 self._boundary = _compile_boundary(entry, expert_mesh, None)
495 else:
496 # Boundary-exit semantics for this branch: the payload is
497 # returned untouched (no DTensor unwrap either).
498 self.noop = True
499 else:
500 self._boundary = _compile_boundary(entry, dense_mesh, active_dim_names)
501 if self._boundary is not None and module is not None:
502 from hyper_parallel.distributed._builder.forward_rewriter import (
503 _bind_input_indices,
504 )
506 _bind_input_indices(self._boundary, module)
508 self._rewrap_plan: Optional[list[tuple[int, str, tuple]]] = None
509 self._rewrap_mesh = None
510 if entry.get("out_src"):
511 rewrap_dims = (
512 tuple(active_dim_names)
513 or tuple(getattr(dense_mesh, "mesh_dim_names", ()) or ())
514 )
515 self._rewrap_plan = _build_rewrap_plan(entry, rewrap_dims)
516 self._rewrap_mesh = dense_mesh
518 def redistribute_inputs(self, payload):
519 """Run the input plan on the ``(args, kwargs)`` pair.
521 Accepts the pair the generated forward builds and takes it as one
522 positional argument, so the lowered call sites stay uniform across
523 boundary forms.
524 """
525 if self.noop:
526 return payload
527 args, kwargs = payload
528 return self._boundary.redistribute_inputs(args, kwargs)
530 def redistribute_outputs(self, outputs):
531 """Run the output plan on the forward's tensor output."""
532 if self.noop:
533 return outputs
534 return self._boundary.redistribute_outputs(outputs)
536 def rewrap_outputs(self, output):
537 """Re-wrap local region outputs into DTensors per ``out_src``.
539 ``_build_rewrap_plan`` pre-resolves the placements at install time; the
540 replay goes through the codegen-side counterpart of the native
541 local-region re-wrap
542 (:func:`~hyper_parallel.distributed._builder.forward_rewriter.rewrap_declared_outputs`).
543 """
544 if self._rewrap_plan is None:
545 return output
546 from hyper_parallel.distributed._builder.forward_rewriter import ( # pylint: disable=import-outside-toplevel
547 rewrap_declared_outputs,
548 )
550 return rewrap_declared_outputs(
551 output,
552 self._rewrap_plan,
553 self._rewrap_mesh,
554 label="InstalledBoundary.rewrap_outputs",
555 )
558class TPOperators:
559 """Bare-operator runtime bound to a statically lowered boundary forward.
561 A ``tp_collective`` template calls ``self._hp_tp.<kind>(...)`` directly
562 (``emit/parallel``'s static form); ``hp_install_boundaries`` binds this
563 object only after re-validating that form against the live mesh, so the
564 baked calls never run unvalidated. Each method is instruction-equivalent
565 to the compiled engine running the same transition (``RedistOp.execute``
566 with the TP lowerer attached): a DTensor input unwraps to its local shard
567 first, ``None`` passes through untouched (the compiled plan's skip), and
568 the collective dispatches through ``TPCollectiveLowerer.execution_op`` —
569 so a gloo ``reduce_scatter`` remaps to all_reduce + local chunk exactly
570 as the placement-driven path does.
571 """
573 def __init__(self, lowerer: Any) -> None:
574 """Capture the install-time TP lowerer every operator dispatches through."""
575 self._lowerer = lowerer
577 def to_local(self, tensor: Any) -> Any:
578 """Identity-transition semantics: a DTensor unwraps to its local shard."""
579 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415
581 if isinstance(tensor, DTensor):
582 return tensor.to_local()
583 return tensor
585 def all_gather(self, tensor: Any, dim: Optional[int] = None) -> Any:
586 """Shard -> Replicate on the tp axis: concat-gather along ``dim``."""
587 return self._execute("all_gather", tensor, dim)
589 def all_reduce(self, tensor: Any) -> Any:
590 """Partial -> Replicate on the tp axis: sum-reduce over the tp group."""
591 return self._execute("all_reduce", tensor)
593 def reduce_scatter(self, tensor: Any, dim: Optional[int] = None) -> Any:
594 """Partial -> Shard on the tp axis: sum-reduce, then shard along ``dim``."""
595 return self._execute("reduce_scatter", tensor, dim)
597 def _execute(self, kind: str, tensor: Any, dim: Optional[int] = None) -> Any:
598 """Unwrap, skip ``None``, and run one lowered collective op."""
599 from hyper_parallel.core.dtensor.dtensor import DTensor # pylint: disable=C0415
601 if tensor is None:
602 return None
603 if isinstance(tensor, DTensor):
604 tensor = tensor.to_local()
605 return self._lowerer.execution_op(kind, tensor_dim=dim).execute(tensor)
608#: Generated class names whose ``forward`` an inline strategy body replaced
609#: (the MoE EP shell): it orchestrates through parallel state bound on the
610#: instance at install time rather than through a compiled boundary. Filled
611#: per generated module from ``meta.external_state_classes`` — the injected
612#: strategy's own resolved spec is the only declaration site.
613_EXTERNAL_STATE_CLASSES_BY_MODULE: dict[str, frozenset[str]] = {}
616def _mesh_group(mesh: Any, name: str) -> Any:
617 """Return a named mesh group, tolerating one-dimensional mesh APIs."""
618 if mesh is None:
619 return None
620 try:
621 return mesh.get_group(name)
622 except TypeError:
623 return mesh.get_group()
626def _install_inline_state_bindings(module: Any, expert_mesh: Any) -> None:
627 """Bind the instance-attribute parallel channel an inlined forward reads.
629 The inlined EP MoE shell orchestrates through the same instance attributes
630 the boundary templates read: ``self._hp_ep_group`` plus the
631 ``self.ep_enable`` guard. The value comes from the live mesh (expert mesh
632 -> EP group) and reproduces the removed module-level state object's verdict
633 exactly: ``ep_enable`` iff the resolved group spans more than one rank.
635 No other axis is bound here. A class whose forward is *not* replaced by an
636 inline body keeps the emitted boundary form, and its channel comes from the
637 shared install paths: the compiled plan plus the TP operators
638 (``hp_install_boundaries``) and the declared inner wrapper
639 (``hp_apply_inner_wrapper``).
640 """
641 ep_group = _mesh_group(expert_mesh, "ep") if expert_mesh is not None else None
642 module._hp_ep_group = ep_group
643 module.ep_enable = ep_group is not None and ep_group.size() > 1
646def hp_bind_inline_state(
647 model: Any,
648 mesh_context: Any,
649 mesh_dim_names: Optional[tuple[str, ...]] = None,
650 *,
651 generated_module: Any = None,
652) -> None:
653 """Bind the parallel channel on every inlined component module instance.
655 A class declared in ``meta.external_state_classes`` has its ``forward``
656 replaced by an inline strategy body, so the boundary / compute /
657 inner-wrapper installers all skip it: its orchestration reads instance
658 attributes instead. This walks the live model once and binds them on every
659 instance of a declared class — the same classes the installers skip, so an
660 inlined forward can never run with a partially bound channel.
662 The mesh facts are resolved lazily and only when a declared class is
663 actually present, so an artifact without inlined components never touches
664 the mesh.
665 """
666 if model is None or generated_module is None:
667 return
668 targets = [
669 module
670 for _, module in model.named_modules()
671 if _is_external_state_inline_module(module, generated_module)
672 ]
673 if not targets:
674 return
675 _, _, expert_mesh, _ = _resolve_runtime_meshes(mesh_context, mesh_dim_names)
676 for module in targets:
677 _install_inline_state_bindings(module, expert_mesh)
680def register_external_state_classes(generated_module: Any, class_names: Any) -> None:
681 """Record which generated classes take their parallel state externally."""
682 if generated_module is None:
683 return
684 _EXTERNAL_STATE_CLASSES_BY_MODULE[generated_module.__name__] = frozenset(
685 class_names or ()
686 )
689def _is_external_state_inline_module(module: Any, generated_module: Any = None) -> bool:
690 """Whether this generated class gets its state from instance attributes.
692 True for an instance of a class the artifact declared in
693 ``meta.external_state_classes``: codegen replaced that class's forward with
694 an inline strategy body and it reads the attributes
695 ``hp_bind_inline_state`` binds, so the boundary installers must leave it
696 alone.
697 """
698 if generated_module is None:
699 return False
700 module_name = getattr(generated_module, "__name__", None)
701 class_names = _EXTERNAL_STATE_CLASSES_BY_MODULE.get(module_name)
702 if not class_names:
703 return False
704 cls = type(module)
705 return cls.__module__ == module_name and cls.__name__ in class_names
708def _install_deferred_bias_pair(module: Any, spec: Any, *, lowered: bool) -> None:
709 """Install the D-22 bias-suppression / deferred-re-add pair on one boundary.
711 Reuses the native Phase C primitives (``_make_bias_free_forward`` /
712 ``_maybe_add_deferred_biases`` from ``forward_rewriter``) so the hiding
713 semantics cannot drift from the trainer's path; the exit hook is bound as
714 ``module._hp_deferred_bias``, which the emitted forward templates call
715 after the output redistribution.
717 The one codegen-specific adjustment vs the native
718 ``_install_bias_suppression`` loop: on a source-lowered boundary
719 (``lowered=True``) a deferred param path with no dot — the boundary's OWN
720 bias — wraps the extracted ``_forward_impl`` (the compute forward the
721 rewritten ``forward`` calls), not ``forward`` itself: wrapping
722 ``forward`` would keep the bias hidden while the emitted
723 ``self._hp_deferred_bias(outputs)`` reads it at the exit. Child-owned
724 paths and imported-class boundaries keep the native placement exactly.
725 """
726 from hyper_parallel.distributed._builder.forward_rewriter import ( # pylint: disable=C0415
727 _make_bias_free_forward,
728 _maybe_add_deferred_biases,
729 )
731 for param_path in spec._deferred_bias_params: # pylint: disable=protected-access
732 owner_path = param_path.rpartition(".")[0]
733 if owner_path:
734 owner = module.get_submodule(owner_path)
735 owner.forward = _make_bias_free_forward(owner, owner.forward)
736 elif lowered:
737 module._forward_impl = _make_bias_free_forward(
738 module, module._forward_impl
739 )
740 else:
741 module.forward = _make_bias_free_forward(module, module.forward)
742 module._hp_deferred_bias = functools.partial(
743 _maybe_add_deferred_biases, module, spec
744 )
747def _wrap_installed_boundary_forward(
748 module: Any, installed: InstalledBoundary, spec: Any = None
749) -> None:
750 """Install native-style boundary entry/exit redistribution on one module.
752 The compiled plan is captured once at install time instead of being
753 re-derived on every call. When ``spec`` carries D-22 deferred bias
754 params, the wrapper also re-adds them after the exit redistribution —
755 the native Phase C exit order (suppress inside, reduce, add once).
756 """
757 if getattr(module, "_codegen_boundary_wrapped", False):
758 return
759 original_forward = module.forward
760 deferred_bias = spec is not None and bool(spec._deferred_bias_params)
762 @functools.wraps(original_forward)
763 def codegen_boundary_forward(*args: Any, **kwargs: Any) -> Any:
764 redist_args, redist_kwargs = installed.redistribute_inputs((args, kwargs))
765 outputs = original_forward(*redist_args, **redist_kwargs)
766 outputs = installed.redistribute_outputs(outputs)
767 if deferred_bias:
768 outputs = module._hp_deferred_bias(outputs)
769 return outputs
771 module.forward = codegen_boundary_forward
772 module._codegen_boundary_wrapped = True
773 module._codegen_original_forward = original_forward
776def _wrap_static_fallback_forward(module: Any, installed: InstalledBoundary) -> None:
777 """Replace a rejected static ``tp_collective`` forward with the generic engine.
779 The static template's body references ``self._hp_tp``, which this
780 install path refused to bind; replacing ``module.forward`` wholesale
781 keeps the class runnable — the compiled plan carries the same
782 transitions through the generic redistribute engine. A D-22 exit hook
783 already bound on the module (``_hp_deferred_bias``, installed before
784 this fallback wraps) is re-applied after the exit redistribution, so the
785 fallback preserves the deferred-bias semantics the emitted template had.
786 """
787 impl = module._forward_impl
788 deferred_bias = getattr(module, "_hp_deferred_bias", None)
790 def generic_boundary_forward(*args: Any, **kwargs: Any) -> Any:
791 redist_args, redist_kwargs = installed.redistribute_inputs((args, kwargs))
792 outputs = impl(*redist_args, **redist_kwargs)
793 outputs = installed.redistribute_outputs(outputs)
794 if deferred_bias is not None:
795 outputs = deferred_bias(outputs)
796 return outputs
798 module.forward = generic_boundary_forward
801def _install_static_tp_operators(
802 module: Any,
803 entry: dict[str, Any],
804 installed: InstalledBoundary,
805 mesh_dim_names: Optional[tuple[str, ...]],
806 injection: Optional[dict[str, Any]],
807) -> None:
808 """Validate a statically lowered ``tp_collective`` forward against the live mesh.
810 The emitter baked bare-operator calls under an optimistic verdict — the
811 frozen entry classified as ``tp_collective`` under the plan's own axes
812 and the source passed the structural gates (the class marker records
813 that). The live mesh may disagree: active axes the plan did not freeze,
814 a tp rank order the lowerer rejects, a backend without the needed
815 collective. Re-derive the form from the install-time routing facts
816 (``installed.active_dim_names``, the live mesh's mesh-major order) and
817 bind ``module._hp_tp`` only when the live verdict reproduces the
818 emitted ops; on any mismatch replace the forward with the generic
819 engine, so the baked calls never run unvalidated.
820 """
821 from hyper_parallel.codegen.emit.parallel import ( # pylint: disable=C0415
822 TP_FORM_ATTRIBUTE,
823 TP_FORM_MARKER,
824 )
825 from hyper_parallel.codegen.plan.boundary_forms import ( # pylint: disable=C0415
826 FORM_TP_COLLECTIVE,
827 classify_boundary_form,
828 )
829 from hyper_parallel.distributed._builder.tp_collective_lowering import ( # pylint: disable=C0415
830 create_tp_collective_lowerer,
831 )
833 if getattr(module, TP_FORM_ATTRIBUTE, None) != TP_FORM_MARKER:
834 return
835 lowerer = create_tp_collective_lowerer(
836 installed.dense_mesh, installed.active_dim_names
837 )
838 emitted = classify_boundary_form(entry, mesh_dim_names, injection)
839 live = classify_boundary_form(entry, installed.active_dim_names, injection)
840 if (
841 lowerer is not None
842 and emitted.form == FORM_TP_COLLECTIVE
843 and live.form == FORM_TP_COLLECTIVE
844 and live.in_ops == emitted.in_ops
845 and live.out_ops == emitted.out_ops
846 ):
847 module._hp_tp = TPOperators(lowerer)
848 return
849 logger.warning(
850 "hp_install_boundaries: static tp_collective forward on %s does not "
851 "match the live mesh (emitted=%s live=%s lowerer=%s); falling back to "
852 "the generic redistribute engine",
853 type(module).__name__, emitted.form, live.form, lowerer is not None,
854 )
855 _wrap_static_fallback_forward(module, installed)
858def _install_toggle_bindings(
859 module: Any,
860 installed: InstalledBoundary,
861) -> None:
862 """Install the switch/toggle enabler booleans and operators onto one module.
864 A ``toggle``-marked forward (see ``emit.parallel.lower_forward_boundaries_toggle``)
865 calls ``self._hp_tp.<kind>(...)`` inside ``if self.{axis}_enable:``
866 segments. This binds that runtime surface from the live mesh facts the
867 compiled ``InstalledBoundary`` already carries — one codegen path, no second
868 execution engine:
870 - ``module.tp_enable`` / ``module.cp_enable`` / ``module.ep_enable`` are
871 set from ``installed.active_dim_names`` (the boundary's actual parallel
872 dimensions on the live mesh), so a segment only runs when its axis is
873 genuinely active.
874 - ``module._hp_tp`` reuses the SAME ``TPOperators`` + native lowerer that
875 the ``tp_collective`` static form binds, so the toggle body dispatches
876 through the production collectives rather than a duplicated path within
877 the raw operator.
878 """
879 from hyper_parallel.codegen.emit.parallel import ( # pylint: disable=C0415
880 TOGGLE_FORM_MARKER,
881 TP_FORM_ATTRIBUTE,
882 )
883 from hyper_parallel.distributed._builder.tp_collective_lowering import ( # pylint: disable=C0415
884 create_tp_collective_lowerer,
885 )
887 if getattr(module, TP_FORM_ATTRIBUTE, None) != TOGGLE_FORM_MARKER:
888 return
889 active = tuple(installed.active_dim_names or ())
890 module.tp_enable = bool(active and "tp" in active)
891 module.cp_enable = bool(active and "cp" in active)
892 module.ep_enable = bool(active and "ep" in active)
893 if (
894 (module.tp_enable or module.cp_enable)
895 and getattr(installed, "dense_mesh", None) is not None
896 ):
897 lowerer = create_tp_collective_lowerer(installed.dense_mesh, active)
898 if lowerer is not None:
899 module._hp_tp = TPOperators(lowerer)
902def hp_install_boundaries(
903 model: Any,
904 param_plan: dict[str, Any],
905 mesh_context: Any,
906 mesh_dim_names: Optional[tuple[str, ...]] = None,
907 *,
908 injections: Optional[list[dict[str, Any]]] = None,
909 generated_module: Any = None,
910) -> None:
911 """Compile every frozen boundary once and bind it to its module instance.
913 Owns the per-boundary install the removed per-forward wrappers used to do:
914 each boundary's mesh selection (dense vs
915 expert, the empty-active-axes no-op), rank-order/backend lowering checks,
916 and input-index binding run once here, and the compiled
917 ``InstalledBoundary`` is stored as ``module._hp_boundary`` — the
918 generated forwards reference that attribute instead of module globals.
920 Source-lowered boundaries (classes rewritten in the generated file) only
921 get the attribute binding; imported-class boundaries (``nn.Embedding``,
922 ``nn.Linear``) additionally get the entry/exit forward wrapper installed,
923 after the vocab-parallel-embedding wrap, mirroring the previous runtime
924 behavior. An imported-class boundary whose frozen injection declares a
925 ``local_compute_fn`` gets no plain wrapper here: ``hp_bind_compute``
926 installs the native local-region skeleton against the compiled boundary,
927 so the injected compute — not the original forward — runs inside the
928 region. A statically lowered ``tp_collective`` forward additionally gets
929 its bare operators bound (``module._hp_tp``) only after the live-mesh
930 re-validation in :func:`_install_static_tp_operators`; a mismatch replaces
931 the forward with the generic engine.
933 A boundary whose frozen entry carries ``deferred_bias_params`` (D-22)
934 additionally gets the native suppression/restore pair installed
935 (:func:`_install_deferred_bias_pair`): child ``Linear`` forwards run
936 bias-free inside the region and the exit hook
937 ``module._hp_deferred_bias`` re-adds each bias exactly once after the
938 output redistribution — the same order the emitted templates encode. An
939 external-state inline boundary cannot run that exit hook and fails fast
940 instead of silently dropping the bias.
942 Boundaries are installed in post-order (deepest FQN first) via the native
943 ``_boundary_post_order_key`` — the same D-14 invariant 2 ordering as the
944 applier's Phase C driving loop, so inner wrappers exist before an outer
945 boundary's local_compute_fn can cache them.
946 """
947 if not param_plan:
948 return
949 from hyper_parallel.distributed._builder.applier import (
950 _boundary_post_order_key,
951 _get_tp_submesh,
952 )
953 from hyper_parallel.distributed._builder.forward_rewriter import (
954 _is_vocab_parallel_embed,
955 )
956 from hyper_parallel.distributed._builder.parameter_sharding import (
957 _resolve_module,
958 )
960 meshes_resolved = False
961 tp_mesh = None
963 for fqn, entry in sorted(param_plan.items(), key=_boundary_post_order_key):
964 if _is_tp_sequence_parallel_embed_entry(entry):
965 parent_fqn = fqn.rpartition(".")[0]
966 _wrap_default_position_ids(_resolve_module(model, parent_fqn))
967 if not is_boundary_entry(entry):
968 continue
969 module = _resolve_module(model, fqn)
970 if _has_lowered_forward(module, generated_module):
971 injection = _injection_spec(injections or [], fqn)
972 installed = InstalledBoundary(
973 entry, mesh_context, mesh_dim_names, module=module
974 )
975 if _is_external_state_inline_module(module, generated_module):
976 # D-22: an inline-strategy class body never passes through
977 # this boundary exit, so a deferred bias would be suppressed
978 # but never re-added — fail fast instead of silently
979 # dropping the bias from every forward.
980 if entry.get("deferred_bias_params"):
981 raise NotImplementedError(
982 f"hp_install_boundaries: boundary {fqn!r} defers bias "
983 f"params {entry['deferred_bias_params']}, but its class is "
984 "an external-state inline strategy the D-22 exit hook "
985 "cannot run in"
986 )
987 continue
988 module._hp_boundary = installed
989 if entry.get("deferred_bias_params"):
990 # Before _install_static_tp_operators: the static-fallback
991 # wrap captures module._forward_impl at wrap time, so the
992 # suppression must already be in place there.
993 _install_deferred_bias_pair(
994 module, _FrozenBoundarySpec(entry, injection), lowered=True
995 )
996 _install_static_tp_operators(
997 module,
998 entry,
999 installed,
1000 mesh_dim_names,
1001 injection,
1002 )
1003 _install_toggle_bindings(module, installed)
1004 continue
1005 spec = _FrozenBoundarySpec(entry)
1006 if spec.params:
1007 if not meshes_resolved:
1008 dense_mesh, _, _, active_dim_names = _resolve_runtime_meshes(
1009 mesh_context, mesh_dim_names
1010 )
1011 tp_mesh = _get_tp_submesh(dense_mesh, tuple(active_dim_names))
1012 meshes_resolved = True
1013 if _is_vocab_parallel_embed(module, spec, tp_mesh):
1014 _install_vocab_parallel_embedding(module, tp_mesh)
1015 installed = InstalledBoundary(
1016 entry, mesh_context, mesh_dim_names, module=module
1017 )
1018 # Every boundary module carries its compiled plan — the imported-class
1019 # wrapper below closes over it, so introspection and debugging see the
1020 # same object the generated source-lowered forwards would.
1021 module._hp_boundary = installed
1022 if spec._deferred_bias_params:
1023 # Native order: suppression first, so the wrapper below captures
1024 # the suppressed forward as its original (D-22, Phase C identity).
1025 _install_deferred_bias_pair(module, spec, lowered=False)
1026 injection = _injection_spec(injections or [], fqn)
1027 if injection is not None and injection.get("local_compute_fn") is not None:
1028 # A local-region boundary on an imported class: hp_bind_compute
1029 # installs the native local-region skeleton (boundary entry ->
1030 # injected compute -> boundary exit) against this compiled
1031 # boundary. The plain entry/exit wrapper here would run the
1032 # original forward instead of the injected compute and double
1033 # the redistribution.
1034 continue
1035 _wrap_installed_boundary_forward(module, installed, spec=spec)
1038# ---------------------------------------------------------------------------
1039# Compute binding and inner-wrapper installation
1040# ---------------------------------------------------------------------------
1042def _injection_spec(injections: list[dict[str, Any]], fqn: str) -> Optional[dict[str, Any]]:
1043 """Find one frozen boundary's injection declaration by its FQN."""
1044 for rule in injections:
1045 if rule.get("match") == fqn:
1046 return rule
1047 return None
1050def hp_expand_injections(
1051 rules: Optional[list[dict[str, Any]]],
1052) -> list[dict[str, Any]]:
1053 """Expand grouped injection rules into the per-FQN form helpers consume.
1055 The generated ``_HYPER_INJECTIONS`` literal groups rules that differ only
1056 in their ``match`` under one record with a ``match`` **list** (one rule
1057 for every layer's CP wrapper instead of one per layer). This expansion —
1058 called once at the top of the generated ``hp_parallelize`` — restores
1059 the per-FQN shape, so ``hp_install_boundaries`` /
1060 ``hp_bind_compute`` / ``hp_apply_inner_wrapper`` keep their frozen
1061 single-FQN ``match`` contract unchanged.
1063 A scalar ``match`` (the frozen ``meta.injections`` shape and every
1064 pre-grouping artifact) passes through untouched; ``None`` / empty becomes
1065 ``[]``.
1066 """
1067 if not rules:
1068 return []
1069 expanded: list[dict[str, Any]] = []
1070 for rule in rules:
1071 match = rule.get("match") if isinstance(rule, dict) else None
1072 if isinstance(match, (list, tuple)):
1073 for fqn in match:
1074 one = dict(rule)
1075 one["match"] = fqn
1076 expanded.append(one)
1077 else:
1078 expanded.append(rule)
1079 return expanded
1082def _install_runtime_region_forward(
1083 module: Any,
1084 spec: Any,
1085 compute_fn: Any,
1086 mesh: Any,
1087 mesh_dim_names: Any,
1088) -> None:
1089 """Install the native local-region skeleton on an imported-class boundary.
1091 The compiled plan ``hp_install_boundaries`` left on the module
1092 (``module._hp_boundary``) carries the native ``PrecompiledBoundary``, so
1093 the region executes the exact native Phase C path: boundary entry ->
1094 injected compute -> boundary exit. A boundary compiled as a no-op (no
1095 active axes) degenerates to the compute swap plus the D-22 exit hook --
1096 its entry/exit are identity, the only remaining skeleton effects.
1097 """
1098 from hyper_parallel.distributed._builder.forward_rewriter import ( # pylint: disable=C0415
1099 _maybe_add_deferred_biases,
1100 _wrap_local_region_forward,
1101 )
1103 installed = getattr(module, "_hp_boundary", None)
1104 if installed is None:
1105 raise RuntimeError(
1106 "codegen: local-region injection on "
1107 f"{type(module).__name__} has no compiled boundary; "
1108 "hp_install_boundaries must run before hp_bind_compute"
1109 )
1110 if installed.noop:
1112 @functools.wraps(module.forward)
1113 def runtime_region_forward(*args: Any, **kwargs: Any) -> Any:
1114 output = compute_fn(*args, **kwargs)
1115 return _maybe_add_deferred_biases(module, spec, output)
1117 module.forward = runtime_region_forward
1118 return
1119 _wrap_local_region_forward(
1120 module,
1121 installed._boundary, # pylint: disable=protected-access
1122 spec,
1123 mesh,
1124 tuple(mesh_dim_names or ()),
1125 validate_mode=False,
1126 compute_fn=compute_fn,
1127 )
1130def hp_bind_compute(
1131 model,
1132 injections: list[dict[str, Any]],
1133 mesh_context,
1134 mesh_dim_names: Optional[tuple[str, ...]] = None,
1135 *,
1136 param_plan: Optional[dict[str, Any]] = None,
1137 generated_module: Any = None,
1138) -> None:
1139 """Bind ``module.__hp_compute__`` for every local-region boundary.
1141 A generated local-compute forward calls
1142 ``self.__hp_compute__(*args, **kwargs)`` on the redistributed inputs.
1143 The bound callable is the applier's local-compute resolution
1144 (``_resolve_local_compute_fn`` → ``functools.partial(compute_fn, module)``),
1145 so the compute function's first argument is the boundary module and the
1146 remaining arguments are the redistributed forward inputs.
1148 ``injections`` is the frozen ``meta.injections`` list (``freeze_injections``
1149 shape); only entries with an explicit ``local_compute_fn`` are bound (the
1150 ``region_dispatch=False`` gate that would otherwise fall back to
1151 ``module.forward`` is deliberately NOT bound — the applier runs that path
1152 by re-entering ``forward``, which the generated file already does as
1153 ``_forward_impl``). ``param_plan`` (``meta.param_plan``) supplies
1154 ``region_dispatch`` / ``ep_size`` for the factory's mesh context, exactly
1155 as the applier reads them off a live spec.
1157 An imported-class boundary (the class is not defined in the generated
1158 file) has no emitted forward that could call ``__hp_compute__``; for
1159 those the native local-region skeleton is installed instead
1160 (:func:`_install_runtime_region_forward`).
1161 """
1162 from hyper_parallel.distributed._builder.rule_resolver import (
1163 _resolve_local_compute_fn,
1164 )
1165 from hyper_parallel.distributed._builder.parameter_sharding import _resolve_module
1167 dense_mesh, _, expert_mesh, active_dim_names = (
1168 _resolve_runtime_meshes(mesh_context, mesh_dim_names)
1169 )
1170 if dense_mesh is None:
1171 raise ValueError(
1172 "hp_bind_compute requires a DeviceMesh / MeshContext to bind compute against"
1173 )
1174 if not injections:
1175 return
1176 for rule in injections:
1177 if rule.get("local_compute_fn") is None:
1178 continue
1179 fqn = rule.get("match")
1180 if not fqn:
1181 raise ValueError(
1182 f"hp_bind_compute: injection rule {rule!r} has no 'match' FQN"
1183 )
1184 entry = (param_plan or {}).get(fqn, {})
1185 spec = _FrozenBoundarySpec(entry, rule)
1186 module = _resolve_module(model, fqn)
1187 compute_fn = _resolve_local_compute_fn(
1188 module, spec, dense_mesh, active_dim_names, expert_mesh)
1189 # External-state inline modules (e.g. Qwen3MoeSparseMoeBlock) carry an
1190 # inlined forward and must not receive __hp_compute__, but the local
1191 # compute factory above still has to run: it installs companion state
1192 # such as ``experts.local_expert_count`` and the bound experts.forward
1193 # that the inlined forward relies on.
1194 if _is_external_state_inline_module(module, generated_module):
1195 continue
1196 # _resolve_local_compute_fn returns module.forward for the
1197 # region_dispatch=False gate; that path is not a bound compute (the
1198 # generated forward already calls _forward_impl), so it must not be
1199 # bound as __hp_compute__.
1200 if compute_fn is None or compute_fn is module.forward:
1201 continue
1202 if _has_lowered_forward(module, generated_module):
1203 module.__hp_compute__ = compute_fn
1204 continue
1205 # An imported-class boundary has no emitted forward that could call
1206 # ``__hp_compute__``: install the native local-region skeleton
1207 # instead (boundary entry -> injected compute -> boundary exit), the
1208 # same code path the native applier runs for this boundary.
1209 _install_runtime_region_forward(
1210 module, spec, compute_fn, dense_mesh, active_dim_names
1211 )
1214def hp_apply_inner_wrapper(
1215 model,
1216 injections: list[dict[str, Any]],
1217 mesh_context,
1218 mesh_dim_names: Optional[tuple[str, ...]] = None,
1219 *,
1220 param_plan: Optional[dict[str, Any]] = None,
1221 generated_module: Any = None,
1222) -> None:
1223 """Install the declared inner CP wrapper on its target module.
1225 A generated redistribution forward delegates computation to
1226 ``_forward_impl``. The
1227 CP K/V all-gather wrapper is applied to ``target.forward`` IN PLACE by the
1228 applier (``_wrap_inner_attention``), which is not bindable as a method
1229 without changing its invocation contract. This helper repeats that
1230 one-shot installation from the frozen ``meta.injections``, so the
1231 wrapper's ``compute`` runs inside the ``_forward_impl`` call at train time.
1233 Uses the applier's own resolution chain — ``_wrap_inner_attention``
1234 resolves the registry name / Target / callable, resolves the target, and
1235 installs the dual-mode adapter (production pass-through / validate local
1236 re-wrap) via ``_install_inner_adapter``. This is the one CP path for a
1237 generated artifact: the wrapper is applied on the module whose forward is
1238 the emitted boundary form (the rendered fused attention), so the companion
1239 ``attention_interface`` swap and the boundary redistribution compose
1240 instead of being re-implemented per component.
1241 """
1242 from hyper_parallel.distributed._builder.applier import (
1243 _get_cp_submesh,
1244 _get_ep_submesh,
1245 _get_tp_submesh,
1246 )
1247 from hyper_parallel.distributed._builder.forward_rewriter import (
1248 _wrap_inner_attention,
1249 )
1250 from hyper_parallel.distributed._builder.parameter_sharding import _resolve_module
1252 dense_mesh, _, expert_mesh, active_dim_names = (
1253 _resolve_runtime_meshes(mesh_context, mesh_dim_names)
1254 )
1255 if dense_mesh is None:
1256 raise ValueError(
1257 "hp_apply_inner_wrapper requires a DeviceMesh / MeshContext to install against"
1258 )
1259 if not injections:
1260 return
1261 dim_names = tuple(active_dim_names) if active_dim_names else ()
1262 cp_mesh = _get_cp_submesh(dense_mesh, dim_names)
1263 tp_mesh = _get_tp_submesh(dense_mesh, dim_names)
1264 ep_mesh = _get_ep_submesh(dense_mesh, dim_names) or expert_mesh
1265 for rule in injections:
1266 if rule.get("inner_wrapper") is None:
1267 continue
1268 fqn = rule.get("match")
1269 if not fqn:
1270 raise ValueError(
1271 f"hp_apply_inner_wrapper: injection rule {rule!r} has no "
1272 "'match' FQN"
1273 )
1274 entry = (param_plan or {}).get(fqn, {})
1275 spec = _FrozenBoundarySpec(entry, rule)
1276 module = _resolve_module(model, fqn)
1277 if _is_external_state_inline_module(module, generated_module):
1278 continue
1279 # The applier's boundary loop runs _resolve_local_compute_fn even for
1280 # inner-wrap boundaries, and the preflight fails fast on a missing
1281 # region_dispatch. Reuse the exact applier entry here for parity.
1282 _wrap_inner_attention(
1283 module,
1284 cp_mesh,
1285 spec=spec,
1286 mesh=dense_mesh,
1287 mesh_dim_names=dim_names,
1288 tp_mesh=tp_mesh,
1289 ep_mesh=ep_mesh,
1290 )
1293def _has_lowered_forward(module: Any, generated_module: Any = None) -> bool:
1294 """Return whether the generated source already owns this boundary wrapper."""
1295 if _is_external_state_inline_module(module, generated_module):
1296 return True
1297 return hasattr(module, "_forward_impl")
1300def _is_tp_sequence_parallel_embed_entry(entry: dict[str, Any]) -> bool:
1301 """Return whether an embedding output is reduce-scattered over sequence."""
1302 out_src = (entry.get("out_src") or {}).get("output") or {}
1303 out_dst = (entry.get("out_dst") or {}).get("output") or {}
1304 return (
1305 str(out_src.get("tp", "")).startswith("P(")
1306 and out_dst.get("tp") == "S(1)"
1307 )
1310def _wrap_default_position_ids(module: Any) -> None:
1311 """Build default position ids before an embedding sequence reduce-scatter.
1313 Transformers commonly derives missing position ids from ``inputs_embeds``.
1314 A generated model's vocab-parallel embedding may already have changed that
1315 tensor from TP-replicated sequence length to a TP sequence shard. Supplying
1316 the default from the still-replicated ``input_ids`` keeps rotary embeddings
1317 in the attention boundary's declared TP-replicated layout.
1318 """
1319 if getattr(module, "_codegen_position_ids_wrapped", False):
1320 return
1321 original_forward = module.forward
1322 signature = inspect.signature(original_forward)
1323 if (
1324 "input_ids" not in signature.parameters
1325 or "position_ids" not in signature.parameters
1326 ):
1327 return
1328 positional_names = [
1329 name
1330 for name, parameter in signature.parameters.items()
1331 if parameter.kind in (
1332 inspect.Parameter.POSITIONAL_ONLY,
1333 inspect.Parameter.POSITIONAL_OR_KEYWORD,
1334 )
1335 ]
1336 position_index = (
1337 positional_names.index("position_ids")
1338 if "position_ids" in positional_names
1339 else None
1340 )
1342 @functools.wraps(original_forward)
1343 def codegen_position_ids_forward(*args: Any, **kwargs: Any) -> Any:
1344 """Supply replicated default position ids to the original model forward."""
1345 bound = signature.bind_partial(*args, **kwargs)
1346 if bound.arguments.get("position_ids") is not None:
1347 return original_forward(*args, **kwargs)
1348 input_ids = bound.arguments.get("input_ids")
1349 if input_ids is None:
1350 return original_forward(*args, **kwargs)
1352 past_key_values = bound.arguments.get("past_key_values")
1353 past_seen_tokens = (
1354 past_key_values.get_seq_length() if past_key_values is not None else 0
1355 )
1356 import torch # pylint: disable=C0415
1358 position_ids = torch.arange(
1359 int(input_ids.shape[-1]),
1360 dtype=input_ids.dtype,
1361 device="cpu",
1362 )
1363 position_ids = (position_ids + past_seen_tokens).unsqueeze(0).to(input_ids.device)
1364 if position_index is not None and position_index < len(args):
1365 call_args = list(args)
1366 call_args[position_index] = position_ids
1367 return original_forward(*call_args, **kwargs)
1368 call_kwargs = kwargs.copy()
1369 call_kwargs["position_ids"] = position_ids
1370 return original_forward(*args, **call_kwargs)
1372 module.forward = codegen_position_ids_forward
1373 module._codegen_position_ids_wrapped = True
1374 module._codegen_original_position_ids_forward = original_forward
1377def _install_vocab_parallel_embedding(module: Any, tp_mesh: Any) -> None:
1378 """Install native's D-02 masked-embedding wrapper, once per module.
1380 The wrapper itself is native
1381 (``forward_rewriter._wrap_vocab_parallel_embedding``): codegen installs the
1382 boundary before its one-shot parameter unwrap, and native derives the
1383 vocab interval from the local shard, so both paths share one
1384 implementation. Only the re-install guard is codegen-local —
1385 ``hp_install_boundaries`` is documented idempotent, while native Phase C
1386 runs exactly once.
1387 """
1388 if getattr(module, "_codegen_vocab_parallel_wrapped", False):
1389 return
1390 from hyper_parallel.distributed._builder.forward_rewriter import ( # pylint: disable=C0415
1391 _wrap_vocab_parallel_embedding,
1392 )
1394 _wrap_vocab_parallel_embedding(module, tp_mesh)
1395 module._codegen_vocab_parallel_wrapped = True
1398# ---------------------------------------------------------------------------
1399# Generated entry-point dispatch
1400# ---------------------------------------------------------------------------
1402def parallelize_from_generated(
1403 model,
1404 mesh_context,
1405 artifact_dir: str,
1406 *,
1407 hf_config: Any = None,
1408) -> dict[str, tuple[Any, Any]]:
1409 """Run the generated module's ``hp_parallelize`` entry point.
1411 The entry point is taken from ``meta.entrypoints["parallelize"]`` and must
1412 return ``tp_grad_info`` with the same semantics as
1413 ``apply_sharding_plan()``. Infrastructure validates the metadata before
1414 dispatch, so this function repeats the check only when called directly
1415 with ``hf_config``.
1416 """
1417 meta = load_codegen_meta(artifact_dir)
1418 if meta is None:
1419 raise FileNotFoundError(
1420 f"codegen: no meta at {os.path.join(artifact_dir, 'codegen_meta.json')}; "
1421 "cannot parallelize from the generated module"
1422 )
1423 if hf_config is not None:
1424 verify_codegen_signature(meta, hf_config)
1425 module = import_generated_module(artifact_dir)
1426 # The artifact's own record of which generated classes take their parallel
1427 # state externally — declared once by the adapter render spec and carried
1428 # here in ``meta``.
1429 register_external_state_classes(module, meta.external_state_classes)
1430 entry_name = meta.entrypoints.get("parallelize")
1431 entry = getattr(module, entry_name, None) if entry_name else None
1432 if entry is not None:
1433 return entry(model, mesh_context)
1434 return _parallelize_inline_from_meta(model, mesh_context, meta, module)
1437def _parallelize_inline_from_meta(model: Any, mesh_context: Any, meta: CodegenMeta, generated_module: Any) -> dict:
1438 """Parallelize a clean inline artifact whose model file carries no plan globals.
1440 Execution is convergent with the native applier (see
1441 :func:`_parallelize_via_native_apply`): Phase A parameter sharding /
1442 B special handlers / D tied weights run against a live ``ShardingPlan``
1443 rebuilt from the frozen meta; Phase C forward assembly stays codegen-owned
1444 because the generated file pre-rewrites forward.
1445 """
1446 return _parallelize_via_native_apply(model, mesh_context, meta, generated_module)
1449def _rebuild_live_plan_from_meta(meta: CodegenMeta):
1450 """Rebuild a live ``ShardingPlan`` from the frozen meta (meta -> plan).
1452 Proof-of-concept reconstruction: re-parses the frozen literal forms back
1453 into real ``Placement``/``Target`` objects so native builder helpers
1454 (``_shard_planned_parameters`` / ``_apply_plan_special_handlers`` /
1455 ``_replicate_tied_weights``) can consume this plan exactly as they consume
1456 a planner-produced one.
1458 Which fields cross the meta boundary is declared in
1459 :mod:`hyper_parallel.codegen.plan.spec_fields` — freeze and this rebuild
1460 walk the same table, so the two sides cannot drift apart. ``init=False``
1461 fields (``_deferred_bias_params``, D-22) are rebuilt via ``setattr``
1462 after construction (``POST_INIT_SPEC_FIELDS``). Exempt fields (never
1463 frozen):
1464 - ``_is_terminal`` — validate-mode only, not used in production;
1465 - ``_tp_local_attr_plan`` — planner-derived; the rebuilt spec leaves it
1466 None and ``maybe_update_head_counts`` takes its backward-compatible
1467 path (auto attrs via the head-sharded heuristic, user attrs via the
1468 rebuilt ``tp_divide_attrs``);
1469 - ``_needs_cp_attn`` / ``_resolved_inner_*`` — apply-time preflight and
1470 introspection state, not consumed on the rebuild path.
1472 Injection-side fields (``local_compute_fn`` / ``inner_wrapper`` /
1473 ``inner_out_src``) are reconstructed with the same :func:`_injection_target`
1474 the existing ``_FrozenBoundarySpec`` uses, so injection fidelity matches
1475 the status-quo codegen path.
1476 """
1477 from hyper_parallel.distributed.plan import ShardingPlan
1478 from hyper_parallel.distributed.recipe_spec import ModuleShardingSpec
1479 from hyper_parallel.codegen.plan.spec_fields import (
1480 INJECTION_SPEC_FIELDS,
1481 PARAM_PLAN_SPEC_FIELDS,
1482 POST_INIT_SPEC_FIELDS,
1483 )
1485 by_fqn: dict[str, Any] = {}
1486 for rule in meta.injections or []:
1487 if isinstance(rule, dict) and rule.get("match"):
1488 by_fqn[rule["match"]] = rule
1490 modules: dict[str, Any] = {}
1491 for fqn, entry in (meta.param_plan or {}).items():
1492 inj = by_fqn.get(fqn, {})
1493 kwargs: dict[str, Any] = {
1494 field.name: _rebuild_spec_field(field, entry)
1495 for field in PARAM_PLAN_SPEC_FIELDS
1496 if not field.post_init
1497 }
1498 for field in INJECTION_SPEC_FIELDS:
1499 kwargs[field.name] = _rebuild_spec_field(field, inj)
1500 spec = ModuleShardingSpec(**kwargs)
1501 for field in POST_INIT_SPEC_FIELDS:
1502 value = _rebuild_spec_field(field, entry)
1503 if value is not None:
1504 setattr(spec, field.name, value)
1505 modules[fqn] = spec
1507 return ShardingPlan(
1508 modules=modules,
1509 special_handlers=dict(meta.special_handlers or {}),
1510 mesh_dim_names=tuple(meta.mesh_dim_names or ()),
1511 tied_pairs=[tuple(pair) for pair in (meta.tied_pairs or [])],
1512 )
1515def _rebuild_spec_field(field: Any, source: dict[str, Any]) -> Any:
1516 """Deserialize one spec field per its :class:`SpecField` kind contract.
1518 The inverse of freeze's ``_freeze_spec_field``: reads the field's meta
1519 key from the frozen entry (param-plan fields) or injection rule
1520 (injection fields) and restores the live value. Absent keys restore the
1521 dataclass defaults — ``None`` for optional declarations, ``True`` for
1522 ``is_boundary`` (pre-flag metas), ``{}``/``0`` for the EP internals —
1523 matching what ``_rebuild_live_plan_from_meta`` passed positionally
1524 before the table drove both sides.
1525 """
1526 kind, key = field.rebuild, field.key
1527 if kind == "placement_map":
1528 # Key absent/None -> None (not declared); an explicit {} ("shards
1529 # nothing") must round-trip as {} — Phase A iterates params.items()
1530 # and a None there is the crash this rebuild must not reintroduce.
1531 if source.get(key) is None:
1532 return None
1533 return parse_named_placement(source[key])
1534 if kind == "named_placements":
1535 return parse_named_placement(source[key]) if source.get(key) else None
1536 if kind in ("name_list", "attr_list"):
1537 return list(source[key]) if source.get(key) is not None else None
1538 if kind == "bool_default_true":
1539 return bool(source[key]) if source.get(key) is not None else True
1540 if kind == "injection_target":
1541 return _injection_target(source[key]) if source.get(key) is not None else None
1542 if kind == "dict_copy":
1543 return dict(source[key]) if source.get(key) else {}
1544 if kind == "int_truthy":
1545 return source.get(key) or 0
1546 if kind == "tuple_empty":
1547 # D-22 param paths: absent/empty restores the dataclass default ()
1548 # — consumers iterate the tuple, so None would crash.
1549 return tuple(source[key]) if source.get(key) else ()
1550 # opt_scalar
1551 return source.get(key)
1554def _parallelize_via_native_apply(
1555 model: Any,
1556 mesh_context: Any,
1557 meta: CodegenMeta,
1558 generated_module: Any,
1559) -> dict:
1560 """Reuse the native execution body for a codegen artifact.
1562 Rebuilds a live ``ShardingPlan`` from the frozen meta
1563 (:func:`_rebuild_live_plan_from_meta`) and runs native Phase A
1564 (``_shard_planned_parameters``) / B (``_apply_plan_special_handlers``) /
1565 D (``_replicate_tied_weights``) against it. This is the converging
1566 execution path: codegen no longer maintains a second literal
1567 parameter-sharding replay.
1569 Phase C forward assembly is deliberately NOT reused: the generated file
1570 has already rewritten every boundary's ``forward`` into the
1571 redistribute-in / ``_forward_impl`` / redistribute-out form (resolved at
1572 runtime against the ``_hp_boundary`` / ``__hp_compute__`` bound
1573 here). Native ``_apply_phase_c`` captures the *current* forward as its
1574 ``original_forward`` and wraps it in ANOTHER
1575 ``redistribute_inputs -> original -> redistribute_outputs`` — running it
1576 on a codegen artifact double-wraps the redistribution (double TP/EP
1577 communication) because no native branch skips the wrap. So full
1578 ``apply_sharding_plan`` reuse is not viable; A/B/D reuse with codegen-owned
1579 Phase C is the minimum viable convergence.
1581 MoE/EP artifacts are supported: native Phase A applies the same
1582 ``_stack_moe_experts`` step here as it does for a planner-produced plan, so
1583 dense and MoE topologies share this one reuse path.
1584 """
1585 from hyper_parallel.distributed.mesh import MeshContext # pylint: disable=C0415
1586 from hyper_parallel.distributed._builder.applier import ( # pylint: disable=C0415
1587 _apply_plan_special_handlers,
1588 _get_active_mesh,
1589 _get_tp_submesh,
1590 )
1591 from hyper_parallel.distributed._builder.parameter_sharding import ( # pylint: disable=C0415
1592 _replicate_tied_weights,
1593 _resolve_parameter_source_meshes,
1594 _shard_planned_parameters,
1595 detect_tied_weights,
1596 )
1598 plan = _rebuild_live_plan_from_meta(meta)
1599 mesh_dim_names = plan.mesh_dim_names
1600 param_plan = meta.param_plan or {}
1601 injections = hp_expand_injections(meta.injections or [])
1603 mc = mesh_context if isinstance(mesh_context, MeshContext) else None
1604 device_mesh = mc.device_mesh if mc is not None else mesh_context
1605 if device_mesh is None:
1606 raise ValueError(
1607 "native A/B/D reuse requires a DeviceMesh / MeshContext"
1608 )
1609 mesh = _get_active_mesh(device_mesh, mesh_dim_names)
1610 tp_mesh = _get_tp_submesh(mesh, mesh_dim_names)
1611 expert_mesh, dense_source_mesh, expert_source_mesh = (
1612 _resolve_parameter_source_meshes(plan, mesh_context, device_mesh, tp_mesh)
1613 )
1615 # Native A / B / D are executed against the reconstructed live plan.
1616 _shard_planned_parameters([model], plan, mesh, expert_mesh, validate_mode=False)
1617 _apply_plan_special_handlers([model], plan, mesh)
1618 _replicate_tied_weights(model, list(plan.tied_pairs) or detect_tied_weights(model))
1620 # Native C is not reusable (see docstring) — keep the codegen forward
1621 # assembly that matches the pre-rewritten generated file.
1622 hp_install_boundaries(
1623 model,
1624 param_plan,
1625 mesh_context,
1626 mesh_dim_names,
1627 injections=injections,
1628 generated_module=generated_module,
1629 )
1630 hp_bind_inline_state(
1631 model,
1632 mesh_context,
1633 mesh_dim_names,
1634 generated_module=generated_module,
1635 )
1636 hp_bind_compute(
1637 model,
1638 injections,
1639 mesh_context,
1640 mesh_dim_names,
1641 param_plan=param_plan,
1642 generated_module=generated_module,
1643 )
1644 hp_apply_inner_wrapper(
1645 model,
1646 injections,
1647 mesh_context,
1648 mesh_dim_names,
1649 param_plan=param_plan,
1650 generated_module=generated_module,
1651 )
1652 return hp_build_tp_grad_info(
1653 model, plan, dense_source_mesh, expert_source_mesh
1654 )
1657__all__ = [
1658 "hp_apply_inner_wrapper",
1659 "hp_bind_compute",
1660 "hp_bind_inline_state",
1661 "hp_build_tp_grad_info",
1662 "hp_install_boundaries",
1663 "hp_to_local_if_dtensor",
1664 "import_generated_module",
1665 "load_codegen_meta",
1666 "parallelize_from_generated",
1667 "verify_codegen_signature",
1668]