Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / trainer / parallel_dims.py: 97%
155 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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# Adapted from https://github.com/pytorch/torchtitan/blob/main/torchtitan/distributed/parallel_dims.py
18"""ParallelDims — fail-fast parallel configuration validator + mesh builder.
20Centralises parallel-degree validation in a single dataclass so
21misconfigurations are caught before model construction.
23What this provides:
251. **Inference** — auto-fill ``dp`` (or ``dp_shard=-1``) from the product
26 constraint ``dp_replicate * dp_shard * cp * tp * pp == world_size``.
282. **Validation against world_size** — raises ``ValueError`` with a clear
29 message when the product mismatches.
313. **Validation against the model** — checks divisibility constraints that
32 would otherwise crash deep inside ``parallelize_module``:
34 - ``num_attention_heads % tp == 0`` (TP shards heads)
35 - ``num_key_value_heads % tp == 0`` (GQA constraint)
36 - ``num_experts % ep == 0`` (when MoE)
37 - ``ulysses_degree <= cp`` and ``cp % ulysses_degree == 0``
38 - ``seq_len % (cp * tp) == 0`` (sequence parallel + CP)
39 - ``etp == tp or etp == 1`` (expert TP rule)
414. **Mesh building** — returns a ``DeviceMesh`` with the canonical dim order
42 ``dp_replicate → dp_shard → ep → cp → tp → pp``. Backwards compatible with
43 the legacy single ``dp`` field (auto-collapses to ``dp_shard``).
45User experience:
47- Default config (no parallel section) — works on 1 GPU, runs as DDP-1.
48- Set only ``tp=4`` on world_size=8 → ``dp`` auto-inferred to 2.
49- Set ``dp_shard=-1`` → fills remaining cards into FSDP shard dim.
50- Misconfig (heads=12, tp=8) → fails at ``_setup`` with a single readable
51 error before any model parallelization is attempted.
52"""
53from __future__ import annotations
55__all__ = ["ParallelDims"]
57import logging
58from dataclasses import dataclass, field
59from typing import Optional
61from hyper_parallel import init_device_mesh
63logger = logging.getLogger(__name__)
66@dataclass
67class ParallelDims:
68 """Validated parallel degrees + lazy mesh builder.
70 Attributes:
71 dp_replicate: DDP replication degree (HSDP outer dim).
72 dp_shard: FSDP shard degree. ``-1`` means "fill the rest from
73 ``world_size / (dp_replicate * cp * tp * pp)``".
74 cp: Context parallel degree.
75 tp: Tensor parallel degree (dense path).
76 pp: Pipeline parallel degree.
77 ep: Expert parallel degree (MoE only).
78 etp: Expert tensor parallel degree. Must equal ``tp`` or ``1``.
79 moe_token_dispatcher_type: Expert token exchange strategy.
80 npu_nums_per_device: Inner expert-parallel degree for deredundency dispatch.
81 ulysses_degree: Ulysses sub-degree inside ``cp``. ``None`` means
82 "pure Ulysses (degree == cp)".
83 world_size: Total number of ranks.
84 """
86 dp_replicate: int = 1
87 dp_shard: int = 1
88 cp: int = 1
89 tp: int = 1
90 pp: int = 1
91 ep: int = 1
92 etp: int = 1
93 moe_token_dispatcher_type: str = "all_to_all"
94 npu_nums_per_device: int = 8
95 ulysses_degree: Optional[int] = None
96 world_size: int = 1
97 # Cached after build_mesh.
98 _device_mesh: object = field(default=None, repr=False)
100 # ------------------------------------------------------------------
101 # Construction & inference
102 # ------------------------------------------------------------------
103 @classmethod
104 def from_config(cls, parallel_cfg, world_size: int) -> "ParallelDims":
105 """Build from a ``ParallelConfig`` (or any object with the same fields).
107 Accepts the legacy single-``dp`` field. If ``dp`` is set and
108 ``dp_replicate``/``dp_shard`` are at default, ``dp`` is mapped to
109 ``dp_shard`` (FSDP behavior).
110 """
111 dp_replicate = getattr(parallel_cfg, 'dp_replicate', 1)
112 dp_shard = getattr(parallel_cfg, 'dp_shard', None)
113 legacy_dp = getattr(parallel_cfg, 'dp', None)
115 # Backward-compat: legacy ``dp`` maps to ``dp_shard`` when both
116 # dp_replicate/dp_shard fields are at defaults.
117 if dp_shard is None:
118 dp_shard = legacy_dp if legacy_dp is not None else 1
120 return cls(
121 dp_replicate=dp_replicate,
122 dp_shard=dp_shard,
123 cp=getattr(parallel_cfg, 'cp', 1),
124 tp=getattr(parallel_cfg, 'tp', 1),
125 pp=getattr(parallel_cfg, 'pp', 1),
126 ep=getattr(parallel_cfg, 'ep', 1),
127 etp=getattr(parallel_cfg, 'etp', getattr(parallel_cfg, 'tp', 1)),
128 moe_token_dispatcher_type=getattr(parallel_cfg, 'moe_token_dispatcher_type', 'all_to_all'),
129 npu_nums_per_device=getattr(parallel_cfg, 'npu_nums_per_device', 8),
130 ulysses_degree=getattr(parallel_cfg, 'ulysses_degree', None),
131 world_size=world_size,
132 )
134 def __post_init__(self) -> None:
135 self._infer_and_validate()
137 def _infer_and_validate(self) -> None:
138 """Auto-fill ``dp_shard=-1`` then validate ``product == world_size``."""
139 for name, value in (
140 ("dp_replicate", self.dp_replicate),
141 ("cp", self.cp),
142 ("tp", self.tp),
143 ("pp", self.pp),
144 ("ep", self.ep),
145 ("etp", self.etp),
146 ("npu_nums_per_device", self.npu_nums_per_device),
147 ):
148 if value < 1:
149 raise ValueError(f"Parallel degree {name}={value} must be >= 1")
151 if self.moe_token_dispatcher_type not in ("all_to_all", "deredundency"):
152 raise ValueError(
153 "moe_token_dispatcher_type must be 'all_to_all' or 'deredundency', "
154 f"got {self.moe_token_dispatcher_type!r}"
155 )
157 if (
158 self.moe_token_dispatcher_type == "deredundency"
159 and self.ep % self.npu_nums_per_device != 0
160 ):
161 raise ValueError(
162 f"ep={self.ep} must be divisible by "
163 f"npu_nums_per_device={self.npu_nums_per_device} when "
164 "moe_token_dispatcher_type='deredundency'."
165 )
167 if self.dp_shard < -1 or self.dp_shard == 0:
168 raise ValueError(
169 f"dp_shard={self.dp_shard} must be -1 (auto) or a positive int"
170 )
172 # Auto-infer dp_shard when -1. ep is an independent peer mesh dim
173 # (see build_mesh) so it does NOT reduce the dp pool.
174 if self.dp_shard == -1:
175 non_dp = self.dp_replicate * self.cp * self.tp * self.pp * self.ep
176 if self.world_size % non_dp != 0:
177 raise ValueError(
178 f"Cannot auto-infer dp_shard: world_size={self.world_size} "
179 f"is not divisible by dp_replicate*cp*tp*pp*ep={non_dp}"
180 )
181 self.dp_shard = max(self.world_size // non_dp, 1)
182 logger.info_rank0(
183 "Auto-inferred dp_shard=%d (world_size=%d / dp_replicate=%d * "
184 "cp=%d * tp=%d * pp=%d * ep=%d)",
185 self.dp_shard, self.world_size,
186 self.dp_replicate, self.cp, self.tp, self.pp, self.ep,
187 )
189 product = (
190 self.dp_replicate * self.dp_shard
191 * self.cp * self.tp * self.pp * self.ep
192 )
193 if product != self.world_size:
194 raise ValueError(
195 f"Invalid parallel dims: dp_replicate({self.dp_replicate}) * "
196 f"dp_shard({self.dp_shard}) * cp({self.cp}) * tp({self.tp}) * "
197 f"pp({self.pp}) * ep({self.ep}) = {product} != "
198 f"world_size({self.world_size}). Set dp_shard=-1 to auto-infer."
199 )
201 # ep is an independent peer mesh dim alongside dp/cp/tp/pp (see build_mesh).
202 # It does NOT need to divide the dp_shard*cp pool; it occupies its own
203 # mesh axis. We only enforce the expert-TP compatibility rule below.
204 if self.ep > 1:
205 if self.etp not in (self.tp, 1):
206 raise ValueError(
207 f"etp={self.etp} must equal tp={self.tp} or 1 "
208 f"(expert tensor-parallel must align with TP or be inactive)"
209 )
211 # No model has implemented PP wiring yet — fail fast.
212 if self.pp > 1:
213 raise NotImplementedError(
214 f"Pipeline parallel (pp={self.pp}>1) is not yet implemented "
215 "for any model. Set pp=1 or add a per-model PP path in "
216 "models/<name>/parallelize.py before enabling."
217 )
219 # Ulysses must divide cp.
220 if self.ulysses_degree is not None and self.cp > 1:
221 if self.ulysses_degree > self.cp:
222 raise ValueError(
223 f"ulysses_degree={self.ulysses_degree} must be <= "
224 f"cp={self.cp}"
225 )
226 if self.cp % self.ulysses_degree != 0:
227 raise ValueError(
228 f"cp={self.cp} must be divisible by "
229 f"ulysses_degree={self.ulysses_degree}"
230 )
232 # ------------------------------------------------------------------
233 # Validate against an actual model
234 # ------------------------------------------------------------------
235 def validate_against_model(
236 self,
237 model,
238 seq_len: Optional[int] = None,
239 ) -> None:
240 """Cross-check parallel degrees against a built model's hyperparams.
242 Reads ``model.config`` for standard transformer fields. Skips silently
243 if a field is absent. Model-specific validation (e.g. "TP unsupported
244 for linear-attn layers") is inlined at the top of each
245 ``parallelize_<model>()`` function — convention.
247 Args:
248 model: The built ``nn.Module`` (must expose ``.config`` to trigger
249 most checks).
250 seq_len: Optional maximum sequence length used for cp/tp
251 divisibility checks.
253 Raises:
254 ValueError: With a single readable message when a constraint is
255 violated. Stops here so the user sees the real cause instead
256 of a stack trace from inside ``parallelize_module``.
257 """
258 cfg = getattr(model, 'config', None)
259 if cfg is None:
260 return
262 heads = getattr(cfg, 'num_attention_heads', None)
263 if heads is not None and self.tp > 1 and heads % self.tp != 0:
264 raise ValueError(
265 f"num_attention_heads={heads} not divisible by tp={self.tp}. "
266 f"Pick tp from the divisors of {heads}."
267 )
269 kv_heads = getattr(cfg, 'num_key_value_heads', None)
270 if kv_heads is not None and self.tp > 1 and kv_heads % self.tp != 0:
271 raise ValueError(
272 f"num_key_value_heads={kv_heads} not divisible by tp={self.tp} "
273 f"(GQA constraint). Pick tp from the divisors of {kv_heads}."
274 )
276 num_experts = getattr(cfg, 'num_experts', None)
277 if num_experts is not None and self.ep > 1 and num_experts % self.ep != 0:
278 raise ValueError(
279 f"num_experts={num_experts} not divisible by ep={self.ep}. "
280 f"Pick ep from the divisors of {num_experts}."
281 )
283 if seq_len is not None and self.cp * self.tp > 1:
284 divisor = self.cp * self.tp
285 if seq_len % divisor != 0:
286 raise ValueError(
287 f"max_seq_len={seq_len} not divisible by cp*tp={divisor}. "
288 f"Increase seq_len to a multiple of {divisor} or reduce "
289 f"cp/tp."
290 )
292 # ------------------------------------------------------------------
293 # Mesh building
294 # ------------------------------------------------------------------
295 def build_mesh(self, device_type: str):
296 """Build the DeviceMesh with canonical dim order and named flatten aliases.
298 Order of base dims: ``dp_replicate → dp_shard → ep → cp → tp → pp``.
299 For deredundency EP, ``ep`` is materialized as ``oep → iep`` and
300 flattened back under the ``"ep"`` alias.
301 Only base dims with degree > 1 are materialized, except deredundency
302 EP keeps both ``oep`` and ``iep`` axes when ``ep > 1`` so the token
303 dispatcher can form its two communication groups. If all dims are 1,
304 a 1D ``dp_shard`` mesh of the world is created so the FSDP code path
305 runs unchanged on single-card.
307 After construction, the following flatten aliases are registered on
308 the root mesh so callers can reach them with ``mesh["fsdp"]`` /
309 ``mesh["dp"]`` regardless of the underlying parallel composition:
311 ``"fsdp"`` – mesh used for ``fully_shard`` / reduce-scatter.
312 Always equals the ``dp_shard`` axis.
313 ``"dp"`` – combined data-parallel mesh used for loss / token
314 all-reduce. ``dp_replicate × dp_shard`` when both
315 are >1 (HSDP); otherwise the single non-trivial
316 DP axis (or ``dp_shard`` for the 1-card case).
318 Args:
319 device_type: Backend device string (``"npu"`` / ``"cuda"``).
321 Returns:
322 ``DeviceMesh`` instance.
323 """
324 dims = []
325 names = []
326 for name, size in self._mesh_dim_specs():
327 force_materialize = (
328 self.moe_token_dispatcher_type == "deredundency"
329 and self.ep > 1
330 and name in ("oep", "iep")
331 )
332 if size > 1 or force_materialize:
333 dims.append(size)
334 names.append(name)
336 if not dims:
337 dims = [self.world_size]
338 names = ["dp_shard"]
340 self._device_mesh = init_device_mesh(
341 device_type=device_type,
342 mesh_shape=tuple(dims),
343 mesh_dim_names=tuple(names),
344 )
345 self._register_flatten_aliases(names)
346 logger.info_rank0(
347 "DeviceMesh built: shape=%s, names=%s",
348 tuple(dims), tuple(names),
349 )
350 return self._device_mesh
352 def _mesh_dim_specs(self) -> tuple[tuple[str, int], ...]:
353 """Return mesh dimension specs in canonical order."""
354 ep_specs = (("ep", self.ep),)
355 if self.moe_token_dispatcher_type == "deredundency":
356 ep_specs = (
357 ("oep", self.ep // self.npu_nums_per_device),
358 ("iep", self.npu_nums_per_device),
359 )
360 return (
361 ("dp_replicate", self.dp_replicate),
362 ("dp_shard", self.dp_shard),
363 *ep_specs,
364 ("cp", self.cp),
365 ("tp", self.tp),
366 ("pp", self.pp),
367 )
369 def _register_flatten_aliases(self, base_names) -> None:
370 """Register named flatten aliases on the root mesh.
372 These aliases give the rest of the trainer a stable, intent-named
373 handle on combined parallel axes so callers never need to fall back
374 to the whole mesh:
376 ``"fsdp"`` – the axis FSDP shards along (= ``dp_shard``).
377 ``"dp"`` – combined data-parallel mesh (replicate × shard).
378 Used for grad/optimizer-state replication accounting.
379 ``"loss"`` – the mesh over which loss / token counts are
380 all-reduced. Equals ``dp × cp`` when CP is enabled
381 (CP-sharded ranks see different sub-sequences and
382 must contribute their token counts to the global
383 denominator); otherwise equals ``dp``.
385 Reserved names (intentionally not registered yet):
386 ``"efsdp"`` – FSDP mesh for expert layers when EP > 1. Will
387 fold ``dp_shard / ep`` once real EP lands.
388 ``"etp"`` – expert TP mesh (= ``ep × tp`` composition)
389 alongside dense ``tp``. Same gate.
390 ``"batch"`` – per-DP batch dispatch mesh; today identical to
391 ``"dp"``, will diverge if we ever support
392 microbatch-sharded scheduling.
394 Idempotent: every flatten call is gated on whether the alias is
395 already on the root mesh, so repeated ``build_mesh`` calls are
396 safe.
398 Args:
399 base_names: Sequence of base mesh-dim names that were materialized
400 (degree > 1, plus the degenerate ``dp_shard`` of size 1 when
401 no other dim was present).
402 """
403 # pylint: disable=protected-access
404 mesh = self._device_mesh
405 existing = set(mesh.mesh_dim_names or ())
406 flatten_keys = set(mesh._get_root_mesh().get_flatten_mapping().keys())
408 def _flatten_unique(source_dims, alias):
409 if alias in existing or alias in flatten_keys:
410 return
411 mesh[source_dims].flatten(alias)
412 flatten_keys.add(alias)
414 has_replicate = "dp_replicate" in base_names
415 has_shard = "dp_shard" in base_names
416 has_cp = "cp" in base_names
417 has_oep = "oep" in base_names
418 has_iep = "iep" in base_names
420 # Deredundency materializes EP as ``oep`` × ``iep`` but callers keep
421 # using the stable full-EP alias ``mesh["ep"]``.
422 if has_oep and has_iep:
423 _flatten_unique(("oep", "iep"), "ep")
425 # ``fsdp`` — the axis ``fully_shard`` actually shards along.
426 if has_shard:
427 _flatten_unique("dp_shard", "fsdp")
429 # ``dp`` — combined replicate×shard data-parallel mesh.
430 if has_replicate and has_shard:
431 _flatten_unique(("dp_replicate", "dp_shard"), "dp")
432 elif has_replicate:
433 _flatten_unique("dp_replicate", "dp")
434 elif has_shard:
435 _flatten_unique("dp_shard", "dp")
437 # ``loss`` — dp folded with cp when context parallelism is active so
438 # loss / token counts include CP-sharded contributions.
439 if has_cp:
440 if has_replicate and has_shard:
441 _flatten_unique(("dp_replicate", "dp_shard", "cp"), "loss")
442 elif has_replicate:
443 _flatten_unique(("dp_replicate", "cp"), "loss")
444 elif has_shard:
445 _flatten_unique(("dp_shard", "cp"), "loss")
446 else:
447 _flatten_unique("cp", "loss")
448 else:
449 # No CP — ``loss`` and ``dp`` are the same group. Re-flatten
450 # the existing 1D dp mesh under the ``loss`` alias so both
451 # names resolve via ``__getitem__``.
452 if "loss" not in flatten_keys and "dp" in flatten_keys:
453 mesh["dp"].flatten("loss")
454 flatten_keys.add("loss")
456 # ------------------------------------------------------------------
457 # Convenience properties
458 # ------------------------------------------------------------------
459 @property
460 def dp_size(self) -> int:
461 """Combined data-parallel size = dp_replicate * dp_shard."""
462 return self.dp_replicate * self.dp_shard
464 @property
465 def non_dp_size(self) -> int:
466 """Product of model-side parallel dims (tp*cp*pp*ep)."""
467 return self.tp * self.cp * self.pp * self.ep
469 @property
470 def tp_enabled(self) -> bool:
471 """Return True if tensor parallelism is enabled (tp > 1)."""
472 return self.tp > 1
474 @property
475 def cp_enabled(self) -> bool:
476 """Return True if context parallelism is enabled (cp > 1)."""
477 return self.cp > 1
479 @property
480 def ep_enabled(self) -> bool:
481 """Return True if expert parallelism is enabled (ep > 1)."""
482 return self.ep > 1
484 @property
485 def pp_enabled(self) -> bool:
486 """Return True if pipeline parallelism is enabled (pp > 1)."""
487 return self.pp > 1
489 @property
490 def fsdp_enabled(self) -> bool:
491 """FSDP is on whenever there's a shard dim or HSDP outer dim."""
492 return self.dp_shard > 1 or self.dp_replicate > 1
494 def summary(self) -> str:
495 """Compact one-line summary for logging."""
496 return (
497 f"dp_replicate={self.dp_replicate} dp_shard={self.dp_shard} "
498 f"cp={self.cp} tp={self.tp} pp={self.pp} ep={self.ep} "
499 f"etp={self.etp} moe_token_dispatcher_type={self.moe_token_dispatcher_type} "
500 f"npu_nums_per_device={self.npu_nums_per_device} | dp={self.dp_size} world={self.world_size}"
501 )