Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / codegen / plan / spec_fields.py: 100%
21 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"""Single source of truth: which ``ModuleShardingSpec`` fields cross the meta boundary.
17``freeze_param_plan`` / ``freeze_injections`` (plan -> meta, in
18``codegen/plan/freeze.py``) and ``_rebuild_live_plan_from_meta``
19(meta -> plan, in ``codegen/runtime.py``) both walk this table, so a spec
20field can no longer be frozen-but-never-rebuilt (or the reverse): adding a
21field to ``ModuleShardingSpec`` without classifying it here fails the
22reflection test in ``tests/codegen/test_spec_field_table.py``.
24The table is pure data (no torch, no imports from freeze/runtime — both
25consume it, so an import here would be circular). Each entry names the
26spec attribute, the meta key it serializes to, the meta section it belongs
27to, and the serializer/deserializer *kind* the two sides dispatch on. The
28kind also fixes the write condition — ``is not None`` (an explicit empty
29dict/list crosses the boundary; this is the params={} case).
31Divergences this table replaces (found 2026-09-22, all closed here):
32- ``tp_divide_attrs`` was never frozen, so a rebuilt spec silently lost the
33 user's YAML declaration and ``maybe_update_head_counts`` fell back to the
34 head-sharded heuristic only;
35- freeze wrote ``params`` under a truthy check, dropping an explicit empty
36 dict and leaving the rebuilt spec with ``params=None`` where Phase A
37 (``_shard_planned_parameters``) expects a mapping;
38- freeze read ``getattr(spec, "needs_cp_attn")`` — the real field is
39 ``_needs_cp_attn``, so the lookup was dead code; the field is now
40 explicitly EXEMPT (fixing the name would add a key to every existing
41 artifact's meta and shift bytes; see the table comment).
42"""
43from __future__ import annotations
45from dataclasses import dataclass
47#: Frozen into ``meta.param_plan[fqn]`` entries (freeze_param_plan).
48GROUP_PARAM_PLAN = "param_plan"
49#: Frozen into ``meta.injections`` rules (freeze_injections).
50GROUP_INJECTIONS = "injections"
51#: Never crosses the meta boundary (planner-internal / apply-time only).
52GROUP_EXEMPT = "exempt"
55@dataclass(frozen=True)
56class SpecField:
57 """One ``ModuleShardingSpec`` field's freeze/rebuild contract.
59 Attributes:
60 name: Attribute on ``ModuleShardingSpec`` (``"_ep_stack"``).
61 key: Key in the frozen meta (``"ep_stack"``); ``""`` for exempt.
62 group: Meta section — one of the ``GROUP_*`` constants.
63 freeze: Freeze-side serializer kind; ``""`` for exempt fields.
64 rebuild: Rebuild-side deserializer kind; ``""`` for exempt fields.
65 post_init: The dataclass field is ``init=False`` (planner-internal
66 output like ``_deferred_bias_params``) — the rebuild sets it
67 via ``setattr`` after construction instead of a constructor
68 kwarg.
70 Serializer kinds (freeze side, dispatched in freeze.py):
72 - ``placement_map``: per-param -> per-axis placements, each via
73 ``placement_to_string``. Written when the value ``is not None`` —
74 an explicit ``params={}`` ("shards nothing") survives the boundary.
75 - ``named_placements``: ``named_placement_to_dict``. Written when the
76 value ``is not None``.
77 - ``name_list`` / ``attr_list``: ``list(value)`` copy. ``name_list``
78 is written only when truthy (legacy out_names behavior);
79 ``attr_list`` when ``is not None`` (an explicit ``[]`` clears an
80 inherited glob declaration and must not be lost).
81 - ``scalar``: value passes through unchanged. Written when the value
82 ``is not None``.
83 - ``value_to_dict``: ``_value_to_dict`` (Target/Placement/callable ->
84 JSON-safe). Written when the value ``is not None``.
85 - ``plain``: value passes through unchanged (string fields).
86 - ``dict_copy``: ``dict(value)``. Written when truthy.
87 - ``int_truthy``: value passes through. Written when truthy (0 is the
88 "inactive" default).
90 Deserializer kinds (rebuild side, dispatched in runtime.py):
91 ``placement_map`` / ``named_placements`` / ``name_list`` / ``attr_list``
92 mirror their freeze kinds; ``bool_default_true`` restores ``is_boundary``
93 with a ``True`` default for pre-flag metas; ``opt_scalar`` reads
94 ``entry.get(key)``; ``injection_target`` re-resolves a serialized
95 Target through ``_injection_target``.
96 """
98 name: str
99 key: str
100 group: str
101 freeze: str = ""
102 rebuild: str = ""
103 post_init: bool = False
106#: Every ``ModuleShardingSpec`` field, classified. Order is load-bearing:
107#: it is the insertion order of keys in frozen entries, so it must match
108#: the legacy hand-written freeze order (params -> in/out placements ->
109#: out_names -> tp_divide_attrs -> flags; injections keep the legacy
110#: inner_* -> ep_* order) to keep regenerated artifacts byte-stable.
111SPEC_FIELDS: tuple[SpecField, ...] = (
112 # ── Frozen into meta.param_plan[fqn] ──
113 SpecField("params", "params", GROUP_PARAM_PLAN,
114 freeze="placement_map", rebuild="placement_map"),
115 SpecField("in_src", "in_src", GROUP_PARAM_PLAN,
116 freeze="named_placements", rebuild="named_placements"),
117 SpecField("in_dst", "in_dst", GROUP_PARAM_PLAN,
118 freeze="named_placements", rebuild="named_placements"),
119 SpecField("out_src", "out_src", GROUP_PARAM_PLAN,
120 freeze="named_placements", rebuild="named_placements"),
121 SpecField("out_dst", "out_dst", GROUP_PARAM_PLAN,
122 freeze="named_placements", rebuild="named_placements"),
123 # out_names: the declared output *order* is separate from the out_src
124 # dict's serialized key order — the runtime resolves tuple indices
125 # against the spec's declared order (declared_out_names prefers
126 # entry["out_names"]), not a sorted order that could swap e.g.
127 # hidden/aux and redistribute the wrong output.
128 SpecField("out_names", "out_names", GROUP_PARAM_PLAN,
129 freeze="name_list", rebuild="name_list"),
130 SpecField("tp_divide_attrs", "tp_divide_attrs", GROUP_PARAM_PLAN,
131 freeze="attr_list", rebuild="attr_list"),
132 SpecField("is_boundary", "is_boundary", GROUP_PARAM_PLAN,
133 freeze="scalar", rebuild="bool_default_true"),
134 SpecField("region_dispatch", "region_dispatch", GROUP_PARAM_PLAN,
135 freeze="scalar", rebuild="opt_scalar"),
136 # _deferred_bias_params: D-22 (rowwise bias defer) — the planner-computed
137 # bias param paths (e.g. ("o_proj.bias",)) whose addition is deferred
138 # until AFTER the boundary exit TP reduction (Megatron
139 # RowParallelLinear semantics). Frozen so the emitted forward can
140 # call ``self._hp_deferred_bias(outputs)`` at the exit and the
141 # install path can run the native suppression/restore pair
142 # (``_install_bias_suppression`` / ``_maybe_add_deferred_biases``).
143 # ``init=False`` on the dataclass — rebuilt via setattr (post_init).
144 SpecField("_deferred_bias_params", "deferred_bias_params", GROUP_PARAM_PLAN,
145 freeze="name_list", rebuild="tuple_empty", post_init=True),
146 # skip_source_lowering: codegen opt-out of the class-level source
147 # forward rewrite — set by meta normalization after freeze
148 # (inline/meta_plan._mark_shared_class_boundaries) when the class is
149 # shared between boundary and non-boundary instances. The
150 # int_truthy freeze kind keeps unmarked artifacts byte-stable (only
151 # an explicit True is written); opt_scalar rebuild restores it onto
152 # the spec so freeze(rebuild(meta)) stays symmetric.
153 SpecField("skip_source_lowering", "skip_source_lowering", GROUP_PARAM_PLAN,
154 freeze="int_truthy", rebuild="opt_scalar"),
155 # ── Frozen into meta.injections rules ──
156 SpecField("inner_wrapper", "inner_wrapper", GROUP_INJECTIONS,
157 freeze="value_to_dict", rebuild="injection_target"),
158 SpecField("inner_target", "inner_target", GROUP_INJECTIONS,
159 freeze="plain", rebuild="opt_scalar"),
160 SpecField("inner_out_src", "inner_out_src", GROUP_INJECTIONS,
161 freeze="value_to_dict", rebuild="injection_target"),
162 SpecField("local_compute_fn", "local_compute_fn", GROUP_INJECTIONS,
163 freeze="value_to_dict", rebuild="injection_target"),
164 SpecField("_ep_stack", "ep_stack", GROUP_INJECTIONS,
165 freeze="dict_copy", rebuild="dict_copy"),
166 SpecField("_ep_size", "ep_size", GROUP_INJECTIONS,
167 freeze="int_truthy", rebuild="int_truthy"),
168 # ── Exempt: never crosses the meta boundary ──
169 # _tp_local_attr_plan: planner-derived from tp_divide_attrs + module
170 # structure (TpLocalAttrPlan, not JSON-safe). A rebuilt spec leaves
171 # it None and maybe_update_head_counts takes its backward-compatible
172 # path: auto attrs via the _is_head_sharded heuristic, user attrs via
173 # the rebuilt tp_divide_attrs — both sources survive.
174 SpecField("_tp_local_attr_plan", "", GROUP_EXEMPT),
175 # _is_terminal: validate-mode propagation marker, not production state.
176 SpecField("_is_terminal", "", GROUP_EXEMPT),
177 # _needs_cp_attn: apply-time preflight metadata (attention boundary
178 # under cp>1 without inner_wrapper fails fast in apply_sharding_plan
179 # — a path codegen's rebuild does not take). The legacy freeze loop
180 # read "needs_cp_attn" (no underscore), which never exists, so no
181 # artifact meta carries the key; fixing the name would shift bytes
182 # for every existing artifact, so the field stays exempt.
183 SpecField("_needs_cp_attn", "", GROUP_EXEMPT),
184 # _resolved_inner_wrapper / _resolved_inner_target: written back by the
185 # applier for introspection; meaningless at generation time.
186 SpecField("_resolved_inner_wrapper", "", GROUP_EXEMPT),
187 SpecField("_resolved_inner_target", "", GROUP_EXEMPT),
188)
191def fields_for_group(group: str) -> tuple[SpecField, ...]:
192 """Return the table entries for one ``GROUP_*`` section, in table order."""
193 return tuple(field for field in SPEC_FIELDS if field.group == group)
196#: The param_plan section (``meta.param_plan[fqn]`` entry keys).
197PARAM_PLAN_SPEC_FIELDS: tuple[SpecField, ...] = fields_for_group(GROUP_PARAM_PLAN)
199#: The injections section (``meta.injections`` rule keys, minus "match").
200INJECTION_SPEC_FIELDS: tuple[SpecField, ...] = fields_for_group(GROUP_INJECTIONS)
202#: ``init=False`` dataclass fields in the frozen groups — the rebuild sets
203#: them via ``setattr`` after constructing the spec.
204POST_INIT_SPEC_FIELDS: tuple[SpecField, ...] = tuple(
205 field for field in SPEC_FIELDS if field.post_init
206)
208#: Every meta key that may appear in a frozen entry/rule (both groups).
209FROZEN_SPEC_KEYS: frozenset[str] = frozenset(
210 field.key for field in SPEC_FIELDS if field.key
211)
214__all__ = [
215 "FROZEN_SPEC_KEYS",
216 "GROUP_EXEMPT",
217 "GROUP_INJECTIONS",
218 "GROUP_PARAM_PLAN",
219 "INJECTION_SPEC_FIELDS",
220 "PARAM_PLAN_SPEC_FIELDS",
221 "POST_INIT_SPEC_FIELDS",
222 "SPEC_FIELDS",
223 "SpecField",
224 "fields_for_group",
225]