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