Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / applier.py: 88%
161 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"""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.components.modules.module_family import get_family_build_info
30from hyper_parallel.core.dtensor.device_mesh import DeviceMesh, init_device_mesh
31from hyper_parallel.distributed._builder.forward_rewriter import (
32 _bind_input_indices,
33 _descendant_boundary_fqns,
34 _install_bias_suppression,
35 _is_vocab_parallel_embed,
36 _keep_loss_parallel_output,
37 _wrap_inner_attention,
38 _wrap_local_region_forward,
39 _wrap_production_forward,
40 _wrap_validate_forward,
41 _wrap_vocab_parallel_embedding,
42)
43from hyper_parallel.distributed._builder.parameter_sharding import (
44 _resolve_module,
45)
46from hyper_parallel.distributed._builder.planner import (
47 _get_architecture,
48)
49from hyper_parallel.distributed._builder.precompiled_boundary import (
50 PrecompiledBoundary,
51)
52from hyper_parallel.distributed._builder.rule_resolver import (
53 _require_region_dispatch,
54 _resolve_local_compute_fn,
55)
56from hyper_parallel.distributed._builder.special_handlers import (
57 SPECIAL_HANDLERS,
58)
59from hyper_parallel.distributed._builder.tp_collective_lowering import (
60 create_tp_collective_lowerer,
61)
62from hyper_parallel.distributed.tensor_parallel.head_count import (
63 maybe_update_head_counts,
64)
66logger = logging.getLogger(__name__)
69# ────────────────────────────────────────────────────────────────────────────
70# Injection preflight (any mutation happens only after all validators pass)
71# ────────────────────────────────────────────────────────────────────────────
73def _validate_inner_wrapper_injections(plan, cp_mesh, transformers_version):
74 """Validate the structural contract of every inner wrapper."""
75 # Lazy import: wrappers imports forward_rewriter at module level;
76 # keeping it lazy decouples the applier's import graph.
77 from hyper_parallel.distributed.context_parallel.wrappers import ( # pylint: disable=C0415
78 INNER_WRAPPER_REQUIREMENTS,
79 )
80 for fqn, spec in plan.modules.items():
81 wrapper = getattr(spec, "inner_wrapper", None)
82 target = getattr(spec, "inner_target", None)
83 region_dispatch = getattr(spec, "region_dispatch", None)
84 if wrapper is None and target is not None:
85 raise ValueError(
86 f"Invalid inner-wrapper plan at boundary {fqn!r} "
87 f"(transformers={transformers_version}): inner_target={target!r} "
88 "locates a module but inner_wrapper is missing. Declare both "
89 "fields or remove inner_target."
90 )
91 if wrapper is not None and target is None:
92 raise ValueError(
93 f"Invalid inner-wrapper plan at boundary {fqn!r} "
94 f"(transformers={transformers_version}): "
95 f"inner_wrapper={wrapper!r}, inner_target=None, "
96 f"region_dispatch={region_dispatch!r}. Declare "
97 "inner_target='self' for the boundary module or an explicit "
98 "child attribute name."
99 )
100 requirements = (
101 INNER_WRAPPER_REQUIREMENTS.get(wrapper)
102 if isinstance(wrapper, str) else None
103 )
104 if requirements is not None:
105 required_dispatch = requirements["region_dispatch"]
106 if region_dispatch is not required_dispatch:
107 raise ValueError(
108 f"Invalid built-in CP wrapper plan at boundary {fqn!r} "
109 f"(transformers={transformers_version}): "
110 f"inner_wrapper={wrapper!r} contains CP communication and "
111 f"requires region_dispatch={required_dispatch}, got "
112 f"{region_dispatch!r}. Set region_dispatch on the "
113 f"match={fqn + '.forward'!r} method replacement."
114 )
115 if requirements["requires_cp"] and (
116 cp_mesh is None or cp_mesh.size() <= 1):
117 raise ValueError(
118 f"Invalid built-in CP wrapper plan at boundary {fqn!r}: "
119 f"inner_wrapper={wrapper!r} requires an active cp mesh, "
120 "but the plan has no cp axis with size > 1"
121 )
122 if wrapper is not None and target != "self" and getattr(
123 spec, "inner_out_src", None) is None:
124 raise ValueError(
125 f"Invalid inner-wrapper plan at boundary {fqn!r}: wrapper "
126 f"{wrapper!r} targets child {target!r}, but inner_out_src is "
127 "missing. Declare inner_out_src='first_input' for a "
128 "layout-preserving single output, provide explicit output "
129 "placements, or use inner_target='self' to reuse boundary "
130 "out_src."
131 )
134def _validate_cp_compute_injections(plan, cp_mesh):
135 """Require an explicit forward replacement for CP attention boundaries."""
136 if cp_mesh is None or cp_mesh.size() <= 1:
137 return
138 for fqn, spec in plan.modules.items():
139 if fqn in plan.parallel_modules:
140 continue
141 if not spec.is_boundary or not getattr(spec, "_needs_cp_attn", False):
142 continue
143 if spec.inner_wrapper is not None or spec.local_compute_fn is not None:
144 continue
145 raise ValueError(
146 f"cp_size={cp_mesh.size()} is active, but attention boundary "
147 f"{fqn!r} has no CP-aware forward replacement. Declare:\n"
148 " plan_overrides:\n"
149 f" - match: {(fqn + '.forward')!r}\n"
150 " when: cp\n"
151 " replace:\n"
152 " kind: method\n"
153 " implementation: hyper_parallel.distributed.context_parallel."
154 "wrappers.sdpa_hf_cp_wrapper\n"
155 "Choose an implementation compatible with this attention's forward signature."
156 )
159def _validate_ep_compute_injections(plan, model):
160 """Require a compute source for each active EP-sharded boundary."""
161 # Lazy import: see _validate_inner_wrapper_injections.
162 from hyper_parallel.distributed.expert_parallel.recipes import ( # pylint: disable=C0415
163 EP_ARCHETYPE_SUGGESTIONS,
164 )
165 for fqn, spec in plan.modules.items():
166 if fqn in plan.parallel_modules:
167 continue
168 if not spec.is_boundary or not getattr(spec, "_ep_size", 0):
169 continue
170 if spec.local_compute_fn is not None or spec.region_dispatch is False:
171 continue
172 suggestion = ""
173 if model is not None:
174 architecture = _get_architecture(model)
175 archetype = EP_ARCHETYPE_SUGGESTIONS.get(architecture)
176 if archetype is not None:
177 suggestion = f"Candidate for {architecture!r}: {archetype!r}. "
178 raise ValueError(
179 f"ep_size={spec._ep_size} is active, but boundary {fqn!r} has "
180 "no EP-aware compute source. "
181 f"{suggestion}Declare one of:\n"
182 " - match: '*.mlp.forward'\n"
183 " when: ep\n"
184 " replace:\n"
185 " kind: method\n"
186 " implementation: <EP-aware factory import path>\n"
187 "or use a rule with region_dispatch: false when the module's "
188 "original forward already owns the EP all-to-all."
189 )
192def _preflight_compute_injection(plan, mesh, model=None):
193 """Fail fast before mutating a plan with incomplete CP/EP injection."""
194 transformers_version = importlib.metadata.version("transformers")
195 cp_mesh = _get_cp_submesh(mesh, plan.mesh_dim_names)
196 _validate_inner_wrapper_injections(plan, cp_mesh, transformers_version)
197 _validate_cp_compute_injections(plan, cp_mesh)
198 _validate_ep_compute_injections(plan, model)
201def _log_injection_choice(module_fqn, spec):
202 """Observability enhancement: the resolution chosen for an injection
203 boundary is immediately visible (one INFO line per boundary), closing
204 the "declare → see the consequence" feedback loop. Must be called after
205 _require_region_dispatch (the injection + region_dispatch combination is
206 already validated as legal at this point)."""
207 rd = getattr(spec, "region_dispatch", None)
208 has_fn = getattr(spec, "local_compute_fn", None) is not None
209 has_wrap = getattr(spec, "inner_wrapper", None) is not None
210 if not (has_fn or has_wrap or rd is not None):
211 return # ordinary boundary: no injection, axiomatic default dispatch-through — do not spam the log
212 candidates = (("local_compute_fn", has_fn), ("inner_wrapper", has_wrap))
213 enabled = []
214 for name, is_enabled in candidates:
215 if is_enabled:
216 enabled.append(name)
217 what = "+".join(enabled) or "the module's own forward (no fn injected, region_dispatch=False declared)"
218 if rd is True:
219 effect = ("validate dispatch-through true validation enabled "
220 "(in-region strategy propagation + true out_src "
221 "validation)")
222 else: # False (None + injection was already intercepted by _require_region_dispatch)
223 effect = ("black-box hosting (local execution inside the region, "
224 "propagation checks skipped, declarative rewrap)")
225 logger.info(
226 "boundary %s: injection[%s], region_dispatch=%s → %s",
227 module_fqn, what, rd, effect)
230# ────────────────────────────────────────────────────────────────────────────
231# Phase B: special handlers (registry dispatch only; concrete handlers live
232# in special_handlers.py)
233# ────────────────────────────────────────────────────────────────────────────
235def _apply_plan_special_handlers(models, plan, mesh):
236 """Run parameter handlers declared by the sharding plan."""
237 for model in models:
238 for param_ref, handler_name in plan.special_handlers.items():
239 handler = SPECIAL_HANDLERS.get(handler_name)
240 if handler is None:
241 logger.warning(
242 "SPECIAL_HANDLERS has no registered handler: %s",
243 handler_name,
244 )
245 continue
246 module_fqn, param_name = param_ref.rsplit(".", 1)
247 handler(_resolve_module(model, module_fqn), param_name, mesh)
250# ────────────────────────────────────────────────────────────────────────────
251# Apply-stage mesh selection
252# ────────────────────────────────────────────────────────────────────────────
254def _get_active_mesh(mesh, mesh_dim_names):
255 """Return the active sub-mesh aligned with plan.mesh_dim_names (the dimension set after stripping size=1 axes)."""
256 names = tuple(getattr(mesh, "mesh_dim_names", ()) or ())
257 if names == tuple(mesh_dim_names):
258 return mesh
259 if mesh_dim_names and names and all(n in names for n in mesh_dim_names):
260 return mesh[tuple(mesh_dim_names)]
261 return mesh
264def _get_tp_submesh(mesh, mesh_dim_names):
265 if "tp" not in mesh_dim_names:
266 return None
267 return mesh["tp"]
270def _get_cp_submesh(mesh, mesh_dim_names):
271 if "cp" not in mesh_dim_names:
272 return None
273 return mesh["cp"]
276def _get_ep_submesh(mesh, mesh_dim_names):
277 if "ep" not in mesh_dim_names:
278 return None
279 return mesh["ep"]
282def _expert_mesh_layout(mesh, mesh_dim_names, ep_size):
283 """(shape, dim_names, rank_list) of the derived expert mesh (pure mapping; no process group is created).
285 D-10 TP-extend-EP (05 §6.4.8 / 06 §4.5.1): the expert domain = the full
286 dense region (all ranks on the non-pp axes of the mesh, i.e.
287 dp_replicate x dp_cp x tp). After a row-major flatten in mesh axis order,
288 it is re-sliced as (edp = D/ep_size, ep = ep_size):
289 - EP groups (inner, the a2a communication domain): ep_size consecutive
290 ranks in flatten order -- tp is usually the innermost axis, so an EP
291 group first spans the entire TP group and then extends to adjacent
292 dp/cp ranks (isomorphic to MindSpeed TP-extend-EP / Megatron etp=1 with
293 ep spanning TP).
294 Example: mesh (dp=4, tp=2), ep_size=4 -> EP groups {0,1,2,3} / {4,5,6,7};
295 - edp groups (outer): expert data-parallel degree = D/ep_size.
296 Expert weights are only Shard(0) on the ep axis (the expert dim); there is
297 no second-axis sharding.
298 """
299 if "pp" in mesh_dim_names:
300 raise NotImplementedError(
301 "D-10 TP-extend-EP v1 does not yet support the pp axis (call after splitting the mesh by stage)")
302 arr = np.array(mesh.rank_list).reshape(mesh.mesh_shape)
303 domain = int(np.prod(arr.shape))
304 if ep_size <= 0 or domain % ep_size != 0:
305 raise ValueError(
306 f"ep_size ({ep_size}) must divide the dense region ({domain})"
307 )
308 edp = domain // ep_size
309 derived = arr.reshape(edp, ep_size)
310 return (edp, ep_size), ("edp", "ep"), tuple(
311 int(r) for r in derived.flatten())
314def _build_expert_mesh(mesh, mesh_dim_names, ep_size):
315 """D-10 (05 §6.4.8 / 06 §4.5.1): repartition the full dense region into the derived expert mesh (edp, ep)."""
316 shape, names, rank_list = _expert_mesh_layout(mesh, mesh_dim_names, ep_size)
317 # Propagate the no-backend (metadata-only) mode of the source mesh —
318 # a meta mesh has no _dim_group_names (no process groups were created).
319 return init_device_mesh(mesh.device_type, shape, mesh_dim_names=names,
320 rank_list=rank_list,
321 init_backend=hasattr(mesh, "_dim_group_names"))
324def build_expert_mesh(mesh: DeviceMesh, ep_size: int) -> DeviceMesh:
325 """Public: derive the D-10 TP-extend-EP expert mesh (edp, ep) from *mesh*
326 (the FULL mesh — the dense region must include dp/cp axes).
328 Standalone helper for introspection and for custom code that needs the
329 expert domain outside the injection path. Injected factories/wrappers do
330 NOT need to call this: the framework derives the expert mesh once at
331 apply time (shared by parameter sharding and injected compute) and hands
332 it to them as the ``ep_mesh`` context.
334 Args:
335 mesh: The full device mesh (dense region including dp/cp axes).
336 ep_size: The expert-parallel group size.
338 Returns:
339 The derived ``(edp, ep)`` expert mesh.
340 """
341 return _build_expert_mesh(mesh, tuple(mesh.mesh_dim_names), ep_size)
344# ────────────────────────────────────────────────────────────────────────────
345# Phase C: forward wrapping driving loop (05 §4.4)
346# ────────────────────────────────────────────────────────────────────────────
348def _boundary_post_order_key(item):
349 """Sort key for boundary installation: deepest FQN first (post-order).
351 D-14 invariant 2 (05 §13.3): boundaries are wrapped in post-order — an
352 outer boundary's local_compute_fn may cache inner forwards, and the
353 unpack-scope exclusion (invariant 3) requires inner wrappers to be
354 installed first.
355 """
356 return -item[0].count(".")
359def _apply_phase_c(model, plan, mesh, validate_mode, expert_mesh=None):
360 """Phase C: wrap forward (production/validate/moe/cp/vocab_embed, five paths).
362 D-14 invariant 2 (05 §13.3): boundaries are wrapped in post-order (deepest
363 FQN first) — an outer boundary's local_compute_fn may cache inner forwards,
364 and the unpack-scope exclusion (invariant 3) requires inner wrappers to be
365 installed first.
366 """
367 mesh_dim_names = plan.mesh_dim_names
368 cp_mesh = _get_cp_submesh(mesh, mesh_dim_names)
369 tp_mesh = _get_tp_submesh(mesh, mesh_dim_names)
370 boundary_op_lowerer = (
371 None
372 if validate_mode
373 else create_tp_collective_lowerer(mesh, mesh_dim_names)
374 )
375 for module_fqn, spec in sorted(
376 plan.modules.items(), key=_boundary_post_order_key):
377 if module_fqn in plan.parallel_modules:
378 continue
379 if not spec.is_boundary:
380 continue
381 module = _resolve_module(model, module_fqn)
382 family = get_family_build_info(module)
383 if family is not None and family.layout is not None:
384 continue
385 boundary = PrecompiledBoundary(
386 spec,
387 mesh,
388 mesh_dim_names,
389 op_lowerer=boundary_op_lowerer,
390 )
391 _bind_input_indices(boundary, module)
392 # Injection discipline: a declared injection requires an explicit
393 # region_dispatch; declaring True without any injection is redundant
394 _require_region_dispatch(spec, source=f"boundary {module_fqn!r}")
395 # Observability: the injection choice is immediately visible (the
396 # "declare → consequence" feedback loop)
397 _log_injection_choice(module_fqn, spec)
399 # Step 0.5 (D-22): rowwise bias defer — install the bias-free region
400 # forward BEFORE inner-wrap/forward wrapping, in BOTH modes (the
401 # in-region computation stays instruction-for-instruction identical).
402 # The bias Parameter itself is untouched; it is added once after the
403 # boundary exit reduction by _maybe_add_deferred_biases.
404 if spec._deferred_bias_params: # pylint: disable=protected-access
405 _install_bias_suppression(module, spec)
407 # Step 1: inner-wrap — the generic "weave into / replace the inner
408 # forward" mechanism
409 # (D-01'': production and validate inject the same wrapper, so the
410 # in-region computation is instruction-for-instruction identical).
411 # NOT gated on cp_mesh since the generalization: the derived gate is
412 # the resolution chain itself (explicit inner_wrapper declaration ->
413 # applied; nothing declared -> None -> no-op). cp_mesh may be None
414 # (no cp axis) — the four shipped CP wrappers self-guard and fail
415 # fast then; custom callables/Targets receive None and own their
416 # semantics.
417 # (_preflight_compute_injection has already failed fast when a CP
418 # attention boundary declares nothing)
419 _wrap_inner_attention(
420 module, cp_mesh, spec=spec, mesh=mesh,
421 mesh_dim_names=mesh_dim_names, tp_mesh=tp_mesh,
422 ep_mesh=expert_mesh, validate_mode=validate_mode,
423 module_fqn=module_fqn,
424 )
426 # Step 2: forward wrapping
427 # local region path (D-03'): the gate is derived from the compute_fn
428 # resolution chain (non-None means take the skeleton) — the remaining
429 # sources after the explicit-injection rework are user
430 # local_compute_fn (callable or factory Target) and the
431 # region_dispatch=False gate; the built-in EP auto-injection link
432 # was REMOVED (05 §4.4.3)
433 compute_fn = _resolve_local_compute_fn(
434 module, spec, mesh, mesh_dim_names, expert_mesh)
435 if compute_fn is not None:
436 if validate_mode:
437 # D-17: inside the local region the module sees local
438 # tensors in both modes -- validate rewrites cached head
439 # counts only for local-region modules (boundary modules
440 # keep global counts for DTensor dispatch)
441 maybe_update_head_counts(
442 module, spec, module_fqn, mesh, mesh_dim_names)
443 # D-14 invariant 3 (05 §13.3): the region's temp-unwrap scope
444 # excludes nested-boundary subtrees — their parameters must stay
445 # DTensors for the inner validate islands (dispatch needs
446 # __torch_function__)
447 nested = _descendant_boundary_fqns(plan, module_fqn)
448 if nested:
449 # E5 (accuracy_fix_plan.md §3): make the nested-boundary call
450 # contract visible at apply time
451 logger.info(
452 "local-region boundary %s: nested boundaries "
453 "(communication_owner: nested_boundary): %s — their "
454 "return values are already in the declared out_dst "
455 "layout with TP communication sealed inside; the region "
456 "compute MUST NOT apply compensating collectives "
457 "(all-reduce/reduce-scatter/all-gather) to them",
458 module_fqn, nested)
459 _wrap_local_region_forward(
460 module, boundary, spec, mesh, mesh_dim_names,
461 validate_mode=validate_mode, compute_fn=compute_fn,
462 exclude_subtrees=nested)
463 elif validate_mode:
464 _wrap_validate_forward(
465 module,
466 boundary,
467 spec,
468 mesh_dim_names,
469 keep_output_dtensor=_keep_loss_parallel_output(plan, spec),
470 )
471 else:
472 # D-02: production vocab-parallel embedding masked wrapper
473 if _is_vocab_parallel_embed(module, spec, tp_mesh):
474 _wrap_vocab_parallel_embedding(module, tp_mesh)
475 _wrap_production_forward(
476 module,
477 boundary,
478 spec,
479 keep_output_dtensor=_keep_loss_parallel_output(plan, spec),
480 )