Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / rule_resolver.py: 97%
306 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +0800
1# Copyright 2025-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"""rule_resolver: plan_overrides merge/insert machinery (05 §3.6.7 + the
16unification rework), contract-field normalization, and injection context
17filling (``fill_context_kwargs``).
19Everything here is a pure (spec, model) → spec transformation: no mesh
20objects, no runtime state.
21"""
23import copy
24import fnmatch
25import functools
26import inspect
27import logging
28from typing import Any, Dict, List, Optional, Tuple
30from hyper_parallel.core.dtensor.placement_types import Placement
31from hyper_parallel.distributed.plan import ShardingPlan
32from hyper_parallel.distributed.recipe_spec import (
33 INNER_WRAPPER,
34 LOCAL_COMPUTE,
35 InjectionMeta,
36 MeshAxisName,
37 ModuleShardingSpec,
38 _normalize_out_fields,
39 require_injection_meta,
40 resolve_placements,
41)
43logger = logging.getLogger(__name__)
46def _last_segment(fqn: str) -> str:
47 return fqn.rsplit(".", 1)[-1].lower() if fqn else ""
50_GLOB_CHARS = ("*", "?", "[")
53def _is_glob_key(key: str) -> bool:
54 return any(c in key for c in _GLOB_CHARS)
57def _can_insert_glob(user_spec: ModuleShardingSpec) -> bool:
58 """Return whether a glob declares a concrete boundary contract."""
59 return any(
60 isinstance(getattr(user_spec, attr), dict)
61 for attr in ("params", "in_src", "in_dst", "out_src", "out_dst")
62 )
65def _merge_plan_overrides(plan_overrides, plan: ShardingPlan, model, *,
66 derive: bool = True) -> None:
67 """Unified override pass, executed before Phase 5 and the D-14 checks.
69 Three modes:
71 - **merge** (key hits an existing boundary — derived or previously
72 inserted): UNSET contract fields (``None``: ``params`` /
73 ``in_src`` / ``in_dst`` / ``out_src`` / ``out_dst`` /
74 ``out_names``) INHERIT, set ones replace at field granularity —
75 **an explicit empty dict {} is a SET value** (explicit "no
76 sharding / no contract", 2026-08-05 "unset inherits, written is
77 honored"); the
78 sentinels ``"auto"``/``"none"`` mean inherit/clear explicitly;
79 injection fields
80 (``local_compute_fn`` / ``inner_target`` / ``inner_wrapper``
81 non-None, ``region_dispatch=False``) always win; internal flags
82 (``_ep_size`` / ``_ep_stack`` / ``_needs_cp_attn`` /
83 ``_is_terminal``) always inherit — they are
84 planner/applier-owned metadata, not user contracts. This is how
85 CP/EP compute injection is declared: an injection-fields-only
86 spec (``ModuleShardingSpec(local_compute_fn=...)``) inherits the
87 whole derived contract;
88 - **insert** (exact key misses every boundary): the spec is
89 deep-copied and inserted as-is; at least one contract field must
90 be declared (an explicit {} counts — the pure I/O-stitch
91 boundary) — an override with EVERYTHING unset fails fast
92 ("no template/boundary matched"). Sentinels are rejected (nothing
93 to inherit/clear). Nesting (ancestor/descendant
94 FQNs) is **allowed** since D-14 (05 §13), subject only to the
95 param-uniqueness invariant (``_check_param_uniqueness``);
96 - **glob keys** (containing ``*``/``?``/``[``): merge-applied to
97 every matching derived/previously inserted boundary (fnmatchcase,
98 ``*`` spans dots). A glob declaring at least one concrete contract
99 dict additionally expands against the final model's
100 ``named_modules`` and inserts every unmatched module using the exact
101 insert validation above. Injection-only/attribute-only globs remain
102 merge-only for existing boundaries and fail fast if they match a real
103 model module without a boundary, so a broad partial override cannot
104 silently ignore or create communication boundaries.
106 Notes:
107 - exact keys must exist in the model's ``named_modules`` (typo
108 fail-fast; PP scenarios plan each single-part model separately);
109 - ``out_src``/``out_dst`` scalar shorthand is normalized;
110 - user spec objects are never mutated (merge reads them, insert
111 deep-copies them) — plan() can be called repeatedly.
112 """
113 entries: List[Tuple[str, ModuleShardingSpec, str]] = [
114 (fqn, spec, "plan_overrides")
115 for fqn, spec in plan_overrides.items()
116 ]
117 if not entries:
118 return
120 module_names = [name for name, _ in model.named_modules()]
121 _validate_override_entries(entries, plan, set(module_names))
123 for key, user_spec, source in entries:
124 if _is_glob_key(key):
125 _apply_glob_override(
126 plan, key, user_spec, source, model, module_names,
127 derive=derive,
128 )
129 elif key in plan.modules:
130 _warn_dropped_params(
131 source, key, key, plan.modules[key], user_spec)
132 _merge_into(plan.modules[key], user_spec)
133 logger.info("%s: merge into the spec of module %s",
134 source, key)
135 else:
136 _insert_spec(plan, key, user_spec, source, model,
137 derive=derive)
140def _validate_override_entries(entries, plan, module_name_set) -> None:
141 """Fail fast on wrong spec types, typo'd axes, and unknown exact FQNs."""
142 for key, user_spec, source in entries:
143 if not isinstance(user_spec, ModuleShardingSpec):
144 raise TypeError(
145 f"{source}[{key!r}] must be a ModuleShardingSpec, "
146 f"got {type(user_spec).__name__}"
147 )
148 _validate_override_axes(key, user_spec, source, plan)
149 if not _is_glob_key(key) and key not in module_name_set:
150 raise ValueError(
151 f"{source} FQN not found in the model's "
152 f"named_modules: {key!r} (check spelling; in PP "
153 f"scenarios plan each single-part model separately)"
154 )
157def _raise_merge_only_glob_error(key, user_spec, source, model, model_hits) -> None:
158 """Fail fast when an insert-incapable glob matches unbounded modules."""
159 configured_fields = [
160 attr
161 for attr in (
162 "local_compute_fn", "inner_target", "inner_wrapper",
163 "inner_out_src", "region_dispatch", "out_names",
164 "tp_divide_attrs",
165 )
166 if getattr(user_spec, attr) is not None
167 ]
168 preview = model_hits[:8]
169 omitted = len(model_hits) - len(preview)
170 matched = f"{preview}"
171 if omitted:
172 matched += f" (and {omitted} more)"
173 example_fqn = model_hits[0]
174 raise ValueError(
175 f"{source}[{key!r}] matched {len(model_hits)} module(s) "
176 "in the final model that have no planner-derived "
177 f"boundary: {matched}. The override declares only "
178 f"merge-only fields {configured_fields or ['none']} and "
179 "no concrete params/in_src/in_dst/out_src/out_dst "
180 "contract, so it cannot create those boundaries. This "
181 "usually means the modules did not match a built-in "
182 "template while the YAML assumed there was a boundary "
183 "to merge into. Fix it by either: (1) declaring the "
184 "module's complete parameter and I/O contract in this "
185 "glob (a concrete contract makes it insert-capable); "
186 "(2) adding a preceding full-contract glob and keeping "
187 "this entry as a later injection-only merge; or (3) if "
188 "the module should be built in, fixing its sharding role/"
189 "template recognition. Narrow the glob if these modules "
190 "were not intended targets.\nSuggested draft for the "
191 f"first unmatched module {example_fqn!r} (placements are "
192 "placeholders; replace the exact match with the original "
193 "glob only when all matched modules share this contract):\n"
194 + _suggest_insert_skeleton(model, example_fqn)
195 )
198def _apply_glob_override(plan: ShardingPlan, key: str,
199 user_spec: ModuleShardingSpec, source: str,
200 model, module_names, *,
201 derive: bool = True) -> None:
202 """Merge-apply one glob override and insert matching unbounded modules."""
203 hits = [fqn for fqn in plan.modules
204 if fnmatch.fnmatchcase(fqn, key)]
205 model_hits = [
206 fqn for fqn in module_names
207 if fqn not in plan.modules and fnmatch.fnmatchcase(fqn, key)
208 ]
209 inserted = []
210 can_insert = _can_insert_glob(user_spec)
211 if model_hits and not can_insert:
212 _raise_merge_only_glob_error(key, user_spec, source, model, model_hits)
213 if can_insert:
214 for fqn in model_hits:
215 _insert_spec(
216 plan,
217 fqn,
218 user_spec,
219 source,
220 model,
221 derive=derive,
222 )
223 inserted.append(fqn)
224 if not hits and not inserted:
225 logger.warning(
226 "%s match=%r hit no boundary spec or model module — check the "
227 "spelling (plan boundaries: %s)",
228 source, key, sorted(plan.modules)[:8],
229 )
230 for fqn in hits:
231 _warn_dropped_params(
232 source, key, fqn, plan.modules[fqn], user_spec)
233 _merge_into(plan.modules[fqn], user_spec)
234 logger.info("%s: merge into %s (glob %r)",
235 source, fqn, key)
238_CONTRACT_FIELDS = ("params", "in_src", "in_dst",
239 "out_src", "out_dst", "out_names")
240# String sentinels accepted on the plan_overrides input side for contract
241# fields (resolved at merge time only; they never reach the plan output):
242# "auto" — explicitly inherit the derived value (derive per template;
243# synonymous with the default unset value, self-documenting)
244# "none" — explicitly clear (params/in_src/in_dst → {}, out_* → None)
245_CONTRACT_SENTINELS = ("auto", "none")
248def _iter_named_placements(spec: ModuleShardingSpec):
249 """Yield (attr, name, named) for every concrete NamedPlacement in an
250 override spec (skips sentinel strings/None/empty; out_* scalar
251 shorthand yields the whole field as one NamedPlacement)."""
252 for attr in ("params", "in_src", "in_dst", "out_src", "out_dst"):
253 value = getattr(spec, attr)
254 if not value or isinstance(value, str):
255 continue
256 if not all(isinstance(v, dict) for v in value.values()):
257 yield attr, "output", value # out_* scalar shorthand
258 else:
259 for name, named in value.items():
260 yield attr, name, named
263def _validate_override_axes(key, user_spec, source, plan) -> None:
264 """Fail fast on typo'd placement axes / non-Placement values.
266 ``resolve_placements`` fills missing axes with Replicate() — so a
267 typo'd axis (``{"tp2": Shard(0)}``) would otherwise be silently
268 IGNORED. Allowed axes = the plan's mesh dims ∪ the canonical
269 ``MeshAxisName`` values (canonical-but-absent axes, e.g. CP
270 placements declared on a tp-only mesh, are tolerated — templates
271 declare all canonical dims and resolve_placements picks the mesh's
272 subset; "ep" is the virtual TP-extend-EP axis). Anything outside
273 both sets is a typo → fail fast. Placement values must already be
274 Placement objects (the YAML string DSL is parsed at desugar time).
275 """
276 allowed = ({str(a) for a in plan.mesh_dim_names}
277 | {axis.value for axis in MeshAxisName})
278 for attr, name, named in _iter_named_placements(user_spec):
279 for axis, placement in named.items():
280 if not isinstance(placement, Placement):
281 raise TypeError(
282 f"{source}[{key!r}]: the value of axis "
283 f"{axis!r} in contract field {attr}[{name!r}] must "
284 f"be a Placement object (Shard(N)/"
285 f"Replicate()/Partial(); the YAML string DSL is "
286 f"parsed into objects at desugar time), "
287 f"got {type(placement).__name__} "
288 f"{placement!r}")
289 # MeshAxisName is a str subclass (hash/eq consistent with a
290 # plain str), so use the raw value for the membership test —
291 # do NOT str(axis) (the enum __str__ would produce
292 # "MeshAxisName.TP")
293 if axis not in allowed:
294 raise ValueError(
295 f"{source}[{key!r}] contract field "
296 f"{attr}[{name!r}] uses the unknown axis {axis!r} — "
297 f"legal axes = mesh axes "
298 f"{sorted(str(a) for a in plan.mesh_dim_names)} ∪ "
299 f"canonical axes "
300 f"{sorted(axis.value for axis in MeshAxisName)}. "
301 f"An unknown axis is silently ignored by "
302 f"resolve_placements, so fail fast (suspected typo)")
305def _merge_contract_field(derived: ModuleShardingSpec,
306 user_spec: ModuleShardingSpec, attr: str) -> None:
307 """Merge one contract field: "unset inherits, written is honored" (2026-08-05).
309 Precedence: ``None`` (unset) / ``"auto"`` → inherit derived;
310 ``"none"`` → explicit clear (a readable alias for the empty value);
311 a concrete dict — **including the empty dict {}** — replaces at field
312 granularity ({} = explicit "no sharding / no contract", the ViT
313 ``params={}`` pattern).
314 """
315 value = getattr(user_spec, attr)
316 if isinstance(value, str):
317 if value == "auto":
318 return # explicit inherit (same as the default unset)
319 if value == "none":
320 setattr(derived, attr,
321 {} if attr in ("params", "in_src", "in_dst") else None)
322 return # explicit clear
323 raise ValueError(
324 f"the string value of plan_overrides contract field {attr} "
325 f"only accepts the sentinels "
326 f"{_CONTRACT_SENTINELS} ('auto'=inherit the "
327 f"template-derived value, 'none'=explicit clear), "
328 f"got {value!r}")
329 if value is None: # undeclared → inherit
330 return
331 setattr(derived, attr, copy.deepcopy(value)) # including {}: explicit empty (clear)
334def _warn_dropped_params(source, key, fqn, derived, user_spec) -> None:
335 """Visibility safeguard: a field-granularity replacement of ``params``
336 during merge strips the derived sharding from every parameter not
337 covered by the override (they stay replicated) — possibly an
338 unintended typo, so a WARNING lists the dropped entries."""
339 user_params = user_spec.params
340 if (isinstance(user_params, dict) and user_params
341 and derived.params):
342 dropped = sorted(set(derived.params) - set(user_params))
343 if dropped:
344 logger.warning(
345 "%s[%r] merge into %s: the field-granularity params "
346 "replacement strips the derived sharding from %d "
347 "parameter(s), which will stay replicated: %s — if "
348 "unintended, write the derived values in too "
349 "(field-granularity replacement, no per-key merge); if "
350 "the de-sharding is intentional, ignore this warning "
351 "(params={} or 'none' explicitly clears all)",
352 source, key, fqn, len(dropped), dropped)
355def _merge_into(derived: ModuleShardingSpec,
356 user_spec: ModuleShardingSpec) -> None:
357 """Merge one user spec into an existing boundary spec (in place).
359 Contract fields: None/"auto" → inherit, "none" → clear, concrete
360 dict (including {}) → replace (field granularity); injection fields
361 win when set; internal flags always inherit from *derived*.
362 """
363 for attr in _CONTRACT_FIELDS:
364 _merge_contract_field(derived, user_spec, attr)
365 for attr in ("local_compute_fn", "inner_target", "inner_wrapper",
366 "inner_out_src", "region_dispatch", "tp_divide_attrs"):
367 value = getattr(user_spec, attr)
368 if value is not None:
369 setattr(derived, attr, value)
370 _normalize_out_fields(derived)
373def _normalize_contract_fields(plan: ShardingPlan) -> None:
374 """Plan output normalization: None (undeclared) of
375 params/in_src/in_dst → {}.
377 "Unset inherits, written is honored" is input-side semantics; specs
378 inside the plan always hold concrete values, so downstream consumers
379 (the applier, the D-14 checks, etc.) need no None branches.
380 """
381 for spec in plan.modules.values():
382 for attr in ("params", "in_src", "in_dst"):
383 if getattr(spec, attr) is None:
384 setattr(spec, attr, {})
387def _is_required_keyword_parameter(parameter: inspect.Parameter) -> bool:
388 """Return whether a parameter is a required named input."""
389 return (parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD,
390 inspect.Parameter.KEYWORD_ONLY)
391 and parameter.default is inspect.Parameter.empty)
394def _suggest_insert_skeleton(model, fqn: str) -> str:
395 """Derive a draft contract skeleton from the module's forward
396 signature (input names) and direct parameters — turns "write a
397 contract from scratch" into "edit a draft". Best-effort: degrades
398 to a generic skeleton when the module/signature is unavailable."""
399 try:
400 module = dict(model.named_modules()).get(fqn)
401 except Exception: # pylint: disable=broad-exception-caught # best-effort hint only
402 module = None
403 in_names = ["hidden_states"]
404 param_names: List[str] = []
405 if module is not None:
406 try:
407 sig = inspect.signature(module.forward)
408 names = [
409 p.name for p in sig.parameters.values()
410 if _is_required_keyword_parameter(p)
411 ]
412 if names:
413 in_names = names
414 except (TypeError, ValueError):
415 pass
416 param_names = [n for n, _ in module.named_parameters(recurse=False)]
417 axis = "tp" # draft placeholder axis — adjust to the actual topology
418 lines = [f' - match: "{fqn}"']
419 if param_names:
420 lines.append(" params: # choose a shard dim per parameter "
421 "(explicit {} = this boundary shards no params)")
422 lines.extend(f" {n}: {{{axis}: \"shard(0)\"}}"
423 for n in param_names)
424 else:
425 lines.append(" params: {} # no direct parameters / pure "
426 "I/O-stitch boundary")
427 in_entries = ", ".join(f'{n}: {{{axis}: "shard(1)"}}'
428 for n in in_names)
429 lines.append(f" in_src: {{{in_entries}}} "
430 "# entry status quo = upstream exit layout")
431 lines.append(f" in_dst: {{{in_entries}}} "
432 "# differing from in_src inserts communication")
433 lines.append(f' out_src: {{output: {{{axis}: "replicate"}}}}'
434 " # multi-output modules: use multiple keys and add "
435 "out_names")
436 lines.append(f' out_dst: {{output: {{{axis}: "replicate"}}}}')
437 return "\n".join(lines)
440def _insert_spec(plan: ShardingPlan, fqn: str,
441 user_spec: ModuleShardingSpec, source: str,
442 model=None, derive: bool = True) -> None:
443 """Insert a fully self-declared spec for a non-derived boundary."""
444 for attr in _CONTRACT_FIELDS:
445 value = getattr(user_spec, attr)
446 if isinstance(value, str):
447 reason = (
448 "derive=False: template derivation is disabled "
449 "entirely, so the plan holds no derived values — "
450 if not derive else
451 "insert (no derived boundary was hit) — ")
452 raise ValueError(
453 f"{source}[{fqn!r}] contract field {attr}={value!r} is "
454 f"meaningless: {reason}"
455 "the 'auto' (inherit the derivation) / 'none' (clear "
456 "the inherited value) sentinels only apply to derived "
457 "boundaries hit by a merge; with no derived value there "
458 "is nothing to inherit/clear. Declare explicitly: a "
459 "concrete dict for the sharding/contract, an explicit "
460 "empty {} for a boundary that shards no params (params "
461 "stay replicated) / has no such contract")
462 contract_values = (getattr(user_spec, attr) for attr in ("params", "in_src", "in_dst", "out_src", "out_dst"))
463 if all(value is None for value in contract_values):
464 hint = ""
465 if model is not None:
466 hint = (
467 "\nSuggested draft (derived from the module forward "
468 "signature/direct parameters; placements are "
469 "placeholders — fix them per the layout semantics):\n"
470 + _suggest_insert_skeleton(model, fqn))
471 derive_note = (
472 "derive=False has disabled template derivation — every "
473 "override is an insert and there are no derived values to "
474 "inherit; " if not derive else "")
475 raise ValueError(
476 f"{source}[{fqn!r}] hit no planner-derived boundary and "
477 "leaves params and the I/O contract all undeclared — "
478 + derive_note +
479 "empty-field inheritance (merge) only applies to derived "
480 "boundaries; inserting a new boundary requires a fully "
481 "self-declared contract (05 §3.6.7 / D-14; an explicit "
482 "empty {} is also a valid declaration, e.g. the pure "
483 "I/O-stitch boundary with params={}), or check the fqn "
484 "spelling"
485 + hint)
486 spec = copy.deepcopy(user_spec)
487 _normalize_out_fields(spec)
488 logger.info("%s: insert the spec of module %s", source, fqn)
489 plan.modules[fqn] = spec
492def fill_context_kwargs(meta: InjectionMeta,
493 context: Dict[str, Any],
494 configured: Optional[Dict[str, Any]] = None,
495 *,
496 source: str = "") -> Dict[str, Any]:
497 """Collect the framework-filled kwargs for the context keys declared in
498 ``meta`` (minimal, no hidden behavior).
500 - The filled set equals the declared set, no more and no less (each fill
501 is logged at INFO);
502 - Context keys are framework-reserved names: a user configuring a
503 same-named key in Target/YAML triggers fail-fast (context parameters
504 have no defaults and the YAML resolver never back-fills them, so any
505 reserved key appearing in ``configured`` was definitely written
506 explicitly by the user).
508 Args:
509 meta: The injection metadata whose declared context keys are filled.
510 context: The framework-provided context values keyed by name.
511 configured: User-configured keys, checked against reserved names.
512 source: Call-site label prepended to log and error messages.
514 Returns:
515 The kwargs dict mapping each declared context key to its
516 framework-provided value.
518 Raises:
519 ValueError: If ``configured`` contains framework-reserved context
520 keys.
521 """
522 configured = configured or {}
523 reserved = sorted(set(configured) & meta.context)
524 if reserved:
525 raise ValueError(
526 f"{source}: framework-reserved context keys {reserved} were "
527 "configured -- context is filled by the framework according to "
528 "the declaration and is not configurable; your config keys must "
529 "not collide with reserved names")
530 kwargs = {}
531 for key in sorted(meta.context):
532 kwargs[key] = context[key]
533 logger.info("%s: context key %s filled by framework (%s)",
534 source, key, type(context[key]).__name__)
535 return kwargs
538# ────────────────────────────────────────────────────────────────────────────
539# Injection resolution: spec/Target/registry → concrete callables and output
540# rules (pure resolution — no module mutation)
541# ────────────────────────────────────────────────────────────────────────────
543def _is_delayed_target(obj) -> bool:
544 """Duck-typed check for a config Target (avoids a components -> trainer
545 import): a Target carries ``build()`` and ``_target_`` and is itself NOT
546 the compute/wrapper callable — it must be built at apply time."""
547 return hasattr(obj, "build") and hasattr(obj, "_target_")
550def _check_target_config_keys(target, kind):
551 """Fail fast on configured Target kwargs the callable would never bind.
553 Target kwargs are bound BY NAME at build time
554 (``{**configured, **runtime}`` -> ``fn(**kwargs)``). When the target
555 callable accepts VAR_KEYWORD (e.g. ``**_context``, there to tolerate the
556 framework-filled generic context), a misspelled configured key is
557 swallowed SILENTLY — the user's value never takes effect and the
558 framework may even auto-fill the intended parameter instead. Guard: any
559 configured key that is not an explicitly declared (keyword-bindable)
560 parameter fails fast with the valid parameter names. A callable whose
561 signature cannot be introspected skips the check (the call itself will
562 surface any mismatch).
563 """
564 configured = getattr(target, "_kwargs", None)
565 if not configured:
566 return
567 fn = getattr(target, "_target_", None)
568 try:
569 params = inspect.signature(fn).parameters
570 except (TypeError, ValueError):
571 return
572 bindable = set()
573 for name, parameter in params.items():
574 if parameter.kind in (inspect.Parameter.POSITIONAL_OR_KEYWORD,
575 inspect.Parameter.KEYWORD_ONLY):
576 bindable.add(name)
577 unknown = sorted(set(configured) - bindable)
578 if unknown:
579 raise ValueError(
580 f"{kind} Target {getattr(target, '_target_path', fn)!r} "
581 f"configures undeclared keys {unknown} — Target kwargs are "
582 f"bound by name to the keyword parameters of the target "
583 f"callable; these keys are not in the explicit parameter list "
584 f"of {getattr(fn, '__name__', fn)} and would be silently "
585 f"swallowed by **kwargs without taking effect (suspected "
586 f"typo). Valid parameters: {sorted(bindable) or '(none)'}")
589def _require_region_dispatch(spec, *, source):
590 """Injection discipline: declaring an injection (local_compute_fn /
591 inner_wrapper) requires an explicit region_dispatch (no default — pass
592 True for a dispatchable pure-ops injection, False for one containing
593 communication primitives / custom kernels; the tutorials and examples
594 explain the reasons case by case)."""
595 rd = getattr(spec, "region_dispatch", None)
596 has_injection = (getattr(spec, "local_compute_fn", None) is not None
597 or getattr(spec, "inner_wrapper", None) is not None)
598 if rd is None:
599 if has_injection:
600 raise ValueError(
601 f"{source}: an injection is declared but region_dispatch "
602 "is not explicitly declared (no default) — if the "
603 "injected fn is pure standard ops that validate can "
604 "dispatch through (fused-op/scripting-style optimizations) "
605 "→ region_dispatch=True (in-region strategy propagation + "
606 "true out_src validation); if it contains communication "
607 "primitives/custom kernels (CP K/V all-gather, EP "
608 "all-to-all, quantized GEMM, etc.) → "
609 "region_dispatch=False (skeleton/adapter black-box hosted "
610 "local execution + declarative rewrap)")
611 return
612 if rd is True and not has_injection:
613 raise ValueError(
614 f"{source}: region_dispatch=True but no injection is declared "
615 "— an ordinary boundary's forward dispatches through natively "
616 "(the axiomatic default), so this declaration is redundant; "
617 "remove it")
620def _build_local_compute_factory(factory, module, mesh, mesh_dim_names,
621 expert_mesh, *, configured=None, source):
622 """Build a ``@local_compute`` factory into the region compute_fn (apply-time).
624 The factory is invoked ONCE with the framework context filled BY NAME per
625 its declared context (``meta.context``): the mesh family ``mesh`` /
626 ``tp_mesh`` / ``cp_mesh`` / ``ep_mesh`` (mandatory declarations; None when
627 the axis is inactive), plus the optional anchor ``module``. Context keys
628 are RESERVED — a user-configured same-name key fails fast
629 (fill_context_kwargs); every fill is logged at INFO. Behavior choices
630 (routing, layouts, ...) are written INTO the factory — config keys carry
631 data only, never functions.
632 The returned compute fn is validated against the module's forward
633 signature (validate_local_compute_signature: params must match the
634 original forward) and bound to *module* by the caller.
635 """
636 # Lazy imports: applier and forward_rewriter both import rule_resolver
637 # at module level — importing them here at module level would cycle.
638 from hyper_parallel.distributed._builder.applier import ( # pylint: disable=C0415
639 _get_cp_submesh,
640 _get_tp_submesh,
641 )
642 from hyper_parallel.distributed._builder.forward_rewriter import ( # pylint: disable=C0415
643 validate_local_compute_signature,
644 )
645 meta = require_injection_meta(factory, LOCAL_COMPUTE, source=source)
646 context = {
647 "module": module,
648 "mesh": mesh,
649 "tp_mesh": _get_tp_submesh(mesh, mesh_dim_names),
650 "cp_mesh": _get_cp_submesh(mesh, mesh_dim_names),
651 "ep_mesh": expert_mesh,
652 }
653 build_kwargs = fill_context_kwargs(
654 meta, context, configured or {}, source=source)
655 # {**configured, **context}: same binding order as Target.build (context
656 # keys are reserved names; fill_context_kwargs already rejects same-name
657 # keys in configured)
658 compute_fn = factory(**{**(configured or {}), **build_kwargs})
659 if not callable(compute_fn):
660 raise TypeError(
661 f"{source} returned {type(compute_fn).__name__}, not callable — "
662 "the @local_compute factory must return the region compute fn "
663 "fn(module, *local_args) (e.g. hyper_parallel.distributed.expert_parallel."
664 "recipes.qwen2moe_ep_compute_fn)")
665 validate_local_compute_signature(
666 compute_fn, module.forward,
667 owner=f"{source!r} on {type(module).__name__}")
668 return compute_fn
671def _resolve_local_compute_fn(module, spec, mesh, mesh_dim_names,
672 expert_mesh):
673 """Resolve the compute_fn of the local region (**single resolution chain**, 05 §4.4.3).
675 Whether a module takes the local-region skeleton is derived by this chain
676 (a non-None return means it does) -- the gate is not a stored bool but the
677 resolution result. Since the explicit-injection rework the built-in EP
678 auto-injection link is REMOVED; the remaining sources:
679 1. spec.local_compute_fn: a user-defined FACTORY — single form since
680 2026-08-10 (the direct compute-fn form was retired: every injection
681 fn is invoked once at apply time with the mesh family filled by the
682 framework, use-it-or-not being the user's choice — same discipline as
683 @inner_wrapper). Accepted shapes: a ``@local_compute``-decorated
684 factory callable (programmatic direct pass), or a Target wrapping one
685 (config keys / YAML ``_target_`` reference); both are built at apply
686 time by _build_local_compute_factory (e.g. a shipped EP archetype
687 factory such as recipes.qwen2moe_ep_compute_fn) — undecorated
688 functions fail fast
689 (injection discipline); the returned compute fn's params are
690 validated against the module's forward (params must match the
691 original function). Declaring it REQUIRES an explicit
692 ``region_dispatch`` (no default — True: dispatchable pure-ops
693 injection, validate dispatches through it and truly validates
694 out_src; False: comm/custom-kernel injection, the skeleton runs it
695 as a black box on local tensors);
696 2. spec.region_dispatch=False without inner_wrapper: the module's own
697 forward cannot dispatch (it IS the data-dependent logic, e.g. an
698 EP-aware in-house MoE with the a2a already inside forward) — the
699 skeleton runs it on local tensors. An explicit inner_wrapper owns the
700 local computation instead and therefore does not select this path;
701 3. none of the above -> None (ordinary module; takes the
702 validate/production path — and an EP-sharded boundary hitting this
703 was already failed fast by _preflight_compute_injection).
704 """
705 custom = getattr(spec, "local_compute_fn", None)
706 if custom is not None:
707 _require_region_dispatch(spec, source="spec.local_compute_fn")
708 if _is_delayed_target(custom):
709 _check_target_config_keys(custom, "local_compute_fn")
710 factory = getattr(custom, "_target_", None)
711 source = (f"local_compute_fn factory "
712 f"{getattr(custom, '_target_path', custom)}")
713 configured = getattr(custom, "_kwargs", {})
714 else:
715 factory = custom
716 source = "spec.local_compute_fn"
717 configured = None
718 compute_fn = _build_local_compute_factory(
719 factory, module, mesh, mesh_dim_names, expert_mesh,
720 configured=configured, source=source)
721 return functools.partial(compute_fn, module)
722 if spec.region_dispatch is False and spec.inner_wrapper is None:
723 return module.forward
724 return None
727def _resolve_inner_target(module, spec=None):
728 """Resolve the inner-wrap target -- pure location resolution.
730 ``inner_target`` is MANDATORY when ``inner_wrapper`` is declared (the
731 pairing is enforced in _resolve_inner_wrapper before this is called):
732 "self" means the boundary module itself, otherwise resolved by attribute
733 name -- fail-fast if the attribute does not exist or has no forward (a
734 typo must not silently degrade). The attention-domain auto-location
735 heuristic (inner_attention/attn/attention attributes, class-name
736 matching, q/k/v_proj structural fallback) was REMOVED (2026-08-10): the
737 inner-wrap mechanism is generic (any module, any submodule), and a
738 silently located target is a silent wrong-target hazard.
739 """
740 explicit = getattr(spec, "inner_target", None) if spec is not None else None
741 if explicit is None:
742 raise ValueError(
743 "inner_target is not declared — declaring inner_wrapper "
744 "requires an explicit paired inner_target (\"self\" or a "
745 "submodule attribute name; the auto-location heuristic was "
746 "removed)")
747 if explicit == "self":
748 return module
749 inner = getattr(module, explicit, None)
750 if inner is not None and hasattr(inner, "forward"):
751 return inner
752 raise ValueError(
753 f"spec.inner_target={explicit!r} did not match anything on "
754 f"{type(module).__name__} (attribute missing or has no forward) "
755 f"-- check the spelling in plan_overrides")
758def _resolve_delayed_inner_wrapper(custom, target, context,
759 apply_custom_inner_wrapper,
760 classify_rewrite_result):
761 """Resolve a delayed Target inner-wrapper declaration."""
762 fn = getattr(custom, "_target_", None)
763 source = (f"inner_wrapper Target "
764 f"{getattr(custom, '_target_path', custom)}")
765 meta = require_injection_meta(fn, INNER_WRAPPER, source=source)
766 _check_target_config_keys(custom, "inner_wrapper")
768 def _apply_target():
769 build_kwargs = fill_context_kwargs(
770 meta, context, getattr(custom, "_kwargs", {}), source=source)
771 result = custom.build(**build_kwargs)
772 if result is None:
773 return None # in-place forward replacement (registry-style fn)
774 if callable(result):
775 if getattr(result, "_injection_meta", None) is not None:
776 # An @inner_wrapper-stamped custom wrapper: invoke it with
777 # its declared context (it replaces in place or returns
778 # its own replacement).
779 return apply_custom_inner_wrapper(result, context)
780 # 4d contract: the built-in wrapper already ran and RETURNED
781 # the replacement forward — hand it to the rewriter.
782 # Validate every supported rewriter result here as well, so a
783 # delayed Target accepts atomic _ForwardRewriteRequest objects
784 # while retaining the resolver's fail-fast behavior for invalid
785 # return values.
786 classify_rewrite_result(result, target, source)
787 return result
789 name = getattr(custom, "_target_path", None) or "custom"
790 return (name, target, _apply_target)
793def _resolve_inner_wrapper(module, spec, cp_mesh, mesh, tp_mesh=None,
794 ep_mesh=None):
795 """Resolve one explicit inner-wrapper declaration without mutating modules.
797 Resolution accepts a delayed Target, a decorated callable, or a registry
798 name and returns ``(name, target, apply_fn)``. ``inner_target`` and
799 ``inner_wrapper`` must be declared together; absent declarations return
800 ``None``. The mesh family is framework-filled for every wrapper.
801 """
802 # Lazy import: forward_rewriter imports rule_resolver at module level.
803 from hyper_parallel.distributed._builder.forward_rewriter import ( # pylint: disable=C0415
804 _apply_custom_inner_wrapper,
805 _classify_rewrite_result,
806 )
807 # Lazy import: wrappers imports forward_rewriter, which imports this
808 # module at module level (circular import).
809 from hyper_parallel.distributed.context_parallel.wrappers import ( # pylint: disable=C0415
810 INNER_WRAPPER_REGISTRY,
811 )
812 custom = getattr(spec, "inner_wrapper", None) if spec is not None else None
813 inner_target = getattr(spec, "inner_target", None) if spec is not None else None
814 if custom is None:
815 if inner_target is not None:
816 raise ValueError(
817 f"spec.inner_target={inner_target!r} only locates "
818 "(replace whom) — after the rework the wrapping scheme is "
819 "no longer chosen heuristically, so inner_wrapper must be "
820 "declared explicitly too: a registry name "
821 f"{sorted(INNER_WRAPPER_REGISTRY)}, an @inner_wrapper "
822 "decorated callable, or a Target pointing to an in-repo "
823 "reference implementation "
824 "(hyper_parallel.distributed.context_parallel.wrappers.*)")
825 return None
827 _require_region_dispatch(spec, source="spec.inner_wrapper")
828 if inner_target is None:
829 raise ValueError(
830 f"inner_wrapper={custom!r} is declared but inner_target is "
831 "not — the attention-domain auto-location heuristic was "
832 "removed (inner-wrap is a generic mechanism; silently "
833 "locating risks wrapping the wrong target). The two fields "
834 "must be declared as an explicit pair: to wrap the boundary "
835 "module itself → inner_target=\"self\"; to wrap a submodule "
836 "→ inner_target=\"<attribute name>\"")
837 target = _resolve_inner_target(module, spec)
838 context = {
839 "target_module": target,
840 "mesh": mesh,
841 "tp_mesh": tp_mesh,
842 "cp_mesh": cp_mesh,
843 "ep_mesh": ep_mesh,
844 }
845 if _is_delayed_target(custom):
846 return _resolve_delayed_inner_wrapper(
847 custom, target, context,
848 _apply_custom_inner_wrapper, _classify_rewrite_result)
850 if callable(custom):
851 require_injection_meta(
852 custom, INNER_WRAPPER, source="spec.inner_wrapper")
853 return ("custom", target,
854 lambda: _apply_custom_inner_wrapper(custom, context))
856 if isinstance(custom, str):
857 fn = INNER_WRAPPER_REGISTRY.get(custom)
858 if fn is None:
859 raise ValueError(
860 f"inner_wrapper={custom!r} is not registered in "
861 f"INNER_WRAPPER_REGISTRY (available: {sorted(INNER_WRAPPER_REGISTRY)})"
862 f" -- check the spelling, or first register "
863 f"INNER_WRAPPER_REGISTRY[{custom!r}] = your_fn")
864 meta = require_injection_meta(
865 fn, INNER_WRAPPER, source=f"INNER_WRAPPER_REGISTRY[{custom!r}]")
866 context["target_module"] = target
867 return (custom, target,
868 lambda: fn(**{k: context[k] for k in sorted(meta.context)}))
870 raise TypeError(
871 f"inner_wrapper must be a registry name (str), an @inner_wrapper "
872 f"decorated callable, or a Target — got {type(custom).__name__}")
875def _resolve_inner_output_placements(
876 spec,
877 boundary_module,
878 target,
879 mesh_dim_names,
880 wrapper_name,
881):
882 """Resolve the declared placement rule for an injected inner forward."""
883 if target is boundary_module:
884 if spec is None or not spec.out_src:
885 return False, None
886 out_names = list(spec.out_names or spec.out_src.keys())
887 missing = [name for name in out_names if name not in spec.out_src]
888 if missing:
889 raise ValueError(
890 f"inner_wrapper {wrapper_name!r}: out_names {missing} "
891 "have no declaration in out_src — a multi-output contract "
892 "must be declared name by name"
893 )
894 return False, [
895 tuple(resolve_placements(spec.out_src[name], mesh_dim_names))
896 for name in out_names
897 ]
899 declared = getattr(spec, "inner_out_src", None) if spec is not None else None
900 if declared is None:
901 raise ValueError(
902 f"inner_wrapper {wrapper_name!r} targets an inner submodule "
903 f"of {type(boundary_module).__name__}, but inner_out_src is "
904 "not declared — the framework infers and guesses nothing "
905 "about inner output layouts. Choose one of: "
906 "① a layout-preserving wrapper sets "
907 "inner_out_src: \"first_input\"; ② declare the placements "
908 "explicitly; ③ or use inner_target=\"self\" to reuse the "
909 "boundary out_src contract"
910 )
911 if isinstance(declared, str):
912 if declared != "first_input":
913 raise ValueError(
914 "the string value of inner_out_src only accepts the "
915 f"sentinel 'first_input', got {declared!r}"
916 )
917 return True, None
918 if all(isinstance(value, Placement) for value in declared.values()):
919 return False, [tuple(resolve_placements(declared, mesh_dim_names))]
920 return False, [
921 tuple(resolve_placements(named, mesh_dim_names))
922 for named in declared.values()
923 ]
926def _inner_target_name(attn_module, target) -> str:
927 """Readable name of the located inner-wrap target: child attribute name,
928 or "self" when the boundary module itself is the target."""
929 if target is attn_module:
930 return "self"
931 for child_name, child in attn_module.named_children():
932 if child is target:
933 return child_name
934 return type(target).__name__