Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / emit / parallel.py: 84%
274 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
15"""Lower the frozen parallel plan into explicit forward code.
17The lowerer sinks forward wrappers into the source. Which form a boundary
18class's ``forward`` takes is decided by
19:func:`hyper_parallel.codegen.plan.boundary_forms.classify_boundary_form` —
20the same pure classifier the runtime re-runs at install time and preflight
21re-runs at train time, so the three stages can never disagree about a
22boundary's form.
24Class vs FQN: the frozen plan is keyed by boundary FQN
25(``model.layers.0.self_attn``), but a ``class.forward`` is shared by many
26instances. The lowerer groups boundary FQNs by their class (meta's
27``boundary_classes``), enforces that every FQN of a class carries the SAME
28boundary contract, and emits one rewritten ``forward`` per class.
30The original ``forward`` is kept
31verbatim as the class-private ``_forward_impl`` (extracted by a zero-width
32insert before the ``def forward`` — a fresh-source operation, so the offsets
33stay valid when combined with the body replacement), and the public
34``forward`` is rewritten in one of the classified forms:
36- ``identity`` (pruned): every declared transition is identity on the plan's
37 active axes — the class forward is NOT rewritten at all. The runtime's
38 install path covers the exact no-op / to-local semantics with its generic
39 wrapper, so pruning is byte-for-byte equivalent to the previous rewrite.
40- ``tp_collective`` (static bare operators): every declared transition is
41 identity or TP-lowerable, and the source-side structural gates pass
42 (single-value return, no ``*args``/``**kwargs``, declared inputs bind to
43 forward parameters). The forward names its form with the class attribute
44 ``_hp_boundary_form = "tp_collective"`` and calls bare operators on
45 ``self._hp_tp`` (bound at install time after the runtime re-validates
46 the compiled plan against the live mesh; on mismatch the runtime replaces
47 the forward with the generic engine — the baked calls never run unvalidated).
48- ``generic``: redistribute in, ``_forward_impl``, redistribute out, through
49 the instance-bound compiled plan (``self._hp_boundary``, bound once by
50 ``hp_install_boundaries``), with a declarative comment block annotating
51 what each side's transitions lower to.
52- ``region`` (EP MoE local region): redistribute in, call the explicit
53 factory-named attribute materialized by source emission, redistribute out.
55A CP inner-wrapper boundary is ``generic`` in the generated code: the wrapper
56is installed at runtime by ``hp_apply_inner_wrapper`` (it mutates
57``target.forward`` in place — not a bindable callable), and runs inside the
58``_forward_impl`` call.
60The lowerer is a pure function over the source text + the frozen plan. It
61returns the patched text. Structural work (finding a class / its forward /
62the import block) is delegated to :mod:`astkit`.
63"""
64from __future__ import annotations
66import json
67import logging
68from typing import Any, Iterator, Optional
70from hyper_parallel.codegen.astkit.edits import (
71 TextEdit,
72 apply_edits,
73 replace_function_body,
74)
75from hyper_parallel.codegen.astkit.index import (
76 FunctionInfo,
77 SourceIndex,
78 build_source_index,
79 returns_any_value,
80 returns_single_value,
81)
82from hyper_parallel.codegen.plan.boundary_forms import (
83 FORM_GENERIC,
84 FORM_IDENTITY,
85 FORM_REGION,
86 FORM_TP_COLLECTIVE,
87 BoundaryForm,
88 TransitionOp,
89 classify_boundary_form,
90 is_boundary_entry,
91)
92from hyper_parallel.codegen.emit.patches import local_compute_binding_name
94#: Injection keys that describe how a boundary's compute is regioned.
95_INJECTION_KEYS = ("inner_wrapper", "inner_target", "inner_out_src", "local_compute_fn")
98def lower_forward_boundaries(
99 source_text: str,
100 frozen_plan: Any,
101 *,
102 boundary_classes: Optional[dict[str, str]] = None,
103 module_name: str = "",
104) -> str:
105 """Rewrite each boundary class's ``forward`` to call the codegen runtime.
107 ``frozen_plan`` is the meta/param-plan dict or a ``CodegenMeta``; the
108 plan entries are read via ``_plan_field``. ``boundary_classes``
109 is the FQN -> class name map (``meta.boundary_classes``); when omitted it
110 degrades to the class name looked up from the source index itself, which is
111 only correct if every boundary FQN's leaf class name equals the key's tail —
112 ``_class_for_fqn``. The plan's active axes (``mesh_dim_names``) drive
113 :func:`classify_boundary_form`; the classified form plus the source-side
114 structural gates decide which template each class gets (see the module
115 docstring).
117 Returns the patched source text. This function only rewrites the
118 ``forward`` bodies; the rewritten bodies reference the ``_hp_boundary``
119 / ``_hp_tp`` instance attributes that ``hp_install_boundaries`` binds
120 at runtime — no per-boundary global constant is emitted.
121 """
122 if _plan_field(frozen_plan, "status") == "not_required":
123 raise ValueError("codegen emit: not_required snapshot has no parallel boundaries")
124 edits: list[TextEdit] = []
125 for class_name, form, emitted, func, injection, deferred_bias in iter_emitted_forms(
126 source_text,
127 frozen_plan,
128 boundary_classes=boundary_classes,
129 module_name=module_name,
130 ):
131 if emitted == FORM_IDENTITY:
132 # Pruned: the runtime's install path covers the exact no-op /
133 # to-local semantics with its generic wrapper, so leaving the
134 # original forward untouched is behavior-preserving.
135 if deferred_bias:
136 # D-22: an unrewritten forward never calls the exit hook, so
137 # the suppressed bias would silently vanish — fail fast
138 # instead (the planner defers only rowwise-bias exits that
139 # reduce, which classify as non-identity, so this is a
140 # mis-declared plan, not a supported shape).
141 raise NotImplementedError(
142 f"lower_forward_boundaries: boundary class {class_name!r} "
143 "prunes to identity but declares deferred bias params; an "
144 "identity forward has no exit hook to re-add them"
145 )
146 continue
147 body = _build_forward_body(
148 source_text, injection, func, form, emitted, deferred_bias
149 )
150 # Keep the original forward as ``_forward_impl``. The
151 # extracted method is inserted BEFORE the ``def forward`` (zero-width
152 # edit at ``def_offset``), and the rewritten body replaces the original
153 # body at ``body_start`` — the two spans never overlap, so the
154 # incremental-shift validation in ``_apply_edits`` passes. The static
155 # form additionally prefixes the class marker naming the form.
156 marker = _TP_FORM_MARKER_LINE if emitted == FORM_TP_COLLECTIVE else None
157 impl = _build_forward_impl_edit(source_text, func, prefix=marker)
158 if impl is not None:
159 edits.append(impl)
160 edits.extend(replace_function_body(source_text, func, body, indent=""))
161 if not edits:
162 return source_text
164 nested = _validate_no_circular(edits)
165 if nested:
166 raise ValueError(
167 "lower_forward_boundaries: edits overlap (circular class nesting); "
168 "cannot safely rewrite: " + ", ".join(nested)
169 )
170 return apply_edits(source_text, edits)
173#: Class attribute a ``tp_collective`` forward carries, naming its form so the
174#: runtime install path and preflight can recognize the static template.
175TP_FORM_ATTRIBUTE = "_hp_boundary_form"
176#: The marker value naming the statically lowered TP-collective form.
177TP_FORM_MARKER = "tp_collective"
178#: The emitted marker statement (method indent, followed by a blank line).
179_TP_FORM_MARKER_LINE = f' {TP_FORM_ATTRIBUTE} = "{TP_FORM_MARKER}"\n\n'
182def iter_emitted_forms(
183 source_text: str,
184 frozen_plan: Any,
185 *,
186 boundary_classes: Optional[dict[str, str]] = None,
187 module_name: str = "",
188) -> Iterator[tuple[str, BoundaryForm, str, FunctionInfo, Any, bool]]:
189 """Yield every boundary class's emit decision over ``source_text``.
191 Single decision path shared by the emitter
192 (:func:`lower_forward_boundaries`) and the train-time preflight
193 (``check/preflight.verify_boundary_forms``): classification
194 (:func:`classify_boundary_form`) plus the structural gates
195 (:func:`resolve_emitted_form`), per boundary class. Yields
196 ``(class_name, form, emitted, func, injection, deferred_bias)`` —
197 ``form`` the classifier's verdict, ``emitted`` the form actually
198 rendered after the gates, ``func`` the class's ``forward``
199 :class:`FunctionInfo`, ``injection`` the class's frozen injection rule
200 (``None`` when absent), and ``deferred_bias`` whether the class's
201 canonical entry carries D-22 deferred bias params (the emitted exit
202 then calls ``self._hp_deferred_bias(outputs)`` after the output
203 redistribution).
205 Boundary classes are iterated in sorted ``class_name`` order. The
206 grouping walks ``param_plan``, whose insertion order is NOT stable
207 across the write/reload cycle: ``meta.param_plan`` is persisted with
208 ``json.dump(..., sort_keys=True)``, which reorders the keys
209 alphabetically, so reload-ing the meta yields a different class order
210 than the in-memory freeze did. Sorting here makes the emitted byte
211 stream a pure function of which classes are boundaries, independent of
212 how ``param_plan`` reached us.
213 """
214 index = build_source_index(source_text)
215 param_plan = _plan_field(frozen_plan, "param_plan") or {}
216 injections = _plan_field(frozen_plan, "injections") or []
217 injection_by_class = _index_injections_by_class(
218 injections, boundary_classes, index, param_plan
219 )
220 mesh_dim_names = tuple(_plan_field(frozen_plan, "mesh_dim_names") or ())
222 groups = _group_boundaries_by_class(
223 param_plan, boundary_classes, index, module_name=module_name
224 )
225 for class_name in sorted(groups):
226 group = groups[class_name]
227 # Fail fast when one class's FQNs disagree on the contract: one shared
228 # ``forward`` cannot carry two boundary plans. The canonical entry is
229 # the per-class contract every FQN agreed on.
230 canonical_entry = _verify_class_contract(group, class_name)
231 func = _require_forward(index, class_name)
232 injection = injection_by_class.get(class_name)
233 form = classify_boundary_form(canonical_entry, mesh_dim_names, injection)
234 emitted = resolve_emitted_form(source_text, func, form)
235 deferred_bias = bool(canonical_entry.get("deferred_bias_params"))
236 yield class_name, form, emitted, func, injection, deferred_bias
239def resolve_emitted_form(
240 source_text: str,
241 func: FunctionInfo,
242 form: BoundaryForm,
243) -> str:
244 """The form the emitter actually renders for one classified boundary.
246 The classifier's ``tp_collective`` is an *optimistic* verdict: the static
247 template additionally requires the source-side structural gates (no
248 ``*args``, value returns, declared inputs binding to forward parameters,
249 output ops targeting the first output). A boundary that fails a gate
250 renders as ``generic`` instead. This resolver is the single decision
251 shared by the emitter and preflight, so the emitted file and the train-time
252 check can never disagree about which template a class got.
253 """
254 if form.form in (FORM_IDENTITY, FORM_REGION, FORM_GENERIC):
255 return form.form
256 if _static_gates_pass(source_text, func, form):
257 return FORM_TP_COLLECTIVE
258 return FORM_GENERIC
261def _static_gates_pass(source_text: str, func: FunctionInfo, form: BoundaryForm) -> bool:
262 """Whether the static bare-operator template is admissible for ``func``.
264 Gates (each mirrors a concrete way the generic engine is more general
265 than a static rewrite):
267 1. No ``*args`` — a positional variable tail cannot be enumerated
268 statically. ``**kwargs`` is allowed and is re-passed by name.
269 2. Every ``return`` carries a value. Single-value returns receive output
270 ops directly; tuple returns receive index-0 output ops and are
271 rebuilt, matching the compiled plan's output-index behavior.
272 3. Every declared input name binds to a forward parameter — the generic
273 engine skips an ``in_src`` name the caller never passed (``_get_arg``
274 default ``None``); the static template would reference an unbound
275 variable instead, so a declared name outside the signature disqualifies
276 the static form.
277 4. Every non-identity output op targets the first declared output — with a
278 single-value return the compiled plan applies index-0 ops only; an op
279 mapped to a later output position cannot be expressed statically.
280 """
281 if func.vararg_name is not None:
282 return False
283 if not returns_any_value(source_text, func):
284 return False
285 signature_names = set(func.param_names) | set(func.kwonly_names)
286 if not set(form.in_declared) <= signature_names:
287 return False
288 out_names = form.out_names
289 for name in form.out_ops:
290 # An op name missing from the declared order compiles to index 0
291 # (``_compile_output_plan``'s ``name_to_idx.get(name, 0)``), which the
292 # single-value template covers; a later position does not.
293 if name in out_names and out_names.index(name) != 0:
294 return False
295 return True
298# ---------------------------------------------------------------------------
299# grouping / consistency
300# ---------------------------------------------------------------------------
302def _group_boundaries_by_class(
303 param_plan: dict[str, Any],
304 boundary_classes: Optional[dict[str, str]],
305 index: SourceIndex,
306 *,
307 module_name: str,
308) -> dict[str, list[tuple[str, dict[str, Any]]]]:
309 """Group boundary FQNs by the class that owns their ``forward``.
311 Returns ``{class_name: [(fqn, entry), ...]}`` in a stable order (FQN sorted
312 for determinism).
314 A boundary whose class is NOT defined in this source file — ``index`` is
315 the source index, so ``index.find_class(name) is None`` means the class is
316 imported (``nn.Embedding`` / ``nn.Linear`` / ``nn.LayerNorm`` for
317 ``embed_tokens`` / ``lm_head`` / ``norm``) — has no ``def forward`` span
318 to rewrite, so its FQN is dropped from the grouping and the boundary is
319 left un-rewritten here. Its boundary contract is still executed at runtime
320 by ``hp_install_boundaries``: the in/out
321 redistribution for an imported class is installed as a forward wrapper on
322 the live module, so the contract is not dropped. Skipping the rewrite here
323 is exactly what lets the real plan (which marks ``embed_tokens`` /
324 ``lm_head`` as boundaries) emit without failing on a class that has no
325 source ``forward``.
327 The same drop applies to a class that IS defined in this source file but
328 inherits its ``forward`` from an imported base (an attention subclassing
329 a component base class): the class body has no ``def forward`` span to
330 rewrite, and live instances carry no ``_forward_impl``, so the runtime
331 install path takes its imported-class branch all the same.
333 An FQN whose param-plan entry carries the explicit ``skip_source_lowering``
334 marker (set by meta normalization when the class is shared between
335 boundary and non-boundary instances — DeepSeek-V3 ``DeepseekV3MLP``:
336 dense decoder layers and shared experts) is skipped: a class-level
337 ``forward`` rewrite is inherited by every instance, so the non-boundary
338 instance would run a boundary body without the attributes
339 ``hp_install_boundaries`` binds only on boundary instances. The
340 per-instance runtime wrapper is installed instead and executes the same
341 boundary contract.
343 An FQN whose class CANNOT be resolved fails fast — a boundary whose forward
344 is never rewritten would silently drop its redistribution.
345 """
346 logger = logging.getLogger(__name__)
347 groups: dict[str, list[tuple[str, dict[str, Any]]]] = {}
348 for fqn, entry in param_plan.items():
349 if not is_boundary_entry(entry):
350 continue
351 if entry.get("skip_source_lowering"):
352 logger.info(
353 "lower_forward_boundaries: boundary %s is marked "
354 "skip_source_lowering (class shared with a non-boundary "
355 "instance); the per-instance runtime wrapper will be "
356 "installed instead",
357 fqn,
358 )
359 continue
360 class_name = _class_for_fqn(fqn, boundary_classes, index, module_name)
361 if class_name is None:
362 raise ValueError(
363 f"lower_forward_boundaries: cannot resolve the class for boundary "
364 f"{fqn!r} — the frozen plan has no class mapping and the source "
365 "index does not contain a class whose name matches; a boundary "
366 "whose forward is not rewritten would silently drop its "
367 "redistribution"
368 )
369 if index.find_class(class_name) is None:
370 # Imported class (Embedding/Linear/LayerNorm): no source ``forward``
371 # to rewrite; runtime ``hp_install_boundaries`` installs the
372 # in/out boundary wrapper. Not added to ``groups``.
373 continue
374 if index.find_forward(class_name) is None:
375 # In-source class with an inherited forward (an attention
376 # subclassing an imported base): no class-body ``def forward``
377 # span to rewrite, the same shape an imported class has. Live
378 # instances carry no ``_forward_impl``, so hp_install_boundaries
379 # takes its imported-class branch and installs the per-instance
380 # entry/exit wrapper.
381 logger.info(
382 "lower_forward_boundaries: boundary %s is a %s instance whose "
383 "forward is inherited (no class-body ``def forward``); the "
384 "per-instance runtime wrapper will be installed instead",
385 fqn,
386 class_name,
387 )
388 continue
389 groups.setdefault(class_name, []).append((fqn, entry))
390 for class_name, items in groups.items():
391 items.sort(key=lambda item: item[0])
392 return groups
395def _class_for_fqn(
396 fqn: str,
397 boundary_classes: Optional[dict[str, str]],
398 index: SourceIndex,
399 module_name: str,
400) -> Optional[str]:
401 """Resolve a boundary FQN to a top-level class name.
403 Resolution order:
404 1. the explicit ``boundary_classes`` map (freeze-time truth);
405 2. a source-index class whose name matches the FQN's leaf segment;
406 3. a source-index class whose name matches after stripping a known module
407 prefix (``module_name``) — remote-code modeling files carry the bare
408 class name, and the FQN leaf is usually it;
409 4. ``None`` (caller fails fast).
410 """
411 if boundary_classes:
412 cls = boundary_classes.get(fqn)
413 if cls:
414 return cls
415 leaf = fqn.rsplit(".", 1)[-1]
416 if leaf in index.classes:
417 return leaf
418 # A remote-code file may name the class without the module-name prefix on
419 # the FQN leaf; only accept it when the leaf matches a class AND the leaf
420 # is a plausible modeling class (no lambdas / generics).
421 if module_name:
422 candidate = module_name.rsplit(".", 1)[-1]
423 if candidate in index.classes and _name_matches(candidate, leaf):
424 return candidate
425 return None
428def _name_matches(class_name: str, leaf: str) -> bool:
429 """Loose match between a module-name class and an FQN leaf.
431 The FQN leaf is the module attribute (``self_attn`` / ``mlp``); the class
432 is its type (``Qwen3MoeAttention``). They usually differ, so this only
433 accepts the trivial identity. The strong mapping is ``boundary_classes``;
434 this fallback is deliberately conservative.
435 """
436 return class_name == leaf
439def _verify_class_contract(
440 group: list[tuple[str, dict[str, Any]]], class_name: str
441) -> dict[str, Any]:
442 """Fail fast unless every FQN of ``class_name`` carries the same contract.
444 One ``class.forward`` is shared by every instance of the class, so all the
445 class's boundaries must be describable by the SAME frozen entry (modulo the
446 ``params`` tree, which is instance-specific and drives the native Phase A
447 parameter sharding, not the forward rewrite). Returns the canonical entry.
448 """
449 canonical: Optional[str] = None
450 canonical_entry: dict[str, Any] = {}
451 for fqn, entry in group:
452 entry = _without_params(entry)
453 blob = json.dumps(entry, sort_keys=True)
454 if canonical is None:
455 canonical = blob
456 canonical_entry = entry
457 elif blob != canonical:
458 raise ValueError(
459 f"lower_forward_boundaries: class {class_name!r} has FQNs with "
460 f"differing boundary contracts ({fqn!r}); a shared forward "
461 "cannot carry two contracts — split the plan or align the "
462 "per-rule contracts"
463 )
464 return canonical_entry
467def _without_params(entry: dict[str, Any]) -> dict[str, Any]:
468 """Drop the instance-specific ``params`` tree from a contract comparison."""
469 if "params" not in entry:
470 return entry
471 return {k: v for k, v in entry.items() if k != "params"}
474def _index_injections_by_class(
475 injections: list[dict[str, Any]],
476 boundary_classes: Optional[dict[str, str]],
477 index: SourceIndex,
478 param_plan: Optional[dict[str, Any]] = None,
479) -> dict[str, dict[str, Any]]:
480 """Fold a class's injection declarations (inner wrapper / local compute).
482 Returns ``{class_name: {key: value}}`` — the union of every injection rule
483 whose ``match`` lands on a FQN of that class. A class with conflicting
484 values for one key (two FQNs declared different local compute fns) fails
485 fast, for the same reason as ``_verify_class_contract``.
486 """
487 by_class: dict[str, dict[str, Any]] = {}
488 for rule in injections:
489 match = rule.get("match")
490 if not isinstance(match, str) or not match:
491 continue
492 entry = (param_plan or {}).get(match) or {}
493 if entry.get("skip_source_lowering"):
494 continue
495 class_name = _class_for_fqn(match, boundary_classes, index, "")
496 if class_name is None:
497 continue
498 payload = {k: rule[k] for k in _INJECTION_KEYS if k in rule}
499 if not payload:
500 continue
501 existing = by_class.setdefault(class_name, {})
502 for key, value in payload.items():
503 if key in existing and _json_stable(existing[key]) != _json_stable(value):
504 raise ValueError(
505 f"lower_forward_boundaries: class {class_name!r} has "
506 f"conflicting injection values for {key!r}"
507 )
508 existing[key] = value
509 return by_class
512def _json_stable(value: Any) -> str:
513 return json.dumps(value, sort_keys=True, default=str)
516# ---------------------------------------------------------------------------
517# forward body construction
518# ---------------------------------------------------------------------------
520def _require_forward(index: SourceIndex, class_name: str) -> FunctionInfo:
521 """Return the ``forward`` ``FunctionInfo`` for a boundary class or raise."""
522 func = index.find_forward(class_name)
523 if func is None:
524 raise ValueError(
525 f"lower_forward_boundaries: class {class_name!r} has no forward "
526 "method; a boundary forward must exist to be rewritten"
527 )
528 if func.body_start is None or func.body_end is None:
529 raise ValueError(
530 f"lower_forward_boundaries: class {class_name!r}.forward has no "
531 "replaceable body (single-line/pass); cannot sink a boundary"
532 )
533 return func
536def _build_forward_impl_edit(
537 source_text: str,
538 func: FunctionInfo,
539 prefix: Optional[str] = None,
540) -> Optional[TextEdit]:
541 """Build an edit that extracts a method as ``_forward_impl``.
543 The original method (signature + body) is copied out of the source and
544 inserted as the private ``_forward_impl`` right before the rewritten
545 ``def forward``. The copy keeps the original signature and body byte-for-
546 byte, re-indented to the method's own indent, with the ``def forward``
547 renamed to ``def _forward_impl``. ``prefix`` (the static form's class
548 marker) is prepended to the same insertion — one zero-width edit at
549 ``def_offset``, so the edit set never carries two edits at one offset.
551 Returns ``None`` when the forward has no replaceable body (single-line/
552 pass) — there is nothing meaningful to extract, and ``forward`` itself was
553 already rejected by ``_require_forward`` for the same reason, so this is a
554 defensive fallback.
555 """
556 if func.body_start is None or func.body_end is None:
557 return None
558 # The extracted method owns the signature + body; DECORATORS stay attached to
559 # the original ``forward`` (which is being rewritten) — copying them onto
560 # ``_forward_impl`` too would duplicate a side effect (e.g.
561 # ``@torch.no_grad()``) on a method that must be pure. ``func.def_offset``
562 # anchors the FIRST line of the method (a decorator when one is present), so
563 # scan forward from it past decorator lines and blanks to the ``def`` keyword
564 # — this finds the signature line for both the plain and the decorated case
565 # without offset arithmetic (a naive backward scan from ``body_start`` finds
566 # the newline that *ends* the ``def`` line, landing on the first body line).
567 def_line_start = func.def_offset
568 for line in source_text[def_line_start:func.body_start].splitlines(keepends=True):
569 stripped = line.lstrip()
570 if stripped.startswith("def "):
571 break
572 def_line_start += len(line)
573 span = source_text[def_line_start:func.body_end]
574 lines = span.splitlines(keepends=True)
575 if not lines:
576 return None
577 # ``_forward_impl`` is a VERBATIM copy of the original method (signature +
578 # body, byte-for-byte, keeping the method indent on the ``def`` and the body
579 # nesting) with only the name changed -- so nothing here needs dedenting or
580 # re-indenting, and the copy can never drift from the original source.
581 for i, line in enumerate(lines):
582 stripped = line.lstrip()
583 if stripped.startswith(f"def {func.name}("):
584 lines[i] = line.replace(f"def {func.name}(", "def _forward_impl(", 1)
585 break
586 impl = "".join(lines)
587 # The copy ends at ``body_end``, which is clamped to EOF for a final
588 # statement with no trailing newline — so the extracted method can lose its
589 # terminating newline and land joined onto the ``def forward`` that follows.
590 # A method must be separated from the next one by a newline, so pad.
591 if not impl.endswith("\n"):
592 impl += "\n"
593 if prefix:
594 impl = prefix + impl
595 return TextEdit(func.def_offset, func.def_offset, impl)
598def _build_forward_body(
599 source_text: str,
600 injection: Optional[dict[str, Any]],
601 func: FunctionInfo,
602 form: BoundaryForm,
603 emitted: str,
604 deferred_bias: bool = False,
605) -> str:
606 """Render the rewritten forward body for one boundary class.
608 ``emitted`` is :func:`resolve_emitted_form`'s verdict (the classifier's
609 form after the static gates). The redistribute calls go through the
610 instance-bound compiled plan (``self._hp_boundary``, installed by
611 ``hp_install_boundaries``); the static ``tp_collective`` form calls bare
612 operators on ``self._hp_tp`` instead; the factory-named compute attribute
613 is bound by ``hp_bind_compute`` for the local-region shape. Either way the class
614 stays runtime-importable without a mesh.
616 ``deferred_bias`` (D-22) appends the exit hook call
617 ``self._hp_deferred_bias(outputs)`` after the output redistribution —
618 the runtime binds that hook when it installs the native
619 suppression/restore pair.
620 """
621 if emitted == FORM_TP_COLLECTIVE:
622 return _render_tp_collective_forward(source_text, func, form, deferred_bias)
624 input_redist = _render_input_redistribute(func)
625 out_redist = _render_output_redistribute(deferred_bias)
627 local = injection.get("local_compute_fn") if injection else None
628 if local is not None:
629 raise RuntimeError(
630 "codegen: local_compute_fn reached source boundary lowering; "
631 "builder recording must mark it skip_source_lowering so runtime "
632 "can install the trainer's native Phase C wrapper"
633 )
634 # A generic boundary (including an inner-wrapper one — the CP wrapper is
635 # installed at runtime by ``hp_apply_inner_wrapper`` and weaves into
636 # this forward) redistributes through the compiled plan, annotated with
637 # what each side's transitions lower to.
638 return _render_plain_redistribute(input_redist, out_redist, form)
641def _render_input_redistribute(func: FunctionInfo) -> str:
642 """Render the input-side redistribute call on the compiled boundary plan.
644 The generated forward re-binds args/kwargs against its own signature, so a
645 positional ``(args, kwargs)`` pair reaches the compiled plan as the
646 boundary's input side. Index binding happens once at install time
647 (``hp_install_boundaries`` → ``_bind_input_indices``); the pair is still
648 built from the signature's own parameter names (not the plan's) — for a
649 ``def forward(self, hidden_states)`` that is ``((hidden_states,), {})``.
651 The full parameter profile is preserved on the re-pass: positional names go
652 in the args tuple (a trailing ``*args`` expands into it), keyword-only
653 parameters and a fixed ``**kwargs`` go into the kwargs dict. The boundary
654 plan rewrites only the args/kwargs it owns and forwards the rest through
655 unchanged, so ``self._forward_impl(*args, **kwargs)`` receives every
656 argument class the original call site used — a keyword-only ``required``
657 or an HF ``**kwargs`` pass-through no longer drops out.
658 """
659 args = list(func.param_names)
660 if func.vararg_name:
661 args.append(f"*{func.vararg_name}")
662 kwargs_parts = [f'"{name}": {name}' for name in func.kwonly_names]
663 if func.kwarg_name:
664 kwargs_parts.append(f"**{func.kwarg_name}")
665 kwargs_literal = "{" + ", ".join(kwargs_parts) + "}" if kwargs_parts else "{}"
666 # The args body is spliced into a tuple literal; a single element needs a
667 # trailing comma (also for a lone ``*vararg``) so it stays a tuple rather
668 # than parenthesized scalars / a generator expression.
669 args_body = ", ".join(args) + ("," if len(args) == 1 else "")
670 return (
671 " args, kwargs = self._hp_boundary.redistribute_inputs(\n"
672 f" (({args_body}), {kwargs_literal}),\n"
673 " )"
674 )
677def _render_output_redistribute(deferred_bias: bool = False) -> str:
678 """Render the output-side boundary exit (no signature to bind).
680 Mirrors the native production path: the compute output remains local and
681 ``PrecompiledBoundary`` executes a lowered collective directly when
682 possible. An unlowerable transition performs any required temporary
683 DTensor wrapping internally and still returns a local tensor.
685 ``deferred_bias`` (D-22) appends the exit hook after redistribution: the
686 rowwise bias was suppressed inside the region so the exit reduction sees
687 a pure matmul contribution, and the hook re-adds it exactly once on the
688 reduced local output — Megatron RowParallelLinear semantics. The hook is
689 bound at install time (``module._hp_deferred_bias``); the emitted call
690 is conditional on the plan declaring deferred params, so the two sides
691 cannot disagree about whether a boundary defers.
692 """
693 lines = [
694 " outputs = self._hp_boundary.redistribute_outputs(outputs)",
695 ]
696 if deferred_bias:
697 lines.append(" outputs = self._hp_deferred_bias(outputs)")
698 return "\n".join(lines)
701def _ordered_notes(notes: dict[str, str]) -> list[str]:
702 """Declarative annotation lines for one side, in deterministic name order."""
703 return [f" # {notes[name]}" for name in sorted(notes)]
706def _render_plain_redistribute(
707 input_redist: str, out_redist: str, form: BoundaryForm
708) -> str:
709 """Generic boundary: annotated redistribute around the original forward.
711 The comment blocks are the plan's own declaration: one line per
712 non-identity transition (``name: axis S(1) -> R ==> all_gather(dim=1)`` or
713 ``==> dtensor redistribute`` for transitions the compiled plan's DTensor
714 fallback owns), so a reader sees what the boundary communicates without
715 re-deriving the plan.
716 """
717 return "\n".join([
718 " # [HYPER BOUNDARY] input redistribution",
719 *_ordered_notes(form.in_notes),
720 input_redist,
721 "",
722 " # [HYPER BOUNDARY] original forward body (now ``_forward_impl``)",
723 " outputs = self._forward_impl(*args, **kwargs)",
724 "",
725 " # [HYPER BOUNDARY] output redistribution",
726 *_ordered_notes(form.out_notes),
727 out_redist,
728 "",
729 " return outputs",
730 ])
733# ---------------------------------------------------------------------------
734# static tp_collective template
735# ---------------------------------------------------------------------------
737def _render_tp_collective_forward(
738 source_text: str,
739 func: FunctionInfo,
740 form: BoundaryForm,
741 deferred_bias: bool = False,
742) -> str:
743 """Static boundary: bare TP operators, semantics of ``RedistOp.execute``.
745 The rendered body is instruction-equivalent to the generic engine with the
746 TP lowerer attached:
748 - identity inputs — ``self._hp_tp.to_local(name)`` (a DTensor unwraps to
749 its local shard, anything else passes through);
750 - TP-lowerable inputs/outputs — ``self._hp_tp.<kind>(...)`` (the bound
751 operator unwraps a DTensor input and runs the collective on the tp
752 group; ``None`` passes through, like the compiled plan's ``None`` skip);
753 - ``_forward_impl`` is called with every declared parameter in signature
754 order (the static forward re-binds names, so no ``*args`` re-pass).
756 ``self._hp_tp`` is bound by ``hp_install_boundaries`` only after the
757 runtime re-validates the emitted form against the live mesh; on mismatch
758 the runtime replaces this forward with the generic engine, so the baked
759 calls never run unvalidated.
761 ``deferred_bias`` (D-22) appends ``self._hp_deferred_bias(outputs)``
762 after the output ops — the Partial(sum) -> R all_reduce above is exactly
763 the exit reduction the deferral waits for, so the bias re-add lands once,
764 on the reduced output.
765 """
766 lines = [
767 " # [HYPER TP-COLLECTIVE] statically lowered boundary — the runtime",
768 " # re-validates this form at install time and replaces the forward",
769 " # with the generic engine on any mismatch.",
770 ]
771 for name in form.in_declared:
772 op = form.in_ops.get(name)
773 if op is None:
774 note = form.in_notes.get(name) or f"{name}: identity ==> to_local"
775 lines.append(f" # {note}")
776 lines.append(f" {name} = self._hp_tp.to_local({name})")
777 else:
778 lines.append(f" # {form.in_notes[name]}")
779 lines.append(_render_tp_call(name, op))
780 lines.append("")
781 lines.append(
782 " # [HYPER TP-COLLECTIVE] original forward body (now ``_forward_impl``)"
783 )
784 lines.append(_render_forward_call(func))
785 tuple_output = not returns_single_value(source_text, func)
786 for name, op in form.out_ops.items():
787 lines.append(f" # {form.out_notes[name]}")
788 lines.extend(_render_output_tp_call(name, op, form, tuple_output))
789 if deferred_bias:
790 lines.append("")
791 lines.append(
792 " # [HYPER D-22] deferred bias re-added once, after the exit reduction"
793 )
794 lines.append(" outputs = self._hp_deferred_bias(outputs)")
795 lines.append("")
796 lines.append(" return outputs")
797 return "\n".join(lines)
800def _render_tp_call_expr(target: str, op: TransitionOp) -> str:
801 """One bare-operator call expression (``self._hp_tp.<kind>(...)``)."""
802 if op.kind == "all_reduce":
803 return f"self._hp_tp.all_reduce({target})"
804 return f"self._hp_tp.{op.kind}({target}, dim={op.tensor_dim})"
807def _render_tp_call(name: str, op: TransitionOp) -> str:
808 """One input-side bare-operator statement (rebinding the named input)."""
809 return f" {name} = {_render_tp_call_expr(name, op)}"
812def _render_output_tp_call(
813 name: str, op: TransitionOp, form: BoundaryForm, tuple_output: bool
814) -> list[str]:
815 """Output-side bare-operator statements.
817 The common single-tensor case rewrites ``outputs`` directly. For tuple
818 outputs, the frozen plan names the first declared output as index 0;
819 rebuild the original container after applying the collective to that item.
820 """
821 out_names = form.out_names
822 if tuple_output and (name not in out_names or out_names.index(name) == 0):
823 return [
824 f" _hp_output_0 = {_render_tp_call_expr('outputs[0]', op)}",
825 " outputs = (_hp_output_0, *outputs[1:])",
826 ]
827 return [f" outputs = {_render_tp_call_expr('outputs', op)}"]
830def _render_forward_call(func: FunctionInfo) -> str:
831 """The ``self._forward_impl(...)`` call, every parameter re-passed by name.
833 Positional parameters go positionally, keyword-only ones by keyword — the
834 call is a pure re-pass of the static forward's own signature, so defaults
835 and call-site keyword usage both survive the rewrite.
836 """
837 args = list(func.param_names) + [f"{name}={name}" for name in func.kwonly_names]
838 if func.kwarg_name:
839 args.append(f"**{func.kwarg_name}")
840 if not args:
841 return " outputs = self._forward_impl()"
842 single = f" outputs = self._forward_impl({', '.join(args)})"
843 if len(single) <= 100:
844 return single
845 lines = [" outputs = self._forward_impl("]
846 lines.extend(f" {arg}," for arg in args)
847 lines.append(" )")
848 return "\n".join(lines)
851def _render_local_compute(
852 local: Any, input_redist: str, out_redist: str, form: BoundaryForm
853) -> str:
854 """Local-region boundary (shape 2, EP MoE): compute fn on local tensors, then exit.
856 The factory-named attribute is ``functools.partial(compute_fn, module)``
857 bound by the emitted constructor patch. The applier calls it as ``compute_fn(*args,
858 **kwargs)`` (``_wrap_local_region_forward``), so the generated forward
859 forwards its own re-distributed ``*args, **kwargs`` through. The local
860 Production keeps the compute output local so the lowered boundary
861 operations can run on plain tensors. If a transition cannot be lowered, the
862 installed boundary owns the necessary fallback wrapping.
864 The per-transition declarative notes (``form.in_notes`` / ``out_notes``)
865 annotate the redistribution steps exactly like the generic and
866 tp-collective templates; identity transitions carry no note, so an
867 identity in/out region emits no noise.
868 """
869 compute_name = _callable_name(local)
870 binding_name = local_compute_binding_name(local)
871 return "\n".join([
872 " # [HYPER LOCAL REGION] boundary entry",
873 *_ordered_notes(form.in_notes),
874 input_redist,
875 " # [HYPER LOCAL REGION] run the region compute fn on local tensors",
876 f" # {compute_name} is bound to self by generated source.",
877 f" outputs = self.{binding_name}(*args, **kwargs)",
878 "",
879 " # [HYPER LOCAL REGION] boundary exit",
880 *_ordered_notes(form.out_notes),
881 out_redist,
882 "",
883 " return outputs",
884 ])
887def _callable_name(value: Any) -> str:
888 """Human name for an injection's callable (dict target / string / fn)."""
889 if isinstance(value, dict):
890 return (
891 value.get("_target_")
892 or value.get("name")
893 or value.get("target")
894 or repr(value)
895 )
896 if isinstance(value, str):
897 return value
898 return getattr(value, "__name__", repr(value))
901# ---------------------------------------------------------------------------
902# field reads / validation
903# ---------------------------------------------------------------------------
905def _plan_field(plan: Any, name: str) -> Any:
906 """Read a frozen-plan field from a dict or a ``CodegenMeta``."""
907 if isinstance(plan, dict):
908 return plan.get(name)
909 return getattr(plan, name, None)
912def _validate_no_circular(edits: list[TextEdit]) -> list[str]:
913 """Detect edits authored against different texts (a soft sanity check).
915 The lowerer computes all edits against one ``source_text``; the index is
916 built once, so a second-pass author (an inner wrapper that itself carries a
917 body rewrite) would need to re-index. We fail fast rather than splice
918 stale offsets.
919 """
920 problems: list[str] = []
921 seen: dict[tuple[int, int], str] = {}
922 for edit in edits:
923 key = (edit.start, edit.end)
924 if key in seen:
925 problems.append(
926 f"[{edit.start},{edit.end}) (was {seen[key]}, now {edit.replacement[:30]!r})"
927 )
928 else:
929 seen[key] = edit.replacement[:30]
930 return problems
933__all__ = [
934 "TP_FORM_ATTRIBUTE",
935 "TP_FORM_MARKER",
936 "iter_emitted_forms",
937 "lower_forward_boundaries",
938 "resolve_emitted_form",
939]