Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / distributed / mesh.py: 39%
134 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-25 05:45 +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"""
24__all__ = [
25 "DistributedSetup",
26 "MeshContext",
27]
29import inspect
30import logging
31import math
32from dataclasses import dataclass, field
33from typing import Any
35import torch.distributed as dist
37from hyper_parallel import init_device_mesh
38from hyper_parallel.models.build_options import ( # pylint: disable=syntax-error
39 get_device_type,
40)
42logger = logging.getLogger(__name__)
45def _init_topology_mesh(mesh_kwargs: dict[str, Any]) -> Any:
46 """Create a mesh while preserving compatibility with lightweight test fakes."""
47 supported_arguments = inspect.signature(init_device_mesh).parameters
48 filtered_kwargs = {
49 name: value
50 for name, value in mesh_kwargs.items()
51 if name in supported_arguments
52 }
53 return init_device_mesh(**filtered_kwargs)
56@dataclass
57class MeshContext:
58 """Runtime topology and the mesh domains consumed by dual-mode Trainer."""
59 device_mesh: Any = None
60 fsdp_non_moe_mesh: Any = None
61 fsdp_moe_mesh: Any = None
62 dp_size: int = 1
63 dp_replicate_size: int = 1
64 dp_shard_size: int = 1
65 edp_shard_size: int = 1
66 tp_size: int = 1
67 cp_size: int = 1
68 pp_size: int = 1
69 ep_size: int = 1
70 dp_rank: int = 0
71 tp_rank: int = 0
72 cp_rank: int = 0
73 pp_rank: int = 0
74 ep_rank: int = 0
75 sequence_parallel: bool = False
76 loss_parallel: bool = False
78 @property
79 def pp_enabled(self) -> bool:
80 """Whether pipeline parallelism is enabled (pp_size > 1)."""
81 return self.pp_size > 1
83 @property
84 def dp_cp_mesh(self) -> Any:
85 """DP+CP joint mesh for all-reduce.
87 Mirrors the derivation in ``FinetuneRecipe.setup()``: take the DP
88 dimension together with CP (if present) and flatten to 1-D.
89 """
90 if self.device_mesh is None:
91 return None
92 dim_names = self.device_mesh.mesh_dim_names or ()
93 selected_dims = tuple(dim_name for dim_name in ("dp", "cp") if dim_name in dim_names)
94 if not selected_dims:
95 selected_dims = tuple(
96 dim_name
97 for dim_name in ("dp_replicate", "dp_shard", "cp")
98 if dim_name in dim_names
99 )
100 if not selected_dims:
101 return None
102 sub = (
103 self.device_mesh[selected_dims[0]]
104 if len(selected_dims) == 1
105 else self.device_mesh[selected_dims]
106 )
107 if sub.ndim > 1:
108 sub = sub.flatten("dp_cp")
109 return sub
111 @property
112 def cp_mesh(self) -> Any:
113 """CP submesh for CP utilities (e.g. shard_batch_for_cp).
115 Returns ``None`` when context parallelism is disabled (cp_size == 1)
116 or no DeviceMesh is available.
117 """
118 if self.device_mesh is None:
119 return None
120 dim_names = self.device_mesh.mesh_dim_names or ()
121 if "cp" in dim_names:
122 return self.device_mesh["cp"]
123 return None
125 def build_meshs(self, device_type: str, world_size: int) -> None:
126 """Build the device, dense-FSDP and expert-FSDP mesh domains.
128 The device mesh preserves the existing (dp, cp, tp) rank order. Dense
129 FSDP flattens the DP+CP domain and reshapes it according to the
130 configured shard size, while expert FSDP derives an (edp, ep)
131 topology from the flattened device domain.
133 Args:
134 device_type: Device type accepted by init_device_mesh.
135 world_size: Number of ranks in the process group.
137 Raises:
138 ValueError: If the configured topology cannot cover world_size or
139 cannot form the requested expert mesh.
140 """
141 dense_world_size = self.dp_size * self.cp_size * self.tp_size
142 expected_world_size = dense_world_size * self.pp_size
143 if expected_world_size != world_size:
144 raise ValueError(
145 "dp_size * cp_size * tp_size * pp_size "
146 f"must equal world_size ({expected_world_size} != {world_size})"
147 )
148 fsdp_data_parallel_size = self.dp_size * self.cp_size
149 configured_fsdp_size = self.dp_replicate_size * self.dp_shard_size
150 if configured_fsdp_size != fsdp_data_parallel_size:
151 raise ValueError(
152 "dp_replicate_size * dp_shard_size must equal dp_size * cp_size "
153 f"({configured_fsdp_size} != {fsdp_data_parallel_size})"
154 )
156 init_backend = dist.is_initialized()
157 stage_rank_start = 0
158 if init_backend and self.pp_size > 1:
159 stage_rank_start = (dist.get_rank() // dense_world_size) * dense_world_size
160 stage_rank_list = tuple(
161 range(stage_rank_start, stage_rank_start + dense_world_size)
162 )
163 device_mesh_kwargs = {
164 "device_type": device_type,
165 "mesh_shape": (self.dp_size, self.cp_size, self.tp_size),
166 "mesh_dim_names": ("dp", "cp", "tp"),
167 "init_backend": init_backend,
168 }
169 if init_backend:
170 device_mesh_kwargs["rank_list"] = stage_rank_list
171 self.device_mesh = _init_topology_mesh(device_mesh_kwargs)
173 dense_rank_list = (
174 tuple(self.device_mesh.rank_list)
175 if hasattr(self.device_mesh, "rank_list")
176 else None
177 )
178 dense_mesh_kwargs = {
179 "device_type": device_type,
180 "mesh_shape": (
181 self.dp_replicate_size,
182 self.dp_shard_size,
183 self.tp_size,
184 ),
185 "mesh_dim_names": ("fsdp_replicate", "fsdp_shard", "tp"),
186 "init_backend": init_backend,
187 }
188 if init_backend and dense_rank_list is not None:
189 dense_mesh_kwargs["rank_list"] = dense_rank_list
190 self.fsdp_non_moe_mesh = _init_topology_mesh(dense_mesh_kwargs)
191 _validate_dense_tp_rank_layout(self.device_mesh, self.fsdp_non_moe_mesh)
193 self._build_expert_parallel_mesh(device_type, init_backend, dense_rank_list)
195 def _build_expert_parallel_mesh(
196 self,
197 device_type: str,
198 init_backend: bool,
199 dense_rank_list: tuple[int, ...] | None,
200 ) -> None:
201 """Build the expert-parallel FSDP mesh when expert parallelism is enabled."""
202 self.fsdp_moe_mesh = None
203 if self.ep_size <= 1:
204 return
205 expert_domain_size = math.prod(self.device_mesh.mesh_shape)
206 if expert_domain_size % self.ep_size != 0:
207 raise ValueError(
208 f"expert domain size ({expert_domain_size}) must be divisible by "
209 f"ep_size ({self.ep_size})"
210 )
211 edp_size = expert_domain_size // self.ep_size
212 if edp_size % self.edp_shard_size != 0:
213 raise ValueError(
214 f"expert data-parallel size ({edp_size}) must be divisible by "
215 f"edp_shard_size ({self.edp_shard_size})"
216 )
217 edp_replicate_size = edp_size // self.edp_shard_size
218 expert_mesh_shape = (
219 (edp_replicate_size, self.edp_shard_size, self.ep_size)
220 if edp_replicate_size > 1
221 else (self.edp_shard_size, self.ep_size)
222 )
223 expert_mesh_names = (
224 ("edp_replicate", "edp_shard", "ep")
225 if edp_replicate_size > 1
226 else ("edp_shard", "ep")
227 )
228 expert_mesh_kwargs = {
229 "device_type": device_type,
230 "mesh_shape": expert_mesh_shape,
231 "mesh_dim_names": expert_mesh_names,
232 "init_backend": init_backend,
233 }
234 if init_backend and dense_rank_list is not None:
235 expert_mesh_kwargs["rank_list"] = dense_rank_list
236 self.fsdp_moe_mesh = _init_topology_mesh(expert_mesh_kwargs)
239def _validate_dense_tp_rank_layout(device_mesh: Any, fsdp_non_moe_mesh: Any) -> None:
240 """Ensure both TP child meshes describe the same groups and local rank."""
241 if not hasattr(device_mesh, "rank_list") or not hasattr(fsdp_non_moe_mesh, "rank_list"):
242 return
243 device_tp_mesh = device_mesh["tp"]
244 fsdp_tp_mesh = fsdp_non_moe_mesh["tp"]
245 if tuple(device_tp_mesh.rank_list) != tuple(fsdp_tp_mesh.rank_list):
246 raise ValueError("device_mesh['tp'] and fsdp_non_moe_mesh['tp'] rank groups differ")
247 if device_tp_mesh.get_local_rank() != fsdp_tp_mesh.get_local_rank():
248 raise ValueError("device_mesh['tp'] and fsdp_non_moe_mesh['tp'] local ranks differ")
251@dataclass
252class DistributedSetup:
253 """Unified distributed configuration container.
255 Following design doc 06 §3.
256 """
257 mesh_context: MeshContext = field(default_factory=MeshContext)
258 strategy_config: Any = None
259 pipeline_config: Any = None
260 moe_parallel_config: Any = None
261 activation_checkpointing: Any = None
262 # Normalized ``{match: ModuleShardingSpec}`` mapping (or None) — the
263 # Trainer desugars YAML PlanOverride entries via
264 # entries_to_plan_overrides (when-filtering against the built mesh
265 # included) BEFORE the AutoModels build pipeline runs; AutoModels never
266 # sees the raw Trainer YAML DTO list (05 §15.2.6).
267 plan_overrides: Any = None
268 # Normalized module replacement rules (tuple[ModuleReplacementSpec])
269 # desugared Trainer-side via entries_to_module_replacements.
270 module_replacements: Any = None
271 # TrainingConfig.training.low_precision. Kept on the setup so model
272 # construction and sharding consume one resolved policy object.
273 low_precision_config: Any = None
274 fp32_main_params: bool = False
277def _build_device_mesh_from_accelerator(
278 accel: Any,
279 dp_shard_size: int,
280 dp_replicate_size: int,
281 world_size: int,
282 edp_shard_size: int = 1,
283) -> tuple[MeshContext, tuple[str, ...]]:
284 """Build the explicit MeshContext domains from accelerator topology."""
285 tp_size = max(1, accel.tp_size)
286 cp_size = max(1, accel.cp_size)
287 pp_size = max(1, accel.pp_size)
288 ep_size = max(1, accel.ep_size)
290 dp_size = world_size // (tp_size * cp_size * pp_size)
291 mesh_context = MeshContext(
292 dp_size=dp_size,
293 dp_replicate_size=dp_replicate_size,
294 dp_shard_size=dp_shard_size,
295 edp_shard_size=edp_shard_size,
296 tp_size=tp_size,
297 cp_size=cp_size,
298 pp_size=pp_size,
299 ep_size=ep_size,
300 sequence_parallel=bool(accel.sequence_parallel),
301 loss_parallel=bool(accel.loss_parallel),
302 )
303 mesh_context.build_meshs(get_device_type(), world_size)
304 device_mesh_shape = (
305 mesh_context.device_mesh.mesh_shape
306 if hasattr(mesh_context.device_mesh, "mesh_shape")
307 else None
308 )
309 dense_shape = (
310 mesh_context.fsdp_non_moe_mesh.mesh_shape
311 if hasattr(mesh_context.fsdp_non_moe_mesh, "mesh_shape")
312 else None
313 )
314 expert_shape = (
315 mesh_context.fsdp_moe_mesh.mesh_shape
316 if mesh_context.fsdp_moe_mesh is not None and hasattr(mesh_context.fsdp_moe_mesh, "mesh_shape")
317 else None
318 )
319 logger.info(
320 "Built device mesh %s, dense FSDP mesh %s and expert FSDP mesh %s",
321 device_mesh_shape,
322 dense_shape,
323 expert_shape,
324 )
325 return mesh_context, mesh_context.device_mesh.mesh_dim_names