Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / forward_rewriter.py: 79%
478 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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# pylint: disable=missing-apache-license-header,forbidden-backend-import
16"""forward_rewriter: forward-mutation machinery for the sharding applier.
18All boundary forward generation, validation and installation lives here:
19the injection signature validators (4b seed), the bias-suppression /
20deferred-bias forwards, the production/validate boundary wrappers, the
21MoE local-region wrapper, the CP inner-attention adapter, and the D-02
22vocab-parallel embedding wrapper.
24This module is the ONLY place in repo-owned code that assigns
25``module.forward`` (AST gate, 05 §15.2.3): every other module returns a
26replacement callable or a private :class:`_ForwardRewriteRequest`, and the
27rewriter performs signature validation, installation and failure rollback.
28The local-compute descendant adapters
29(``install_local_compute_forward_adapters``) also live here;
30``local_compute_context.py`` keeps only the region state. Qwen model-family
31forward mutations remain in their transitional homes until changesets
32M2/M3 (auto_models_adjust.md §5.3).
33"""
35import functools
36import inspect
37import logging
38import types
39import weakref
40from typing import Any, Callable, List, Sequence
42import torch
43from torch import nn
45from hyper_parallel.core.dtensor.dtensor import DTensor
46from hyper_parallel.core.dtensor.placement_types import Shard
47from hyper_parallel.distributed.recipe_spec import (
48 INNER_WRAPPER,
49 PlacementMismatchError,
50 require_injection_meta,
51 resolve_placements,
52)
53from hyper_parallel.distributed._builder.local_compute_context import (
54 _LOCAL_COMPUTE_ACTIVE, # pylint: disable=protected-access
55 local_compute_context,
56)
57from hyper_parallel.distributed._builder.parameter_sharding import (
58 _temp_local_params,
59)
60from hyper_parallel.distributed._builder.rule_resolver import (
61 _inner_target_name,
62 _resolve_inner_output_placements,
63 _resolve_inner_wrapper,
64)
65from hyper_parallel.distributed.tensor_parallel.head_count import (
66 maybe_update_head_counts,
67)
69logger = logging.getLogger(__name__)
72# ────────────────────────────────────────────────────────────────────────────
73# Runtime-layer signature validation (principle 1: the injected function's
74# arguments must match the original function's)
75# ────────────────────────────────────────────────────────────────────────────
77def _is_subsequence(sub: List[str], full: List[str]) -> bool:
78 """Return True if ``sub`` is an in-order subsequence of ``full``."""
79 it = iter(full)
80 return all(name in it for name in sub)
83def _compute_parameters(compute_fn: Callable[..., Any], owner: str) -> List[inspect.Parameter]:
84 """Return the declared compute parameters after its module argument."""
85 try:
86 params = list(inspect.signature(compute_fn).parameters.values())
87 except (TypeError, ValueError) as exc:
88 raise TypeError(f"{owner}: cannot introspect the signature of the injected compute fn") from exc
89 if not params:
90 raise TypeError(
91 f"{owner}: a compute fn needs at least the module first "
92 "parameter -- the contract is fn(module, *forward_args)"
93 )
94 return params[1:]
97def _validate_explicit_compute_parameters(params: List[inspect.Parameter], owner: str) -> None:
98 """Reject variadic compute parameters that hide signature mismatches."""
99 for param in params:
100 if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
101 raise TypeError(
102 f"{owner}: a compute fn must not use *args/**kwargs "
103 f"(parameter {param.name!r}) -- its arguments must explicitly "
104 "match the original forward"
105 )
108def _validate_compute_parameter_names(
109 fn_params: List[inspect.Parameter],
110 fwd_params: List[inspect.Parameter],
111 owner: str,
112) -> None:
113 """Validate compute parameter names and positional ordering."""
114 fwd_names = {
115 param.name
116 for param in fwd_params
117 if param.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
118 }
119 for param in fn_params:
120 if param.name not in fwd_names:
121 raise TypeError(
122 f"{owner}: compute-fn parameter {param.name!r} has no "
123 f"same-named entry among the original forward parameters {sorted(fwd_names)} -- "
124 "the injected function's arguments must match the original function's "
125 "(the skeleton forwards the actual forward arguments, so a name mismatch "
126 "becomes a runtime TypeError)"
127 )
128 positional_kinds = (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.POSITIONAL_OR_KEYWORD)
129 fwd_pos = [param.name for param in fwd_params if param.kind in positional_kinds]
130 fn_pos = [param.name for param in fn_params if param.kind in positional_kinds]
131 if not _is_subsequence(fn_pos, fwd_pos):
132 raise TypeError(
133 f"{owner}: the compute fn's positional parameters {fn_pos} are not an in-order subsequence "
134 f"of the original forward's positional parameters {fwd_pos} -- the skeleton forwards "
135 "arguments positionally, so a reordering would bind them wrong"
136 )
139def _validate_required_compute_parameters(
140 fn_params: List[inspect.Parameter],
141 fwd_params: List[inspect.Parameter],
142 owner: str,
143) -> None:
144 """Require the compute function to accept every required forward parameter."""
145 fn_names = {param.name for param in fn_params}
146 required = []
147 for param in fwd_params:
148 if param.default is not inspect.Parameter.empty:
149 continue
150 if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
151 continue
152 required.append(param.name)
153 missing = [name for name in required if name not in fn_names]
154 if missing:
155 raise TypeError(
156 f"{owner}: the original forward's required parameters {missing} are not accepted by the "
157 "compute fn (the compute fn's declared arguments must cover the original function's required ones)"
158 )
161def validate_local_compute_signature(compute_fn: Callable[..., Any],
162 forward: Callable[..., Any],
163 *,
164 owner: str) -> None:
165 """Validate that a regional compute fn's arguments match the original
166 forward (fail-fast at apply time).
168 Rules (compute_fn's first parameter is module and forward's first
169 parameter is self; both are skipped):
170 1. The compute fn must not have *args/**kwargs (the skeleton forwards
171 the actual forward arguments, and swallowing them would hide
172 mismatches);
173 2. Every compute-fn parameter must have a **same-named** parameter in
174 forward (the skeleton forwards by kwarg, so a name mismatch is a
175 runtime TypeError);
176 3. The compute fn's positional-parameter sequence must be an in-order
177 **subsequence** of forward's positional-parameter sequence
178 (positional forwarding must not reorder);
179 4. All required (defaultless) parameters of forward must be accepted by
180 the compute fn.
182 Args:
183 compute_fn: The compute fn returned by the @local_compute factory.
184 forward: The original module forward being replaced.
185 owner: Call-site label prepended to error messages.
187 Raises:
188 TypeError: If any of the rules above is violated.
189 """
190 fn_params = _compute_parameters(compute_fn, owner)
191 _validate_explicit_compute_parameters(fn_params, owner)
192 fwd_params = list(inspect.signature(forward).parameters.values())
193 _validate_compute_parameter_names(fn_params, fwd_params, owner)
194 _validate_required_compute_parameters(fn_params, fwd_params, owner)
197def validate_wrapped_forward(orig_forward: Callable[..., Any],
198 new_forward: Callable[..., Any],
199 *,
200 owner: str) -> None:
201 """Validate that the forward replaced by inner_wrapper can receive all
202 arguments of the original forward.
204 A dummy-bind probe is constructed from the original signature
205 (*args/**kwargs parameters cannot be forged and are skipped); the
206 replacement side may pass through tolerantly with *args/**kwargs (it
207 must forward arguments it does not care about to the original
208 implementation), but every named parameter of the original forward must
209 remain bindable by name.
211 Args:
212 orig_forward: The original module forward.
213 new_forward: The replacement forward installed by the wrapper.
214 owner: Call-site label prepended to error messages.
216 Raises:
217 TypeError: If the replacement signature cannot be introspected or
218 cannot bind the original arguments.
219 """
220 orig_sig = inspect.signature(orig_forward)
221 try:
222 new_sig = inspect.signature(new_forward, follow_wrapped=False)
223 except (TypeError, ValueError) as exc:
224 raise TypeError(
225 f"{owner}: cannot introspect the signature of the forward "
226 "replaced by the injected wrapper") from exc
227 args, kwargs = [], {}
228 for p in orig_sig.parameters.values():
229 if p.kind in (inspect.Parameter.VAR_POSITIONAL,
230 inspect.Parameter.VAR_KEYWORD):
231 continue
232 if p.kind is inspect.Parameter.POSITIONAL_ONLY:
233 args.append(None)
234 else:
235 # Probe POSITIONAL_OR_KEYWORD by name (kwarg) -- the
236 # name-dimension contract is what the injection discipline
237 # cares about; purely positional passing of optional positional
238 # arguments is not required
239 kwargs[p.name] = None
240 try:
241 new_sig.bind(*args, **kwargs)
242 except TypeError as exc:
243 raise TypeError(
244 f"{owner}: the forward replaced by the injected wrapper is "
245 f"incompatible with the original forward's arguments ({exc}) -- "
246 f"original signature {orig_sig}, replacement signature {new_sig}; "
247 "the replacement forward must be able to receive all arguments "
248 "of the original forward (it may pass through tolerantly with "
249 "*args/**kwargs)") from exc
252# ────────────────────────────────────────────────────────────────────────────
253# Boundary forward wrapping: bias suppression / deferred bias (05 §4.4)
254# ────────────────────────────────────────────────────────────────────────────
256def _make_bias_free_forward(owner: nn.Module, original: Callable[..., Any]) -> Callable[..., Any]:
257 """Create a forward wrapper that temporarily hides one module bias."""
259 @functools.wraps(original)
260 def bias_free_forward(*args: Any, **kwargs: Any) -> Any:
261 """Run the owner's forward with its bias temporarily hidden."""
262 bias = owner.bias
263 owner._parameters["bias"] = None # pylint: disable=protected-access
264 try:
265 return original(*args, **kwargs)
266 finally:
267 owner._parameters["bias"] = bias # pylint: disable=protected-access
269 return bias_free_forward
272def _install_bias_suppression(module, spec):
273 """D-22: make each defer-listed Linear run bias-free inside the region.
275 The bias Parameter itself is never touched (same FQN / same object /
276 state_dict / optimizer unchanged); the wrapper only hides it from
277 ``F.linear`` for the duration of the region forward, so the boundary
278 exit's Partial reduction sees a pure matmul contribution and the bias is
279 added exactly once afterwards by :func:`_maybe_add_deferred_biases`.
280 Both modes install the same wrapper (instruction-for-instruction
281 identity, D-01''); nesting the suppression is idempotent.
282 """
283 for param_path in spec._deferred_bias_params: # pylint: disable=protected-access
284 owner_path = param_path.rpartition(".")[0]
285 owner = module.get_submodule(owner_path) if owner_path else module
286 original = owner.forward
287 owner.forward = _make_bias_free_forward(owner, original)
290def _add_bias_to_primary_output(output, bias, module_name):
291 """Add a deferred bias to the primary Tensor while preserving output structure."""
292 if isinstance(output, torch.Tensor):
293 primary_output = output
294 rebuild = None
295 elif isinstance(output, (tuple, list)):
296 if not output or not isinstance(output[0], torch.Tensor):
297 primary_type = type(output[0]).__name__ if output else "missing"
298 raise TypeError(
299 f"{module_name}: deferred bias requires output index 0 to be a Tensor, "
300 f"got {primary_type}"
301 )
302 primary_output = output[0]
303 rebuild = list(output)
304 else:
305 raise TypeError(
306 f"{module_name}: deferred bias requires a Tensor, tuple, or list output, "
307 f"got {type(output).__name__}"
308 )
310 if not isinstance(primary_output, DTensor) and isinstance(bias, DTensor):
311 bias = bias.to_local()
312 biased_output = primary_output + bias
313 if rebuild is None:
314 return biased_output
315 rebuild[0] = biased_output
316 return tuple(rebuild) if isinstance(output, tuple) else rebuild
319def _maybe_add_deferred_biases(module, spec, output):
320 """D-22: add each deferred bias exactly once AFTER the boundary exit.
322 The bias is read at forward time (never captured in a closure), so
323 production sees the unwrapped local tensor and validate sees the DTensor
324 — whichever form the boundary output currently has: nested validate keeps
325 the DTensor (dispatch add, Shard(1)+Replicate→Shard(1)); the outermost /
326 production / local-region exits are local, and a DTensor bias is
327 unwrapped to match. For structured attention outputs, index 0 is the
328 primary hidden-state Tensor; metadata/cache entries are preserved.
329 """
330 if not spec._deferred_bias_params: # pylint: disable=protected-access
331 return output
332 for param_path in spec._deferred_bias_params: # pylint: disable=protected-access
333 owner_path = param_path.rpartition(".")[0]
334 owner = module.get_submodule(owner_path) if owner_path else module
335 bias = owner.bias
336 if bias is None:
337 continue
338 output = _add_bias_to_primary_output(output, bias, type(module).__name__)
339 return output
342def _unwrap_dtensor_outputs(output):
343 """Return local tensors while preserving tuple/list output structure."""
344 if isinstance(output, DTensor):
345 return output.to_local()
346 if isinstance(output, tuple):
347 return tuple(_unwrap_dtensor_outputs(item) for item in output)
348 if isinstance(output, list):
349 return [_unwrap_dtensor_outputs(item) for item in output]
350 return output
353# ────────────────────────────────────────────────────────────────────────────
354# Production/validate boundary wrappers and output-contract validation
355# ────────────────────────────────────────────────────────────────────────────
357def _keep_loss_parallel_output(plan, spec):
358 """Return whether a terminal vocab-sharded output must remain a DTensor."""
359 if not plan.loss_parallel or not spec._is_terminal: # pylint: disable=protected-access
360 return False
361 output_placements = (spec.out_dst or {}).get("output", {})
362 return any(
363 isinstance(placement, Shard) and placement.dim == -1
364 for placement in output_placements.values()
365 )
368def _descendant_boundary_fqns(plan, module_fqn):
369 """Relative FQNs of boundaries nested inside *module_fqn* (D-14).
371 Returned relative to module_fqn (matching the name space of
372 module.named_parameters(recurse=True)); the root spec (fqn "") treats
373 every other boundary as a descendant.
374 """
375 boundaries = [fqn for fqn, spec in plan.modules.items() if spec.is_boundary]
376 if module_fqn == "":
377 return [f for f in boundaries if f]
378 prefix = module_fqn + "."
379 return [f[len(prefix):] for f in boundaries if f.startswith(prefix)]
382def _bind_input_indices(boundary, module):
383 """Bind the arg_name of in_plan to the positional index of the forward signature.
385 Inter-module calls are mostly positional (self.mlp(x) inside a layer), so
386 RedistOp's kwargs lookup by name would miss -- the signature index is bound
387 at compile time, and at runtime _get_arg checks kwargs first, then args.
388 """
389 try:
390 sig = inspect.signature(module.forward)
391 except (TypeError, ValueError):
392 sig = None
393 if sig is not None:
394 positional = [
395 name for name, p in sig.parameters.items()
396 if p.kind in (p.POSITIONAL_ONLY, p.POSITIONAL_OR_KEYWORD)
397 ]
398 name_to_idx = {name: i for i, name in enumerate(positional)}
399 for op in boundary.in_plan:
400 if op.arg_index is None and op.arg_name in name_to_idx:
401 op.arg_index = name_to_idx[op.arg_name]
402 # Positional fallback: a single-input contract (in_plan has only 1 op) is
403 # bound to the first positional parameter -- covers cases where the
404 # template key (e.g. "hidden_states") differs from the leaf module's
405 # signature (nn.Linear.forward(input)).
406 if len(boundary.in_plan) == 1 and boundary.in_plan[0].arg_index is None:
407 boundary.in_plan[0].arg_index = 0
410def _wrap_production_forward(
411 module, boundary, spec=None, keep_output_dtensor=False):
412 """Production mode: pure local tensor computation + precompiled boundary communication (05 §4.4.1).
414 _local_params_context was already invoked at the Phase C entry (parameters
415 permanently unpacked).
416 """
417 original_forward = module.forward
419 @functools.wraps(original_forward)
420 def production_forward(*args: Any, **kwargs: Any) -> Any:
421 """Production boundary forward: precompiled entry communication,
422 local computation, exit redistribution, then D-22 deferred biases."""
423 args, kwargs = boundary.redistribute_inputs(args, kwargs)
424 outputs = original_forward(*args, **kwargs)
425 outputs = boundary.redistribute_outputs(
426 outputs,
427 as_dtensor_input=keep_output_dtensor,
428 )
429 # D-22: deferred rowwise biases — added once after the exit reduction
430 return _maybe_add_deferred_biases(module, spec, outputs) \
431 if spec is not None else outputs
433 module.forward = production_forward
436def _wrap_validate_forward(
437 module, boundary, spec, mesh_dim_names, keep_output_dtensor=False):
438 """Validate mode: DTensor propagation end to end -> validate out_src (core) + out_dst (terminal modules only).
440 The in-house DTensor is forward-only: validation covers only forward
441 placement propagation; backward is local autograd in both modes (05 §1.0),
442 and gradient equivalence is guaranteed by tests/ut/auto_models/distributed/grad_equiv.py.
443 """
444 original_forward = module.forward
445 module_name = type(module).__name__
447 @functools.wraps(original_forward)
448 def validate_forward(*args: Any, **kwargs: Any) -> Any:
449 """Validate boundary forward: wrap inputs as DTensors, propagate
450 placements through the original forward, validate out_src, then
451 redistribute to out_dst."""
452 # D-14 nesting (05 §13.4): detect whether the call arrives from an
453 # outer DTensor-propagating boundary BEFORE Step 1 wraps everything
454 # into DTensors (Step 1 would make the check useless).
455 nested = any(isinstance(a, DTensor) for a in args) or any(
456 isinstance(v, DTensor) for v in kwargs.values())
457 # Step 1: inputs -> DTensor
458 args, kwargs = boundary.redistribute_inputs(args, kwargs, as_dtensor=True)
460 # Step 2: parameters stay DTensors; placement propagates via __torch_function__ dispatch
461 outputs = original_forward(*args, **kwargs)
463 # Step 3: [core validation] out_src -- native DTensor-propagated output vs declaration
464 if spec.out_src is not None:
465 _validate_out_src(outputs, spec, mesh_dim_names, module_name)
467 # Step 4: redistribute to out_dst
468 outputs = boundary.redistribute_outputs(outputs, as_dtensor_input=True)
470 # Step 5: [defensive validation] out_dst -- terminal modules only
471 if spec._is_terminal and spec.out_dst is not None: # pylint: disable=protected-access
472 _validate_out_dst(outputs, spec, mesh_dim_names, module_name)
474 # D-22: deferred biases are added after the exit reduction while the
475 # output is still a DTensor (Shard/Replicate + Replicate bias dispatch
476 # add — the Partial reduction is already done, so every rank adds the
477 # bias exactly once); Step 6 then unwraps as usual.
478 outputs = _maybe_add_deferred_biases(module, spec, outputs)
480 # Step 6: return local (isomorphic to production boundary outputs) --
481 # but under an outer DTensor-propagating boundary (D-14 nesting, 05
482 # §13.4) keep the DTensor so the outer forward's dispatch chain is
483 # unbroken; the outermost boundary exit unwraps.
484 if nested or keep_output_dtensor:
485 return outputs
486 return _unwrap_dtensor_outputs(outputs)
488 module.forward = validate_forward
491def _validate_out_src(outputs, spec, mesh_dim_names, module_name):
492 _validate_outputs(outputs, spec, mesh_dim_names, module_name, "out_src")
495def _validate_out_dst(outputs, spec, mesh_dim_names, module_name):
496 _validate_outputs(outputs, spec, mesh_dim_names, module_name, "out_dst")
499def _normalize_placements_ndim(placements, ndim):
500 """Normalize negative dims like Shard(-1) against the tensor ndim (Shard(-1) == Shard(ndim-1))."""
501 out = []
502 for p in placements:
503 if isinstance(p, Shard) and p.dim < 0:
504 out.append(Shard(p.dim + ndim))
505 else:
506 out.append(p)
507 return tuple(out)
510def _validate_output_placement(tensor, expected_named, mesh_dim_names, module_name, stage):
511 """Validate one DTensor output's declared placement."""
512 ndim = len(tensor.shape)
513 expected = _normalize_placements_ndim(
514 tuple(resolve_placements(expected_named, mesh_dim_names)), ndim)
515 actual = _normalize_placements_ndim(tuple(tensor.placements), ndim)
516 if expected != actual:
517 raise PlacementMismatchError(module_name, expected, actual, stage)
520def _validate_outputs(outputs, spec, mesh_dim_names, module_name, stage):
521 """Placement validation for single/multi outputs (shared by out_src / out_dst).
523 Multi outputs are mapped to tuple positions via spec.out_names (falling
524 back to declaration key order); outputs that are not returned or are not
525 DTensors are skipped. Declared and actual placements are
526 negative-dim-normalized before comparison.
527 """
528 declared = getattr(spec, stage)
529 if isinstance(outputs, (tuple, list)):
530 out_names = getattr(spec, "out_names", None) or list(declared.keys())
531 name_to_idx = {name: i for i, name in enumerate(out_names)}
532 items = list(outputs)
533 else:
534 name_to_idx = {name: 0 for name in declared}
535 items = [outputs]
536 for out_name, expected_named in declared.items():
537 idx = name_to_idx.get(out_name)
538 if idx is None or idx >= len(items):
539 continue
540 tensor = items[idx]
541 if not isinstance(tensor, DTensor):
542 continue
543 suffix = f"[{out_name}]" if len(declared) > 1 else ""
544 _validate_output_placement(
545 tensor, expected_named, mesh_dim_names, module_name, f"{stage}{suffix}")
548# ────────────────────────────────────────────────────────────────────────────
549# MoE EP local region (05 §4.4.3 + D-03')
550# ────────────────────────────────────────────────────────────────────────────
552def _rewrap_local_outputs(output, spec, mesh, mesh_dim_names, module_name):
553 """Wrap every declared local Tensor output with its ``out_src`` layout."""
554 declared = spec.out_src or {}
555 if not declared:
556 return output
558 is_sequence = isinstance(output, (tuple, list))
559 items = list(output) if is_sequence else [output]
560 out_names = list(getattr(spec, "out_names", None) or declared.keys())
561 name_to_idx = {name: index for index, name in enumerate(out_names)}
563 for out_name, named_placement in declared.items():
564 index = name_to_idx.get(out_name)
565 if index is None:
566 raise ValueError(
567 f"{module_name}: out_src declares output {out_name!r}, but "
568 f"out_names={out_names!r} does not contain it"
569 )
570 if index >= len(items):
571 raise ValueError(
572 f"{module_name}: out_src maps output {out_name!r} to index "
573 f"{index}, but forward returned only {len(items)} output(s)"
574 )
575 item = items[index]
576 if item is None:
577 continue
578 if isinstance(item, DTensor):
579 continue
580 if not isinstance(item, torch.Tensor):
581 raise TypeError(
582 f"{module_name}: declared output {out_name!r} at index "
583 f"{index} must be a Tensor or None, got {type(item).__name__}"
584 )
585 placements = tuple(resolve_placements(named_placement, mesh_dim_names))
586 items[index] = DTensor.from_local(item, mesh, placements)
588 if isinstance(output, tuple):
589 return tuple(items)
590 if isinstance(output, list):
591 return items
592 if len(items) != 1:
593 raise ValueError(
594 f"{module_name}: scalar forward output cannot satisfy "
595 f"{len(declared)} declared out_src entries"
596 )
597 return items[0]
600def _wrap_local_region_forward(module, boundary, spec, mesh, mesh_dim_names,
601 *, validate_mode=False, compute_fn=None,
602 exclude_subtrees=()):
603 """Generic local-region forward wrapper (D-03', formerly the _wrap_moe_forward skeleton).
605 Structure: boundary entry -> local region -> boundary exit. Applies to any
606 module containing data-dependent logic that DTensor dispatch cannot express
607 (e.g. MoE all-to-all) -- injected by _apply_phase_c when
608 _resolve_local_compute_fn resolves to non-None (derived gate, 05 §4.4.3).
610 production: parameters were permanently unpacked at build time and inputs
611 are local (boundary passthrough), and the output stays local for lowered
612 boundary operations; validate: inputs are DTensors -> to_local ->
613 temporarily unwrap parameters -> local computation -> re-wrap the output
614 per the declared out_src via from_local (for data-dependent modules out_src
615 is declarative validation -- the data dependence of all-to-all makes the
616 placement underivable; this is an inherent limitation). Both modes share
617 the same wrapper code (local_region tolerant passthrough semantics).
619 compute_fn: the function actually executed inside the region; defaults to
620 the module's own forward. Resolved uniformly by
621 _resolve_local_compute_fn (user local_compute_fn / region_dispatch=False
622 gate), independent of the original forward.
624 spec.region_dispatch=True (dispatchable pure-ops injection): validate
625 feeds the DTensors straight into compute_fn — strategy propagation runs
626 THROUGH the injected fn and out_src is TRULY validated (propagated vs
627 declared); production is unchanged (always local passthrough).
628 """
629 original_forward = module.forward
630 if compute_fn is None:
631 compute_fn = original_forward
633 dispatch_through = bool(getattr(spec, "region_dispatch", None))
634 if validate_mode and not dispatch_through:
635 install_local_compute_forward_adapters(module, exclude=exclude_subtrees)
637 @functools.wraps(original_forward)
638 def local_region_forward(*args: Any, **kwargs: Any) -> Any:
639 """Local-region forward: boundary entry → local region computation →
640 validate re-wrap → boundary exit (both modes share this wrapper)."""
641 # Step 1: PrecompiledBoundary entry (e.g. TP all-gather; identity passthrough)
642 args, kwargs = boundary.redistribute_inputs(
643 args, kwargs, as_dtensor=validate_mode)
645 # Step 2: local region -- the data-dependent computation (e.g. EP
646 # dispatch/combine) executes on local tensors; region_dispatch=True
647 # instead dispatches THROUGH the injected fn (pure standard ops) so
648 # validate's strategy propagation covers it
649 if validate_mode and dispatch_through:
650 try:
651 output = compute_fn(*args, **kwargs)
652 except Exception as exc:
653 raise type(exc)(
654 f"{exc}\n[region_dispatch=True] dispatch failed while "
655 "validate was dispatching through the injected "
656 "function — does the injection contain "
657 "non-dispatchable communication primitives/custom "
658 "kernels? Declare region_dispatch=False instead "
659 "(skeleton black-box hosting)"
660 ) from exc
661 # True validation: propagated result vs the out_src declaration
662 # (a wrong declaration fails fast)
663 if spec.out_src is not None:
664 _validate_out_src(output, spec, mesh_dim_names,
665 type(module).__name__)
666 elif validate_mode:
667 local_args = tuple(
668 a.to_local() if isinstance(a, DTensor) else a for a in args)
669 local_kwargs = {
670 k: (v.to_local() if isinstance(v, DTensor) else v)
671 for k, v in kwargs.items()
672 }
673 with local_compute_context():
674 with _temp_local_params(module, exclude=exclude_subtrees):
675 output = compute_fn(*local_args, **local_kwargs)
676 else:
677 output = compute_fn(*args, **kwargs)
679 # Step 3: validate restores the DTensor metadata broken by all-to-all.
680 # Production keeps the output local so lowered boundary operations can
681 # execute directly; an unlowerable transition wraps only in its
682 # necessary DTensor fallback inside RedistOp.
683 if validate_mode and not isinstance(output, DTensor):
684 output = _rewrap_local_outputs(
685 output, spec, mesh, mesh_dim_names, type(module).__name__)
687 # Step 4: PrecompiledBoundary exit (e.g. TP reduce-scatter)
688 output = boundary.redistribute_outputs(
689 output, as_dtensor_input=validate_mode)
690 # The final boundary exit is always local.
691 output = _unwrap_dtensor_outputs(output)
692 # D-22: deferred rowwise biases — added once after the exit reduction
693 return _maybe_add_deferred_biases(module, spec, output)
695 module.forward = local_region_forward
698# ────────────────────────────────────────────────────────────────────────────
699# CP inner attention wrapper (05 §4.4.2 + D-01'' + D-04)
700# ────────────────────────────────────────────────────────────────────────────
702def _apply_custom_inner_wrapper(custom_fn, context):
703 """Apply a user-defined inner_wrapper (@inner_wrapper callable).
705 Contract: custom_fn must be decorated ``@inner_wrapper``; it is invoked
706 with its declared context (anchor target_module + the mandatory mesh
707 family, filled by name). It either replaces target.forward in place and
708 returns None (the external contract) or RETURNS the replacement — a
709 forward callable or _ForwardRewriteRequest(s) — for the rewriter to
710 validate and install (the in-repo discipline, 05 §15.2.3). The
711 replacement forward runs in a LOCAL-TENSOR world — the dual-mode
712 adapter installed by _wrap_inner_attention owns all DTensor
713 conversion (to_local / _temp_local_params / from_local rewrap per the
714 declared placements), so custom wrappers never touch DTensor.
715 """
716 meta = require_injection_meta(
717 custom_fn, INNER_WRAPPER, source="spec.inner_wrapper")
718 return custom_fn(**{k: context[k] for k in sorted(meta.context)})
721def _find_dtensor_argument(args, kwargs):
722 """Return the first DTensor positional or keyword argument."""
723 return next(
724 (
725 value
726 for value in (*args, *kwargs.values())
727 if isinstance(value, DTensor)
728 ),
729 None,
730 )
733def _validate_inner_dispatch_output(
734 output,
735 expected_placements,
736 boundary_module,
737 wrapper_name,
738):
739 """Validate DTensor outputs produced through a dispatchable inner wrapper."""
740 if expected_placements is None:
741 return
742 outputs = list(output) if isinstance(output, (tuple, list)) else [output]
743 tensor_outputs = [tensor for tensor in outputs if isinstance(tensor, torch.Tensor)]
744 if len(tensor_outputs) != len(expected_placements):
745 raise RuntimeError(
746 f"inner_wrapper {wrapper_name!r}: tensor output count "
747 f"{len(tensor_outputs)} does not match the "
748 f"{len(expected_placements)} declared placements — a "
749 "multi-output contract must be declared name by name with "
750 "matching counts"
751 )
752 for tensor, placements in zip(tensor_outputs, expected_placements):
753 if not isinstance(tensor, DTensor):
754 raise RuntimeError(
755 f"inner_wrapper {wrapper_name!r} [region_dispatch=True]: "
756 f"the dispatch-propagated output is not a DTensor "
757 f"({type(tensor).__name__}) — the injection appears to "
758 "have left the dispatch chain, so true validation cannot "
759 "complete"
760 )
761 if tuple(tensor.placements) != tuple(placements):
762 raise PlacementMismatchError(
763 f"{type(boundary_module).__name__} "
764 f"(inner_wrapper {wrapper_name!r})",
765 tuple(placements),
766 tuple(tensor.placements),
767 "inner_out_src",
768 )
771def _rewrap_inner_tensor(tensor, placements, mesh):
772 """Wrap one local tensor with its declared layout and pass other values through."""
773 if not isinstance(tensor, torch.Tensor) or isinstance(tensor, DTensor):
774 return tensor
775 return DTensor.from_local(tensor, mesh, placements)
778def _rewrap_inner_outputs(output, out_placements, mesh, wrapper_name):
779 """Rewrap tensor outputs while preserving auxiliary non-tensor outputs."""
780 if not isinstance(output, (tuple, list)):
781 return _rewrap_inner_tensor(output, out_placements[0], mesh)
782 tensor_output_count = sum(isinstance(value, torch.Tensor) for value in output)
783 if tensor_output_count != len(out_placements):
784 raise RuntimeError(
785 f"inner_wrapper {wrapper_name!r}: the replacement forward "
786 f"returned {tensor_output_count} tensor outputs, which does "
787 f"not match the {len(out_placements)} declared placements — "
788 "a multi-output contract must be declared name by name with "
789 "matching counts"
790 )
791 placement_iter = iter(out_placements)
792 wrapped = [
793 _rewrap_inner_tensor(value, next(placement_iter), mesh)
794 if isinstance(value, torch.Tensor)
795 else value
796 for value in output
797 ]
798 return tuple(wrapped) if isinstance(output, tuple) else wrapped
801def _install_inner_adapter(target, user_fwd, boundary_module, spec, mesh,
802 mesh_dim_names, wrapper_name, validate_mode=False):
803 """Unified dual-mode adapter: rewrap rules are resolved at install time,
804 with zero runtime decisions (05 §4.4.2 + D-01'').
806 The replacement forward of a user wrapper only faces local tensors. The
807 adapter handles:
808 - validate (any input is a DTensor): every DTensor input is to_local'd
809 (non-tensors pass through) + ``_temp_local_params(target)`` temporarily
810 unwraps the parameters → call the user forward → rewrap the outputs
811 back into DTensors per the declaration (the propagation chain is
812 re-linked and boundary validation continues);
813 - production (no DTensor inputs): straight passthrough, zero conversion
814 overhead.
816 Source of the rewrap placements (the framework infers and guesses
817 nothing — everything is explicitly declared):
818 - case A (target is the boundary module itself): the boundary
819 ``spec.out_src`` declaration (multiple outputs positionally per
820 ``out_names``/declaration key order);
821 - case B (an inner submodule): the explicit ``spec.inner_out_src``
822 declaration — the sentinel ``"first_input"`` (output layout == the
823 runtime layout of the first DTensor input; for layout-preserving
824 wrappers, single output only) or NamedPlacement /
825 {name: NamedPlacement} (multiple outputs mapped to tuple positions per
826 declaration key order); undeclared → fail-fast at install time.
827 """
828 first_input_rule, out_placements = _resolve_inner_output_placements(
829 spec,
830 boundary_module,
831 target,
832 mesh_dim_names,
833 wrapper_name,
834 )
835 if validate_mode and not getattr(spec, "region_dispatch", None):
836 install_local_compute_forward_adapters(target)
838 @functools.wraps(user_fwd)
839 def adapted(*args: Any, **kwargs: Any) -> Any:
840 """Dual-mode adapter around the user's replacement forward.
842 production (no DTensor input): straight passthrough; validate:
843 unwrap DTensor inputs/parameters to local, run the user forward, and
844 rewrap outputs per the declaration resolved at install time (or, with
845 region_dispatch=True, dispatch through and truly validate instead).
846 """
847 source_dtensor = _find_dtensor_argument(args, kwargs)
848 if source_dtensor is None:
849 return user_fwd(*args, **kwargs) # production: straight passthrough
850 if getattr(spec, "region_dispatch", None):
851 # Validate dispatch-through (region_dispatch=True): the injection
852 # is pure standard ops — DTensors go straight in and dispatch
853 # propagation runs through the inner region; the declared rewrap
854 # rules are promoted to the true-validation baseline (propagated
855 # result vs declaration, a mismatch fails fast).
856 try:
857 out = user_fwd(*args, **kwargs)
858 except Exception as exc:
859 raise type(exc)(
860 f"{exc}\n[region_dispatch=True] dispatch failed "
861 f"while validate was dispatching through "
862 f"inner_wrapper {wrapper_name!r} — does the injection "
863 "contain non-dispatchable communication "
864 "primitives/custom kernels? Declare "
865 "region_dispatch=False instead (adapter black-box "
866 "hosting)") from exc
867 expected = (
868 [tuple(source_dtensor.placements)]
869 if first_input_rule
870 else out_placements
871 )
872 _validate_inner_dispatch_output(
873 out,
874 expected,
875 boundary_module,
876 wrapper_name,
877 )
878 return out
879 # validate: uniformly convert to local (parameters temporarily
880 # unwrapped, restored on exit)
881 local_args = tuple(
882 a.to_local() if isinstance(a, DTensor) else a for a in args)
883 local_kwargs = {k: (v.to_local() if isinstance(v, DTensor) else v)
884 for k, v in kwargs.items()}
885 with local_compute_context():
886 with _temp_local_params(target):
887 out = user_fwd(*local_args, **local_kwargs)
889 if first_input_rule:
890 if isinstance(out, (tuple, list)):
891 raise RuntimeError(
892 f"inner_wrapper {wrapper_name!r}: inner_out_src="
893 "'first_input' only supports a single output — for "
894 "multiple outputs declare inner_out_src explicitly in "
895 "the {name: {axis: placement}} form")
896 return _rewrap_inner_tensor(
897 out,
898 tuple(source_dtensor.placements),
899 source_dtensor.device_mesh,
900 )
901 if out_placements is None:
902 return out # case A with no out_src declaration: no rewrap
903 return _rewrap_inner_outputs(out, out_placements, mesh, wrapper_name)
905 target.forward = adapted
908def _wrap_inner_attention(module, cp_mesh, *, spec=None, mesh=None,
909 mesh_dim_names=(), tp_mesh=None, ep_mesh=None,
910 validate_mode=False, module_fqn=""):
911 """Inject an inner forward wrapper (one-shot replacement at apply time, 05 §4.4.2).
913 General "weave into / replace the inner forward" mechanism — **not gated
914 on CP**: whenever ``spec.inner_wrapper`` is declared the wrapper is
915 applied (declaration == application; the resolution chain is the derived
916 gate). The shipped CP wrappers (INNER_WRAPPER_REGISTRY) are its first-class
917 built-in use case and still require an active cp axis (fail-fast in the
918 resolution chain otherwise); custom callables/Targets receive
919 ``cp_mesh=None`` when no cp axis exists and own their semantics.
921 Resolution (_resolve_inner_wrapper, a pure function of the chain) is
922 separated from application; when resolution returns None (no explicit
923 declaration), this returns None and injects nothing. Returns the resolved
924 wrapper name (or None).
926 D-01'': production and validate inject **the same** wrapper, so the
927 in-region computation is instruction-for-instruction identical
928 (kernel-level equivalence).
929 D-04: when is_causal and CP is active, replace it with an offset-aware
930 explicit mask.
931 Nothing is located silently: after injection an INFO log records the
932 target/wrapper/source, and spec._resolved_inner_wrapper +
933 spec._resolved_inner_target are written back for plan introspection.
934 """
935 resolved = _resolve_inner_wrapper(
936 module, spec, cp_mesh, mesh, tp_mesh=tp_mesh, ep_mesh=ep_mesh)
937 if resolved is None:
938 return None
939 name, target = _apply_resolved_inner_wrapper(
940 module, resolved, spec, mesh, mesh_dim_names, validate_mode, module_fqn)
941 target_name = _inner_target_name(module, target)
942 if spec is not None:
943 spec._resolved_inner_wrapper = name # pylint: disable=protected-access
944 spec._resolved_inner_target = target_name # pylint: disable=protected-access
945 if name == "custom":
946 source = "custom callable"
947 elif spec is not None and isinstance(
948 getattr(spec, "inner_wrapper", None), str):
949 source = "explicitly specified (registry)"
950 else:
951 source = "explicitly specified (Target)"
952 logger.info("inner-wrap: %s target=%s <- wrapper %r (%s)",
953 type(module).__name__, target_name, name, source)
954 return name
957def _rollback_inner_rewrite(target, saved_state, committed_secondaries):
958 """Restore primary and secondary targets after a failed inner rewrite."""
959 for secondary_target, secondary_state in reversed(committed_secondaries):
960 for attr, saved in secondary_state.items():
961 _restore_attr(secondary_target, attr, saved)
962 for attr, saved in saved_state.items():
963 _restore_attr(target, attr, saved)
966def _prepare_inner_rewrite(target, apply_fn, name):
967 """Apply an inner wrapper and retain the state needed for rollback."""
968 orig_forward = target.forward
969 # Record the pre-rewrite state for failure rollback: an in-place wrapper
970 # writes an instance attribute, so __dict__ holds the full mutation.
971 saved_state = {"forward": target.__dict__.get("forward", _MISSING)}
972 committed_secondaries = []
973 try:
974 primary, secondaries = _classify_rewrite_result(
975 apply_fn(), target, name)
976 for request in secondaries:
977 committed_secondaries.append(
978 (request.target, _commit_forward_rewrite(request))
979 )
980 if primary is None:
981 # In-place contract (external @inner_wrapper wrappers). Detect
982 # "a replacement really happened": attribute access on a bound
983 # method creates a new object each time, so an `is` comparison is
984 # always true — the underlying function objects (__func__) must
985 # be compared, otherwise a pure-probe wrapper (which does not
986 # replace forward) would also get an adapter installed by mistake
987 # and be forced to declare inner_out_src
988 new_forward = target.forward
989 replaced = (getattr(new_forward, "__func__", new_forward)
990 is not getattr(orig_forward, "__func__", orig_forward))
991 else:
992 # In-repo discipline: the wrapper returned its replacement, the
993 # rewriter commits the companion attributes and installs.
994 saved_state.update({
995 attr: target.__dict__.get(attr, _MISSING)
996 for attr in primary.companion_attrs
997 })
998 for attr, value in primary.companion_attrs.items():
999 setattr(target, attr, value)
1000 new_forward = primary.forward
1001 replaced = True
1002 except Exception:
1003 _rollback_inner_rewrite(
1004 target, saved_state, committed_secondaries)
1005 raise
1006 return (orig_forward, new_forward, replaced, saved_state,
1007 committed_secondaries)
1010def _apply_resolved_inner_wrapper(module, resolved, spec, mesh, mesh_dim_names,
1011 validate_mode, module_fqn):
1012 """Apply, validate, and atomically install a resolved inner wrapper."""
1013 name, target, apply_fn = resolved
1014 if validate_mode and getattr(spec, "region_dispatch", None) is False:
1015 # A black-box inner wrapper receives local tensors and local parameter
1016 # shards in validate mode. Cached head counts must therefore match the
1017 # TP-local projection widths, just as they already do in production.
1018 maybe_update_head_counts(
1019 target,
1020 spec,
1021 module_fqn or type(module).__name__,
1022 mesh,
1023 mesh_dim_names,
1024 )
1025 (orig_forward, new_forward, replaced, saved_state,
1026 committed_secondaries) = _prepare_inner_rewrite(
1027 target, apply_fn, name)
1028 try:
1029 if replaced:
1030 # Principle 1: the replaced forward must accept all inputs of the
1031 # original forward
1032 validate_wrapped_forward(
1033 orig_forward, new_forward,
1034 owner=f"inner_wrapper {name!r} on {type(module).__name__}")
1035 # Uniformly install the dual-mode adapter: the user wrapper only
1036 # faces local tensors; DTensor conversion and declarative rewrap
1037 # are managed by the adapter (local_map semantics; validate skips
1038 # propagation checks for the inner region — the safety net is at
1039 # the boundary layer)
1040 _install_inner_adapter(
1041 target, new_forward, module, spec, mesh, mesh_dim_names, name,
1042 validate_mode=validate_mode)
1043 except Exception:
1044 _rollback_inner_rewrite(
1045 target, saved_state, committed_secondaries)
1046 raise
1047 return name, target
1050# ────────────────────────────────────────────────────────────────────────────
1051# D-02 vocab-parallel embedding wrapper
1052# ────────────────────────────────────────────────────────────────────────────
1054def _is_vocab_parallel_embed(module, spec, tp_mesh) -> bool:
1055 """production embed boundary check: nn.Embedding + weight Shard(0) on TP + TP>1."""
1056 if tp_mesh is None or tp_mesh.size() <= 1:
1057 return False
1058 if not isinstance(module, nn.Embedding):
1059 return False
1060 weight_named = spec.params.get("weight", {})
1061 return weight_named.get("tp") == Shard(0)
1064def _wrap_vocab_parallel_embedding(module, tp_mesh):
1065 """D-02: Megatron-style masked embedding (injected at the production embed boundary).
1067 The vocab-range mask logic of DTensor dispatch is lost after the
1068 parameter unwrap -- HF native F.embedding would index out of range when
1069 given global token ids. The wrapper: tokens outside the local vocab
1070 interval [lo, hi) are zeroed and indices are shifted by the offset, so the
1071 output is naturally a Partial contribution and the boundary exit's
1072 Partial->Shard(1) reduction is unchanged.
1074 The interval must come from the *local* vocab shard. The applier reaches
1075 this after the Phase C one-shot ``_local_params_context`` unwrap (plain
1076 local parameter), but a caller may install the boundary while the weight
1077 is still a DTensor -- ``DTensor.shape`` is the global shape and would give
1078 a wrong interval -- so the local shard is read explicitly.
1079 """
1080 original_forward = module.forward
1081 weight = module.weight
1082 to_local = getattr(weight, "to_local", None)
1083 local_weight = to_local() if callable(to_local) else weight
1084 v_local = local_weight.shape[0]
1085 lo = tp_mesh.get_local_rank() * v_local
1086 hi = lo + v_local
1088 @functools.wraps(original_forward)
1089 def masked_embedding_forward(
1090 input_ids: torch.Tensor, *args: Any, **kwargs: Any,
1091 ) -> torch.Tensor:
1092 """Masked vocab-parallel embedding forward: zero out-of-range token
1093 ids, shift indices by the local vocab offset, and mask the output so
1094 it is a pure Partial contribution.
1096 Args:
1097 input_ids: Global token IDs at the vocab-parallel boundary.
1098 *args: Additional original embedding arguments.
1099 **kwargs: Additional original embedding keyword arguments.
1100 """
1101 mask = (input_ids >= lo) & (input_ids < hi)
1102 local_ids = torch.where(mask, input_ids - lo, torch.zeros_like(input_ids))
1103 out = original_forward(local_ids, *args, **kwargs)
1104 return out * mask.unsqueeze(-1).to(out.dtype)
1106 module.forward = masked_embedding_forward
1108# ────────────────────────────────────────────────────────────────────────────
1109# Install primitives — the ONLY forward-assignment sites in repo-owned code
1110# (AST gate, 05 §15.2.3). Everywhere else, a forward mutation is expressed as
1111# a returned callable or _ForwardRewriteRequest and committed here.
1112# ────────────────────────────────────────────────────────────────────────────
1114_MISSING = object()
1117class _ForwardRewriteRequest:
1118 """An atomic forward-rewrite request returned by wrapper factories.
1120 Carries the target module, the replacement forward and the companion
1121 attribute updates (e.g. ``attention_interface``/context flags) that must
1122 commit together with the forward swap; the rewriter validates, commits
1123 and rolls back as one unit (05 §15.2.3, P0 atomicity requirement).
1124 """
1126 __slots__ = ("target", "forward", "companion_attrs")
1128 def __init__(self, target: nn.Module, forward: Callable[..., Any], companion_attrs: Any = ()) -> None:
1129 """Keep one forward replacement and its companion attribute updates."""
1130 self.target = target
1131 self.forward = forward
1132 self.companion_attrs = dict(companion_attrs)
1135def _restore_attr(module, name, saved):
1136 """Restore one attribute to its pre-rewrite state (_MISSING deletes it)."""
1137 if saved is _MISSING:
1138 vars(module).pop(name, None)
1139 else:
1140 setattr(module, name, saved)
1143def _commit_forward_rewrite(request):
1144 """Commit one rewrite request atomically, verbatim (no dual-mode adapter).
1146 Used for secondary targets of a wrapper (e.g. the text-model input
1147 sharding of the MLA/DSA Ulysses scheme); the primary target of a plan
1148 injection goes through ``_wrap_inner_attention`` instead, which validates
1149 the signature and installs the dual-mode adapter.
1150 """
1151 target = request.target
1152 saved = {
1153 "forward": target.__dict__.get("forward", _MISSING),
1154 **{name: target.__dict__.get(name, _MISSING)
1155 for name in request.companion_attrs},
1156 }
1157 try:
1158 for attr_name, value in request.companion_attrs.items():
1159 setattr(target, attr_name, value)
1160 target.forward = request.forward
1161 except Exception:
1162 for attr_name, value in saved.items():
1163 _restore_attr(target, attr_name, value)
1164 raise
1165 return saved
1168def _classify_rewrite_result(returned, target, wrapper_name):
1169 """Normalize a wrapper's return into (primary request, secondary requests).
1171 Return contract for @inner_wrapper callables: None for in-place
1172 replacement (the external contract — target.forward was assigned by the
1173 wrapper itself); a replacement forward callable for the resolved target;
1174 or _ForwardRewriteRequest(s) — the entry targeting the resolved module is
1175 the primary (validated + dual-mode adapter installed), the rest are
1176 committed verbatim.
1177 """
1178 if returned is None:
1179 return None, []
1180 entries = (list(returned)
1181 if isinstance(returned, (list, tuple)) else [returned])
1182 requests = []
1183 for entry in entries:
1184 if isinstance(entry, _ForwardRewriteRequest):
1185 requests.append(entry)
1186 elif callable(entry):
1187 requests.append(_ForwardRewriteRequest(target, entry))
1188 else:
1189 raise TypeError(
1190 f"inner_wrapper {wrapper_name!r} returned "
1191 f"{type(entry).__name__} — expected None (in-place forward "
1192 "replacement), a replacement forward callable, or a "
1193 "_ForwardRewriteRequest")
1194 primaries = [request for request in requests if request.target is target]
1195 if len(primaries) > 1:
1196 raise TypeError(
1197 f"inner_wrapper {wrapper_name!r} returned multiple rewrite "
1198 "requests for the resolved target — exactly one primary "
1199 "replacement is allowed")
1200 secondaries = [request for request in requests
1201 if request.target is not target]
1202 return (primaries[0] if primaries else None), secondaries
1205def _install_bound_forward(module, unbound_forward):
1206 """Install ``unbound_forward(module, *args, **kwargs)`` as a bound forward.
1208 The repo's only ``types.MethodType`` forward installation: for entry-point
1209 binders whose replacement is written as a plain function taking the module
1210 explicitly (e.g. the EP local-expert entry point); companion attributes
1211 (``local_expert_count`` and friends) are set by the binder beforehand.
1212 """
1213 module.forward = types.MethodType(unbound_forward, module)
1216# ────────────────────────────────────────────────────────────────────────────
1217# Local-compute descendant adapters (moved from local_compute_context.py,
1218# which keeps only the region state)
1219# ────────────────────────────────────────────────────────────────────────────
1221_ADAPTED_MODULES = weakref.WeakSet()
1224def install_local_compute_forward_adapters(
1225 module: Any,
1226 exclude: Sequence[str] = (),
1227) -> None:
1228 """Make descendant forwards expose their directly owned local parameters.
1230 A descendant FSDP pre-hook executes before its wrapped ``forward``, so the
1231 wrapper observes and temporarily unwraps the newly installed unsharded
1232 DTensor. Its restoration finishes before the FSDP post-hook. The same
1233 wrapper is a no-op when FSDP is disabled or execution is outside a local
1234 compute region.
1236 The root is intentionally omitted: the local-region skeleton already
1237 handles all parameters visible when its forward starts. Descendant
1238 installation is idempotent across overlapping local regions.
1240 Args:
1241 module: Root module of a trainer-managed local compute region.
1242 exclude: Relative descendant FQNs that remain DTensor dispatch islands.
1243 """
1244 excluded = tuple(name.rstrip(".") for name in exclude)
1245 descendants = tuple(module.named_modules())[1:]
1246 for relative_fqn, target in descendants:
1247 if any(
1248 relative_fqn == name or relative_fqn.startswith(name + ".")
1249 for name in excluded
1250 ):
1251 continue
1252 if target in _ADAPTED_MODULES:
1253 continue
1254 _wrap_module_forward(target)
1255 _ADAPTED_MODULES.add(target)
1258def _wrap_module_forward(module: Any) -> None:
1259 """Wrap one descendant module without changing its public forward contract."""
1260 original_forward = module.forward
1262 @functools.wraps(original_forward)
1263 def local_param_forward(*args: Any, **kwargs: Any) -> Any:
1264 """Run forward with directly owned DTensor parameters unwrapped locally."""
1265 if not _LOCAL_COMPUTE_ACTIVE.get():
1266 return original_forward(*args, **kwargs)
1267 saved = []
1268 for name, parameter in module.named_parameters(recurse=False):
1269 if isinstance(parameter, DTensor):
1270 saved.append((name, parameter))
1271 module._parameters[name] = parameter.to_local() # pylint: disable=protected-access
1272 try:
1273 return original_forward(*args, **kwargs)
1274 finally:
1275 for name, parameter in saved:
1276 module._parameters[name] = parameter # pylint: disable=protected-access
1278 module.forward = local_param_forward