Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / applier.py: 80%
159 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"""applier: apply-phase orchestration — injection preflight (fail-fast before
16any mutation), apply-stage mesh selection, special-handler dispatch, and the
17Phase C per-boundary driving loop.
19The forward writers themselves live in ``forward_rewriter.py``; Target/spec
20resolution lives in ``rule_resolver.py``; parameter sharding lives in
21``parameter_sharding.py``.
22"""
24import importlib.metadata
25import logging
27import numpy as np
29from hyper_parallel.core.dtensor.device_mesh import DeviceMesh, init_device_mesh
30from hyper_parallel.distributed._builder.forward_rewriter import (
31 _bind_input_indices,
32 _descendant_boundary_fqns,
33 _install_bias_suppression,
34 _is_vocab_parallel_embed,
35 _keep_loss_parallel_output,
36 _wrap_inner_attention,
37 _wrap_local_region_forward,
38 _wrap_production_forward,
39 _wrap_validate_forward,
40 _wrap_vocab_parallel_embedding,
41)
42from hyper_parallel.distributed._builder.parameter_sharding import (
43 _resolve_module,
44)
45from hyper_parallel.distributed._builder.planner import (
46 _get_architecture,
47)
48from hyper_parallel.distributed._builder.precompiled_boundary import (
49 PrecompiledBoundary,
50)
51from hyper_parallel.distributed._builder.rule_resolver import (
52 _require_region_dispatch,
53 _resolve_local_compute_fn,
54)
55from hyper_parallel.distributed._builder.special_handlers import (
56 SPECIAL_HANDLERS,
57)
58from hyper_parallel.distributed._builder.tp_collective_lowering import (
59 create_tp_collective_lowerer,
60)
61from hyper_parallel.distributed.tensor_parallel.head_count import (
62 maybe_update_head_counts,
63)
65logger = logging.getLogger(__name__)
68# ────────────────────────────────────────────────────────────────────────────
69# Injection preflight (any mutation happens only after all validators pass)
70# ────────────────────────────────────────────────────────────────────────────
72def _validate_inner_wrapper_injections(plan, cp_mesh, transformers_version):
73 """Validate the structural contract of every inner wrapper."""
74 # Lazy import: wrappers imports forward_rewriter at module level;
75 # keeping it lazy decouples the applier's import graph.
76 from hyper_parallel.distributed.context_parallel.wrappers import ( # pylint: disable=C0415
77 INNER_WRAPPER_REQUIREMENTS,
78 )
79 for fqn, spec in plan.modules.items():
80 wrapper = getattr(spec, "inner_wrapper", None)
81 target = getattr(spec, "inner_target", None)
82 region_dispatch = getattr(spec, "region_dispatch", None)
83 if wrapper is None and target is not None:
84 raise ValueError(
85 f"Invalid inner-wrapper plan at boundary {fqn!r} "
86 f"(transformers={transformers_version}): inner_target={target!r} "
87 "locates a module but inner_wrapper is missing. Declare both "
88 "fields or remove inner_target."
89 )
90 if wrapper is not None and target is None:
91 raise ValueError(
92 f"Invalid inner-wrapper plan at boundary {fqn!r} "
93 f"(transformers={transformers_version}): "
94 f"inner_wrapper={wrapper!r}, inner_target=None, "
95 f"region_dispatch={region_dispatch!r}. Declare "
96 "inner_target='self' for the boundary module or an explicit "
97 "child attribute name."
98 )
99 requirements = (
100 INNER_WRAPPER_REQUIREMENTS.get(wrapper)
101 if isinstance(wrapper, str) else None
102 )
103 if requirements is not None:
104 required_dispatch = requirements["region_dispatch"]
105 if region_dispatch is not required_dispatch:
106 raise ValueError(
107 f"Invalid built-in CP wrapper plan at boundary {fqn!r} "
108 f"(transformers={transformers_version}): "
109 f"inner_wrapper={wrapper!r} contains CP communication and "
110 f"requires region_dispatch={required_dispatch}, got "
111 f"{region_dispatch!r}. Suggested YAML:\n"
112 f" - match: {fqn!r}\n"
113 " when: cp\n"
114 f" region_dispatch: {str(required_dispatch).lower()}\n"
115 f" inner_target: {target!r}\n"
116 f" inner_wrapper: {wrapper}"
117 )
118 if requirements["requires_cp"] and (
119 cp_mesh is None or cp_mesh.size() <= 1):
120 raise ValueError(
121 f"Invalid built-in CP wrapper plan at boundary {fqn!r}: "
122 f"inner_wrapper={wrapper!r} requires an active cp mesh, "
123 "but the plan has no cp axis with size > 1"
124 )
125 if wrapper is not None and target != "self" and getattr(
126 spec, "inner_out_src", None) is None:
127 raise ValueError(
128 f"Invalid inner-wrapper plan at boundary {fqn!r}: wrapper "
129 f"{wrapper!r} targets child {target!r}, but inner_out_src is "
130 "missing. Declare inner_out_src='first_input' for a "
131 "layout-preserving single output, provide explicit output "
132 "placements, or use inner_target='self' to reuse boundary "
133 "out_src."
134 )
137def _validate_cp_compute_injections(plan, cp_mesh):
138 """Require an explicit CP wrapper for each active attention boundary."""
139 # Lazy import: see _validate_inner_wrapper_injections.
140 from hyper_parallel.distributed.context_parallel.wrappers import ( # pylint: disable=C0415
141 INNER_WRAPPER_REGISTRY,
142 )
143 if cp_mesh is not None and cp_mesh.size() > 1:
144 for fqn, spec in plan.modules.items():
145 if (spec.is_boundary and getattr(spec, "_needs_cp_attn", False)
146 and getattr(spec, "inner_wrapper", None) is None):
147 raise ValueError(
148 f"cp_size={cp_mesh.size()} is active, so attention "
149 f"boundary {fqn!r} needs a CP-aware inner forward "
150 "(K/V all-gather), but no inner_wrapper is declared — "
151 "the framework no longer picks one heuristically. "
152 "Inject explicitly:\n"
153 " plan_overrides:\n"
154 " - match: \"*.self_attn\"\n"
155 " when: cp\n"
156 " region_dispatch: false # the wrapper contains "
157 "communication; must not dispatch\n"
158 " inner_wrapper:\n"
159 " _target_: hyper_parallel.distributed.context_parallel."
160 "wrappers.sdpa_hf_cp_wrapper\n"
161 f"(the registry {sorted(INNER_WRAPPER_REGISTRY)} can be "
162 "referenced by str name; use sdpa_qkv for the "
163 "NeMo-style (q,k,v) signature, sdpa_hf for the "
164 "HF-style forward(hidden_states); or provide a custom "
165 "callable/Target implementation)")
168def _validate_ep_compute_injections(plan, model):
169 """Require an explicit compute source for each active EP boundary."""
170 # Lazy import: see _validate_inner_wrapper_injections.
171 from hyper_parallel.distributed.expert_parallel.recipes import ( # pylint: disable=C0415
172 EP_ARCHETYPE_SUGGESTIONS,
173 )
174 for fqn, spec in plan.modules.items():
175 if not spec.is_boundary:
176 continue
177 if (getattr(spec, "_ep_size", 0) # pylint: disable=protected-access
178 and getattr(spec, "local_compute_fn", None) is None
179 and getattr(spec, "region_dispatch", None) is not False):
180 suggestion = ""
181 if model is not None:
182 arch = _get_architecture(model)
183 archetype = EP_ARCHETYPE_SUGGESTIONS.get(arch)
184 if archetype is not None:
185 suggestion = (
186 f"Detected model architecture {arch!r}; the "
187 f"matching archetype is likely {archetype!r} (a "
188 "suggestion, not an automatic choice — confirm and "
189 "configure it explicitly; a wrong pick fails the "
190 "apply-time interface assertion and lists the "
191 "module's actual submodule names).\n"
192 )
193 raise ValueError(
194 f"ep_size={spec._ep_size} is active (parameters explicitly " # pylint: disable=protected-access
195 f"declared with EP sharding use the virtual-EP source mesh), but boundary "
196 f"{fqn!r} has no local-region compute source — nothing "
197 "executes the corresponding local compute and all-to-all, and the "
198 "framework no longer injects any implementation "
199 "automatically. Choose one of:\n"
200 f"{suggestion}"
201 " ① Routed MoE: pick a built-in archetype factory per the model's "
202 "behavior (full semantics: router / shared expert / gate / "
203 "merge all implemented cohesively; see the recipes.py "
204 "module docstring for the available archetypes and the "
205 "module interface each one expects):\n"
206 " plan_overrides:\n"
207 " - match: \"*.mlp\"\n"
208 " when: ep\n"
209 " region_dispatch: false # a2a is inside the "
210 "region; must not dispatch\n"
211 " local_compute_fn:\n"
212 " _target_: hyper_parallel.distributed.expert_parallel."
213 "recipes.qwen2moe_ep_compute_fn # pick per the "
214 "archetype table\n"
215 " ② Model-specific MoE or another EP-sharded module (for "
216 "example a sparse table): provide a local_compute_fn from "
217 "models/<family>/adapter and validate its interface at build time\n"
218 " ③ A module whose forward already owns all-to-all: declare "
219 "region_dispatch: false")
222def _preflight_compute_injection(plan, mesh, model=None):
223 """Fail fast before mutating a plan with incomplete CP/EP injection."""
224 transformers_version = importlib.metadata.version("transformers")
225 cp_mesh = _get_cp_submesh(mesh, plan.mesh_dim_names)
226 _validate_inner_wrapper_injections(plan, cp_mesh, transformers_version)
227 _validate_cp_compute_injections(plan, cp_mesh)
228 _validate_ep_compute_injections(plan, model)
231def _log_injection_choice(module_fqn, spec):
232 """Observability enhancement: the resolution chosen for an injection
233 boundary is immediately visible (one INFO line per boundary), closing
234 the "declare → see the consequence" feedback loop. Must be called after
235 _require_region_dispatch (the injection + region_dispatch combination is
236 already validated as legal at this point)."""
237 rd = getattr(spec, "region_dispatch", None)
238 has_fn = getattr(spec, "local_compute_fn", None) is not None
239 has_wrap = getattr(spec, "inner_wrapper", None) is not None
240 if not (has_fn or has_wrap or rd is not None):
241 return # ordinary boundary: no injection, axiomatic default dispatch-through — do not spam the log
242 candidates = (("local_compute_fn", has_fn), ("inner_wrapper", has_wrap))
243 enabled = []
244 for name, is_enabled in candidates:
245 if is_enabled:
246 enabled.append(name)
247 what = "+".join(enabled) or "the module's own forward (no fn injected, region_dispatch=False declared)"
248 if rd is True:
249 effect = ("validate dispatch-through true validation enabled "
250 "(in-region strategy propagation + true out_src "
251 "validation)")
252 else: # False (None + injection was already intercepted by _require_region_dispatch)
253 effect = ("black-box hosting (local execution inside the region, "
254 "propagation checks skipped, declarative rewrap)")
255 logger.info(
256 "boundary %s: injection[%s], region_dispatch=%s → %s",
257 module_fqn, what, rd, effect)
260# ────────────────────────────────────────────────────────────────────────────
261# Phase B: special handlers (registry dispatch only; concrete handlers live
262# in special_handlers.py)
263# ────────────────────────────────────────────────────────────────────────────
265def _apply_plan_special_handlers(models, plan, mesh):
266 """Run parameter handlers declared by the sharding plan."""
267 for model in models:
268 for param_ref, handler_name in plan.special_handlers.items():
269 handler = SPECIAL_HANDLERS.get(handler_name)
270 if handler is None:
271 logger.warning(
272 "SPECIAL_HANDLERS has no registered handler: %s",
273 handler_name,
274 )
275 continue
276 module_fqn, param_name = param_ref.rsplit(".", 1)
277 handler(_resolve_module(model, module_fqn), param_name, mesh)
280# ────────────────────────────────────────────────────────────────────────────
281# Apply-stage mesh selection
282# ────────────────────────────────────────────────────────────────────────────
284def _get_active_mesh(mesh, mesh_dim_names):
285 """Return the active sub-mesh aligned with plan.mesh_dim_names (the dimension set after stripping size=1 axes)."""
286 names = tuple(getattr(mesh, "mesh_dim_names", ()) or ())
287 if names == tuple(mesh_dim_names):
288 return mesh
289 if mesh_dim_names and names and all(n in names for n in mesh_dim_names):
290 return mesh[tuple(mesh_dim_names)]
291 return mesh
294def _get_active_mesh_with_names(mesh, mesh_dim_names):
295 """Return ``(active_sub_mesh, active_dim_names)`` sliced in mesh-major order.
297 A variant of :func:`_get_active_mesh` for callers whose axis names may
298 arrive in an order that differs from the mesh's declared order (the
299 codegen runtime's frozen plans record axes in offline-mesh order while a
300 live ``DeviceMesh`` declares e.g. ``("dp", "cp", "tp")``).
301 ``DeviceMesh.__getitem__`` only accepts a slice whose axis indices are
302 ascending, so the names are reordered into the mesh's own declared order
303 before slicing, and the reordered tuple is returned alongside the mesh:
304 the caller must resolve placements against the *returned* names, because
305 the sliced sub-mesh carries exactly that axis order.
307 A mesh without the given axes (e.g. an ``OfflineMesh`` built without
308 ``ep``) returns itself with the intersecting names; a mesh that HAS the
309 axes but cannot be sliced raises — never silently returns the unsliced
310 mesh (that would shard along a wrong axis).
311 """
312 if not mesh_dim_names:
313 return mesh, ()
314 names = tuple(getattr(mesh, "mesh_dim_names", ()) or ())
315 if not names:
316 return mesh, tuple(mesh_dim_names)
317 # Mesh-major order: keep only the axes the plan shards on, reordered to
318 # the mesh's own declared order (so `mesh[tuple(...)]` indices ascend).
319 active = tuple(n for n in names if n in mesh_dim_names)
320 if not active:
321 return mesh, ()
322 # Only skip slicing when the mesh declares EXACTLY the plan's axes (in the
323 # same order). Comparing against ``names`` (not the filtered ``active``)
324 # prevents a mesh that also carries non-plan axes -- e.g. a ('dp', 'tp')
325 # mesh vs a ('tp',) frozen plan -- from returning the unsliced 2-D mesh
326 # alongside a 1-D placement axis. When the mesh carries any extra axis the
327 # plan does not shard on, it must be sliced down to ``active`` so the dense
328 # mesh dimension count always matches the placement tuple.
329 if tuple(names) == tuple(mesh_dim_names):
330 return mesh, tuple(mesh_dim_names)
331 # The plan may name an axis the mesh does not carry (e.g. a plan frozen for
332 # an ``ep`` axis on a mesh that derives EP from the dense region). Never
333 # silently return the unsliced mesh with a dim name that does not exist on
334 # it — that would resolve placements against a phantom axis and corrupt
335 # every downstream DTensor layout. Slice to the overlapping active axes
336 # and keep the same mesh-major names.
337 return mesh[active], active
340def _get_tp_submesh(mesh, mesh_dim_names):
341 if "tp" not in mesh_dim_names:
342 return None
343 return mesh["tp"]
346def _get_cp_submesh(mesh, mesh_dim_names):
347 if "cp" not in mesh_dim_names:
348 return None
349 return mesh["cp"]
352def _get_ep_submesh(mesh, mesh_dim_names):
353 if "ep" not in mesh_dim_names:
354 return None
355 return mesh["ep"]
358def _expert_mesh_layout(mesh, mesh_dim_names, ep_size):
359 """(shape, dim_names, rank_list) of the derived expert mesh (pure mapping; no process group is created).
361 D-10 TP-extend-EP (05 §6.4.8 / 06 §4.5.1): the expert domain = the full
362 dense region (all ranks on the non-pp axes of the mesh, i.e.
363 dp_replicate x dp_cp x tp). After a row-major flatten in mesh axis order,
364 it is re-sliced as (edp = D/ep_size, ep = ep_size):
365 - EP groups (inner, the a2a communication domain): ep_size consecutive
366 ranks in flatten order -- tp is usually the innermost axis, so an EP
367 group first spans the entire TP group and then extends to adjacent
368 dp/cp ranks (isomorphic to MindSpeed TP-extend-EP / Megatron etp=1 with
369 ep spanning TP).
370 Example: mesh (dp=4, tp=2), ep_size=4 -> EP groups {0,1,2,3} / {4,5,6,7};
371 - edp groups (outer): expert data-parallel degree = D/ep_size.
372 Expert weights are only Shard(0) on the ep axis (the expert dim); there is
373 no second-axis sharding.
374 """
375 if "pp" in mesh_dim_names:
376 raise NotImplementedError(
377 "D-10 TP-extend-EP v1 does not yet support the pp axis (call after splitting the mesh by stage)")
378 arr = np.array(mesh.rank_list).reshape(mesh.mesh_shape)
379 domain = int(np.prod(arr.shape))
380 if ep_size <= 0 or domain % ep_size != 0:
381 raise ValueError(
382 f"ep_size ({ep_size}) must divide the dense region ({domain})"
383 )
384 edp = domain // ep_size
385 derived = arr.reshape(edp, ep_size)
386 return (edp, ep_size), ("edp", "ep"), tuple(
387 int(r) for r in derived.flatten())
390def _build_expert_mesh(mesh, mesh_dim_names, ep_size):
391 """D-10 (05 §6.4.8 / 06 §4.5.1): repartition the full dense region into the derived expert mesh (edp, ep)."""
392 shape, names, rank_list = _expert_mesh_layout(mesh, mesh_dim_names, ep_size)
393 # Propagate the no-backend (metadata-only) mode of the source mesh —
394 # a meta mesh has no _dim_group_names (no process groups were created).
395 return init_device_mesh(mesh.device_type, shape, mesh_dim_names=names,
396 rank_list=rank_list,
397 init_backend=hasattr(mesh, "_dim_group_names"))
400def build_expert_mesh(mesh: DeviceMesh, ep_size: int) -> DeviceMesh:
401 """Public: derive the D-10 TP-extend-EP expert mesh (edp, ep) from *mesh*
402 (the FULL mesh — the dense region must include dp/cp axes).
404 Standalone helper for introspection and for custom code that needs the
405 expert domain outside the injection path. Injected factories/wrappers do
406 NOT need to call this: the framework derives the expert mesh once at
407 apply time (shared by parameter sharding and injected compute) and hands
408 it to them as the ``ep_mesh`` context.
410 Args:
411 mesh: The full device mesh (dense region including dp/cp axes).
412 ep_size: The expert-parallel group size.
414 Returns:
415 The derived ``(edp, ep)`` expert mesh.
416 """
417 return _build_expert_mesh(mesh, tuple(mesh.mesh_dim_names), ep_size)
420# ────────────────────────────────────────────────────────────────────────────
421# Phase C: forward wrapping driving loop (05 §4.4)
422# ────────────────────────────────────────────────────────────────────────────
424def _boundary_post_order_key(item):
425 """Sort key for boundary installation: deepest FQN first (post-order).
427 D-14 invariant 2 (05 §13.3): boundaries are wrapped in post-order — an
428 outer boundary's local_compute_fn may cache inner forwards, and the
429 unpack-scope exclusion (invariant 3) requires inner wrappers to be
430 installed first. Shared with the codegen runtime's install loop so both
431 paths keep a single ordering contract.
432 """
433 return -item[0].count(".")
436def _apply_phase_c(model, plan, mesh, validate_mode, expert_mesh=None):
437 """Phase C: wrap forward (production/validate/moe/cp/vocab_embed, five paths).
439 D-14 invariant 2 (05 §13.3): boundaries are wrapped in post-order (deepest
440 FQN first) — an outer boundary's local_compute_fn may cache inner forwards,
441 and the unpack-scope exclusion (invariant 3) requires inner wrappers to be
442 installed first.
443 """
444 mesh_dim_names = plan.mesh_dim_names
445 cp_mesh = _get_cp_submesh(mesh, mesh_dim_names)
446 tp_mesh = _get_tp_submesh(mesh, mesh_dim_names)
447 boundary_op_lowerer = (
448 None
449 if validate_mode
450 else create_tp_collective_lowerer(mesh, mesh_dim_names)
451 )
452 for module_fqn, spec in sorted(
453 plan.modules.items(), key=_boundary_post_order_key):
454 if not spec.is_boundary:
455 continue
456 module = _resolve_module(model, module_fqn)
457 boundary = PrecompiledBoundary(
458 spec,
459 mesh,
460 mesh_dim_names,
461 op_lowerer=boundary_op_lowerer,
462 )
463 _bind_input_indices(boundary, module)
464 # Injection discipline: a declared injection requires an explicit
465 # region_dispatch; declaring True without any injection is redundant
466 _require_region_dispatch(spec, source=f"boundary {module_fqn!r}")
467 # Observability: the injection choice is immediately visible (the
468 # "declare → consequence" feedback loop)
469 _log_injection_choice(module_fqn, spec)
471 # Step 0.5 (D-22): rowwise bias defer — install the bias-free region
472 # forward BEFORE inner-wrap/forward wrapping, in BOTH modes (the
473 # in-region computation stays instruction-for-instruction identical).
474 # The bias Parameter itself is untouched; it is added once after the
475 # boundary exit reduction by _maybe_add_deferred_biases.
476 if spec._deferred_bias_params: # pylint: disable=protected-access
477 _install_bias_suppression(module, spec)
479 # Step 1: inner-wrap — the generic "weave into / replace the inner
480 # forward" mechanism
481 # (D-01'': production and validate inject the same wrapper, so the
482 # in-region computation is instruction-for-instruction identical).
483 # NOT gated on cp_mesh since the generalization: the derived gate is
484 # the resolution chain itself (explicit inner_wrapper declaration ->
485 # applied; nothing declared -> None -> no-op). cp_mesh may be None
486 # (no cp axis) — the four shipped CP wrappers self-guard and fail
487 # fast then; custom callables/Targets receive None and own their
488 # semantics.
489 # (_preflight_compute_injection has already failed fast when a CP
490 # attention boundary declares nothing)
491 _wrap_inner_attention(
492 module, cp_mesh, spec=spec, mesh=mesh,
493 mesh_dim_names=mesh_dim_names, tp_mesh=tp_mesh,
494 ep_mesh=expert_mesh, validate_mode=validate_mode,
495 module_fqn=module_fqn,
496 )
498 # Step 2: forward wrapping
499 # local region path (D-03'): the gate is derived from the compute_fn
500 # resolution chain (non-None means take the skeleton) — the remaining
501 # sources after the explicit-injection rework are user
502 # local_compute_fn (callable or factory Target) and the
503 # region_dispatch=False gate; the built-in EP auto-injection link
504 # was REMOVED (05 §4.4.3)
505 compute_fn = _resolve_local_compute_fn(
506 module, spec, mesh, mesh_dim_names, expert_mesh)
507 if compute_fn is not None:
508 if validate_mode:
509 # D-17: inside the local region the module sees local
510 # tensors in both modes -- validate rewrites cached head
511 # counts only for local-region modules (boundary modules
512 # keep global counts for DTensor dispatch)
513 maybe_update_head_counts(
514 module, spec, module_fqn, mesh, mesh_dim_names)
515 # D-14 invariant 3 (05 §13.3): the region's temp-unwrap scope
516 # excludes nested-boundary subtrees — their parameters must stay
517 # DTensors for the inner validate islands (dispatch needs
518 # __torch_function__)
519 nested = _descendant_boundary_fqns(plan, module_fqn)
520 if nested:
521 # E5 (accuracy_fix_plan.md §3): make the nested-boundary call
522 # contract visible at apply time
523 logger.info(
524 "local-region boundary %s: nested boundaries "
525 "(communication_owner: nested_boundary): %s — their "
526 "return values are already in the declared out_dst "
527 "layout with TP communication sealed inside; the region "
528 "compute MUST NOT apply compensating collectives "
529 "(all-reduce/reduce-scatter/all-gather) to them",
530 module_fqn, nested)
531 _wrap_local_region_forward(
532 module, boundary, spec, mesh, mesh_dim_names,
533 validate_mode=validate_mode, compute_fn=compute_fn,
534 exclude_subtrees=nested)
535 elif validate_mode:
536 _wrap_validate_forward(
537 module,
538 boundary,
539 spec,
540 mesh_dim_names,
541 keep_output_dtensor=_keep_loss_parallel_output(plan, spec),
542 )
543 else:
544 # D-02: production vocab-parallel embedding masked wrapper
545 if _is_vocab_parallel_embed(module, spec, tp_mesh):
546 _wrap_vocab_parallel_embedding(module, tp_mesh)
547 _wrap_production_forward(
548 module,
549 boundary,
550 spec,
551 keep_output_dtensor=_keep_loss_parallel_output(plan, spec),
552 )