Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / inline / meta_plan.py: 14%
83 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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 fnmatch import fnmatch
21from typing import Any
23from hyper_parallel.codegen.inline.ir import InlineRule
24from hyper_parallel.codegen.inline.specs import strategy_spec
26logger = logging.getLogger(__name__)
29def normalize_inline_meta(meta: Any, rules: tuple[InlineRule, ...], model_type: str | None = None) -> None:
30 """Align the frozen plan with the inline-generated source shape.
32 Every stage is structure-proven: the strategy each strategy target
33 resolves to names the class whose forward is inlined and the sub-patterns
34 that forward owns, and ``meta.external_state_classes`` (recorded from the
35 same structure at freeze time) names the classes whose forward an inline
36 strategy body *replaces* (the MoE EP shell) and which therefore read their
37 parallel state from install-time instance attributes. A rendered component
38 class (the fused attention) keeps the component's own forward and is not
39 stripped: its boundary is emitted and compiled like any other.
40 """
42 # EP strategy forwards inline the entire routed+shared path, so any
43 # boundary entries for sub-modules of an EP target (e.g. routed experts
44 # and shared_experts) must be stripped — otherwise the shared class used
45 # in both EP and TP contexts fails the single-contract invariant.
46 # Capture the complete class map BEFORE stripping, so we can detect
47 # classes shared between boundary and non-boundary instances (e.g.
48 # DeepseekV3MLP used as both dense decoder layers and shared experts
49 # inside MoE). Source lowering rewrites class.forward, which ALL
50 # instances inherit -- a non-boundary instance would get a boundary
51 # body it cannot satisfy (no _hp_boundary installed).
52 full_class_map = dict(getattr(meta, "boundary_classes", None) or {})
53 _strip_ep_sub_boundaries(meta, rules, model_type)
54 _strip_external_state_boundaries(meta)
55 _mark_shared_class_boundaries(meta, full_class_map)
58def _strip_external_state_boundaries(meta: Any) -> None:
59 """Keep parameter sharding without duplicating inline-body-owned collectives.
61 Only a class whose forward an inline strategy body replaced is stripped:
62 that body carries its own orchestration, so a boundary wrapper for the same
63 class would rewrite the very forward the inline patch installed. A
64 rendered component class is deliberately kept — the emitted boundary form
65 wraps its ``_forward_impl`` and is compiled at install time.
66 """
67 external_classes = set(getattr(meta, "external_state_classes", None) or ())
68 boundary_classes = getattr(meta, "boundary_classes", None) or {}
69 param_plan = getattr(meta, "param_plan", None) or {}
70 for fqn, class_name in list(boundary_classes.items()):
71 if class_name not in external_classes:
72 continue
73 entry = param_plan.get(fqn)
74 if isinstance(entry, dict):
75 entry["is_boundary"] = False
76 boundary_classes.pop(fqn)
79def _strip_ep_sub_boundaries(
80 meta: Any, rules: tuple[InlineRule, ...], model_type: str | None,
81) -> None:
82 """Mark boundary entries owned by inline EP strategy forwards as non-boundary.
84 The EP strategy forward replaces the target class's ``forward`` and calls
85 sub-modules directly (routed experts via all-to-all, shared experts via
86 ``self.shared_experts(...)``). Boundary entries for the EP target itself
87 and for routed-expert / shared-expert sub-modules are marked
88 ``is_boundary: false`` so ``lower_forward_boundaries`` does not double-wrap
89 them. The param_plan entries are kept so FSDP / EP param sharding still
90 applies and the preflight coverage check passes.
91 """
92 param_plan = getattr(meta, "param_plan", None) or {}
93 boundary_classes = getattr(meta, "boundary_classes", None) or {}
94 ep_fqns: set[str] = set()
95 strip_patterns: tuple[str, ...] = ()
96 for rule in rules:
97 target = rule.local_compute_target
98 if target is None:
99 continue
100 spec = strategy_spec(target, model_type)
101 if spec is None or spec.target_class is None:
102 continue
103 if spec.strip_boundary_subpatterns:
104 strip_patterns = spec.strip_boundary_subpatterns
105 for pattern in rule.match:
106 for fqn, cls in boundary_classes.items():
107 if fnmatch(fqn, pattern) and cls == spec.target_class:
108 ep_fqns.add(fqn)
109 if not ep_fqns:
110 return
111 mark_fqns: set[str] = set()
112 for ep_fqn in ep_fqns:
113 mark_fqns.add(ep_fqn)
114 if not strip_patterns:
115 continue
116 sub_prefix = ep_fqn + "."
117 for fqn in list(param_plan):
118 if not fqn.startswith(sub_prefix):
119 continue
120 local_suffix = fqn[len(sub_prefix):]
121 for pat in strip_patterns:
122 if fnmatch(local_suffix, pat):
123 mark_fqns.add(fqn)
124 break
125 for fqn in mark_fqns:
126 entry = param_plan.get(fqn)
127 if isinstance(entry, dict):
128 entry["is_boundary"] = False
129 boundary_classes.pop(fqn, None)
132def _mark_shared_class_boundaries(meta: Any, full_class_map: dict[str, str]) -> None:
133 """Mark source lowering as skipped for classes shared with a non-boundary instance.
135 Source lowering rewrites class.forward, which ALL instances of the class
136 inherit -- including instances that are NOT boundaries. A non-boundary
137 instance would inherit a boundary body that references
138 self._hp_boundary, but hp_install_boundaries skips non-boundary FQNs
139 (is_boundary=False), so the attribute is never installed and the
140 forward raises AttributeError at train time.
142 The classic trigger is DeepSeek-V3 hybrid stack: DeepseekV3MLP is
143 used both as the dense decoder MLP (a TP boundary) and as the shared
144 expert inside the MoE layer (stripped to is_boundary=False by
145 _strip_ep_sub_boundaries). Source lowering the dense-layer boundary
146 rewrites the class forward; the shared-expert instance inherits it;
147 the runtime never installs _hp_boundary on the shared expert -> crash.
149 The decision travels as an explicit ``skip_source_lowering`` marker on
150 the param-plan entry (``is_boundary`` stays True; ``boundary_classes``
151 is kept intact) rather than by removing the class mapping: the
152 emitter's fail-fast on an unresolvable boundary class is a real safety
153 contract (a silently dropped redistribution would change the model's
154 numeric behavior), so "skip source lowering" must be stated, not
155 implied by missing data. ``hp_install_boundaries`` installs the
156 per-instance runtime wrapper for the marked boundary, which executes
157 the same contract.
158 """
159 param_plan = getattr(meta, "param_plan", None) or {}
160 boundary_classes = getattr(meta, "boundary_classes", None) or {}
162 non_boundary_classes: set[str] = set()
163 for fqn, entry in param_plan.items():
164 if not isinstance(entry, dict):
165 continue
166 if entry.get("is_boundary") is False:
167 cls = full_class_map.get(fqn)
168 if cls is not None:
169 non_boundary_classes.add(cls)
171 if not non_boundary_classes:
172 return
174 for fqn, class_name in list(boundary_classes.items()):
175 if class_name not in non_boundary_classes:
176 continue
177 entry = param_plan.get(fqn)
178 if isinstance(entry, dict):
179 entry["skip_source_lowering"] = True
180 logger.info(
181 "codegen: boundary %s (class %s) marked skip_source_lowering "
182 "-- the class is shared with a non-boundary instance; the "
183 "runtime per-instance wrapper is used instead",
184 fqn, class_name,
185 )
188__all__ = ["normalize_inline_meta"]