Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / tensor_parallel / param_role.py: 93%
61 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"""param_role: ParamRole enum + ParameterClassifier (05 §3.6 Phase 1).
17ParamRole is the bridge between naming rules and ShardingTemplate:
18- Phase 1 classifies named_parameters() into ParamRole by naming rules;
19- Phase 2 aggregates communication boundaries by ParamRole;
20- Phase 4 fills the placements of spec.params from the Template by ParamRole.
22ParamRole does not decide the I/O contract — that is decided by the Template's
23semantic role (attention/mlp/...).
24"""
26import logging
27from enum import Enum, auto
28from typing import Any, Dict, List, Optional, Tuple
30logger = logging.getLogger(__name__)
33class ParamRole(Enum):
34 """Semantic roles of parameters (14 enum values, 05 §3.6)."""
35 COLWISE = auto() # column-sharded linear layers: q/k/v/gate/up proj → Shard(0)
36 ROWWISE = auto() # row-sharded linear layers: o/down proj → Shard(1)
37 NORM = auto() # RMSNorm/LayerNorm weight → Replicate
38 EMBED = auto() # embedding weight → Shard(0) (vocab dim)
39 LM_HEAD = auto() # lm_head weight → Shard(0) (vocab dim)
40 MOE_GATE = auto() # MoE router/gate → Replicate
41 MOE_EXPERT = auto() # MoE routed expert → EP Shard(0) + TP colwise/rowwise
42 SHARED_EXPERT = auto() # MoE shared expert → EP Replicate + TP colwise/rowwise
43 FUSED_QKV = auto() # fused QKV → Shard(0) (a later SpecialHandler may adjust)
44 FUSED_GATE_UP = auto() # fused gate/up → Shard(0)
45 BIAS = auto() # unmatched bias → Replicate; Linear bias follows its weight role
46 REPLICATED = auto() # linear-layer weights forced to be replicated
47 # (e.g. MLA q_a/kv_a down projections)
48 # → Replicate on all dims; assigned only explicitly via
49 # the family's sharding_rules (ModelAdapterSpec),
50 # never produced by the default naming rules
51 SPECIAL = auto() # special parameters (gated_delta etc.) → Phase 6 SpecialHandler
52 SKIP = auto() # frozen / not sharded → excluded from spec.params
55def _match_any(name: str, patterns: List[str]) -> bool:
56 """Substring matching: name contains any of the patterns."""
57 return any(p in name for p in patterns)
60# Match modes for default naming rules (accuracy fix F1 — segment-aware
61# matching; the accuracy_problem.md 10.1 misclassification came from matching
62# "shared_expert" as a raw substring of the whole FQN, which also hits
63# "shared_expert_gate").
64SEGMENT_EXACT = "segment_exact" # a full '.'-segment must equal the pattern
65SEGMENT_SUBSTRING = "segment_substring" # pattern is a substring OF ONE segment
68def _match_pattern(name_lower: str, pattern: str, mode: str) -> bool:
69 """Match one pattern against a lower-cased parameter FQN.
71 Dotted patterns (e.g. ".mlp.gate.", "embed_tokens.weight", ".bias")
72 always keep the legacy dotted-path substring semantics. Dot-less patterns
73 are matched per segment according to the rule's declared mode:
75 - SEGMENT_EXACT: a whole segment must equal the pattern — for container
76 names ("shared_expert"/"shared_experts") where a longer sibling name
77 ("shared_expert_gate") must NOT match;
78 - SEGMENT_SUBSTRING: the pattern may be a fragment of one segment — for
79 leaf-name fragments ("norm" hits "input_layernorm"; "q_proj" hits
80 "q_proj"), while never spanning a '.' boundary.
81 """
82 if "." in pattern:
83 return pattern in name_lower
84 segments = name_lower.split(".")
85 if mode == SEGMENT_EXACT:
86 return any(seg == pattern for seg in segments)
87 return any(pattern in seg for seg in segments)
90def _match_rule(name_lower: str, patterns: List[str], mode: str) -> bool:
91 """Return True when any pattern of one rule matches the lower-cased FQN."""
92 return any(_match_pattern(name_lower, p, mode) for p in patterns)
95def _build_default_rules() -> List[Tuple[List[str], ParamRole, str]]:
96 """Default naming rules: list[(patterns, ParamRole, match_mode)],
97 first match wins.
99 Ordering principle: more specific rules come first (shared_expert(s)
100 before experts; dotted patterns of the MoE gate before the bare "gate"
101 word). Colwise Linear-name rules intentionally precede the generic bias
102 fallback so their bias follows the output-channel shard. The fallback
103 precedes rowwise rules: rowwise bias is a full output vector and remains
104 replicated; adding it exactly once after TP reduction is an execution-
105 layer concern. The "ln"/"norm"-style patterns do not misfire on
106 "linear"/"kernel" (neither contains them within one segment).
108 Match modes (F1): MoE container names use SEGMENT_EXACT so that e.g.
109 "shared_expert_gate" is NOT classified SHARED_EXPERT (it is a scalar
110 gate Linear, not the shared expert body — see accuracy_fix_plan.md §2);
111 leaf-name fragments use SEGMENT_SUBSTRING; dotted patterns keep the
112 legacy path-substring semantics regardless of the declared mode.
113 """
114 return [
115 (["embed_tokens.weight", "wte.weight", "tok_embeddings.weight",
116 "embed_in.weight", "word_embeddings.weight"], ParamRole.EMBED, SEGMENT_SUBSTRING),
117 (["lm_head.weight", "embed_out.weight", "output_layer.weight"], ParamRole.LM_HEAD, SEGMENT_SUBSTRING),
118 (["shared_expert", "shared_experts"], ParamRole.SHARED_EXPERT, SEGMENT_EXACT),
119 (["experts"], ParamRole.MOE_EXPERT, SEGMENT_SUBSTRING),
120 ([".mlp.gate.", ".router.", "moe_gate", "mlp.router"], ParamRole.MOE_GATE, SEGMENT_SUBSTRING),
121 (["fused_qkv", "linear_qkv", "qkv_proj", "query_key_value"],
122 ParamRole.FUSED_QKV, SEGMENT_SUBSTRING),
123 (["gate_up_proj", "fused_gate_up", ".w13."], ParamRole.FUSED_GATE_UP, SEGMENT_SUBSTRING),
124 (["a_log", "dt_bias", "gated_delta"], ParamRole.SPECIAL, SEGMENT_SUBSTRING),
125 (["norm", "layernorm", "rmsnorm", "ln_"], ParamRole.NORM, SEGMENT_SUBSTRING),
126 (["q_proj", "k_proj", "v_proj", "gate_proj", "up_proj", ".w1.", ".w3."],
127 ParamRole.COLWISE, SEGMENT_SUBSTRING),
128 ([".bias"], ParamRole.BIAS, SEGMENT_SUBSTRING),
129 (["o_proj", "down_proj", ".w2."], ParamRole.ROWWISE, SEGMENT_SUBSTRING),
130 ]
133class ParameterClassifier:
134 """Classifies named_parameters into ParamRole by naming rules + arch overrides (05 §3.6.6).
136 Rule sources (in decreasing priority):
137 1. ``arch_overrides[arch]`` — explicit (pattern | [patterns], ParamRole)
138 overrides (legacy full-name substring semantics — they are explicit
139 user intent);
140 2. default naming rules (first match; each rule declares its match mode,
141 see ``_build_default_rules``);
142 3. no match → ``ParamRole.SKIP``.
144 User-supplied ``name_rules`` may use either the current 3-tuple form
145 ``(patterns, role, mode)`` or the legacy 2-tuple form ``(patterns, role)``
146 (treated as full-name substring, the pre-F1 semantics).
147 """
149 def __init__(
150 self,
151 name_rules: Optional[List[tuple]] = None,
152 arch_overrides: Optional[Dict[str, list]] = None,
153 ) -> None:
154 """Initialize the classifier.
156 Args:
157 name_rules: Custom naming rules; falls back to
158 ``_build_default_rules()`` when None.
159 arch_overrides: ``{arch: [(pattern | [patterns], ParamRole)]}``
160 per-architecture overrides; empty when None.
161 """
162 self._name_rules = (
163 name_rules if name_rules is not None else _build_default_rules()
164 )
165 self._arch_overrides = arch_overrides if arch_overrides is not None else {}
167 def classify(self, model: Any, arch: str = "") -> Dict[str, ParamRole]:
168 """Iterate over all named parameters and return {param_fqn: ParamRole}.
170 Args:
171 model: The model whose ``named_parameters()`` are classified.
172 arch: Architecture key selecting the ``arch_overrides`` entry.
174 Returns:
175 Mapping of parameter FQN to its classified ParamRole.
176 """
177 roles: Dict[str, ParamRole] = {}
178 overrides = self._arch_overrides.get(arch, [])
179 for name, _ in model.named_parameters():
180 roles[name] = self.classify_param(name, overrides)
181 return roles
183 def classify_param(
184 self,
185 name: str,
186 overrides: Optional[List[Tuple[Any, ParamRole]]] = None,
187 ) -> ParamRole:
188 """Classify a single parameter (arch overrides are not applied when overrides is omitted).
190 Args:
191 name: The parameter FQN to classify.
192 overrides: Explicit ``(pattern | [patterns], ParamRole)`` override
193 rules checked before the default naming rules.
195 Returns:
196 The matched ParamRole; ``ParamRole.SKIP`` when nothing matches.
197 """
198 name_lower = name.lower()
199 # 1. Explicit arch overrides (three forms: exact FQN / substring /
200 # list-of-patterns)
201 for pattern, forced_role in (overrides or []):
202 patterns = [pattern] if isinstance(pattern, str) else list(pattern)
203 if _match_any(name_lower, [p.lower() for p in patterns]):
204 return forced_role
205 # 2. Default naming rules (first match)
206 for rule in self._name_rules:
207 if len(rule) == 3:
208 patterns, default_role, mode = rule
209 else: # legacy 2-tuple: full-name substring
210 patterns, default_role = rule
211 mode = None
212 if mode is None:
213 if _match_any(name_lower, patterns):
214 return default_role
215 elif _match_rule(name_lower, patterns, mode):
216 return default_role
217 # 3. Fallback
218 return ParamRole.SKIP