Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / meta_plan.py: 53%
110 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +0800
1# Copyright 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"""Normalize frozen meta so it matches the inline-generated source shape."""
17from __future__ import annotations
19import logging
20from copy import copy
21from fnmatch import fnmatch
22from typing import Any
24from hyper_parallel.codegen.inline.ir import SourceObservation
25from hyper_parallel.codegen.inline.framework_spec import strategy_spec_for as strategy_spec
26from hyper_parallel.codegen.inline.strategy_pass import rule_block_classes
28logger = logging.getLogger(__name__)
31def normalize_inline_meta(
32 meta: Any,
33 rules: tuple[SourceObservation, ...],
34 model_type: str | None = None,
35) -> Any:
36 """Align the frozen plan with the inline-generated source shape.
38 Every stage is structure-proven: the strategy each strategy target
39 resolves to names the class whose forward is inlined and the sub-patterns
40 that forward owns, and ``meta.external_state_classes`` (recorded from the
41 same structure at freeze time) names the classes whose forward an inline
42 strategy body *replaces* (the MoE EP shell) and which therefore read their
43 parallel state from install-time instance attributes. A rendered component
44 class (the fused attention) keeps the component's own forward and is not
45 stripped: its boundary is emitted and compiled like any other.
46 """
48 # EP strategy forwards inline the entire routed+shared path, so any
49 # boundary entries for sub-modules of an EP target (e.g. routed experts
50 # and shared_experts) must be stripped — otherwise the shared class used
51 # in both EP and TP contexts fails the single-contract invariant.
52 # Capture the complete class map BEFORE stripping, so we can detect
53 # classes shared between boundary and non-boundary instances (e.g.
54 # DeepseekV3MLP used as both dense decoder layers and shared experts
55 # inside MoE). Source lowering rewrites class.forward, which ALL
56 # instances inherit -- a non-boundary instance would get a boundary
57 # body it cannot satisfy (no _hp_boundary installed).
58 normalized = _copy_inline_meta(meta)
59 full_class_map = dict(getattr(normalized, "boundary_classes", None) or {})
60 _strip_ep_sub_boundaries(normalized, rules, model_type)
61 _strip_external_state_boundaries(normalized)
62 _mark_shared_class_boundaries(normalized, full_class_map)
63 return normalized
66def _copy_inline_meta(meta: Any) -> Any:
67 """Copy only fields normalized by source emission, preserving all other context."""
68 normalized = copy(meta)
69 normalized.param_plan = dict(getattr(meta, "param_plan", None) or {})
70 normalized.boundary_classes = dict(
71 getattr(meta, "boundary_classes", None) or {}
72 )
73 normalized.external_state_classes = list(
74 getattr(meta, "external_state_classes", None) or ()
75 )
76 return normalized
79def _strip_external_state_boundaries(meta: Any) -> None:
80 """Keep parameter sharding without duplicating inline-body-owned collectives.
82 Only a class whose forward an inline strategy body replaced is stripped:
83 that body carries its own orchestration, so a boundary wrapper for the same
84 class would rewrite the very forward the inline patch installed. A
85 rendered component class is deliberately kept — the emitted boundary form
86 wraps its ``_forward_impl`` and is compiled at install time.
87 """
88 external_classes = set(getattr(meta, "external_state_classes", None) or ())
89 boundary_classes = getattr(meta, "boundary_classes", None) or {}
90 param_plan = getattr(meta, "param_plan", None) or {}
91 for fqn, class_name in list(boundary_classes.items()):
92 if class_name not in external_classes:
93 continue
94 entry = param_plan.get(fqn)
95 if isinstance(entry, dict):
96 param_plan[fqn] = {**entry, "is_boundary": False}
97 boundary_classes.pop(fqn)
100def _strip_ep_sub_boundaries(
101 meta: Any, rules: tuple[SourceObservation, ...], model_type: str | None,
102) -> None:
103 """Mark boundary entries owned by inline EP strategy forwards as non-boundary.
105 The EP strategy forward replaces the target class's ``forward`` and calls
106 sub-modules directly (routed experts via all-to-all, shared experts via
107 ``self.shared_experts(...)``). Boundary entries for the EP target itself
108 and for routed-expert / shared-expert sub-modules are marked
109 ``is_boundary: false`` so ``lower_forward_boundaries`` does not double-wrap
110 them. The param_plan entries are kept so FSDP / EP param sharding still
111 applies and the preflight coverage check passes.
112 """
113 param_plan = getattr(meta, "param_plan", None) or {}
114 boundary_classes = getattr(meta, "boundary_classes", None) or {}
115 ep_fqns, strip_patterns = _collect_ep_boundaries(
116 meta, rules, model_type, boundary_classes
117 )
118 if not ep_fqns:
119 return
120 mark_fqns = _expanded_ep_boundaries(ep_fqns, strip_patterns, param_plan)
121 for fqn in mark_fqns:
122 entry = param_plan.get(fqn)
123 if isinstance(entry, dict):
124 param_plan[fqn] = {**entry, "is_boundary": False}
125 boundary_classes.pop(fqn, None)
128def _collect_ep_boundaries(
129 meta: Any,
130 rules: tuple[SourceObservation, ...],
131 model_type: str | None,
132 boundary_classes: dict[str, str],
133) -> tuple[set[str], tuple[str, ...]]:
134 """Collect EP host FQNs and sub-boundary patterns from inline rules."""
135 ep_fqns: set[str] = set()
136 strip_patterns: tuple[str, ...] = ()
137 for rule in rules:
138 resolved = _ep_rule_hosts(rule, model_type, boundary_classes)
139 if resolved is None:
140 continue
141 spec, host_classes = resolved
142 if spec.strip_boundary_subpatterns:
143 strip_patterns = spec.strip_boundary_subpatterns
144 if host_classes:
145 # Persist the derived classes so a re-emission of this
146 # already-normalized meta (whose boundary entries the strip below
147 # consumes) can still resolve its patch hosts.
148 recorded = set(getattr(meta, "external_state_classes", None) or ())
149 meta.external_state_classes = sorted(recorded | host_classes)
150 for fqn, cls in boundary_classes.items():
151 if cls not in host_classes:
152 continue
153 if any(fnmatch(fqn, pattern) for pattern in rule.match):
154 ep_fqns.add(fqn)
155 return ep_fqns, strip_patterns
158def _expanded_ep_boundaries(
159 ep_fqns: set[str], strip_patterns: tuple[str, ...], param_plan: dict[str, Any]
160) -> set[str]:
161 """Expand EP hosts with matching routed/shared expert sub-boundaries."""
162 mark_fqns: set[str] = set()
163 for ep_fqn in ep_fqns:
164 mark_fqns.add(ep_fqn)
165 if not strip_patterns:
166 continue
167 sub_prefix = ep_fqn + "."
168 for fqn in list(param_plan):
169 if not fqn.startswith(sub_prefix):
170 continue
171 local_suffix = fqn[len(sub_prefix):]
172 for pat in strip_patterns:
173 if fnmatch(local_suffix, pat):
174 mark_fqns.add(fqn)
175 break
176 return mark_fqns
179def _ep_rule_hosts(
180 rule: SourceObservation, model_type: str | None, boundary_classes: dict[str, str]
181) -> tuple[Any, set[str]] | None:
182 """Resolve an inline EP rule and the source classes that host it."""
183 if rule.local_compute_target is None:
184 return None
185 spec = strategy_spec(rule.local_compute_target, model_type)
186 if spec is None or spec.body_template is None:
187 return None
188 if spec.target_class is not None:
189 return spec, {spec.target_class}
190 return spec, rule_block_classes(rule, boundary_classes)
193def _mark_shared_class_boundaries(meta: Any, full_class_map: dict[str, str]) -> None:
194 """Mark source lowering as skipped for classes shared with a non-boundary instance.
196 Source lowering rewrites class.forward, which ALL instances of the class
197 inherit -- including instances that are NOT boundaries. A non-boundary
198 instance would inherit a boundary body that references
199 self._hp_boundary, but hp_install_boundaries skips non-boundary FQNs
200 (is_boundary=False), so the attribute is never installed and the
201 forward raises AttributeError at train time.
203 The classic trigger is DeepSeek-V3 hybrid stack: DeepseekV3MLP is
204 used both as the dense decoder MLP (a TP boundary) and as the shared
205 expert inside the MoE layer (stripped to is_boundary=False by
206 _strip_ep_sub_boundaries). Source lowering the dense-layer boundary
207 rewrites the class forward; the shared-expert instance inherits it;
208 the runtime never installs _hp_boundary on the shared expert -> crash.
210 The decision travels as an explicit ``skip_source_lowering`` marker on
211 the param-plan entry (``is_boundary`` stays True; ``boundary_classes``
212 is kept intact) rather than by removing the class mapping: the
213 emitter's fail-fast on an unresolvable boundary class is a real safety
214 contract (a silently dropped redistribution would change the model's
215 numeric behavior), so "skip source lowering" must be stated, not
216 implied by missing data. ``hp_install_boundaries`` installs the
217 per-instance runtime wrapper for the marked boundary, which executes
218 the same contract.
219 """
220 param_plan = getattr(meta, "param_plan", None) or {}
221 boundary_classes = getattr(meta, "boundary_classes", None) or {}
223 non_boundary_classes: set[str] = set()
224 for fqn, entry in param_plan.items():
225 if not isinstance(entry, dict):
226 continue
227 if entry.get("is_boundary") is False:
228 cls = full_class_map.get(fqn)
229 if cls is not None:
230 non_boundary_classes.add(cls)
232 if not non_boundary_classes:
233 return
235 for fqn, class_name in list(boundary_classes.items()):
236 if class_name not in non_boundary_classes:
237 continue
238 entry = param_plan.get(fqn)
239 if isinstance(entry, dict):
240 param_plan[fqn] = {**entry, "skip_source_lowering": True}
241 logger.info(
242 "codegen: boundary %s (class %s) marked skip_source_lowering "
243 "-- the class is shared with a non-boundary instance; the "
244 "runtime per-instance wrapper is used instead",
245 fqn, class_name,
246 )
249__all__ = ["normalize_inline_meta"]