Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / mesh.py: 77%
193 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-11-15 07:18 +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# ============================================================================
15"""Mesh domains and the distributed-setup container for AutoModels.
17Split out of the former ``components/distributed/infrastructure.py`` in
18stage 7 (05 §15.11 step 3): the mesh topology objects are AutoModels-side
19model-construction inputs, while process-group lifecycle helpers moved to
20``hyper_parallel.trainer.runtime.distributed`` /
21``hyper_parallel.trainer.runtime.logging``.
22"""
24from __future__ import annotations
26__all__ = [
27 "DistributedSetup",
28 "LogicalDeviceMesh",
29 "MeshContext",
30 "ParallelLayout",
31 "ParallelTopologySpec",
32 "infer_minimum_world_size",
33 "resolve_parallel_layout",
34]
36import inspect
37import math
38from dataclasses import dataclass, field
39from typing import Any, Optional, Tuple, Union
41import torch.distributed as dist
43from hyper_parallel import init_device_mesh
46def _compatible_parallel_degree(value: Any) -> Any:
47 """Preserve the Trainer's historical minimum parallel degree of one."""
48 return max(1, value)
51def _positive_world_size(value: Any) -> int:
52 """Validate the process-group or inferred world size."""
53 if isinstance(value, bool) or not isinstance(value, int) or value < 1:
54 raise ValueError(f"world_size must be a positive integer, got {value!r}")
55 return value
58@dataclass(frozen=True)
59class ParallelTopologySpec:
60 """Normalized parallel degrees shared by runtime and offline planning."""
62 tp_size: int
63 cp_size: int
64 pp_size: int
65 ep_size: int
66 dp_shard_size: int
67 edp_shard_size: int
68 sequence_parallel: bool
69 loss_parallel: bool
71 @classmethod
72 def from_config(cls, config: Any) -> "ParallelTopologySpec":
73 """Read parallel degrees using the Trainer's compatibility semantics."""
74 accelerator = config.accelerator
75 fsdp_config = config.fsdp_config
76 return cls(
77 tp_size=_compatible_parallel_degree(accelerator.tp_size),
78 cp_size=_compatible_parallel_degree(accelerator.cp_size),
79 pp_size=_compatible_parallel_degree(accelerator.pp_size),
80 ep_size=_compatible_parallel_degree(accelerator.ep_size),
81 dp_shard_size=_compatible_parallel_degree(fsdp_config.dp_shard_size),
82 edp_shard_size=_compatible_parallel_degree(fsdp_config.edp_shard_size),
83 sequence_parallel=bool(accelerator.sequence_parallel),
84 loss_parallel=bool(accelerator.loss_parallel),
85 )
88@dataclass(frozen=True)
89class ParallelLayout:
90 """Resolved parallel topology independent of devices and process groups."""
92 spec: ParallelTopologySpec
93 world_size: int
94 dp_size: int
95 dp_replicate_size: int
97 @property
98 def dense_world_size(self) -> int:
99 """Return the rank count in one pipeline stage."""
100 return self.dp_size * self.spec.cp_size * self.spec.tp_size
103def resolve_parallel_layout(
104 spec: ParallelTopologySpec, world_size: int
105) -> ParallelLayout:
106 """Validate a world size and resolve its DP/FSDP/EP layout."""
107 world_size = _positive_world_size(world_size)
108 non_dp_size = spec.tp_size * spec.cp_size * spec.pp_size
109 if world_size % non_dp_size != 0:
110 raise ValueError(
111 f"world_size {world_size} is not divisible by non-DP size "
112 f"{non_dp_size}"
113 )
114 dp_size = world_size // non_dp_size
115 fsdp_data_parallel_size = dp_size * spec.cp_size
116 if fsdp_data_parallel_size % spec.dp_shard_size != 0:
117 raise ValueError(
118 "DP+CP size "
119 f"{fsdp_data_parallel_size} is not divisible by FSDP shard size "
120 f"{spec.dp_shard_size}"
121 )
122 dense_world_size = dp_size * spec.cp_size * spec.tp_size
123 if dense_world_size % spec.ep_size != 0:
124 raise ValueError(
125 f"expert domain size ({dense_world_size}) must be divisible by "
126 f"ep_size ({spec.ep_size})"
127 )
128 expert_data_parallel_size = dense_world_size // spec.ep_size
129 if expert_data_parallel_size % spec.edp_shard_size != 0:
130 raise ValueError(
131 f"expert data-parallel size ({expert_data_parallel_size}) must be "
132 f"divisible by edp_shard_size ({spec.edp_shard_size})"
133 )
134 return ParallelLayout(
135 spec=spec,
136 world_size=world_size,
137 dp_size=dp_size,
138 dp_replicate_size=fsdp_data_parallel_size // spec.dp_shard_size,
139 )
142def infer_minimum_world_size(spec: ParallelTopologySpec) -> int:
143 """Return the smallest world size satisfying every parallel contract."""
144 dense_world_size = math.lcm(
145 spec.tp_size * spec.cp_size,
146 spec.ep_size,
147 spec.tp_size * spec.dp_shard_size,
148 spec.ep_size * spec.edp_shard_size,
149 )
150 return dense_world_size * spec.pp_size
153@dataclass(frozen=True)
154class LogicalDeviceMesh:
155 """Process-group-free subset of the DeviceMesh planning interface."""
157 mesh_dim_names: tuple[str, ...]
158 mesh_shape: tuple[int, ...]
160 @property
161 def shape(self) -> tuple[int, ...]:
162 """Return the logical mesh shape."""
163 return self.mesh_shape
165 def __getitem__(
166 self, mesh_dim_names: Union[str, Tuple[str, ...]]
167 ) -> "LogicalDeviceMesh":
168 """Return a logical submesh selected by dimension name."""
169 names = (
170 (mesh_dim_names,)
171 if isinstance(mesh_dim_names, str)
172 else tuple(mesh_dim_names)
173 )
174 missing = tuple(name for name in names if name not in self.mesh_dim_names)
175 if missing:
176 raise KeyError(f"logical mesh has no dimensions {missing!r}")
177 shape_by_name = dict(zip(self.mesh_dim_names, self.mesh_shape))
178 return LogicalDeviceMesh(
179 mesh_dim_names=names,
180 mesh_shape=tuple(shape_by_name[name] for name in names),
181 )
183 def size(self, mesh_dim: Optional[Union[int, str]] = None) -> int:
184 """Return the total or selected logical dimension size."""
185 if mesh_dim is None:
186 return math.prod(self.mesh_shape)
187 index = (
188 self.mesh_dim_names.index(mesh_dim)
189 if isinstance(mesh_dim, str)
190 else mesh_dim
191 )
192 return self.mesh_shape[index]
195def _init_topology_mesh(mesh_kwargs: dict[str, Any]) -> Any:
196 """Create a mesh while preserving compatibility with lightweight test fakes."""
197 supported_arguments = inspect.signature(init_device_mesh).parameters
198 filtered_kwargs = {
199 name: value
200 for name, value in mesh_kwargs.items()
201 if name in supported_arguments
202 }
203 return init_device_mesh(**filtered_kwargs)
206@dataclass
207class MeshContext:
208 """Runtime topology and the mesh domains consumed by dual-mode Trainer."""
209 device_mesh: Any = None
210 fsdp_non_moe_mesh: Any = None
211 fsdp_moe_mesh: Any = None
212 dp_size: int = 1
213 dp_replicate_size: int = 1
214 dp_shard_size: int = 1
215 edp_shard_size: int = 1
216 tp_size: int = 1
217 cp_size: int = 1
218 pp_size: int = 1
219 ep_size: int = 1
220 dp_rank: int = 0
221 tp_rank: int = 0
222 cp_rank: int = 0
223 pp_rank: int = 0
224 ep_rank: int = 0
225 sequence_parallel: bool = False
226 loss_parallel: bool = False
228 @classmethod
229 def from_layout(cls, layout: ParallelLayout) -> "MeshContext":
230 """Create common mesh metadata from a resolved parallel layout."""
231 spec = layout.spec
232 return cls(
233 dp_size=layout.dp_size,
234 dp_replicate_size=layout.dp_replicate_size,
235 dp_shard_size=spec.dp_shard_size,
236 edp_shard_size=spec.edp_shard_size,
237 tp_size=spec.tp_size,
238 cp_size=spec.cp_size,
239 pp_size=spec.pp_size,
240 ep_size=spec.ep_size,
241 sequence_parallel=spec.sequence_parallel,
242 loss_parallel=spec.loss_parallel,
243 )
245 @property
246 def pp_enabled(self) -> bool:
247 """Whether pipeline parallelism is enabled (pp_size > 1)."""
248 return self.pp_size > 1
250 @property
251 def dp_cp_mesh(self) -> Any:
252 """DP+CP joint mesh for all-reduce.
254 Mirrors the derivation in ``FinetuneRecipe.setup()``: take the DP
255 dimension together with CP (if present) and flatten to 1-D.
256 """
257 if self.device_mesh is None:
258 return None
259 dim_names = self.device_mesh.mesh_dim_names or ()
260 selected_dims = tuple(dim_name for dim_name in ("dp", "cp") if dim_name in dim_names)
261 if not selected_dims:
262 selected_dims = tuple(
263 dim_name
264 for dim_name in ("dp_replicate", "dp_shard", "cp")
265 if dim_name in dim_names
266 )
267 if not selected_dims:
268 return None
269 sub = (
270 self.device_mesh[selected_dims[0]]
271 if len(selected_dims) == 1
272 else self.device_mesh[selected_dims]
273 )
274 if sub.ndim > 1:
275 sub = sub.flatten("dp_cp")
276 return sub
278 @property
279 def cp_mesh(self) -> Any:
280 """CP submesh for CP utilities (e.g. shard_batch_for_cp).
282 Returns ``None`` when context parallelism is disabled (cp_size == 1)
283 or no DeviceMesh is available.
284 """
285 if self.device_mesh is None:
286 return None
287 dim_names = self.device_mesh.mesh_dim_names or ()
288 if "cp" in dim_names:
289 return self.device_mesh["cp"]
290 return None
292 def build_meshs(self, device_type: str, world_size: int) -> None:
293 """Build the device, dense-FSDP and expert-FSDP mesh domains.
295 The device mesh preserves the existing (dp, cp, tp) rank order. Dense
296 FSDP flattens the DP+CP domain and reshapes it according to the
297 configured shard size, while expert FSDP derives an (edp, ep)
298 topology from the flattened device domain.
300 Args:
301 device_type: Device type accepted by init_device_mesh.
302 world_size: Number of ranks in the process group.
304 Raises:
305 ValueError: If the configured topology cannot cover world_size or
306 cannot form the requested expert mesh.
307 """
308 dense_world_size = self.dp_size * self.cp_size * self.tp_size
309 expected_world_size = dense_world_size * self.pp_size
310 if expected_world_size != world_size:
311 raise ValueError(
312 "dp_size * cp_size * tp_size * pp_size "
313 f"must equal world_size ({expected_world_size} != {world_size})"
314 )
315 fsdp_data_parallel_size = self.dp_size * self.cp_size
316 configured_fsdp_size = self.dp_replicate_size * self.dp_shard_size
317 if configured_fsdp_size != fsdp_data_parallel_size:
318 raise ValueError(
319 "dp_replicate_size * dp_shard_size must equal dp_size * cp_size "
320 f"({configured_fsdp_size} != {fsdp_data_parallel_size})"
321 )
323 init_backend = dist.is_initialized()
324 stage_rank_start = 0
325 if init_backend and self.pp_size > 1:
326 stage_rank_start = (dist.get_rank() // dense_world_size) * dense_world_size
327 stage_rank_list = tuple(
328 range(stage_rank_start, stage_rank_start + dense_world_size)
329 )
330 device_mesh_kwargs = {
331 "device_type": device_type,
332 "mesh_shape": (self.dp_size, self.cp_size, self.tp_size),
333 "mesh_dim_names": ("dp", "cp", "tp"),
334 "init_backend": init_backend,
335 }
336 if init_backend:
337 device_mesh_kwargs["rank_list"] = stage_rank_list
338 self.device_mesh = _init_topology_mesh(device_mesh_kwargs)
340 dense_rank_list = (
341 tuple(self.device_mesh.rank_list)
342 if hasattr(self.device_mesh, "rank_list")
343 else None
344 )
345 dense_mesh_kwargs = {
346 "device_type": device_type,
347 "mesh_shape": (
348 self.dp_replicate_size,
349 self.dp_shard_size,
350 self.tp_size,
351 ),
352 "mesh_dim_names": ("fsdp_replicate", "fsdp_shard", "tp"),
353 "init_backend": init_backend,
354 }
355 if init_backend and dense_rank_list is not None:
356 dense_mesh_kwargs["rank_list"] = dense_rank_list
357 self.fsdp_non_moe_mesh = _init_topology_mesh(dense_mesh_kwargs)
358 _validate_dense_tp_rank_layout(self.device_mesh, self.fsdp_non_moe_mesh)
360 self._build_expert_parallel_mesh(device_type, init_backend, dense_rank_list)
362 def _build_expert_parallel_mesh(
363 self,
364 device_type: str,
365 init_backend: bool,
366 dense_rank_list: Optional[Tuple[int, ...]],
367 ) -> None:
368 """Build the expert-parallel FSDP mesh when expert parallelism is enabled."""
369 self.fsdp_moe_mesh = None
370 if self.ep_size <= 1:
371 return
372 expert_domain_size = math.prod(self.device_mesh.mesh_shape)
373 if expert_domain_size % self.ep_size != 0:
374 raise ValueError(
375 f"expert domain size ({expert_domain_size}) must be divisible by "
376 f"ep_size ({self.ep_size})"
377 )
378 edp_size = expert_domain_size // self.ep_size
379 if edp_size % self.edp_shard_size != 0:
380 raise ValueError(
381 f"expert data-parallel size ({edp_size}) must be divisible by "
382 f"edp_shard_size ({self.edp_shard_size})"
383 )
384 edp_replicate_size = edp_size // self.edp_shard_size
385 expert_mesh_shape = (
386 (edp_replicate_size, self.edp_shard_size, self.ep_size)
387 if edp_replicate_size > 1
388 else (self.edp_shard_size, self.ep_size)
389 )
390 expert_mesh_names = (
391 ("edp_replicate", "edp_shard", "ep")
392 if edp_replicate_size > 1
393 else ("edp_shard", "ep")
394 )
395 expert_mesh_kwargs = {
396 "device_type": device_type,
397 "mesh_shape": expert_mesh_shape,
398 "mesh_dim_names": expert_mesh_names,
399 "init_backend": init_backend,
400 }
401 if init_backend and dense_rank_list is not None:
402 expert_mesh_kwargs["rank_list"] = dense_rank_list
403 self.fsdp_moe_mesh = _init_topology_mesh(expert_mesh_kwargs)
406def _validate_dense_tp_rank_layout(device_mesh: Any, fsdp_non_moe_mesh: Any) -> None:
407 """Ensure both TP child meshes describe the same groups and local rank."""
408 if not hasattr(device_mesh, "rank_list") or not hasattr(fsdp_non_moe_mesh, "rank_list"):
409 return
410 device_tp_mesh = device_mesh["tp"]
411 fsdp_tp_mesh = fsdp_non_moe_mesh["tp"]
412 if tuple(device_tp_mesh.rank_list) != tuple(fsdp_tp_mesh.rank_list):
413 raise ValueError("device_mesh['tp'] and fsdp_non_moe_mesh['tp'] rank groups differ")
414 if device_tp_mesh.get_local_rank() != fsdp_tp_mesh.get_local_rank():
415 raise ValueError("device_mesh['tp'] and fsdp_non_moe_mesh['tp'] local ranks differ")
418@dataclass
419class DistributedSetup:
420 """Unified distributed configuration container.
422 Following design doc 06 §3.
423 """
424 mesh_context: MeshContext = field(default_factory=MeshContext)
425 strategy_config: Any = None
426 # Normalized ``{match: ModuleShardingSpec}`` mapping (or None) — the
427 # Trainer desugars YAML PlanOverride entries via
428 # entries_to_plan_overrides (when-filtering against the built mesh
429 # included) BEFORE the AutoModels build pipeline runs; AutoModels never
430 # sees the raw Trainer YAML DTO list (05 §15.2.6).
431 plan_overrides: Any = None
432 # Normalized module replacement rules (tuple[ModuleReplacementSpec])
433 # desugared Trainer-side via entries_to_module_replacements.
434 module_replacements: Any = None
435 module_replacement_requests: Any = None
436 # Non-forward method patches installed after module replacement and before planning.
437 method_replacements: Any = None
438 # TrainingConfig.training.low_precision. Kept on the setup so model
439 # construction and sharding consume one resolved policy object.
440 low_precision_config: Any = None
441 fp32_main_params: bool = False
443 @classmethod
444 def from_config(
445 cls, config: Any, mesh_context: MeshContext, *, strategy_enabled: bool
446 ) -> "DistributedSetup":
447 """Assemble model-build configuration independently of mesh execution.
449 Args:
450 config: Resolved training configuration.
451 mesh_context: Topology prepared by the caller.
452 strategy_enabled: Whether to retain the configured FSDP strategy.
454 Returns:
455 Setup with common model-build policies populated.
456 """
457 return cls(
458 mesh_context=mesh_context,
459 strategy_config=config.fsdp_config if strategy_enabled else None,
460 low_precision_config=getattr(
461 getattr(config, "training", None), "low_precision", None
462 ),
463 fp32_main_params=config.optimizer.fp32_main_params,
464 )