Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / parallel_dims.py: 97%
183 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-13 05:07 +0800
1# Copyright 2026 Huawei Technologies Co., Ltd
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7# http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14# ============================================================================
16"""ParallelDims — fail-fast parallel configuration validator + mesh builder.
18Centralises parallel-degree validation in a single dataclass so
19misconfigurations are caught before model construction.
21What this provides:
231. **Inference** — auto-fill ``dp`` (or ``dp_shard=-1``) from the product
24 constraint ``dp_replicate * dp_shard * cp * tp * pp == world_size``.
262. **Validation against world_size** — raises ``ValueError`` with a clear
27 message when the product mismatches.
293. **Validation against the model** — checks divisibility constraints that
30 would otherwise crash deep inside ``parallelize_module`` or the model's
31 own parallel setup:
33 - ``num_attention_heads % tp == 0`` (TP shards heads)
34 - ``num_key_value_heads % tp == 0`` (GQA constraint)
35 - ``num_experts % ep == 0`` (when MoE)
36 - context-parallel modules validate ``ulysses_degree``
37 - ``seq_len % (cp * tp) == 0`` (sequence parallel + CP)
38 - ``etp == tp or etp == 1`` (expert TP rule)
404. **Mesh building** — returns a ``DeviceMesh`` with the canonical dim order
41 ``pp → dp_replicate → dp_shard → ep → cp → tp``. Backwards compatible with
42 the legacy single ``dp`` field (auto-collapses to ``dp_shard``).
44User experience:
46- Default config (no parallel section) — works on 1 GPU, runs as DDP-1.
47- Set only ``tp=4`` on world_size=8 → ``dp`` auto-inferred to 2.
48- Set ``dp_shard=-1`` → fills remaining cards into FSDP shard dim.
49- Misconfig (heads=12, tp=8) → fails at ``_setup`` with a single readable
50 error before any model parallelization is attempted.
51"""
52from __future__ import annotations
54__all__ = ["ParallelDims"]
56import logging
57from dataclasses import dataclass, field
58from typing import Optional
60from hyper_parallel import init_device_mesh
62logger = logging.getLogger(__name__)
65@dataclass
66class ParallelDims:
67 """Validated parallel degrees + lazy mesh builder.
69 Attributes:
70 dp_replicate: DDP replication degree (HSDP outer dim).
71 dp_shard: FSDP shard degree. ``-1`` means "fill the rest from
72 ``world_size / (dp_replicate * cp * tp * pp)``".
73 cp: Context parallel degree.
74 tp: Tensor parallel degree (dense path).
75 pp: Pipeline parallel degree.
76 ep: Expert parallel degree (MoE only).
77 etp: Expert tensor parallel degree. Must equal ``tp`` or ``1``.
78 moe_token_dispatcher_type: Expert token exchange strategy.
79 npu_nums_per_device: Inner expert-parallel degree for deredundency dispatch.
80 ulysses_degree: Legacy direct-construction override. The YAML trainer
81 path leaves this to the context-parallel module, which owns the
82 actual Ulysses strategy validation.
83 world_size: Total number of ranks.
84 _allow_pp: Whether a trainer/model path has opted into pipeline
85 parallelism. Direct ``ParallelDims`` construction keeps the legacy
86 fail-fast guard used by the existing UTs.
87 """
89 dp_replicate: int = 1
90 dp_shard: int = 1
91 cp: int = 1
92 tp: int = 1
93 pp: int = 1
94 ep: int = 1
95 etp: int = 1
96 moe_token_dispatcher_type: str = "all_to_all"
97 npu_nums_per_device: int = 8
98 ulysses_degree: Optional[int] = None
99 world_size: int = 1
100 _allow_pp: bool = False
101 # Cached after build_mesh.
102 _device_mesh: object = field(default=None, repr=False)
104 # ------------------------------------------------------------------
105 # Construction & inference
106 # ------------------------------------------------------------------
107 @classmethod
108 def from_config(cls, parallel_cfg, world_size: int) -> "ParallelDims":
109 """Build from a ``ParallelConfig`` (or any object with the same fields).
111 Accepts the legacy single-``dp`` field. If ``dp`` is set and
112 ``dp_replicate``/``dp_shard`` are at default, ``dp`` is mapped to
113 ``dp_shard`` (FSDP behavior).
114 """
115 dp_replicate = getattr(parallel_cfg, 'dp_replicate', 1)
116 dp_shard = getattr(parallel_cfg, 'dp_shard', None)
117 legacy_dp = getattr(parallel_cfg, 'dp', None)
119 # Backward-compat: legacy ``dp`` maps to ``dp_shard`` when both
120 # dp_replicate/dp_shard fields are at defaults.
121 if dp_shard is None:
122 dp_shard = legacy_dp if legacy_dp is not None else 1
124 return cls(
125 dp_replicate=dp_replicate,
126 dp_shard=dp_shard,
127 cp=getattr(parallel_cfg, 'cp', 1),
128 tp=getattr(parallel_cfg, 'tp', 1),
129 pp=getattr(parallel_cfg, 'pp', 1),
130 ep=getattr(parallel_cfg, 'ep', 1),
131 etp=getattr(parallel_cfg, 'etp', getattr(parallel_cfg, 'tp', 1)),
132 moe_token_dispatcher_type=getattr(parallel_cfg, 'moe_token_dispatcher_type', 'all_to_all'),
133 npu_nums_per_device=getattr(parallel_cfg, 'npu_nums_per_device', 8),
134 world_size=world_size,
135 _allow_pp=True,
136 )
138 def __post_init__(self) -> None:
139 self._infer_and_validate()
141 def _infer_and_validate(self) -> None:
142 """Auto-fill ``dp_shard=-1`` then validate ``product == world_size``."""
143 self._validate_positive_degrees()
144 self._validate_moe_dispatcher()
145 self._validate_and_infer_dp_shard()
146 self._validate_parallel_product()
147 self._validate_expert_tensor_parallel()
148 self._validate_pipeline_parallel()
149 self._validate_ulysses_degree()
151 def _validate_positive_degrees(self) -> None:
152 """Require every non-auto parallel degree to be positive."""
153 for name, value in (
154 ("dp_replicate", self.dp_replicate),
155 ("cp", self.cp),
156 ("tp", self.tp),
157 ("pp", self.pp),
158 ("ep", self.ep),
159 ("etp", self.etp),
160 ("npu_nums_per_device", self.npu_nums_per_device),
161 ):
162 if value < 1:
163 raise ValueError(f"Parallel degree {name}={value} must be >= 1")
165 def _validate_moe_dispatcher(self) -> None:
166 """Validate the MoE dispatcher name and deredundency topology."""
167 if self.moe_token_dispatcher_type not in ("all_to_all", "deredundency"):
168 raise ValueError(
169 "moe_token_dispatcher_type must be 'all_to_all' or 'deredundency', "
170 f"got {self.moe_token_dispatcher_type!r}"
171 )
173 if (
174 self.moe_token_dispatcher_type == "deredundency"
175 and self.ep % self.npu_nums_per_device != 0
176 ):
177 raise ValueError(
178 f"ep={self.ep} must be divisible by "
179 f"npu_nums_per_device={self.npu_nums_per_device} when "
180 "moe_token_dispatcher_type='deredundency'."
181 )
183 def _validate_and_infer_dp_shard(self) -> None:
184 """Validate ``dp_shard`` and resolve its auto value when requested."""
185 if self.dp_shard < -1 or self.dp_shard == 0:
186 raise ValueError(
187 f"dp_shard={self.dp_shard} must be -1 (auto) or a positive int"
188 )
190 # Auto-infer dp_shard when -1. ep is an independent peer mesh dim
191 # (see build_mesh) so it does NOT reduce the dp pool.
192 if self.dp_shard == -1:
193 non_dp = self.dp_replicate * self.cp * self.tp * self.pp * self.ep
194 if self.world_size % non_dp != 0:
195 raise ValueError(
196 f"Cannot auto-infer dp_shard: world_size={self.world_size} "
197 f"is not divisible by dp_replicate*cp*tp*pp*ep={non_dp}"
198 )
199 self.dp_shard = max(self.world_size // non_dp, 1)
200 logger.info_rank0(
201 "Auto-inferred dp_shard=%d (world_size=%d / dp_replicate=%d * "
202 "cp=%d * tp=%d * pp=%d * ep=%d)",
203 self.dp_shard, self.world_size,
204 self.dp_replicate, self.cp, self.tp, self.pp, self.ep,
205 )
207 def _validate_parallel_product(self) -> None:
208 """Require the resolved parallel degrees to cover the world exactly."""
209 product = (
210 self.dp_replicate * self.dp_shard
211 * self.cp * self.tp * self.pp * self.ep
212 )
213 if product != self.world_size:
214 raise ValueError(
215 f"Invalid parallel dims: dp_replicate({self.dp_replicate}) * "
216 f"dp_shard({self.dp_shard}) * cp({self.cp}) * tp({self.tp}) * "
217 f"pp({self.pp}) * ep({self.ep}) = {product} != "
218 f"world_size({self.world_size}). Set dp_shard=-1 to auto-infer."
219 )
221 def _validate_expert_tensor_parallel(self) -> None:
222 """Validate expert tensor parallelism when expert parallelism is active."""
223 # ep is an independent peer mesh dim alongside dp/cp/tp/pp (see build_mesh).
224 # It does NOT need to divide the dp_shard*cp pool; it occupies its own
225 # mesh axis. We only enforce the expert-TP compatibility rule below.
226 if self.ep > 1:
227 if self.etp not in (self.tp, 1):
228 raise ValueError(
229 f"etp={self.etp} must equal tp={self.tp} or 1 "
230 f"(expert tensor-parallel must align with TP or be inactive)"
231 )
233 def _validate_pipeline_parallel(self) -> None:
234 """Require model-owned pipeline construction for a non-trivial PP axis."""
235 # PP is opt-in per model: trainer config construction sets the private
236 # ``_allow_pp`` marker and ``ModelSpec.pipelining_fn`` owns the stage split.
237 # Direct construction keeps the legacy fail-fast guard for callers that
238 # do not provide a model-level PP path.
239 if self.pp > 1 and not self._allow_pp:
240 raise NotImplementedError(
241 f"Pipeline parallel (pp={self.pp}>1) requires a model-specific "
242 "pipelining_fn. Use ParallelDims.from_config() from the trainer "
243 "path, or set pp=1 for direct construction."
244 )
246 def _validate_ulysses_degree(self) -> None:
247 """Validate the optional Ulysses degree against context parallelism."""
248 # Ulysses must divide cp.
249 if self.ulysses_degree is not None and self.cp > 1:
250 if self.ulysses_degree > self.cp:
251 raise ValueError(
252 f"ulysses_degree={self.ulysses_degree} must be <= "
253 f"cp={self.cp}"
254 )
255 if self.cp % self.ulysses_degree != 0:
256 raise ValueError(
257 f"cp={self.cp} must be divisible by "
258 f"ulysses_degree={self.ulysses_degree}"
259 )
261 # ------------------------------------------------------------------
262 # Validate against an actual model
263 # ------------------------------------------------------------------
264 def validate_against_model(
265 self,
266 model,
267 seq_len: Optional[int] = None,
268 ) -> None:
269 """Cross-check parallel degrees against a built model's hyperparams.
271 Reads common decoder config fields from ``model.config``. Skips silently
272 if a field is absent. Model-specific validation (e.g. "TP unsupported
273 for linear-attn layers") is inlined at the top of each
274 ``parallelize_<model>()`` function — convention.
276 Args:
277 model: The built ``nn.Module`` (must expose ``.config`` to trigger
278 most checks).
279 seq_len: Optional maximum sequence length used for cp/tp
280 divisibility checks.
282 Raises:
283 ValueError: With a single readable message when a constraint is
284 violated. Stops here so the user sees the real cause instead
285 of a stack trace from inside ``parallelize_module``.
286 """
287 cfg = getattr(model, 'config', None)
288 if cfg is None:
289 return
290 cfg = getattr(cfg, 'text_config', cfg)
292 heads = getattr(cfg, 'num_attention_heads', None)
293 if heads is not None and self.tp > 1 and heads % self.tp != 0:
294 raise ValueError(
295 f"num_attention_heads={heads} not divisible by tp={self.tp}. "
296 f"Pick tp from the divisors of {heads}."
297 )
299 kv_heads = getattr(cfg, 'num_key_value_heads', None)
300 if kv_heads is not None and self.tp > 1 and kv_heads % self.tp != 0:
301 raise ValueError(
302 f"num_key_value_heads={kv_heads} not divisible by tp={self.tp} "
303 f"(GQA constraint). Pick tp from the divisors of {kv_heads}."
304 )
306 num_experts = getattr(cfg, 'num_experts', None)
307 if num_experts is not None and self.ep > 1 and num_experts % self.ep != 0:
308 raise ValueError(
309 f"num_experts={num_experts} not divisible by ep={self.ep}. "
310 f"Pick ep from the divisors of {num_experts}."
311 )
313 if seq_len is not None and self.cp * self.tp > 1:
314 divisor = self.cp * self.tp
315 if seq_len % divisor != 0:
316 raise ValueError(
317 f"max_seq_len={seq_len} not divisible by cp*tp={divisor}. "
318 f"Increase seq_len to a multiple of {divisor} or reduce "
319 f"cp/tp."
320 )
322 # ------------------------------------------------------------------
323 # Mesh building
324 # ------------------------------------------------------------------
325 def build_mesh(self, device_type: str, force_dp_shard: bool = False):
326 """Build the DeviceMesh with canonical dim order and named flatten aliases.
328 Order of base dims: ``pp → dp_replicate → dp_shard → ep → cp → tp``.
329 PP is placed **outermost** (lowest stride is on TP), so pipeline stages
330 consume contiguous ranges of ranks while TP shards stay co-located on
331 the same host.
332 For deredundency EP, ``ep`` is materialized as ``oep → iep`` and
333 flattened back under the ``"ep"`` alias. Only base dims with degree
334 > 1 are materialized, except deredundency EP keeps both ``oep`` and
335 ``iep`` axes when ``ep > 1`` so the token dispatcher can form its two
336 communication groups. If all dims are 1, a 1D ``dp_shard`` mesh of the
337 world is created so the FSDP code path runs unchanged on single-card.
339 A size-1 ``dp_shard`` dim is also materialized whenever
340 ``dp_replicate > 1``. This preserves the explicit two-dimensional
341 HSDP topology for pure replicated data parallelism: sharding over the
342 size-1 inner axis is a no-op and the outer replicate all-reduce gives
343 DDP-equivalent gradient synchronization.
345 ``force_dp_shard=True`` otherwise materializes the ``dp_shard`` dim
346 even at size 1 when neither ``dp_shard`` nor ``cp`` is > 1. Mixed
347 precision is realised entirely through FSDP2's
348 ``MixedPrecisionPolicy``, so a pure-TP / pure-PP mesh (no shardable
349 axis) would otherwise skip the FSDP wrap and silently run the model in
350 full precision; the size-1 axis gives ``fully_shard`` a no-op
351 communication group to hang the dtype policy on.
353 After construction, the following flatten aliases are registered on
354 the root mesh so callers can reach them with ``mesh["fsdp"]`` /
355 ``mesh["dp"]`` regardless of the underlying parallel composition:
357 ``"fsdp"`` – mesh used for ``fully_shard`` / reduce-scatter.
358 Always equals the ``dp_shard`` axis.
359 ``"dp"`` – combined data-parallel mesh used for loss / token
360 all-reduce. ``dp_replicate × dp_shard`` when both
361 are >1 (HSDP); otherwise the single non-trivial
362 DP axis (or ``dp_shard`` for the 1-card case).
364 Args:
365 device_type: Backend device string (``"npu"`` / ``"cuda"``).
367 Returns:
368 ``DeviceMesh`` instance.
369 """
370 # PP is excluded: the schedule drives FSDP unshard/reshard itself and
371 # its world-1 FSDP path is not wired — the trainer rejects pure-PP
372 # low-precision runs instead (use PP x FSDP).
373 force_dp_shard = (
374 force_dp_shard and self.dp_shard == 1 and self.cp == 1
375 and self.pp == 1
376 )
377 dims = []
378 names = []
379 for name, size in self._mesh_dim_specs():
380 force_materialize = (
381 (
382 name == "dp_shard"
383 and (force_dp_shard or self.dp_replicate > 1)
384 )
385 or (
386 self.moe_token_dispatcher_type == "deredundency"
387 and self.ep > 1
388 and name in ("oep", "iep")
389 )
390 )
391 if size > 1 or force_materialize:
392 dims.append(size)
393 names.append(name)
395 if not dims:
396 dims = [self.world_size]
397 names = ["dp_shard"]
399 self._device_mesh = init_device_mesh(
400 device_type=device_type,
401 mesh_shape=tuple(dims),
402 mesh_dim_names=tuple(names),
403 )
404 self._register_flatten_aliases(names)
405 logger.info_rank0(
406 "DeviceMesh built: shape=%s, names=%s",
407 tuple(dims), tuple(names),
408 )
409 return self._device_mesh
411 def _mesh_dim_specs(self) -> tuple[tuple[str, int], ...]:
412 """Return mesh dimension specs in canonical order."""
413 ep_specs = (("ep", self.ep),)
414 if self.moe_token_dispatcher_type == "deredundency":
415 ep_specs = (
416 ("oep", self.ep // self.npu_nums_per_device),
417 ("iep", self.npu_nums_per_device),
418 )
419 return (
420 ("pp", self.pp),
421 ("dp_replicate", self.dp_replicate),
422 ("dp_shard", self.dp_shard),
423 *ep_specs,
424 ("cp", self.cp),
425 ("tp", self.tp),
426 )
428 def _register_flatten_aliases(self, base_names) -> None:
429 """Register named flatten aliases on the root mesh.
431 These aliases give the rest of the trainer a stable, intent-named
432 handle on combined parallel axes so callers never need to fall back
433 to the whole mesh:
435 ``"fsdp"`` – the axis FSDP shards along (= ``dp_shard``).
436 ``"dp"`` – combined data-parallel mesh (replicate × shard).
437 Used for grad/optimizer-state replication accounting.
438 ``"loss"`` – the mesh over which loss / token counts are
439 all-reduced. Equals ``dp × cp`` when CP is enabled
440 (CP-sharded ranks see different sub-sequences and
441 must contribute their token counts to the global
442 denominator); otherwise equals ``dp``.
444 Reserved names (intentionally not registered yet):
445 ``"efsdp"`` – FSDP mesh for expert layers when EP > 1. Will
446 fold ``dp_shard / ep`` once real EP lands.
447 ``"etp"`` – expert TP mesh (= ``ep × tp`` composition)
448 alongside dense ``tp``. Same gate.
449 ``"batch"`` – per-DP batch dispatch mesh; today identical to
450 ``"dp"``, will diverge if we ever support
451 microbatch-sharded scheduling.
453 Idempotent: every flatten call is gated on whether the alias is
454 already on the root mesh, so repeated ``build_mesh`` calls are
455 safe.
457 Args:
458 base_names: Sequence of base mesh-dim names that were materialized
459 (degree > 1, plus the degenerate ``dp_shard`` of size 1 when
460 no other dim was present).
461 """
462 # pylint: disable=protected-access
463 mesh = self._device_mesh
464 existing = set(mesh.mesh_dim_names or ())
465 flatten_keys = set(mesh._get_root_mesh().get_flatten_mapping().keys())
467 def _flatten_unique(source_dims, alias):
468 if alias in existing or alias in flatten_keys:
469 return
470 mesh[source_dims].flatten(alias)
471 flatten_keys.add(alias)
473 self._register_data_flatten_aliases(base_names, _flatten_unique)
474 self._register_loss_flatten_alias(
475 base_names, mesh, flatten_keys, _flatten_unique,
476 )
478 @staticmethod
479 def _register_data_flatten_aliases(base_names, flatten_unique) -> None:
480 """Register EP, FSDP, and DP aliases in canonical order."""
481 has_replicate = "dp_replicate" in base_names
482 has_shard = "dp_shard" in base_names
483 has_oep = "oep" in base_names
484 has_iep = "iep" in base_names
486 # Deredundency materializes EP as ``oep`` × ``iep`` but callers keep
487 # using the stable full-EP alias ``mesh["ep"]``.
488 if has_oep and has_iep:
489 flatten_unique(("oep", "iep"), "ep")
491 # ``fsdp`` — the axis ``fully_shard`` actually shards along.
492 if has_shard:
493 flatten_unique("dp_shard", "fsdp")
495 # ``dp`` — combined replicate×shard data-parallel mesh.
496 if has_replicate and has_shard:
497 flatten_unique(("dp_replicate", "dp_shard"), "dp")
498 elif has_replicate:
499 flatten_unique("dp_replicate", "dp")
500 elif has_shard:
501 flatten_unique("dp_shard", "dp")
503 @staticmethod
504 def _register_loss_flatten_alias(
505 base_names, mesh, flatten_keys, flatten_unique,
506 ) -> None:
507 """Register the loss-reduction alias after the DP aliases."""
508 has_replicate = "dp_replicate" in base_names
509 has_shard = "dp_shard" in base_names
510 has_cp = "cp" in base_names
512 # ``loss`` — dp folded with cp when context parallelism is active so
513 # loss / token counts include CP-sharded contributions.
514 if has_cp:
515 if has_replicate and has_shard:
516 flatten_unique(("dp_replicate", "dp_shard", "cp"), "loss")
517 elif has_replicate:
518 flatten_unique(("dp_replicate", "cp"), "loss")
519 elif has_shard:
520 flatten_unique(("dp_shard", "cp"), "loss")
521 else:
522 flatten_unique("cp", "loss")
523 else:
524 # No CP — ``loss`` and ``dp`` are the same group. Re-flatten
525 # the existing 1D dp mesh under the ``loss`` alias so both
526 # names resolve via ``__getitem__``.
527 if "loss" not in flatten_keys and "dp" in flatten_keys:
528 mesh["dp"].flatten("loss")
529 flatten_keys.add("loss")
531 # ------------------------------------------------------------------
532 # Convenience properties
533 # ------------------------------------------------------------------
534 @property
535 def dp_size(self) -> int:
536 """Combined data-parallel size = dp_replicate * dp_shard."""
537 return self.dp_replicate * self.dp_shard
539 @property
540 def non_dp_size(self) -> int:
541 """Product of model-side parallel dims (tp*cp*pp*ep)."""
542 return self.tp * self.cp * self.pp * self.ep
544 @property
545 def seq_divisor(self) -> int:
546 """Sequence-length divisor the data pipeline must pad to.
548 SequenceParallel TP and context parallel both slice the sequence
549 dim, so variable-length batches pad up to a multiple of ``tp * cp``
550 (the trailing pad rides label ``-100`` and is masked out of the CE).
551 """
552 return self.tp * self.cp
554 @property
555 def tp_enabled(self) -> bool:
556 """Return True if tensor parallelism is enabled (tp > 1)."""
557 return self.tp > 1
559 @property
560 def cp_enabled(self) -> bool:
561 """Return True if context parallelism is enabled (cp > 1)."""
562 return self.cp > 1
564 @property
565 def ep_enabled(self) -> bool:
566 """Return True if expert parallelism is enabled (ep > 1)."""
567 return self.ep > 1
569 @property
570 def pp_enabled(self) -> bool:
571 """Return True if pipeline parallelism is enabled (pp > 1)."""
572 return self.pp > 1
574 @property
575 def fsdp_enabled(self) -> bool:
576 """FSDP is on whenever there's a shard dim or HSDP outer dim."""
577 return self.dp_shard > 1 or self.dp_replicate > 1
579 def summary(self) -> str:
580 """Compact one-line summary for logging."""
581 return (
582 f"dp_replicate={self.dp_replicate} dp_shard={self.dp_shard} "
583 f"cp={self.cp} tp={self.tp} pp={self.pp} ep={self.ep} "
584 f"etp={self.etp} moe_token_dispatcher_type={self.moe_token_dispatcher_type} "
585 f"npu_nums_per_device={self.npu_nums_per_device} | dp={self.dp_size} world={self.world_size}"
586 )