Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / _builder / planner.py: 94%
582 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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# pylint: disable=missing-apache-license-header
16"""planner: ShardingPlanner 6-phase derivation pipeline (05 §3.6 canonical).
18Phase 1 parameter role classification (ParameterClassifier + the family's
19 ModelAdapterSpec.sharding_rules from models/registry.py)
20Phase 2 communication boundary grouping (two passes: first group by owning
21 module, then merge upward depth-first — fixes the flaw in the
22 05 §3.6.6 pseudocode where "single-parameter group inference"
23 misjudges a q_proj leaf module as an mlp boundary)
24Phase 3 semantic role inference (explicit FQN patterns > structural guards
25 > parameter role combinations)
26Phase 4 template lookup to generate spec (_build_spec_from_template)
27Phase 4.5 user plan_overrides merge/insert (_merge_plan_overrides, 05 §3.6.7)
28Phase 5 _is_terminal marking only (D-14, 05 §13: compile-time chain
29 propagation/validation removed — specs are fully self-declared and
30 each module vouches for its own propagation in validate mode)
31Phase 6 special parameter handler collection (SPECIAL_HANDLERS)
33Registries:
34- family sharding rules: ``ModelAdapterSpec.sharding_rules`` providers in
35 ``models/<family>/adapter/registration.py``, resolved lazily through
36 ``models/registry.py::get_model_adapter`` (any spelling — model_type,
37 HF architecture or canonical arch name);
38- ``SPECIAL_HANDLERS``: {handler_name: callable(module, param_name, mesh)} —
39 lives in ``_builder/special_handlers.py``.
40"""
41# pylint: disable=forbidden-backend-import
43import copy
44import logging
45import re
46from typing import Any, Dict, List, Optional, Tuple
48from torch import nn
50from hyper_parallel.components.modules.module_family import FamilyShardingContext, _family_parameter_shapes
51from hyper_parallel.distributed._builder.family_layout import collect_owned_layouts, validate_owned_layouts
52from hyper_parallel.models.registry import resolve_model_adapter
53from hyper_parallel.core.dtensor.device_mesh import DeviceMesh
54from hyper_parallel.core.dtensor.placement_types import (
55 Partial,
56 Replicate,
57 Shard,
58)
59from hyper_parallel.distributed.plan import ShardingPlan
60from hyper_parallel.distributed.recipe_spec import (
61 DP,
62 EP,
63 TP,
64 ModuleShardingSpec,
65)
66from hyper_parallel.distributed.tensor_parallel.param_role import (
67 ParameterClassifier,
68 ParamRole,
69 _match_any,
70)
71from hyper_parallel.distributed._builder.default_templates import (
72 ShardingTemplate,
73 TEMPLATES,
74 _build_spec_from_template,
75 _finalize_tp_local_attr_plans,
76 _moe_expert_tp_placement,
77 _multi_dim,
78)
79from hyper_parallel.distributed._builder.function_module import FunctionModule
80from hyper_parallel.distributed._builder.rule_resolver import (
81 _last_segment,
82 _merge_plan_overrides,
83 _normalize_contract_fields,
84)
85from hyper_parallel.distributed._builder.special_handlers import (
86 _SPECIAL_HANDLER_PATTERNS,
87 _collect_special_handlers,
88)
90logger = logging.getLogger(__name__)
92# Per-family naming-rule overrides no longer live here: each family declares
93# them via ``ModelAdapterSpec.sharding_rules`` in
94# ``models/<family>/adapter/registration.py`` (DeepSeek MLA, Qwen2-MoE
95# shared_expert_gate), discovered through ``models/registry.py``. The
96# planner core stays family-agnostic (05 §15.9 step 1).
99# Leaf-segment guard for projection/container segment names: these segment
100# names are not boundary containers themselves; inference returns unknown
101# and continues upward.
102# NOTE (accuracy fix F3): "shared_experts" was REMOVED from this guard —
103# the shared expert (singular or plural spelling) must surface as its own
104# nested "mlp" boundary (its boundary exit owns the RowWise Partial TP
105# reduction, accuracy_problem.md 10.3 Option A); keeping it in the guard would
106# merge it into the parent MoE boundary and silently leave the TP reduction
107# without an owner.
108_LEAF_SEGMENT_GUARD = frozenset({
109 "q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj",
110 "qkv_proj", "fused_qkv", "linear_qkv", "gate_up_proj", "query_key_value",
111 "experts", "gate", "linear", "proj",
112 "fc1", "fc2", "w1", "w2", "w3", "w13", "dense", "dense_h_to_4h", "dense_4h_to_h",
113})
115_ATTN_PATTERNS = ("attn", "attention")
116_MLP_PATTERNS = ("mlp", "ffn", "feed_forward")
117_MOE_CONTAINER_PATTERNS = ("mlp", "moe", "moe_block", "moe_layer")
120def _get_architecture(model) -> str:
121 """Detect the canonical architecture name:
122 config.architectures[0] > config.model_type > class name;
123 lowercased with ForCausalLM-style suffixes stripped."""
124 cfg = getattr(model, "config", None)
125 arch_str = None
126 archs = getattr(cfg, "architectures", None)
127 if archs:
128 arch_str = archs[0]
129 if not arch_str:
130 arch_str = getattr(cfg, "model_type", None)
131 if not arch_str:
132 arch_str = type(model).__name__
134 s = arch_str.lower()
135 for suffix in ("forcausallm", "forconditionalgeneration",
136 "forsequenceclassification", "forimagetexttotext"):
137 if s.endswith(suffix):
138 s = s[: -len(suffix)]
139 return s
142class ShardingPlanner:
143 """Automatically derive a ShardingPlan from any HF-style model (05 §3.6.6).
145 ``plan_overrides``: {module_fqn | fqn_glob: ModuleShardingSpec} — the
146 SINGLE injection/override interface (05 §3.6.7 + the unification rework):
148 - **merge mode** (key hits a planner-derived boundary): empty contract
149 fields (``params``/``in_src``/``in_dst``/``out_src``/``out_dst``/
150 ``out_names``) INHERIT the derived values, non-empty ones replace
151 (field granularity); the string sentinels ``"auto"`` (explicit
152 inherit — self-documenting "derive per template") and ``"none"``
153 (explicit clear) are resolved at merge time and never reach the plan
154 output; the injection fields (``local_compute_fn`` /
155 ``inner_target`` / ``inner_wrapper`` non-None, ``region_dispatch=False``)
156 always win; internal flags (``_ep_size`` / ``_ep_stack`` /
157 ``_needs_cp_attn``) always inherit. This is how
158 CP/EP compute injection is declared — no need to re-declare contracts;
159 - **insert mode** (exact key misses every derived boundary): the spec is
160 inserted as-is and must be fully self-declared — an override with
161 empty params AND empty contracts fails fast ("no template matched");
162 - **glob keys** (containing ``*``/``?``/``[``): merge-applied to every
163 matching boundary. A glob with a concrete parameter or I/O contract
164 also expands against the final model tree and inserts unmatched modules
165 with the same validation as exact insert mode. Partial globs carrying
166 only injection/attribute changes remain merge-only and fail if they
167 target a model module for which no boundary was derived.
169 The YAML transport (``hyper_parallel.trainer.config.PlanOverride``) is
170 converted to this dict by ``entries_to_plan_overrides()`` on the trainer
171 side BEFORE the
172 planner is constructed — the planner itself has exactly one override
173 interface.
175 ``derive=False`` turns the planner into a pure declaration assembler:
176 template derivation (Phases 2-4) is skipped and the plan contains ONLY
177 the ``plan_overrides`` specs — every key is insert mode and must be
178 fully self-declared. Use it for subtrees where automatic derivation is
179 semantically wrong, e.g. the multimodal encoder_dp ViT bridge (each
180 rank encodes a different image shard, so any derived TP collective
181 inside the ViT would silently mix samples). This replaces the error-
182 prone post-hoc pruning idiom ``plan.modules = {"": plan.modules[""]}``.
183 """
185 def __init__(
186 self,
187 plan_overrides: Optional[Dict[str, ModuleShardingSpec]] = None,
188 *,
189 derive: bool = True,
190 allow_uncovered_params: bool = False,
191 ) -> None:
192 """Create a ShardingPlanner.
194 Args:
195 plan_overrides: ``{module_fqn | fqn_glob: ModuleShardingSpec}`` —
196 the single injection/override interface (05 §3.6.7); see the
197 class docstring for merge/insert/glob semantics.
198 derive: When ``False``, template derivation (Phases 2-4) is
199 skipped and the plan contains ONLY the ``plan_overrides``
200 specs (every key is insert mode, fully self-declared).
201 allow_uncovered_params: When ``True``, downgrade the F4b
202 "every trainable parameter must be covered by the plan"
203 hard error to a warning (exploratory debugging only).
204 """
205 self._classifier = ParameterClassifier()
206 self._templates = TEMPLATES
207 self._special_handler_patterns = dict(_SPECIAL_HANDLER_PATTERNS)
208 self._plan_overrides = dict(plan_overrides or {})
209 # derive=False: skip template derivation entirely — the plan contains
210 # ONLY the plan_overrides specs (every key is insert mode and must be
211 # fully self-declared). This replaces post-hoc pruning of
212 # ``plan.modules`` (e.g. the multimodal encoder_dp ViT bridge, where
213 # each rank encodes different images and ANY derived TP collective
214 # inside the ViT subtree would be a math error).
215 self._derive = derive
216 # F4b (accuracy_fix_plan.md §2): by default every trainable parameter
217 # must be covered by the plan (fail-fast). Set True ONLY for
218 # exploratory debugging to downgrade the hard error to a warning.
219 self._allow_uncovered_params = allow_uncovered_params
221 # ── Main entry point ────────────────────────────────────────────────
223 def _derive_boundary_specs(
224 self,
225 plan: ShardingPlan,
226 boundary_groups,
227 *,
228 sequence_parallel: bool,
229 loss_parallel: bool,
230 mesh_dim_names: Tuple[str, ...],
231 arch: str,
232 ep_extend: int,
233 mesh: DeviceMesh,
234 model: Any,
235 param_ndims: Dict[str, int],
236 ) -> None:
237 """Materialize template-derived boundary specs into a plan."""
238 if not self._derive:
239 logger.info(
240 "derive=False: template derivation skipped — the plan will contain only plan_overrides "
241 "specs (all insert mode, fully self-declared)"
242 )
243 return
244 for boundary_fqn, group in boundary_groups.items():
245 boundary_type = self._infer_boundary_type(boundary_fqn, group)
246 if self._templates.get(boundary_type) is None:
247 logger.warning("No template for boundary_type=%s at %s", boundary_type, boundary_fqn)
248 continue
249 spec = _build_spec_from_template(
250 self._templates,
251 boundary_fqn,
252 group,
253 self._templates.get(boundary_type),
254 sequence_parallel,
255 loss_parallel,
256 mesh_dim_names,
257 param_ndims=param_ndims,
258 )
259 if spec is None:
260 continue
261 if boundary_type == "moe_mlp":
262 self._mark_hf_native_moe(
263 spec, group, boundary_fqn, self._templates.get(boundary_type), mesh_dim_names, arch,
264 ep_extend=ep_extend, mesh=mesh, model=model, param_ndims=param_ndims,
265 )
266 plan.modules[boundary_fqn] = spec
268 def _finalize_boundary_specs(
269 self,
270 plan: ShardingPlan,
271 model: Any,
272 *,
273 tp_size: int,
274 ep_size: int,
275 mesh: DeviceMesh,
276 mesh_dim_names: Tuple[str, ...],
277 ) -> None:
278 """Normalize overrides and finish placement-dependent boundary metadata."""
279 _merge_plan_overrides(self._plan_overrides, plan, model, derive=self._derive)
280 _normalize_contract_fields(plan)
281 self._mark_virtual_ep_parameters(plan, ep_size, mesh)
282 _finalize_tp_local_attr_plans(plan, model, tp_size=tp_size, mesh_dim_names=mesh_dim_names)
283 self._finalize_deferred_biases(plan, model, mesh_dim_names)
284 self._finalize_fused_expert_tp_guard(plan, tp_size=tp_size)
286 @staticmethod
287 def _log_explanation(plan: ShardingPlan, enabled: bool) -> None:
288 """Log the plan explanation when requested."""
289 if enabled:
290 logger.info("ShardingPlan explain:\n%s", plan.explain())
292 @staticmethod
293 def _resolve_isolated_roots(model: Any, arch: str) -> tuple[str, ...]:
294 """Find subtrees whose activation layout is owned by a separate runtime."""
295 from hyper_parallel.models.registry import get_model_adapter # pylint: disable=C0415
297 adapter = (model._hp_model_adapter_spec if hasattr(model, "_hp_model_adapter_spec") # pylint: disable=protected-access
298 else get_model_adapter(arch))
299 if adapter is None or adapter.encoder_parallel is None:
300 return ()
301 spec = adapter.encoder_parallel()
302 names = {name for name, _ in model.named_modules()}
303 return tuple(root for root in spec.isolated_subtrees if root in names)
305 def _classify_boundary_groups(self, model: Any, arch: str, isolated_roots=(), excluded_names=()):
306 """Classify model parameters and group them into candidate boundaries."""
307 param_roles = {name: role for name, role in self._classify_all_params(model, arch).items()
308 if name not in excluded_names}
309 regular_roles = {
310 name: role for name, role in param_roles.items()
311 if not any(name.startswith(root + ".") for root in isolated_roots)
312 }
313 return param_roles, self._group_by_boundary(regular_roles)
315 def plan(
316 self,
317 model: Any,
318 mesh: DeviceMesh,
319 *,
320 tp_size: int = 1,
321 cp_size: int = 1,
322 ep_size: int = 1,
323 sequence_parallel: bool = True,
324 loss_parallel: bool = False,
325 explain: bool = False,
326 ) -> ShardingPlan:
327 """Derive a ShardingPlan from *model* and *mesh* (6-phase pipeline).
329 Args:
330 model: Any HuggingFace ``PreTrainedModel``.
331 mesh: A :class:`~hyper_parallel.core.dtensor.device_mesh.DeviceMesh`
332 whose ``mesh_dim_names`` declare the physical topology axes (e.g.
333 ``("dp", "tp")`` or ``("dp_replicate", "tp", "cp")``). The
334 ``tp`` / ``cp`` axes are managed by DTensor; ``dp*`` / ``pp``
335 axes belong to FSDP2 / pipeline-parallel runtimes and are
336 filtered out.
337 tp_size: Tensor-parallel group size. Must be **equal** to
338 ``mesh[\"tp\"].size()`` when ``mesh_dim_names`` contains ``"tp"``
339 and tp_size > 1; a mismatch raises :class:`ValueError`
340 immediately (fail-first).
341 cp_size: Context-parallel group size. Same fail-first contract as
342 *tp_size* with respect to ``mesh[\"cp\"]``.
343 ep_size: Expert-parallel group size.
345 - **D-10 TP-extend-EP** (the default MoE path, 05 §6.4.8): the
346 mesh does **not** have an ``"ep"`` axis; *ep_size* is the
347 extended EP group size, and the expert mesh ``(edp, ep)`` is
348 derived by flattening the full dense region
349 (``dp × tp × cp``). *ep_size* is **not** validated against
350 the mesh in this case.
351 - **Old-style EP** (when ``mesh_dim_names`` contains ``"ep"``):
352 *ep_size* must equal ``mesh[\"ep\"].size()``, validated
353 fail-first.
354 sequence_parallel: When ``True`` (default), sequence dimension is
355 sharded across TP (Shard(1) on the ``tp`` axis of activations).
356 loss_parallel: When ``True``, lm_head output stays Shard(-1);
357 otherwise gathered to Replicate (cross-entropy compatibility).
358 explain: When ``True``, log the plan introspection report
359 (:meth:`ShardingPlan.explain`) at INFO level after planning —
360 per-boundary param sharding, compiled communication plans,
361 and injection resolutions. The same report is available
362 standalone via ``plan.explain(fqn=None)``.
364 Returns:
365 :class:`ShardingPlan`: module FQN → :class:`ModuleShardingSpec`.
367 Raises:
368 ValueError: If *tp_size* / *cp_size* / *ep_size* do not match the
369 corresponding mesh dimensions (fail-first), if model-level
370 constraints are violated (head divisibility, expert count, etc.),
371 or if any ``plan_overrides`` spec declares a DP placement
372 (05 §3.1.1 coordinate-system convention: the plan is always a
373 single dp slice).
374 """
375 arch = _get_architecture(model)
376 self._check_overrides_no_dp() # fail-first: the plan's coordinate system = a single dp slice
377 mesh_dim_names = self._build_mesh_dim_names(mesh, tp_size, cp_size, ep_size)
378 plan, param_roles, param_metadata = self._build_plan(
379 model,
380 mesh,
381 mesh_dim_names,
382 arch,
383 tp_size=tp_size,
384 cp_size=cp_size,
385 ep_size=ep_size,
386 sequence_parallel=sequence_parallel,
387 loss_parallel=loss_parallel,
388 )
390 # D-14 invariants, Phase 5/6, F4 lints, DX guard, tied weights
391 plan = self._finalize_plan(
392 plan,
393 model,
394 param_roles,
395 param_metadata,
396 tp_size=tp_size,
397 cp_size=cp_size,
398 ep_size=ep_size,
399 )
401 self._log_explanation(plan, explain)
403 return plan
405 def _build_plan(
406 self,
407 model: Any,
408 mesh: DeviceMesh,
409 mesh_dim_names: Tuple[str, ...],
410 arch: str,
411 *,
412 tp_size: int,
413 cp_size: int,
414 ep_size: int,
415 sequence_parallel: bool,
416 loss_parallel: bool,
417 ) -> Tuple[ShardingPlan, Dict[str, ParamRole], Tuple[Dict[str, int], Dict[str, Tuple[int, ...]]]]:
418 """Run Phases 1-4.5: classify, derive, and merge boundary specs."""
419 # D-10 TP-extend-EP (05 §6.4.8): ep_size is the extended EP group
420 # size (the a2a communication domain, extended from the TP group to
421 # neighboring dp/cp ranks; expert weights are sharded only along
422 # the expert dim; no separate etp configuration). Validation
423 # happens when _mark_hf_native_moe actually matches an HF-native
424 # MoE (pre-stacked EP-aware modules use their own dispatcher and
425 # are not subject to this constraint)
426 # Phase 1: parameter role classification
427 layout_context = FamilyShardingContext(
428 mesh_dim_names, {"tp": tp_size, "cp": cp_size, "ep": ep_size},
429 _family_parameter_shapes(model),
430 sequence_parallel, loss_parallel,
431 )
432 owned = collect_owned_layouts(model, layout_context, resolve_model_adapter(model))
433 isolated_roots = self._resolve_isolated_roots(model, arch) if self._derive else ()
434 param_roles, boundary_groups = self._classify_boundary_groups(
435 model, arch, isolated_roots, excluded_names=owned.owned_names,
436 )
438 # Phase 3+4: semantic inference + template-fills I/O
439 # F4a/F4b: full shapes + requires_grad for the plan-time lints
440 # (meta tensors carry shapes too — the zero-memory path is unaffected)
441 param_metadata = (
442 {name: p.ndim for name, p in model.named_parameters()},
443 dict(layout_context.parameter_shapes),
444 )
445 plan = ShardingPlan(
446 mesh_dim_names=mesh_dim_names,
447 sequence_parallel=sequence_parallel,
448 loss_parallel=loss_parallel,
449 )
450 self._derive_boundary_specs(
451 plan,
452 boundary_groups,
453 sequence_parallel=sequence_parallel,
454 loss_parallel=loss_parallel,
455 mesh_dim_names=mesh_dim_names,
456 arch=arch,
457 ep_extend=ep_size if ep_size > 1 else 0,
458 mesh=mesh,
459 model=model,
460 param_ndims=param_metadata[0],
461 )
463 for root in isolated_roots:
464 module = model.get_submodule(root)
465 params = {name: {TP: Replicate()} for name, parameter in module.named_parameters()
466 if parameter.requires_grad}
467 plan.modules[root] = ModuleShardingSpec(params=params, in_src={}, in_dst={}, is_boundary=False)
469 plan.modules.update(owned.specs)
470 plan.parallel_modules.update(owned.parallel)
471 plan.owned_parameter_names.update(owned.owned_names)
472 plan.tied_pairs.extend(owned.aliases)
473 plan.parameter_aliases.extend(owned.aliases)
475 # Phase 4.5: merge derived matches and expand fully declared globs
476 # against the final model tree to insert new boundaries.
477 self._finalize_boundary_specs(
478 plan,
479 model,
480 tp_size=tp_size,
481 ep_size=ep_size,
482 mesh=mesh,
483 mesh_dim_names=mesh_dim_names,
484 )
485 validate_owned_layouts(model, plan, owned, layout_context)
486 return plan, param_roles, param_metadata
488 def _finalize_plan(
489 self,
490 plan: ShardingPlan,
491 model: Any,
492 param_roles: Dict[str, ParamRole],
493 param_metadata: Tuple[Dict[str, int], Dict[str, Tuple[int, ...]]],
494 *,
495 tp_size: int,
496 cp_size: int,
497 ep_size: int,
498 ) -> ShardingPlan:
499 """Run the D-14 invariants, Phase 5/6, the F4 lints, and tied pairs."""
500 # D-14 invariants (05 §13.2/§13.3): full self-declaration + param
501 # uniqueness (the only nesting check that remains)
502 self._check_full_declaration(plan)
503 self._check_param_uniqueness(plan)
505 # Phase 5: _is_terminal marking (D-14: chain propagation removed)
506 plan = self._mark_terminal(plan, model)
508 # Phase 6: special parameter handling
509 plan.special_handlers = _collect_special_handlers(param_roles, self._special_handler_patterns)
511 # F4 plan-time lints (accuracy_fix_plan.md §2 — after Phase 4.5
512 # overrides merge and Phase 6, so hand-written specs and special
513 # handlers are all accounted for):
514 # F4a: every Shard(dim) must divide the parameter shape — an empty
515 # shard at apply time becomes a plan-time teaching error;
516 # F4b: every trainable parameter must be covered by the plan.
517 self._check_shard_divisibility(
518 plan, param_metadata[1], tp_size=tp_size, cp_size=cp_size,
519 ep_size=ep_size)
520 self._check_all_trainable_params_covered(plan, model)
522 # DX guard: FunctionModule instances not covered by any spec run
523 # without any boundary communication — warn instead of silently passing
524 self._warn_uncovered_function_modules(model, plan)
526 # tied-weight detection (embed <-> lm_head sharing storage)
527 plan.tied_pairs = list(dict.fromkeys([*plan.tied_pairs, *self._detect_tied_pairs(model)]))
528 return plan
530 # ── Architecture detection ──────────────────────────────────────────
532 @staticmethod
533 def _get_architecture(model) -> str:
534 """Compatibility wrapper for the module-level architecture helper."""
535 return _get_architecture(model)
537 @staticmethod
538 def _validate_dtensor_axes(
539 mesh, tp_size: int, cp_size: int, ep_size: int,
540 ) -> None:
541 """Fail-first: validate passed tp/cp/ep sizes against the mesh dimensions.
543 The mesh *must* declare ``mesh_dim_names`` (and the corresponding
544 ``mesh_shape``) for validation to proceed; a mesh without names skips
545 validation (backward-compatible fallback).
547 Degenerate single-rank meshes (all dimensions size 1, common in
548 compile-time unit tests) also skip validation — on a single rank every
549 DTensor placement is a no-op.
551 Rules:
552 - tp_size > 1 → mesh must contain a "tp" axis whose size equals tp_size.
553 - cp_size > 1 → mesh must contain a "cp" axis whose size equals cp_size.
554 - ep_size > 1 is validated only when the mesh has an explicit "ep"
555 dimension (old-style EP). In D-10 TP-extend-EP the ep group is
556 *derived* from the dense region, not a native mesh axis, so the
557 absence of "ep" in mesh_dim_names is expected and not an error.
558 """
559 mesh_names = tuple(getattr(mesh, "mesh_dim_names", ()) or ())
560 mesh_shape = tuple(getattr(mesh, "mesh_shape", ()) or ())
561 if not mesh_names or not mesh_shape:
562 return # cannot validate without mesh metadata
564 # Degenerate single-rank mesh (compile-time test fixtures): skip
565 if all(sz == 1 for sz in mesh_shape):
566 return
568 name_to_size = dict(zip(mesh_names, mesh_shape))
570 for ax, size in (("tp", tp_size), ("cp", cp_size)):
571 if size > 1:
572 if ax not in name_to_size:
573 raise ValueError(
574 f"{ax}_size={size} > 1, but the mesh has no '{ax}' dimension. "
575 f"Mesh dimensions: {list(name_to_size.keys())}. "
576 f"Either add '{ax}' to the mesh's mesh_dim_names, or set "
577 f"{ax}_size=1."
578 )
579 mesh_sz = name_to_size[ax]
580 if size != mesh_sz:
581 raise ValueError(
582 f"{ax}_size ({size}) does not match mesh['{ax}'] size "
583 f"({mesh_sz}). They must be equal."
584 )
586 # ep_size: only validate when "ep" is a declared mesh axis (old-style
587 # EP). D-10 TP-extend-EP does not put "ep" in mesh_dim_names — the
588 # ep group is derived from the dense region by _expert_mesh_layout.
589 if ep_size > 1 and "ep" in name_to_size:
590 mesh_sz = name_to_size["ep"]
591 if ep_size != mesh_sz:
592 raise ValueError(
593 f"ep_size ({ep_size}) does not match mesh['ep'] size "
594 f"({mesh_sz}). They must be equal."
595 )
597 def _build_mesh_dim_names(
598 self, mesh, tp_size: int, cp_size: int, ep_size: int,
599 ) -> Tuple[str, ...]:
600 """Filter tp/cp/ep with mesh.mesh_dim_names as the authoritative
601 order; fall back to (tp,cp,ep) when undeclared; drop size=1 axes.
603 dp* axes are ALWAYS stripped (05 §3 coordinate-system convention): the plan's
604 coordinate system is a single dp slice — dp semantics live in the
605 data pipeline (data assignment) and FSDP (weight/grad domain),
606 never in placements. Overrides declaring DP placements are
607 rejected fail-first by :meth:`_check_overrides_no_dp`.
609 Raises :class:`ValueError` when *tp_size* / *cp_size* do not match the
610 corresponding mesh dimension sizes (fail-first before any planning work).
611 """
612 self._validate_dtensor_axes(mesh, tp_size, cp_size, ep_size)
613 mesh_names = tuple(getattr(mesh, "mesh_dim_names", ()) or ())
614 dtensor_axes = ("tp", "cp", "ep")
615 active = {ax for ax, sz in (("tp", tp_size), ("cp", cp_size), ("ep", ep_size))
616 if sz and sz > 1}
617 if mesh_names:
618 return tuple(n for n in mesh_names if n in dtensor_axes and n in active)
619 return tuple(ax for ax in dtensor_axes if ax in active)
621 def _check_overrides_no_dp(self) -> None:
622 """Fail-first: plan_overrides must not declare DP placements.
624 The plan's coordinate system is a single dp slice (05 §3 coordinate-system convention):
625 dp placements never drive boundary communication, and parameter dp
626 layout is owned by FSDP *after* the plan — a DP key here would be
627 silently dropped at resolve time, so reject it with a teaching
628 message instead of letting the misreading happen.
629 """
631 def _declares_dp(node) -> bool:
632 return isinstance(node, dict) and any(
633 k == DP or _declares_dp(v) for k, v in node.items())
635 for fqn, spec in self._plan_overrides.items():
636 if not isinstance(spec, ModuleShardingSpec):
637 continue # wrong types are reported as TypeError by _merge_plan_overrides
638 for field in ("params", "in_src", "in_dst", "out_src", "out_dst"):
639 if _declares_dp(getattr(spec, field)):
640 raise ValueError(
641 f'plan_overrides["{fqn}"].{field}: declaring a DP '
642 "placement is not allowed. The plan's coordinate "
643 "system is a single dp slice (tp/cp/ep, 05 §3 "
644 "coordinate-system convention) — dp data sharding is "
645 "expressed by the data pipeline, parameter/gradient "
646 "sharding by FSDP; for scenarios such as multimodal "
647 "encoder_dp the dp semantics live in vit_mesh + data "
648 "assignment + fully_shard, and the I/O contract only "
649 "needs to declare tp/cp/ep."
650 )
652 # ── Phase 1 ─────────────────────────────────────────────────────────
654 @staticmethod
655 def _warn_uncovered_function_modules(model, plan) -> None:
656 """DX guard: a FunctionModule without a boundary spec gets NO
657 redistribution communication (it is invisible to the planner's
658 derivation) — surface it instead of silently passing."""
659 uncovered = [fqn for fqn, m in model.named_modules()
660 if isinstance(m, FunctionModule) and fqn not in plan.modules]
661 if uncovered:
662 logger.warning(
663 "FunctionModule %s has no boundary spec — no communication "
664 "will be inserted around it. Declare one via "
665 "plan_overrides (params={}, region_dispatch=False).",
666 uncovered,
667 )
669 def _classify_all_params(self, model, arch: str) -> Dict[str, ParamRole]:
670 """Phase 1 entry: default naming rules plus the family's declared
671 sharding rules (``ModelAdapterSpec.sharding_rules``, resolved through
672 the models registry — the planner core carries no per-family
673 knowledge)."""
674 from hyper_parallel.models.registry import ( # pylint: disable=import-outside-toplevel
675 get_model_adapter,
676 )
677 spec = (model._hp_model_adapter_spec if hasattr(model, "_hp_model_adapter_spec") # pylint: disable=protected-access
678 else get_model_adapter(arch))
679 rules = []
680 if spec is not None and spec.sharding_rules is not None:
681 rules = list(spec.sharding_rules())
682 if not rules:
683 return self._classifier.classify(model, arch)
684 classifier = ParameterClassifier(arch_overrides={arch: rules})
685 return classifier.classify(model, arch)
687 # ── Phase 2 ─────────────────────────────────────────────────────────
689 def _group_by_boundary(
690 self, param_roles: Dict[str, ParamRole],
691 ) -> Dict[str, List[Tuple[str, ParamRole]]]:
692 """Single deepest-first sweep over the module hierarchy:
694 Pass 1: group by owning module FQN (strip the leaf parameter name).
695 Pass 2: materialize ALL ancestor modules upfront and sweep them in
696 strictly decreasing depth — every shallower ancestor is
697 therefore inferred only after ALL its descendants' parameters
698 have merged into it. On ``unknown`` the whole group's
699 parameters merge upward into the parent module. If still
700 unknown after backtracking to the root, attribute the group
701 to the parameter's own module (no template will match later
702 → warning and skip).
704 (Fix for the merge-ordering hazard of the former tail-enqueue work
705 queue: a parent enqueued by an early child — e.g. the MoE gate —
706 was inferred before later descendants — e.g. the experts subtree —
707 had merged into it, and the late parameters were then silently
708 dropped into an already-consumed module. Accuracy fix F3 made this
709 observable: a MOE_GATE-only ``mlp`` group must merge upward, which
710 stranded the experts params entirely.)
711 """
712 # Pass 1
713 own: Dict[str, List[Tuple[str, ParamRole]]] = {}
714 for fqn, role in param_roles.items():
715 module_fqn = ".".join(fqn.split(".")[:-1])
716 own.setdefault(module_fqn, []).append((fqn, role))
718 # Pass 2
719 all_modules = set(own)
720 for mfqn in own:
721 parts = mfqn.split(".")
722 for i in range(1, len(parts)):
723 all_modules.add(".".join(parts[:i]))
724 merged: Dict[str, List[Tuple[str, ParamRole]]] = {
725 mfqn: list(own.get(mfqn, [])) for mfqn in all_modules
726 }
727 groups: Dict[str, List[Tuple[str, ParamRole]]] = {}
728 for mfqn in sorted(all_modules, key=lambda f: f.count("."),
729 reverse=True):
730 params = merged[mfqn]
731 if not params:
732 continue
733 if self._infer_boundary_type(mfqn, params) != "unknown":
734 groups[mfqn] = params
735 else:
736 parent = mfqn.rsplit(".", 1)[0] if "." in mfqn else ""
737 if parent:
738 merged[parent].extend(params)
739 else:
740 # Still unknown after backtracking to the root:
741 # attribute each parameter to its own module (no template
742 # will match later → skipped). Unrelated unknown
743 # subtrees may have merged at the root, so anchoring the
744 # whole group to the first parameter would produce invalid
745 # relative names for the remaining parameters.
746 for param_fqn, role in params:
747 origin = param_fqn.rpartition(".")[0]
748 groups.setdefault(origin, []).append((param_fqn, role))
749 return groups
751 # ── Phase 3 ─────────────────────────────────────────────────────────
753 @staticmethod
754 def _explicit_boundary_type(fqn_lower: str, segment: str) -> Optional[str]:
755 """Return an explicit FQN-based boundary type when one matches."""
756 if _match_any(
757 fqn_lower,
758 ["embed_tokens", "wte", ".embed.", "tok_embeddings", "embed_in", "word_embeddings"],
759 ):
760 return "embed"
761 if _match_any(fqn_lower, ["lm_head", "embed_out", "output_layer"]):
762 return "lm_head"
763 if _match_any(fqn_lower, ["norm", "layernorm", "rmsnorm", "ln_"]):
764 return "norm"
765 if _match_any(segment, ["router"]):
766 return "moe_gate"
767 return None
769 @staticmethod
770 def _moe_boundary_type(fqn: str, roles: set[ParamRole]) -> Optional[str]:
771 """Return the boundary type implied by MoE-specific parameter roles."""
772 if ParamRole.MOE_EXPERT in roles:
773 return "moe_mlp" if _match_any(_last_segment(fqn), list(_MOE_CONTAINER_PATTERNS)) else "unknown"
774 if ParamRole.SHARED_EXPERT in roles:
775 return "mlp"
776 if ParamRole.MOE_GATE in roles:
777 return "unknown"
778 return None
780 @staticmethod
781 def _dense_boundary_type(fqn_lower: str, group: List[Tuple[str, ParamRole]]) -> str:
782 """Infer an attention or MLP boundary from dense parameter roles."""
783 has_colwise = any(
784 role in (ParamRole.COLWISE, ParamRole.FUSED_QKV, ParamRole.FUSED_GATE_UP)
785 for _, role in group
786 )
787 has_rowwise = any(role == ParamRole.ROWWISE for _, role in group)
788 if has_colwise and has_rowwise:
789 if _match_any(fqn_lower, list(_ATTN_PATTERNS)):
790 return "attention"
791 if _match_any(fqn_lower, list(_MLP_PATTERNS)):
792 return "mlp"
793 return "attention"
794 if has_colwise and _match_any(fqn_lower, list(_MLP_PATTERNS)):
795 return "mlp"
796 return "unknown"
798 def _infer_boundary_type(self, fqn: str, group: List[Tuple[str, ParamRole]]) -> str:
799 """Identify the semantic role from the module FQN + the group's
800 parameter roles.
802 Priority: explicit FQN patterns > leaf-segment guard > MoE roles >
803 parameter role combinations > default.
804 """
805 fqn_lower = fqn.lower()
806 seg = _last_segment(fqn)
808 explicit_type = self._explicit_boundary_type(fqn_lower, seg)
809 if explicit_type is not None:
810 return explicit_type
812 # 2. Leaf-segment guard: projection/expert leaf modules are not
813 # boundary containers themselves
814 if seg in _LEAF_SEGMENT_GUARD:
815 return "unknown"
816 # Numeric-segment guard: HF per-expert containers (experts.0..N) are
817 # not boundaries; parameters must aggregate upward into the moe
818 # container (D-09, 05 §6.4.7)
819 if seg.isdigit():
820 return "unknown"
822 # 3. MoE roles (accuracy fix F3, accuracy_fix_plan.md §2):
823 # structural lint — a "moe_mlp" boundary must contain routed experts
824 # (MOE_EXPERT); a MOE_GATE-only group (e.g. a scalar gate Linear
825 # misclassified as router) must NOT anchor a MoE boundary. Container
826 # identity is judged on the module's OWN last segment, not a
827 # substring of the whole FQN (a parent path segment like "mlp" must
828 # not qualify a leaf Linear such as shared_expert_gate).
829 roles = {r for _, r in group}
830 moe_type = self._moe_boundary_type(fqn, roles)
831 return moe_type if moe_type is not None else self._dense_boundary_type(fqn_lower, group)
833 # ── Phase 4 post-processing: MoE EP marking (D-09/D-10, 05 §6.4.7/§6.4.8) ──
835 @staticmethod
836 def _validate_ep_domain(ep_extend, mesh) -> None:
837 """Validate that virtual EP evenly partitions the dense rank domain.
839 The dense region = all ranks of the non-pp mesh axes
840 (dp_replicate × dp_cp × tp).
841 """
842 names = tuple(getattr(mesh, "mesh_dim_names", ()) or ())
843 shape = tuple(getattr(mesh, "mesh_shape", ()) or ())
844 domain = 1
845 for name, size in zip(names, shape):
846 if name == "pp" and size > 1:
847 raise NotImplementedError(
848 "virtual EP v1 does not support pp>1 "
849 "(split the mesh by stage before calling)"
850 )
851 if name != "pp":
852 domain *= size
853 if ep_extend > domain or domain % ep_extend != 0:
854 raise ValueError(
855 f"ep_size ({ep_extend}) must not exceed and must divide "
856 f"the dense region (dp_replicate × dp_cp × tp = {domain})"
857 )
859 @classmethod
860 def _validate_ep_extend(cls, ep_extend, mesh, model) -> None:
861 """Validate virtual EP plus the routed-expert count constraint."""
862 cls._validate_ep_domain(ep_extend, mesh)
863 num_experts = (getattr(getattr(model, "config", None), "num_experts", None)
864 or getattr(getattr(model, "config", None), "n_routed_experts", 0))
865 if num_experts and num_experts % ep_extend != 0:
866 raise ValueError(
867 f"num_experts ({num_experts}) must be divisible by ep_size ({ep_extend})"
868 )
870 @classmethod
871 def _mark_virtual_ep_parameters(
872 cls,
873 plan: ShardingPlan,
874 ep_size: int,
875 mesh: DeviceMesh,
876 ) -> None:
877 """Map explicit virtual-EP parameters to the derived expert mesh.
879 This deliberately keys off placement semantics rather than parameter
880 names, so routed experts and sparse lookup tables use one mechanism.
881 The boundary still has to inject its matching all-to-all computation.
882 """
883 if ep_size <= 1 or "ep" in tuple(getattr(mesh, "mesh_dim_names", ()) or ()):
884 return
885 marked = []
886 for fqn, spec in plan.modules.items():
887 if getattr(spec, "_ep_size", 0): # pylint: disable=protected-access
888 continue
889 ep_params = [
890 name
891 for name, named in (spec.params or {}).items()
892 if isinstance((named or {}).get(EP), Shard)
893 ]
894 if ep_params:
895 spec._ep_size = ep_size # pylint: disable=protected-access
896 marked.append((fqn, ep_params))
897 if not marked:
898 return
899 cls._validate_ep_domain(ep_size, mesh)
900 logger.info(
901 "virtual EP: mapped explicitly EP-sharded parameters to the "
902 "derived expert mesh: %s",
903 marked,
904 )
906 # per-expert parameter pattern: experts.<idx>.<proj>.weight (legacy HF /
907 # in-house MoE layouts).
908 _PER_EXPERT_RE = re.compile(r"^experts\.(\d+)\.([^.]+)\.weight$")
909 # batched parameter pattern: experts.<attr> (a single attribute with no
910 # numeric segment, the layout after the HF 2025 refactor —
911 # gate_up_proj [E, 2I, H] / down_proj [E, H, I], natively stacked with
912 # no stacking needed; the automodel names gate_and_up_projs/down_projs
913 # are isomorphic).
914 _BATCHED_EXPERT_RE = re.compile(
915 r"^experts\.(gate_up_proj|gate_and_up_projs|down_proj|down_projs"
916 r"|gate_proj|up_proj)$")
917 # Fused (gate|up merged into ONE parameter, chunk-split inside forward)
918 # expert weight pattern — covers both the batched 3D layout
919 # (``experts.gate_up_proj [E, 2I, H]``) and the per-expert 2D layout
920 # (``experts.<idx>.gate_up_proj.weight [2I, H]``). Such weights are
921 # incompatible with contiguous-block TP Shard on the fused dim; see
922 # _finalize_fused_expert_tp_guard.
923 _FUSED_EXPERT_WEIGHT_RE = re.compile(
924 r"^experts\.(?:\d+\.)?(?:gate_up_proj|gate_and_up_projs"
925 r"|fused_gate_up|w13)(?:\.weight)?$")
927 def _detect_expert_layout(self, group, boundary_fqn: str, param_ndims):
928 """Return expert parameters grouped by per-expert and batched layouts."""
929 expert_params = [fqn for fqn, role in group if role == ParamRole.MOE_EXPERT]
930 if not expert_params:
931 return None
932 stacks: Dict[str, List[Tuple[int, str]]] = {}
933 batched: List[str] = []
934 for param_fqn in expert_params:
935 rel = param_fqn[len(boundary_fqn) + 1:]
936 if "bias" in rel.lower():
937 logger.warning(
938 "%s: MoE expert has bias (%s); not supported in v1, skipping EP marking",
939 boundary_fqn, rel,
940 )
941 return None
942 match = self._PER_EXPERT_RE.match(rel)
943 if match is not None:
944 stacks.setdefault(match.group(2), []).append((int(match.group(1)), rel))
945 elif self._BATCHED_EXPERT_RE.match(rel) is not None and (param_ndims or {}).get(param_fqn, 2) >= 3:
946 batched.append(rel)
947 if stacks and batched:
948 logger.warning(
949 "%s: mixed per-expert and batched layouts (%s ...); skipping EP marking",
950 boundary_fqn, batched[0],
951 )
952 return None
953 return expert_params, stacks, batched
955 @staticmethod
956 def _mark_old_style_experts(
957 spec: ModuleShardingSpec,
958 stacks: Dict[str, List[Tuple[int, str]]],
959 template: ShardingTemplate,
960 ) -> None:
961 """Stack per-expert parameters while retaining explicit TP and EP axes."""
962 for projection, items in stacks.items():
963 items.sort()
964 sources = [rel for _, rel in items]
965 stacked = f"experts.{projection}"
966 tp_placement = _moe_expert_tp_placement(stacked, ndim=3, template=template)
967 for rel in sources:
968 spec.params.pop(rel, None)
969 spec.params[stacked] = _multi_dim(
970 tp=tp_placement, cp=Replicate(), ep=template.moe_expert_placement
971 )
972 spec._ep_stack[stacked] = sources # pylint: disable=protected-access # planner owns the spec DSL internals
974 @staticmethod
975 def _mark_extended_expert_params(
976 spec: ModuleShardingSpec,
977 expert_params: List[str],
978 stacks: Dict[str, List[Tuple[int, str]]],
979 batched: List[str],
980 boundary_fqn: str,
981 template: ShardingTemplate,
982 ) -> None:
983 """Mark expert parameters for TP-extended EP sharding."""
984 if stacks:
985 for projection, items in stacks.items():
986 items.sort()
987 sources = [rel for _, rel in items]
988 stacked = f"experts.{projection}"
989 for rel in sources:
990 spec.params.pop(rel, None)
991 spec.params[stacked] = _multi_dim(
992 tp=None, cp=Replicate(), ep=template.moe_expert_placement
993 )
994 # The planner owns the spec DSL internals.
995 spec._ep_stack[stacked] = sources # pylint: disable=protected-access
996 spec.region_dispatch = None
997 return
998 if batched:
999 for rel in batched:
1000 spec.params[rel] = _multi_dim(
1001 tp=None, cp=Replicate(), ep=template.moe_expert_placement
1002 )
1003 spec.region_dispatch = None
1004 return
1005 for param_fqn in expert_params:
1006 rel = param_fqn[len(boundary_fqn) + 1:]
1007 spec.params[rel] = _multi_dim(
1008 tp=None, cp=Replicate(), ep=template.moe_expert_placement
1009 )
1011 @staticmethod
1012 def _set_extended_ep_contract(spec: ModuleShardingSpec, ep_extend: int) -> None:
1013 """Set the identity boundary contract used by TP-extended EP."""
1014 identity = copy.deepcopy(spec.in_src)
1015 spec.in_dst = copy.deepcopy(identity)
1016 out_layout = copy.deepcopy(next(iter(identity.values())))
1017 spec.out_src = {"output": copy.deepcopy(out_layout)}
1018 spec.out_dst = {"output": copy.deepcopy(out_layout)}
1019 spec._ep_size = ep_extend # pylint: disable=protected-access # planner owns the spec DSL internals
1021 def _mark_hf_native_moe(
1022 self, spec: ModuleShardingSpec, group, boundary_fqn: str,
1023 template: ShardingTemplate, mesh_dim_names: Tuple[str, ...], arch: str,
1024 *, ep_extend: int = 0, mesh=None, model=None, param_ndims=None,
1025 ) -> None:
1026 """Post-process moe_mlp spec for EP (05 §6.4.7/§6.4.8).
1028 The EP **mode** is determined solely by whether the mesh includes an
1029 explicit ``"ep"`` axis — NOT by parameter naming:
1031 * **Old-style EP** (``"ep" in mesh_dim_names``): the mesh already has
1032 an ``"ep"`` axis. :meth:`_build_spec_from_template` produced
1033 ``{TP: Shard(…), EP: Shard(0)}`` — the correct dual-axis
1034 sharding. This method only performs **stacking** when the expert
1035 layout is per-expert 2D (``experts.<idx>.<proj>.weight``); the
1036 stacked entry keeps both TP and EP keys. Batched 3D and custom
1037 (``w1``/``w2``/``w3``) layouts are already ready — nothing to do.
1038 No ``_ep_size`` is set; communication is handled by the module's
1039 own dispatcher or external ``_attach_ep``.
1041 * **D-10 TP-extend-EP** (``"ep" not in mesh_dim_names``, the
1042 default for HF-native models): the ep group is *derived* from the
1043 dense region. Expert weights are rewritten to
1044 ``{EP: Shard(0)}`` (no TP key), the boundary contract is changed
1045 to SP-in identity, and ``_ep_size`` is set.
1047 The expert **layout** (per-expert / batched / custom) only affects
1048 the *stacking strategy* — orthogonal to the EP mode:
1050 - **per-expert**: ``experts.<idx>.<proj>.weight`` → stack into
1051 ``experts.<proj>`` 3D before sharding;
1052 - **batched** (D-11): ``experts.gate_up_proj`` etc., natively 3D;
1053 - **custom**: ``experts.w1`` / ``w2`` / ``w3``, already 3D,
1054 pre-stacked by the module author.
1056 In D-10 mode all three layouts are supported; in old-style EP mode
1057 only per-expert needs stacking (batched and custom are already 3D).
1058 """
1059 # ``arch`` is part of the dispatch-interface signature shared with the
1060 # other boundary post-processors; the EP mode is decided by the mesh
1061 # and layout alone, so it is intentionally unused here.
1062 _ = arch
1063 if not ep_extend:
1064 return
1065 layout = self._detect_expert_layout(group, boundary_fqn, param_ndims)
1066 if layout is None:
1067 return
1068 expert_params, stacks, batched = layout
1069 if "ep" in mesh_dim_names:
1070 if stacks:
1071 self._mark_old_style_experts(spec, stacks, template)
1072 return
1073 self._validate_ep_extend(ep_extend, mesh, model)
1074 self._mark_extended_expert_params(
1075 spec, expert_params, stacks, batched, boundary_fqn, template
1076 )
1077 self._set_extended_ep_contract(spec, ep_extend)
1080 @staticmethod
1081 def _finalize_fused_expert_tp_guard(plan: ShardingPlan, *, tp_size: int) -> None:
1082 """Fail fast on fused expert weights TP-sharded without D-10 EP.
1084 A fused expert weight (gate/up merged into ONE parameter — e.g.
1085 ``experts.gate_up_proj [E, 2I, H]`` — and chunk-split inside the
1086 module's forward) is incompatible with the framework's
1087 contiguous-block TP ``Shard`` on the fused dim: after Phase A each
1088 rank holds a contiguous slice of the fused dim (tp=2 → rank0 all
1089 gate, rank1 all up), so the runtime ``chunk`` result is NOT
1090 equivalent to the non-TP semantics (Megatron solves this with
1091 chunk-aware sharding, which the framework does not support). The
1092 error is silent: the MoE boundary is local-region black-box hosted
1093 in BOTH modes, so production and validate run the same (wrong)
1094 local forward and the dual-mode comparison cannot expose it.
1096 Runs AFTER all plan_overrides are merged (next to
1097 ``_finalize_deferred_biases``) and anchors detection on the FINAL
1098 spec declarations, so an explicit TP-Replicate override on the
1099 fused params legitimately clears the hazard. D-10 TP-extend-EP
1100 (``_ep_size > 0``) removes the TP key from expert weights entirely
1101 (each rank holds complete experts), so it is exempt.
1102 """
1103 if tp_size <= 1:
1104 return
1105 offenders = []
1106 for fqn, spec in plan.modules.items():
1107 if getattr(spec, "_ep_size", 0): # pylint: disable=protected-access
1108 # D-10 TP-extend-EP: expert weights carry no TP shard — the
1109 # fused dim stays complete per rank, chunk is correct.
1110 continue
1111 for param_name, named in (spec.params or {}).items():
1112 if not isinstance(named, dict):
1113 continue
1114 if ShardingPlanner._FUSED_EXPERT_WEIGHT_RE.match(param_name) is None:
1115 continue
1116 if isinstance(named.get(TP), Shard):
1117 offenders.append((fqn, param_name, named[TP]))
1118 if not offenders:
1119 return
1120 listing = "\n".join(
1121 f" - {fqn}.{param_name} → {{TP: {placement}, ...}}"
1122 for fqn, param_name, placement in offenders[:8])
1123 more = (f"\n ... and {len(offenders) - 8} more"
1124 if len(offenders) > 8 else "")
1125 boundary_fqn, example_param, _ = offenders[0]
1126 raise ValueError(
1127 f"Detected {len(offenders)} fused MoE expert weight(s) (gate/up "
1128 "merged into a single parameter, chunk-split inside forward) "
1129 "declared with TP sharding:\n"
1130 f"{listing}{more}\n"
1131 "This pattern is incompatible with the framework's contiguous-"
1132 "block TP Shard: after Phase A shards the fused dim contiguously, "
1133 "each rank holds a slice whose gate/up layout is misordered, so "
1134 "the runtime chunk is NOT equivalent to the non-TP semantics "
1135 "(Megatron handles this with chunk-aware sharding, which the "
1136 "framework does not support). Worse, the error is silent: the "
1137 "MoE boundary is local-region black-box hosted in BOTH modes, so "
1138 "production and validate run the same (wrong) local forward and "
1139 "the dual-mode comparison cannot expose the accuracy corruption. "
1140 "Choose one of:\n"
1141 " (1) Enable ep_size (recommended): "
1142 "ShardingPlanner().plan(..., ep_size=N) with N>1, N dividing the "
1143 "dense region dp*cp*tp and dividing num_experts. Enabling ep_size "
1144 "takes the TP-extend-EP path (design revision D-10): the dense "
1145 "region is re-partitioned into the derived expert mesh "
1146 "(edp_shard, ep), and expert weights are sharded ONLY as "
1147 "{EP: Shard(0)} along the expert dim with NO TP key — every rank "
1148 "holds complete experts, so the gate/up produced by the runtime "
1149 "chunk matches the non-parallel semantics. Note: the MoE "
1150 "boundary must then also declare local_compute_fn explicitly "
1151 "(EP compute injection; the apply-time preflight checks this and "
1152 "prints a YAML example).\n"
1153 " (2) Keep ep=1: override the fused weight to TP Replicate via "
1154 "plan_overrides (it then opts out of TP sharding and saves no "
1155 "memory for that weight), e.g.:\n"
1156 " ShardingPlanner(plan_overrides={\n"
1157 f" {boundary_fqn!r}: ModuleShardingSpec(\n"
1158 " params={\n"
1159 f" {example_param!r}: "
1160 "{TP: Replicate(), CP: Replicate()},\n"
1161 " # ...declare ALL params of this boundary too "
1162 "(experts.down_proj / gate.weight etc.) — params merge replaces "
1163 "the field wholesale (no per-key merge)\n"
1164 " },\n"
1165 " # out_src TP must ALSO be changed from the template-"
1166 "derived Partial to Replicate: with replicated weights each "
1167 "rank's MoE output is already complete (not Partial); keeping "
1168 "Partial would make the boundary exit reduction count the "
1169 "output tp_size times\n"
1170 " out_src={'output': {TP: Replicate(), CP: Replicate()}},"
1171 " # keep the template-derived CP value (Shard(1) when "
1172 "sequence_parallel is on)\n"
1173 " ),\n"
1174 " })\n"
1175 " (3) Switch to separate w1/w3 expert weights (gate_proj and "
1176 "up_proj as independent parameters) — natively supported, no "
1177 "override needed.")
1179 @staticmethod
1180 def _should_defer_rowwise_bias(
1181 module_fqn: str,
1182 module: nn.Module,
1183 owner_path: str,
1184 param_name: str,
1185 bias_tp,
1186 out_src,
1187 partial_outputs: List[str],
1188 ) -> bool:
1189 """Validate a rowwise bias and return whether it must be deferred."""
1190 if not partial_outputs:
1191 return False
1192 if len(out_src) != 1:
1193 raise ValueError(
1194 f"boundary {module_fqn!r}: rowwise bias deferral (D-22) v1 only supports single-output "
1195 f"boundaries — this boundary's out_src declares {len(out_src)} outputs with a TP Partial "
1196 f"reduction, so the framework cannot attribute {param_name!r} to a unique output. Take over "
1197 "the region with local_compute_fn (add the bias yourself after the reduction)"
1198 )
1199 if bias_tp is not None and not isinstance(bias_tp, Replicate):
1200 raise ValueError(
1201 f"boundary {module_fqn!r}: the bias of rowwise Linear {owner_path!r} declares a non-Replicate "
1202 f"TP placement ({bias_tp!r}) — D-22 deferred addition requires the bias to stay Replicate "
1203 f"(added exactly once as a whole after the TP reduction). Remove {param_name!r} from "
1204 "spec.params, or change it to {TP: replicate()}"
1205 )
1206 owner = module.get_submodule(owner_path) if owner_path else module
1207 if isinstance(owner, nn.Linear):
1208 return True
1209 logger.warning(
1210 "boundary %s: rowwise Linear %r carries a bias and the boundary out_src is TP Partial — but "
1211 "the owner type is %s (not nn.Linear), so the framework does not touch its forward semantics: "
1212 "the bias will be counted multiple times by the Partial reduction (production output = correct "
1213 "value + tp_size × bias). Move the bias after the boundary communication, switch to nn.Linear, "
1214 "or take over the region with local_compute_fn",
1215 module_fqn, owner_path, type(owner).__name__,
1216 )
1217 return False
1219 @staticmethod
1220 def _validate_colwise_bias(
1221 module_fqn: str,
1222 weight_path: str,
1223 param_name: str,
1224 param,
1225 bias_named,
1226 bias_tp,
1227 shard_dim: int,
1228 ) -> None:
1229 """Require a colwise bias to follow its weight's output shard."""
1230 bias_dim = None
1231 if isinstance(bias_tp, Shard):
1232 bias_dim = bias_tp.dim if bias_tp.dim >= 0 else bias_tp.dim + param.ndim
1233 if bias_dim == shard_dim:
1234 return
1235 declared = repr(bias_tp) if bias_named else "undeclared"
1236 raise ValueError(
1237 f"boundary {module_fqn!r}: {weight_path!r} is sharded along the output dim as Shard({shard_dim}), "
1238 f"but {param_name!r} is not sharded along the output channels the same way ({declared}) — "
1239 f"template mismatch (typical: lm_head.bias). Declare it explicitly via plan_overrides as "
1240 f"{{'{param_name}': {{TP: shard({shard_dim})}}}}, or remove the bias"
1241 )
1243 @staticmethod
1244 def _bias_owner_path(param_name: str) -> Optional[str]:
1245 """Return a bias owner's module path, or ``None`` for other params."""
1246 if param_name == "bias":
1247 return ""
1248 if param_name.endswith(".bias"):
1249 return param_name[: -len(".bias")]
1250 return None
1252 @staticmethod
1253 def _collect_deferred_biases(module_fqn: str, spec: ModuleShardingSpec, module: nn.Module) -> List[str]:
1254 """Validate one boundary's biases and collect rowwise deferred names."""
1255 named_params = dict(module.named_parameters())
1256 out_src = spec.out_src or {}
1257 partial_outputs = [
1258 out_name for out_name, named in out_src.items() if isinstance(named.get(TP), Partial)
1259 ]
1260 deferred: List[str] = []
1261 for param_name, param in named_params.items():
1262 owner_path = ShardingPlanner._bias_owner_path(param_name)
1263 if owner_path is None:
1264 continue
1265 weight_path = f"{owner_path}.weight" if owner_path else "weight"
1266 weight = named_params.get(weight_path)
1267 tp_placement = (spec.params or {}).get(weight_path)
1268 if weight is None or tp_placement is None:
1269 continue
1270 tp_placement = tp_placement.get(TP)
1271 if not isinstance(tp_placement, Shard):
1272 continue
1273 shard_dim = tp_placement.dim if tp_placement.dim >= 0 else tp_placement.dim + weight.ndim
1274 bias_tp = spec.params.get(param_name)
1275 bias_tp = bias_tp.get(TP) if bias_tp else None
1276 if shard_dim == weight.ndim - 1:
1277 if ShardingPlanner._should_defer_rowwise_bias(
1278 module_fqn, module, owner_path, param_name, bias_tp, out_src, partial_outputs
1279 ):
1280 deferred.append(param_name)
1281 else:
1282 ShardingPlanner._validate_colwise_bias(
1283 module_fqn, weight_path, param_name, param, spec.params.get(param_name), bias_tp, shard_dim
1284 )
1285 return deferred
1287 @staticmethod
1288 def _finalize_deferred_biases(
1289 plan: ShardingPlan, model, mesh_dim_names: Tuple[str, ...],
1290 ) -> None:
1291 """D-22: decide which Linear biases are deferred to after the boundary
1292 exit TP reduction, and validate every bias against its sibling weight's
1293 sharding direction.
1295 Runs AFTER all plan_overrides are merged (next to
1296 ``_finalize_tp_local_attr_plans``) and anchors detection on the FINAL
1297 spec declarations plus the model structure — never on ParamRole — so
1298 derived / merge / insert / derive=False specs (and family sharding-rule
1299 naming like ``wo``/``c_proj``) all share one code path.
1301 Per boundary, for every physical ``X.bias`` parameter (scanned from
1302 the module, regardless of whether spec.params declares it — a
1303 physically present bias is fused by F.linear either way):
1305 - sibling ``X.weight`` TP placement is ``Shard(weight.ndim - 1)``
1306 (contraction-dim = rowwise) AND the boundary out_src reduces on TP
1307 (Partial) → the fused bias would be counted once per TP rank by the
1308 exit reduction → defer it (bias stays Replicate, added exactly once
1309 after the reduction). Fail-fast guards: a declared non-Replicate
1310 bias placement is rejected; a non-nn.Linear owner gets a WARNING
1311 and is skipped (its forward semantics are not touched); a
1312 multi-output boundary cannot attribute the bias to one output →
1313 fail-fast (use local_compute_fn).
1314 - sibling ``X.weight`` TP placement is an output-dim ``Shard(d)``
1315 (colwise / lm_head / embed) → the bias must follow the same output
1316 shard; Replicate / undeclared / wrong-dim → plan-time "template
1317 mismatch"
1318 error (typical: lm_head.bias), instead of a remote runtime
1319 broadcast crash.
1320 - out_src without a TP Partial (no boundary reduction, or the user
1321 reduces inside the region) → the bias is already added exactly once
1322 → no defer, no check.
1323 """
1324 has_tp = "tp" in mesh_dim_names
1325 modules = dict(model.named_modules())
1326 for module_fqn, spec in plan.modules.items():
1327 spec._deferred_bias_params = () # pylint: disable=protected-access
1328 if not has_tp or not spec.is_boundary:
1329 continue
1330 module = modules.get(module_fqn)
1331 if module is None:
1332 continue
1333 deferred = ShardingPlanner._collect_deferred_biases(module_fqn, spec, module)
1334 spec._deferred_bias_params = tuple(deferred) # pylint: disable=protected-access
1335 if deferred:
1336 logger.info(
1337 "boundary %s: deferred bias (D-22, added exactly once "
1338 "after the TP reduction): %s",
1339 module_fqn, list(deferred))
1341 @staticmethod
1342 def _check_param_uniqueness(plan: ShardingPlan) -> None:
1343 """D-14 invariant 1 (05 §13.3): every parameter is sharded by exactly
1344 one boundary. spec.params keys are resolved to full parameter FQNs
1345 (relative to the boundary module); any parameter declared by ≥2 specs
1346 fails fast (double sharding corrupts silently under production)."""
1347 seen: Dict[str, str] = {} # full param fqn -> first declaring boundary
1348 for fqn, spec in plan.modules.items():
1349 prefix = fqn + "." if fqn else ""
1350 for pname in spec.params:
1351 full = prefix + pname
1352 if full in seen:
1353 raise ValueError(
1354 f"param {full!r} is declared by two boundaries: "
1355 f"{seen[full]!r} and {fqn!r} — each parameter must be "
1356 f"sharded by exactly one boundary (D-14, 05 §13.3); "
1357 f"drop it from one of the specs (an outer boundary "
1358 f"may only declare params of its own/intermediate "
1359 f"modules, never of a nested boundary's subtree)"
1360 )
1361 seen[full] = fqn
1363 @staticmethod
1364 def _check_full_declaration(plan: ShardingPlan) -> None:
1365 """D-14 (05 §13.2): chain fill is removed — every boundary spec must
1366 fully declare its I/O contract. A non-empty in_dst with an empty
1367 in_src fails fast (the previous Scenario-1 fill no longer exists)."""
1368 for fqn, spec in plan.modules.items():
1369 if spec.is_boundary and spec.in_dst and not spec.in_src:
1370 raise ValueError(
1371 f"boundary {fqn!r} declares in_dst but an empty in_src — "
1372 f"chain fill was removed (D-14, 05 §13.2); declare in_src "
1373 f"explicitly (keys must mirror in_dst)"
1374 )
1376 # ── F4 plan-time lints (accuracy_fix_plan.md §2) ─────────────────────
1378 @staticmethod
1379 def _check_param_shard_divisibility(
1380 fqn: str,
1381 prefix: str,
1382 pname: str,
1383 placement: Dict[str, Any],
1384 param_shapes: Dict[str, Tuple[int, ...]],
1385 ep_stack: Dict[str, List[str]],
1386 axis_sizes: Dict[str, int],
1387 ) -> None:
1388 """Check every sharded mesh axis for one declared parameter."""
1389 full = prefix + pname
1390 shape = param_shapes.get(full)
1391 if shape is None and pname in ep_stack:
1392 sources = ep_stack[pname]
1393 shape = param_shapes.get(prefix + sources[0]) if sources else None
1394 if shape is not None:
1395 shape = (len(sources), *shape)
1396 if shape is None:
1397 return # created later (inner_target factory) — nothing to check
1398 for axis, p in (placement or {}).items():
1399 if not isinstance(p, Shard):
1400 continue
1401 size = axis_sizes.get(axis, 1)
1402 if size <= 1:
1403 continue
1404 axis_name = getattr(axis, "value", axis) # MeshAxisName → "tp"
1405 dim = p.dim + len(shape) if p.dim < 0 else p.dim
1406 if dim >= len(shape):
1407 raise ValueError(
1408 f"plan-time shard check failed: {full!r} has shape "
1409 f"{tuple(shape)} but boundary {fqn!r} declares "
1410 f"{{{axis_name}: Shard({p.dim})}} — dim {p.dim} is out "
1411 f"of range for a {len(shape)}D parameter; fix the "
1412 f"plan_overrides declaration"
1413 )
1414 if shape[dim] % size != 0:
1415 raise ValueError(
1416 f"plan-time shard check failed: {full!r} has shape "
1417 f"{tuple(shape)} but boundary {fqn!r} declares "
1418 f"{{{axis_name}: Shard({p.dim})}} — shape[{dim}]={shape[dim]} "
1419 f"is not divisible by {axis_name} size {size} (it would "
1420 f"produce empty shards at apply time). This is most "
1421 f"often a parameter-classification error (e.g. a "
1422 f"replicated/gate parameter misclassified into a "
1423 f"sharded role — see accuracy_fix_plan.md §2): fix "
1424 f"the naming rule / family sharding_rules entry, or correct "
1425 f"the plan_overrides declaration"
1426 )
1428 @staticmethod
1429 def _check_shard_divisibility(
1430 plan: ShardingPlan,
1431 param_shapes: Dict[str, Tuple[int, ...]],
1432 *,
1433 tp_size: int,
1434 cp_size: int,
1435 ep_size: int,
1436 ) -> None:
1437 """F4a: every Shard(dim) in every spec.params must divide the
1438 parameter's shape along that dim — fail at plan time with a teaching
1439 error instead of producing an empty shard at apply time
1440 (accuracy_problem.md 10.1: a (1, 64) scalar-gate weight misclassified
1441 SHARED_EXPERT was Shard(0)'d over tp=2 into a (0, 64) empty shard,
1442 surfacing only much later at weight reconstruction).
1444 Runs after the Phase 4.5 override merge, so hand-written specs are
1445 checked too. Stacked expert params (D-09/D-10 ``_ep_stack``) resolve
1446 their shape as ``(num_experts, *source_shape)``; D-10 experts use
1447 ``spec._ep_size`` as the EP axis size. Params not resolvable at plan
1448 time (e.g. created later by an inner_target factory) are skipped.
1449 """
1450 for fqn, spec in plan.modules.items():
1451 prefix = fqn + "." if fqn else ""
1452 ep_stack = getattr(spec, "_ep_stack", None) or {}
1453 # D-10: the EP axis lives on the derived expert mesh whose size
1454 # is spec._ep_size; old-style EP uses the mesh "ep" axis (ep_size).
1455 axis_sizes = {
1456 "tp": tp_size,
1457 "cp": cp_size,
1458 "ep": getattr(spec, "_ep_size", 0) or ep_size,
1459 }
1460 for pname, placement in (spec.params or {}).items():
1461 ShardingPlanner._check_param_shard_divisibility(
1462 fqn, prefix, pname, placement, param_shapes, ep_stack, axis_sizes
1463 )
1465 @staticmethod
1466 def _declared_parameter_names(plan: ShardingPlan) -> set[str]:
1467 """Collect full names covered by placements and special handlers."""
1468 covered = set()
1469 for fqn, spec in plan.modules.items():
1470 prefix = fqn + "." if fqn else ""
1471 for pname in (spec.params or {}):
1472 covered.add(prefix + pname)
1473 for sources in (getattr(spec, "_ep_stack", None) or {}).values():
1474 covered.update(prefix + source for source in sources)
1475 covered.update(plan.special_handlers or {})
1476 return covered
1478 def _check_all_trainable_params_covered(self, plan: ShardingPlan, model) -> None:
1479 """F4b: every ``requires_grad=True`` parameter must appear in some
1480 spec.params (resolved to a full FQN) or in special_handlers —
1481 otherwise its gradient-sync semantics would be decided SILENTLY by
1482 the consumer-side default (the same class of shape-legal/silent-wrong
1483 hazard as accuracy_problem.md 10.1/10.2), so the plan must be a
1484 complete declaration of every trainable parameter's layout.
1486 Covered = spec.params ∪ _ep_stack sources ∪ special_handlers (a
1487 param explicitly declared via plan_overrides therefore counts as
1488 covered). Remediation paths, listed per parameter in the error:
1489 ① classify it explicitly (naming rule / the family's sharding_rules, e.g.
1490 shared_expert_gate → REPLICATED);
1491 ② declare it in a spec via plan_overrides (e.g. all-Replicate);
1492 ③ mark it SPECIAL (special_handlers) or freeze it
1493 (requires_grad=False).
1494 Escape hatch for exploratory debugging:
1495 ``ShardingPlanner(allow_uncovered_params=True)`` downgrades to a
1496 warning. Skipped entirely under derive=False (pure declaration
1497 assembler for subtrees — coverage is the caller's responsibility).
1498 """
1499 if not self._derive:
1500 return
1501 covered = self._declared_parameter_names(plan)
1502 covered.update(plan.owned_parameter_names)
1503 uncovered = [
1504 name for name, p in model.named_parameters()
1505 if p.requires_grad and name not in covered
1506 ]
1507 if not uncovered:
1508 return
1509 lines = "\n".join(f" - {n}" for n in uncovered[:20])
1510 more = (f" ... and {len(uncovered) - 20} more\n"
1511 if len(uncovered) > 20 else "")
1512 msg = (
1513 f"plan-time coverage check failed: {len(uncovered)} trainable "
1514 f"parameter(s) are not covered by any spec.params / "
1515 f"special_handlers:\n{lines}\n{more}"
1516 "An uncovered parameter is never sharded (kept replicated) AND is "
1517 "absent from source_shard_info — its gradient-sync semantics would be "
1518 "decided silently by the consumer-side default. Resolve each "
1519 "parameter explicitly: ① add a naming rule / a family sharding_rules "
1520 "entry "
1521 "(e.g. shared_expert_gate → REPLICATED); ② declare it in a spec "
1522 "via plan_overrides; ③ mark it SPECIAL or freeze it "
1523 "(requires_grad=False). For exploratory debugging only, construct "
1524 "ShardingPlanner(allow_uncovered_params=True) to downgrade this "
1525 "error to a warning (accuracy_fix_plan.md §2 F4b)"
1526 )
1527 if self._allow_uncovered_params:
1528 logger.warning("%s", msg)
1529 else:
1530 raise ValueError(msg)
1532 # ── Phase 5 ─────────────────────────────────────────────────────────
1534 def _mark_terminal(self, plan: ShardingPlan, model) -> ShardingPlan:
1535 """D-14 (05 §13.2): compile-time chain propagation/validation is
1536 removed; only ``_is_terminal`` marking is kept — the last boundary
1537 in forward order (each module vouches for its own propagation in
1538 validate mode; adjacent-contract checks no longer exist)."""
1539 boundaries = [fqn for fqn, spec in plan.modules.items() if spec.is_boundary]
1540 sorted_fqns = self._topological_sort_by_forward_order(boundaries, model)
1541 terminal = sorted_fqns[-1] if sorted_fqns else None
1542 for fqn, spec in plan.modules.items():
1543 # The planner owns the spec DSL internals.
1544 spec._is_terminal = fqn == terminal # pylint: disable=protected-access
1545 return plan
1547 @staticmethod
1548 def _topological_sort_by_forward_order(fqns: List[str], model) -> List[str]:
1549 """Sort by named_modules registration order; unmatched FQNs are
1550 appended at the end with a warning."""
1551 fqn_set = set(fqns)
1552 ordered: List[str] = []
1553 seen: set = set()
1554 for name, _ in model.named_modules():
1555 if name in fqn_set and name not in seen:
1556 ordered.append(name)
1557 seen.add(name)
1558 missing = fqn_set - seen
1559 if missing:
1560 logger.warning(
1561 "_topological_sort_by_forward_order: %d FQNs not found in "
1562 "named_modules; appended at the end: %s",
1563 len(missing), sorted(missing)[:5],
1564 )
1565 ordered.extend(sorted(missing))
1566 return ordered
1568 # ── tied weights ────────────────────────────────────────────────────
1570 @staticmethod
1571 def _detect_tied_pairs(model) -> List[Tuple[str, str]]:
1572 """Detect the embed_tokens.weight <-> lm_head.weight tied pair.
1574 With HF tie_word_embeddings both ends share storage; in PP
1575 scenarios the two ends cannot be detected across stages, so the
1576 user must declare plan.tied_pairs explicitly (05
1577 detect_tied_weights comment).
1578 """
1579 if not getattr(getattr(model, "config", None), "tie_word_embeddings", False):
1580 return []
1581 embed_fqn = lm_head_fqn = None
1582 # remove_duplicate=False: under the default deduplication of
1583 # named_parameters a tied parameter appears only once.
1584 for name, _ in model.named_parameters(remove_duplicate=False):
1585 if name.endswith("embed_tokens.weight"):
1586 embed_fqn = name
1587 elif name.endswith("lm_head.weight"):
1588 lm_head_fqn = name
1589 if embed_fqn and lm_head_fqn:
1590 return [(embed_fqn, lm_head_fqn)]
1591 return []
1594def _validate_tp_model_compatibility(config: Any, tp_size: int, ep_size: int) -> None:
1595 """Check model head and expert width divisibility on the TP axis."""
1596 heads = getattr(config, "num_attention_heads", None)
1597 if heads is not None and heads % tp_size != 0:
1598 raise ValueError(
1599 f"num_attention_heads ({heads}) must be divisible by TP ({tp_size})"
1600 )
1601 kv_heads = getattr(config, "num_key_value_heads", None)
1602 if kv_heads is not None and kv_heads % tp_size != 0:
1603 raise ValueError(
1604 f"num_key_value_heads ({kv_heads}) must be divisible by TP ({tp_size})"
1605 )
1606 moe_inter = getattr(config, "moe_intermediate_size", None)
1607 if ep_size > 1 and moe_inter is not None and moe_inter % tp_size != 0:
1608 raise ValueError(
1609 f"moe_intermediate_size ({moe_inter}) must be divisible by TP ({tp_size})"
1610 )
1613def _validate_ep_model_compatibility(config: Any, ep_size: int) -> None:
1614 """Check that the model has a divisible routed-expert count."""
1615 num_experts = (getattr(config, "num_experts", None)
1616 or getattr(config, "n_routed_experts", None) or 0)
1617 if num_experts <= 0:
1618 raise ValueError("EP>1 requires MoE model (num_experts > 0)")
1619 if num_experts % ep_size != 0:
1620 raise ValueError(
1621 f"num_experts ({num_experts}) must be divisible by EP ({ep_size})"
1622 )
1625def validate_model_compatibility(
1626 model: Any, *, tp_size: int = 1, cp_size: int = 1, ep_size: int = 1,
1627 seq_len: Optional[int] = None,
1628) -> None:
1629 """Validate model config against requested TP, CP, and EP degrees.
1631 Args:
1632 model: Model whose configuration declares attention and expert sizes.
1633 tp_size: Requested decoder tensor-parallel degree.
1634 cp_size: Requested decoder context-parallel degree.
1635 ep_size: Requested expert-parallel degree.
1636 seq_len: Optional global decoder token length.
1637 """
1638 config = getattr(model, "config", None)
1639 if config is None:
1640 return
1642 if tp_size > 1:
1643 _validate_tp_model_compatibility(config, tp_size, ep_size)
1645 if cp_size > 1 and seq_len is not None and seq_len % (cp_size * 2) != 0:
1646 raise ValueError(
1647 f"seq_len ({seq_len}) must be divisible by 2*cp ({2 * cp_size})"
1648 )
1650 if ep_size > 1:
1651 _validate_ep_model_compatibility(config, ep_size)