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