Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / default_templates.py: 99%
112 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"""default_templates: semantic role → placement templates (05 §3.5 canonical).
17``ShardingTemplate`` is the public template type (re-exported from
18``hyper_parallel.distributed``); ``TEMPLATES`` is the private
19default table consumed by the planner. The template-referencing Phase 4
20derivation helpers (``_build_spec_from_template`` / ``_placement_for_role``
21and friends) also live here — they are pure functions of the templates.
22"""
24import copy
25from dataclasses import dataclass, field
26from typing import Dict, List, Optional, Tuple
28from hyper_parallel.core.dtensor.placement_types import (
29 Partial,
30 Placement,
31 Replicate,
32 Shard,
33)
34from hyper_parallel.distributed.plan import ShardingPlan
35from hyper_parallel.distributed.recipe_spec import (
36 CP,
37 EP,
38 TP,
39 ModuleShardingSpec,
40 NamedPlacement,
41 _normalize_out_fields,
42)
43from hyper_parallel.distributed.tensor_parallel.head_count import (
44 build_tp_local_attr_plan,
45)
46from hyper_parallel.distributed.tensor_parallel.param_role import ParamRole
49@dataclass
50class ShardingTemplate:
51 """Semantic role → placement template (05 §3.5).
53 Each I/O field declares placements for all active mesh dimensions
54 (TP+CP+EP); the ShardingPlanner filters out dimensions not enabled by the
55 actual mesh_dim_names (resolve_placements looks up keys by mesh_dim_names,
56 so extra keys are naturally dropped).
58 Note: sp_out_src / nosp_out_src etc. are scalar NamedPlacement shorthands
59 (single-output modules); they are wrapped into {"output": ...} during
60 normalization in _build_spec_from_template.
61 """
62 # Parameter sharding rules
63 colwise_placement: Placement = field(default_factory=lambda: Shard(0))
64 rowwise_placement: Placement = field(default_factory=lambda: Shard(1))
65 norm_placement: Placement = field(default_factory=Replicate)
66 moe_expert_placement: Placement = field(default_factory=lambda: Shard(0))
68 # SP-mode I/O (full TP+CP+EP, three dims)
69 sp_in_src: Dict[str, NamedPlacement] = field(default_factory=dict)
70 sp_in_dst: Dict[str, NamedPlacement] = field(default_factory=dict)
71 sp_out_src: Optional[NamedPlacement] = None
72 sp_out_dst: Optional[NamedPlacement] = None
74 # non-SP-mode I/O
75 nosp_in_src: Dict[str, NamedPlacement] = field(default_factory=dict)
76 nosp_in_dst: Dict[str, NamedPlacement] = field(default_factory=dict)
77 nosp_out_src: Optional[NamedPlacement] = None
78 nosp_out_dst: Optional[NamedPlacement] = None
80 # Special flags
81 # region_dispatch: template-level declaration that the matched module's own
82 # forward CANNOT dispatch (data-dependent logic inside, e.g. an EP-aware
83 # custom MoE's a2a) — the planner copies False into the spec (derived
84 # metadata, logged); None = ordinary dispatchable module.
85 region_dispatch: Optional[bool] = None
86 needs_cp_attn: bool = False # CP: inner attention needs a CP-aware forward
89def _multi_dim(tp=None, cp=None, ep=None) -> NamedPlacement:
90 """Build multi-dim placement dict, filtering out None dims."""
91 result = {}
92 if tp is not None:
93 result[TP] = tp
94 if cp is not None:
95 result[CP] = cp
96 if ep is not None:
97 result[EP] = ep
98 return result
101def _hid(tp_p, cp_p, ep_p=None) -> Dict[str, NamedPlacement]:
102 """Shorthand for the single-input hidden_states contract."""
103 return {"hidden_states": _multi_dim(tp=tp_p, cp=cp_p, ep=ep_p or Replicate())}
106def _out(tp_p, cp_p, ep_p=None) -> NamedPlacement:
107 """Shorthand for the single-output (scalar shorthand) contract."""
108 return _multi_dim(tp=tp_p, cp=cp_p, ep=ep_p or Replicate())
111# ── TEMPLATES: complete templates for the 7 semantic roles (05 §3.5, declared over TP+CP+EP) ──
112# CP-dim rule: parameters are always Replicate (CP does not shard parameters);
113# input sequence activations use Shard(1) independently of TP sequence_parallel.
114# EP-dim rule: non-MoE modules Replicate; MoE experts Shard(0).
115TEMPLATES: Dict[str, ShardingTemplate] = {
116 # ── Attention (q/k/v Colwise + o Rowwise) ──
117 # The CP dim keeps Shard(1) in in_dst: the K/V all-gather is done by the
118 # inner attention wrapper inside SDPA/FlexAttention (needs_cp_attn=True),
119 # not at the boundary layer.
120 "attention": ShardingTemplate(
121 colwise_placement=Shard(0), # q/k/v: [H/tp, H]
122 rowwise_placement=Shard(1), # o: [H, H/tp]
123 sp_in_src=_hid(Shard(1), Shard(1)),
124 sp_in_dst=_hid(Replicate(), Shard(1)),
125 sp_out_src=_out(Partial(), Shard(1)), # local Q-chunk output → CP Shard(1)
126 sp_out_dst=_out(Shard(1), Shard(1)),
127 nosp_in_src=_hid(Replicate(), Shard(1)),
128 nosp_in_dst=_hid(Replicate(), Shard(1)),
129 nosp_out_src=_out(Partial(), Shard(1)),
130 nosp_out_dst=_out(Replicate(), Shard(1)),
131 needs_cp_attn=True,
132 ),
134 # ── MLP (gate/up Colwise + down Rowwise) ──
135 # The CP dim stays Shard(1) throughout (revision D-06): MLP is pointwise and
136 # CP needs no communication; if in_dst had CP=Replicate, the full-sequence
137 # reduce-scatter under TP×CP would produce a tp-major sequence layout
138 # inconsistent with embed/attention (cp-major).
139 "mlp": ShardingTemplate(
140 colwise_placement=Shard(0),
141 rowwise_placement=Shard(1),
142 sp_in_src=_hid(Shard(1), Shard(1)),
143 sp_in_dst=_hid(Replicate(), Shard(1)),
144 sp_out_src=_out(Partial(), Shard(1)),
145 sp_out_dst=_out(Shard(1), Shard(1)),
146 nosp_in_src=_hid(Replicate(), Shard(1)),
147 nosp_in_dst=_hid(Replicate(), Shard(1)),
148 nosp_out_src=_out(Partial(), Shard(1)),
149 nosp_out_dst=_out(Replicate(), Shard(1)),
150 ),
152 # ── Norm (RMSNorm/LayerNorm: weight replicated, zero communication) ──
153 "norm": ShardingTemplate(
154 norm_placement=Replicate(),
155 sp_in_src=_hid(Shard(1), Shard(1)),
156 sp_in_dst=_hid(Shard(1), Shard(1)), # identity
157 sp_out_src=_out(Shard(1), Shard(1)),
158 sp_out_dst=_out(Shard(1), Shard(1)), # identity
159 nosp_in_src=_hid(Replicate(), Shard(1)),
160 nosp_in_dst=_hid(Replicate(), Shard(1)),
161 nosp_out_src=_out(Replicate(), Shard(1)),
162 nosp_out_dst=_out(Replicate(), Shard(1)),
163 ),
165 # ── Embedding (weight Shard(0) along the vocab dim, output Partial → SP+CP) ──
166 "embed": ShardingTemplate(
167 colwise_placement=Shard(0), # weight: [V/tp, H]
168 sp_in_src={"input": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
169 sp_in_dst={"input": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
170 sp_out_src=_out(Partial(), Shard(1)),
171 sp_out_dst=_out(Shard(1), Shard(1)), # reduce-scatter → SP+CP
172 nosp_in_src={"input": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
173 nosp_in_dst={"input": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
174 nosp_out_src=_out(Partial(), Shard(1)),
175 nosp_out_dst=_out(Replicate(), Shard(1)),
176 ),
178 # ── LM Head (weight Shard(0), output Shard(-1); out_dst is overridden according to loss_parallel) ──
179 # The CP dim stays Shard(1) throughout (revision D-07): R8 — the CP dim is
180 # always identity at the boundary layer (the CP sequence all-gather happens
181 # only inside attention for K/V). lm_head computes logits/loss on the local
182 # CP chunk (the standard Megatron CP practice), with no CP gather.
183 "lm_head": ShardingTemplate(
184 colwise_placement=Shard(0), # weight: [V/tp, H]
185 sp_in_src=_hid(Shard(1), Shard(1)),
186 sp_in_dst=_hid(Replicate(), Shard(1)),
187 sp_out_src=_out(Shard(-1), Shard(1)),
188 sp_out_dst=_out(Shard(-1), Shard(1)), # loss_parallel=true default;
189 # when loss_parallel=false, _build_spec_from_template overrides it to
190 # {TP: Replicate, CP: Shard(1)}
191 nosp_in_src=_hid(Replicate(), Shard(1)),
192 nosp_in_dst=_hid(Replicate(), Shard(1)),
193 nosp_out_src=_out(Shard(-1), Shard(1)),
194 nosp_out_dst=_out(Replicate(), Shard(1)),
195 ),
197 # ── MoE Gate (Router: weight replicated, output redistributes → EP) ──
198 "moe_gate": ShardingTemplate(
199 norm_placement=Replicate(), # router weight/bias: replicated
200 sp_in_src=_hid(Shard(1), Shard(1)),
201 sp_in_dst=_hid(Replicate(), Replicate()),
202 sp_out_src=_out(Replicate(), Replicate()),
203 sp_out_dst=_out(Replicate(), Replicate(), Shard(0)),
204 nosp_in_src=_hid(Replicate(), Shard(1)),
205 nosp_in_dst=_hid(Replicate(), Replicate()),
206 nosp_out_src=_out(Replicate(), Replicate()),
207 nosp_out_dst=_out(Replicate(), Replicate(), Shard(0)),
208 ),
210 # ── MoE MLP (gate + routed experts + optional shared experts) ──
211 # CP dim same as mlp (D-06): pointwise per-token, CP stays Shard(1) throughout.
212 "moe_mlp": ShardingTemplate(
213 colwise_placement=Shard(0), # expert w1/w3: Colwise on TP
214 rowwise_placement=Shard(1), # expert w2: Rowwise on TP
215 norm_placement=Replicate(), # gate/norm: replicated
216 moe_expert_placement=Shard(0), # expert params: Shard(0) on EP
217 sp_in_src={"x_BLD": _multi_dim(tp=Shard(1), cp=Shard(1), ep=Replicate())},
218 sp_in_dst={"x_BLD": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
219 sp_out_src=_out(Partial(), Shard(1)),
220 sp_out_dst=_out(Shard(1), Shard(1)),
221 nosp_in_src={"x_BLD": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
222 nosp_in_dst={"x_BLD": _multi_dim(tp=Replicate(), cp=Shard(1), ep=Replicate())},
223 nosp_out_src=_out(Partial(), Shard(1)),
224 nosp_out_dst=_out(Replicate(), Shard(1)),
225 region_dispatch=False, # MoE forward has its own a2a; dispatch not allowed
226 ),
227}
230# ────────────────────────────────────────────────────────────────────────────
231# Template → spec derivation (planner Phase 4 machinery, 05 §3.5)
232# ────────────────────────────────────────────────────────────────────────────
234def _infer_colwise_vs_rowwise(param_path: str, template: ShardingTemplate):
235 """Infer the TP placement from the parameter name suffix: w2/down → rowwise, everything else → colwise."""
236 name = param_path.lower()
237 if any(k in name for k in ("w2", "down_proj", "down.")):
238 return template.rowwise_placement
239 return template.colwise_placement
242def _moe_expert_tp_placement(param_path: str, ndim: int,
243 template: ShardingTemplate):
244 """TP placement for MOE_EXPERT (revision D-08, ndim-aware per parameter).
246 When expert weights use a batched 3D layout [E, H_out, H_in] (ndim>=3),
247 tensor dim 0 is the expert dim (owned by EP Shard(0)), so the TP
248 colwise/rowwise sharding must apply at dim +1:
249 colwise (shard H_out) → Shard(1); rowwise (shard the contraction dim
250 H_in) → Shard(2).
251 The per-expert 2D layout (experts.N.w1 [H_out, H_in]) keeps the standard
252 Shard(0)/Shard(1) — but then EP Shard(0) would shard H_out, which is
253 semantically invalid: EP must be implemented as "each rank holds a
254 subset of experts" (module level), requiring family sharding_rules /
255 SpecialHandler; it is outside template coverage.
256 """
257 name = param_path.lower()
258 is_rowwise = any(k in name for k in ("w2", "down_proj", "down."))
259 if ndim >= 3:
260 return Shard(2) if is_rowwise else Shard(1)
261 return template.rowwise_placement if is_rowwise else template.colwise_placement
264def _placement_for_role(
265 param_path: str, role: ParamRole, template: ShardingTemplate,
266 has_tp: bool, has_ep: bool, ndim: int = 2,
267) -> Optional[NamedPlacement]:
268 """13 roles → placement mapping (05 §3.5 mapping table + D-08
269 ndim-aware)."""
270 if role in (ParamRole.COLWISE, ParamRole.EMBED, ParamRole.LM_HEAD,
271 ParamRole.FUSED_QKV, ParamRole.FUSED_GATE_UP):
272 return _multi_dim(tp=template.colwise_placement if has_tp else None,
273 cp=Replicate(), ep=Replicate())
274 if role == ParamRole.ROWWISE:
275 return _multi_dim(tp=template.rowwise_placement if has_tp else None,
276 cp=Replicate(), ep=Replicate())
277 if role in (ParamRole.NORM, ParamRole.MOE_GATE):
278 return _multi_dim(tp=template.norm_placement if has_tp else None,
279 cp=Replicate(), ep=Replicate())
280 if role == ParamRole.MOE_EXPERT:
281 # 05 §3.5 NOTE: when has_tp=False, use an explicit Replicate
282 # (rather than omitting the TP key).
283 # D-08: the TP dim of a 3D expert weight [E, H_out, H_in] is
284 # shifted according to ndim.
285 tp_p = (_moe_expert_tp_placement(param_path, ndim, template)
286 if has_tp else Replicate())
287 return _multi_dim(tp=tp_p, cp=Replicate(),
288 ep=template.moe_expert_placement if has_ep else None)
289 if role == ParamRole.SHARED_EXPERT:
290 # replicated along the EP dim; TP per w1/w3(colwise)/w2(rowwise)
291 tp_p = _infer_colwise_vs_rowwise(param_path, template)
292 return _multi_dim(tp=tp_p if has_tp else None,
293 cp=Replicate(), ep=Replicate())
294 if role == ParamRole.BIAS:
295 # Generic bias whose owning Linear role could not be inferred.
296 # Named colwise/rowwise Linear biases are classified with their
297 # weight role before reaching this fallback.
298 return _multi_dim(tp=Replicate(), cp=Replicate(), ep=Replicate())
299 if role == ParamRole.REPLICATED:
300 # MLA down-projections etc. (explicitly assigned via
301 # family sharding_rules): replicate on all dims.
302 # The output latent is identical within the TP group, so the
303 # input contract of the downstream q_b/kv_b (COLWISE), sharded
304 # along the head dim, matches standard attention.
305 return _multi_dim(tp=Replicate(), cp=Replicate(), ep=Replicate())
306 # SPECIAL → Phase 6; SKIP → not sharded
307 return None
310def _build_param_placements(
311 boundary_fqn: str, group: List[Tuple[str, ParamRole]],
312 template: ShardingTemplate, has_tp: bool, has_ep: bool,
313 param_ndims: Optional[Dict[str, int]],
314) -> Dict[str, NamedPlacement]:
315 """Build relative parameter paths and their resolved placements."""
316 params = {}
317 for param_fqn, role in group:
318 param_path = param_fqn[len(boundary_fqn) + 1:]
319 ndim = (param_ndims or {}).get(param_fqn, 2)
320 placement = _placement_for_role(
321 param_path, role, template, has_tp, has_ep, ndim=ndim,
322 )
323 if placement is not None:
324 params[param_path] = placement
325 return params
328def _build_spec_from_template(
329 templates, boundary_fqn: str, group: List[Tuple[str, ParamRole]],
330 template: ShardingTemplate, sequence_parallel: bool, loss_parallel: bool,
331 mesh_dim_names: Tuple[str, ...], param_ndims: Optional[Dict[str, int]] = None,
332) -> Optional[ModuleShardingSpec]:
333 """Template + ParamRole → ModuleShardingSpec (05 §3.5 Template Mapping)."""
334 has_tp = "tp" in mesh_dim_names
335 has_cp = "cp" in mesh_dim_names
336 has_ep = "ep" in mesh_dim_names
337 # The derived spec is materialized directly with concrete dicts (None
338 # is the "undeclared" semantics on the override input side; derived
339 # artifacts always hold concrete values).
340 spec = ModuleShardingSpec(params=_build_param_placements(
341 boundary_fqn, group, template, has_tp, has_ep, param_ndims,
342 ))
344 # Step 2: select the I/O contract per the SP switch (deep copy, so
345 # chain propagation cannot dirty the shared templates)
346 if sequence_parallel:
347 spec.in_src = copy.deepcopy(template.sp_in_src)
348 spec.in_dst = copy.deepcopy(template.sp_in_dst)
349 spec.out_src = copy.deepcopy(template.sp_out_src)
350 spec.out_dst = copy.deepcopy(template.sp_out_dst)
351 else:
352 spec.in_src = copy.deepcopy(template.nosp_in_src)
353 spec.in_dst = copy.deepcopy(template.nosp_in_dst)
354 spec.out_src = copy.deepcopy(template.nosp_out_src)
355 spec.out_dst = copy.deepcopy(template.nosp_out_dst)
357 # Step 2.5: lm_head's out_dst depends on loss_parallel (a runtime
358 # decision).
359 # The CP dim is always Shard(1) (D-07/R8): under CP the loss is
360 # computed on the local chunk; no gather is performed.
361 if template is templates.get("lm_head"):
362 cp_placement = Shard(1) if has_cp else Replicate()
363 for input_contract in (spec.in_src, spec.in_dst):
364 for placements in input_contract.values():
365 placements[CP] = cp_placement
366 spec.out_src = _multi_dim(
367 tp=Shard(-1),
368 cp=cp_placement,
369 ep=Replicate(),
370 )
371 spec.out_dst = _multi_dim(
372 tp=Shard(-1) if loss_parallel else Replicate(),
373 cp=cp_placement,
374 ep=Replicate(),
375 )
377 # Step 2.6: embed's CP contract (revision D-05): the CP data
378 # pipeline (shard_batch_for_cp, 05 §6.3.4) has already sharded
379 # input_ids along CP — the CP dim of in/out is Shard(1) rather than
380 # the template's default Replicate, otherwise the boundary would
381 # scatter the already-sharded chunk a second time (the sequence
382 # would be sharded twice).
383 if template is templates.get("embed") and has_cp:
384 spec.in_src = {"input": _multi_dim(tp=Replicate(), cp=Shard(1),
385 ep=Replicate())}
386 spec.in_dst = {"input": _multi_dim(tp=Replicate(), cp=Shard(1),
387 ep=Replicate())}
388 spec.out_src = _multi_dim(tp=Partial(), cp=Shard(1), ep=Replicate())
390 # Step 3: special flags
391 spec.region_dispatch = template.region_dispatch
392 if template.needs_cp_attn:
393 spec._needs_cp_attn = True # pylint: disable=protected-access # planner owns the spec DSL internals
395 # Step 4: normalize out_src/out_dst scalar shorthand
396 return _normalize_out_fields(spec)
399def _finalize_tp_local_attr_plans(
400 plan: ShardingPlan, model, *, tp_size: int,
401 mesh_dim_names: Tuple[str, ...],
402) -> None:
403 """Build internal auto/user TP-local attribute plans after overrides."""
404 modules = dict(model.named_modules())
405 for module_fqn, spec in plan.modules.items():
406 module = modules.get(module_fqn)
407 if module is None:
408 raise ValueError(
409 f"Cannot finalize TP-local attributes: module "
410 f"{module_fqn!r} is not present in model.named_modules()")
411 spec._tp_local_attr_plan = build_tp_local_attr_plan( # pylint: disable=protected-access # planner owns the spec DSL internals
412 module, spec, module_fqn, tp_size, mesh_dim_names,
413 )