Coverage for / home / jenkins / .local / lib / python3.10 / site-packages / hyper_parallel / core / dtensor / device_mesh.py: 75%
862 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-07-06 05:41 +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"""device mesh"""
17import copy
18import threading
19from types import TracebackType
20from typing import Any, List, Literal, Optional, Sequence, Type, Union
21import numpy as np
23from hyper_parallel.core.dtensor._mesh_layout import IntTuple, _MeshLayout, _contiguous_strides, _is_int
24from hyper_parallel.platform import get_platform
25from hyper_parallel.platform.platform import EXISTING_COMM_GROUPS, Platform, PlatformType
27platform = get_platform()
28Tensor = platform.Tensor
31def _host_tensor_from_numpy(np_array: np.ndarray):
32 """Build a host-resident int tensor from a NumPy array for rank/mesh bookkeeping.
34 A real platform's ``from_numpy`` keeps the tensor off the meta device, so it stays
35 ``asnumpy``-able even when a DeviceMesh is built under ``ms.DeviceCtx("meta")``
36 (``fully_shard`` with ``mesh=None``). Unit tests run with a mocked ``platform``
37 (no real ``from_numpy``, and never under a meta context), so fall back to the plain
38 ``Tensor`` constructor there.
39 """
40 if isinstance(platform, Platform):
41 return platform.from_numpy(np_array)
42 return Tensor(np_array).int()
45class _MeshEnv(threading.local):
46 """Per-thread stack of active :class:`DeviceMesh` (PyTorch ``_mesh_resources`` parity)."""
48 def __init__(self) -> None:
49 super().__init__()
50 self.mesh_stack: List["DeviceMesh"] = []
52 def get_current_mesh(self) -> "DeviceMesh":
53 """Return the innermost active :class:`DeviceMesh` for this thread (PyTorch parity)."""
54 if len(self.mesh_stack) == 0:
55 raise RuntimeError("No device mesh is currently active!")
56 return self.mesh_stack[-1]
59_mesh_resources = _MeshEnv()
61BackendConfig = Optional[str]
62_CP_MESH_DIM_NAMES = {"cp", "co", "ds"}
65def _get_sub_rank_list(mesh_shape, mesh_dim_names, rank_list, sub_mesh_dim_names, current_rank):
66 """
67 Get the sub rank list for a sub mesh.
69 Args:
70 mesh_shape (tuple[int]): The shape of the original mesh.
71 mesh_dim_names (tuple[str]): The mesh dim names of the original mesh dimensions.
72 rank_list (tuple[int]): A tuple of ranks that participate in this mesh.
73 sub_mesh_dim_names (tuple[str]): The mesh dim names of the sub mesh to extract.
74 current_rank (int): The current process rank.
76 Returns:
77 list: The sub rank list for the sub mesh.
78 """
79 mesh_tensor = np.array(rank_list).reshape(mesh_shape)
81 for dim_index, dim_name in enumerate(mesh_dim_names):
82 if dim_name in sub_mesh_dim_names:
83 continue
85 dim_size = mesh_shape[dim_index]
86 sliced_tensors = np.split(mesh_tensor, dim_size, axis=dim_index)
88 for sliced_tensor in sliced_tensors:
89 rank_exists = np.isin(np.array([current_rank]), sliced_tensor).any()
90 if rank_exists:
91 mesh_tensor = sliced_tensor
92 break
94 sub_rank_list = mesh_tensor.reshape(-1).tolist()
95 return sub_rank_list
98def _normalize_backend_value(value: Any) -> BackendConfig:
99 if value is None:
100 return None
101 if isinstance(value, str):
102 return value
103 if isinstance(value, tuple) and len(value) > 0:
104 backend = value[0]
105 if backend is None or isinstance(backend, str):
106 return backend
107 return None
110def _get_cp_pg_options(mesh_dim_names: Optional[tuple[str, ...]], dim: int) -> Optional[dict[str, Any]]:
111 if (
112 platform.platform_type == PlatformType.MINDSPORE
113 and mesh_dim_names
114 and mesh_dim_names[dim] in _CP_MESH_DIM_NAMES
115 ):
116 return {"hccl_config": {"hccl_op_expansion_mode": "AIV"}}
117 return None
120def _normalize_backend_override(
121 backend_override: dict[Union[int, str], Any],
122 ndim: int,
123 mesh_dim_names: Optional[tuple[str, ...]] = None,
124) -> tuple[BackendConfig, ...]:
125 """Normalize backend overrides by dim index/name."""
126 remaining = dict(backend_override)
127 normalized: list[BackendConfig] = []
128 mesh_dim_names = mesh_dim_names or ()
130 for dim_idx in range(ndim):
131 dim_name = mesh_dim_names[dim_idx] if dim_idx < len(mesh_dim_names) else None
132 if dim_name is not None and dim_name in remaining:
133 if dim_idx in remaining:
134 raise RuntimeError(
135 f"Found redundant dim index {dim_idx} and name {dim_name} in backend_override"
136 )
137 normalized.append(_normalize_backend_value(remaining.pop(dim_name)))
138 elif dim_idx in remaining:
139 normalized.append(_normalize_backend_value(remaining.pop(dim_idx)))
140 else:
141 normalized.append(None)
143 if remaining:
144 raise RuntimeError(
145 f"Found invalid keys in backend_override: got {list(remaining.keys())}, "
146 f"expected integers in range [0, {ndim}) or one of {mesh_dim_names}"
147 )
148 return tuple(normalized)
151def _should_defer_group_init(sub_layout: _MeshLayout, backend_override: BackendConfig) -> bool:
152 """Whether this mesh dimension should skip eager process-group creation."""
153 return backend_override == "fake" or sub_layout.numel() == 1
156class DeviceMesh:
157 """
158 Topological abstraction describing cluster devices.
160 Args:
161 device_type (str): Device type. Valid values depend on the active platform:
163 - **PyTorch** (same as ``torch.distributed.device_mesh.DeviceMesh``):
164 ``"cpu"``, ``"cuda"``, ``"npu"``.
165 - **MindSpore** (mapped to the corresponding communication backend):
166 ``"cpu"`` → mccl, ``"gpu"`` → nccl, ``"npu"`` → hccl.
167 mesh (Union[Tensor, list, tuple, np.ndarray, None]): A multi-dimensional array, list, or integer
168 tensor describing the device layout. The IDs in the mesh are global IDs of the
169 default process group, representing the multi-dimensional networking structure
170 of devices in distributed training (e.g., [[0,1],[2,3]] represents a 2x2 device mesh).
171 If a list or non-int32 tensor is provided, it will be automatically converted
172 to an int32 tensor. If None, a 1D mesh containing all ranks
173 (i.e., ``[0, 1, ..., world_size-1]``) will be created automatically.
174 mesh_dim_names (tuple[str]): A tuple[str] of mesh dim names for each dimension of mesh.
175 _init_backend (boolean): Whether initial process group.
177 Attributes:
178 ndim (int): Number of dimensions in the mesh.
179 mesh_shape (tuple[int]): Shape of the device mesh.
180 rank_list (tuple[int]): Flattened list of ranks from the mesh.
181 root_mesh (DeviceMesh): The parent mesh if this is a sub mesh, None otherwise.
182 sub_mesh (list[DeviceMesh]): List of child meshes created from this mesh.
184 Context manager:
185 Use ``with device_mesh:`` to set the **current** mesh for this thread.
186 """
188 device_type: Literal["cpu", "cuda", "gpu", "npu"]
189 mesh: Union[Tensor, list, tuple, np.ndarray]
190 mesh_dim_names: Union[tuple[str, ...], list[str], None]
192 _VALID_DEVICE_TYPES = {
193 PlatformType.PYTORCH: {"cpu", "cuda", "npu"},
194 PlatformType.MINDSPORE: {"cpu", "gpu", "npu"},
195 }
197 def __init__(self,
198 device_type: Literal["cpu", "cuda", "gpu", "npu"],
199 mesh: Union[Tensor, list, tuple, np.ndarray, None] = None,
200 *,
201 mesh_dim_names: Union[tuple[str, ...], list[str], None] = None,
202 _init_backend: bool = True,
203 _layout: Optional[_MeshLayout] = None,
204 _rank_map: Optional[Tensor] = None,
205 _root_mesh: Optional['DeviceMesh'] = None,
206 ):
207 self._validate_device_type(device_type)
208 self.device_type = device_type
210 if _init_backend:
211 platform.init_process_group()
213 self._layout, self._rank_map = self._resolve_layout_and_rank_map(mesh, _layout, _rank_map)
214 self._rank = platform.get_rank()
215 self._root_mesh = _root_mesh
216 self._refresh_mesh_view()
217 self._set_mesh_dim_names(mesh_dim_names)
218 self._initialize_runtime_state(_init_backend)
219 self._coordinate_on_dim = self._compute_coordinate_on_dim()
221 @classmethod
222 def _validate_device_type(cls, device_type: str) -> None:
223 """Validate that the requested device type is supported on the active platform."""
224 valid_device_types = cls._VALID_DEVICE_TYPES.get(platform.platform_type)
225 if valid_device_types is not None and device_type not in valid_device_types:
226 raise ValueError(
227 f"Invalid device_type '{device_type}' for {platform.platform_type.name} platform. "
228 f"Valid device types are: {sorted(valid_device_types)}"
229 )
231 @classmethod
232 def _resolve_layout_and_rank_map(
233 cls,
234 mesh: Union[Tensor, list, tuple, np.ndarray, None],
235 layout: Optional[_MeshLayout],
236 rank_map: Optional[Tensor],
237 ) -> tuple[_MeshLayout, Tensor]:
238 """Build the internal layout and rank map from either public or private constructor inputs."""
239 if mesh is not None and (layout is not None or rank_map is not None):
240 raise TypeError("Cannot provide both explicit mesh and private _layout/_rank_map arguments.")
242 if mesh is None and (layout is None or rank_map is None):
243 world_size = platform.get_world_size()
244 mesh = list(range(world_size))
246 if mesh is not None:
247 mesh_tensor = cls._convert_mesh_to_tensor(mesh)
248 if mesh_tensor.ndim == 0:
249 raise ValueError("mesh must be at least 1-dimensional")
250 return cls._build_layout_from_mesh(mesh_tensor), cls._build_rank_map_from_mesh(mesh_tensor)
252 rank_map_tensor = cls._convert_rank_map_to_tensor(rank_map)
253 if layout is None or rank_map_tensor is None:
254 raise TypeError("The mesh argument is required except for private _layout/_rank_map construction.")
255 if not layout.check_non_overlap():
256 raise ValueError(f"Invalid overlapping layout {layout}.")
257 return layout, rank_map_tensor
259 def _refresh_mesh_view(self) -> None:
260 """Materialize the visible mesh tensor and the derived shape/rank metadata."""
261 # Compute everything in numpy first so the intermediate ops don't need
262 # a real device. Otherwise the call would fail (or SIGSEGV on Ascend)
263 # when DeviceMesh is constructed inside a ``ms.DeviceCtx("meta")``
264 # block — e.g., from ``DeviceMesh.concatenate`` invoked under
265 # ``fully_shard``, which forces fresh ``Tensor()`` constructions onto
266 # the meta device and any subsequent op (asnumpy, nonzero, …) crashes.
267 rank_map_np = platform.tensor_to_numpy(self._rank_map).reshape(-1)
268 full_mesh_np = self._layout.remap_to_numpy(rank_map_np)
269 if full_mesh_np.shape[0] == 1:
270 per_rank_mesh_np = full_mesh_np[0]
271 else:
272 coords = np.argwhere(full_mesh_np == self._rank)
273 if coords.shape[0] == 0:
274 raise RuntimeError(
275 "In order to get the mesh tensor of a DeviceMesh it needs to "
276 "either have all its original dimensions or contain the local rank."
277 )
278 per_rank_mesh_np = full_mesh_np[coords[0, 0]]
279 # Cache the numpy view so ``_compute_coordinate_on_dim`` doesn't need
280 # to operate on ``self.mesh`` (which may be on the meta device).
281 self._per_rank_mesh_np = per_rank_mesh_np
282 self.mesh = Tensor(per_rank_mesh_np.astype(np.int32)).int()
283 self._mesh_shape = tuple(per_rank_mesh_np.shape)
284 self._rank_list = tuple(per_rank_mesh_np.reshape(-1).tolist())
285 self._flatten_rank_map = tuple(rank_map_np.tolist())
286 self._dev_num = np.prod(np.array(self._mesh_shape))
287 self._dev_rank = len(self._mesh_shape)
289 def _set_mesh_dim_names(
290 self,
291 mesh_dim_names: Union[tuple[str, ...], list[str], None],
292 ) -> None:
293 """Validate mesh dim names and build lookup tables for named access."""
294 self.mesh_dim_names = tuple(mesh_dim_names) if mesh_dim_names else None
295 if self.mesh_dim_names is None:
296 return
298 if len(self._mesh_shape) != len(self.mesh_dim_names):
299 raise ValueError(
300 f'mesh dimensions ({len(self._mesh_shape)}) should be equal to '
301 f'mesh_dim_names length ({len(self.mesh_dim_names)})'
302 )
303 if len(set(self.mesh_dim_names)) != len(self.mesh_dim_names):
304 raise ValueError(f'Each element of mesh_dim_names {self.mesh_dim_names} should be different')
305 inter_key = "interleaved_parallel"
306 if inter_key in self.mesh_dim_names and self.mesh_dim_names.index(inter_key) != len(self.mesh_dim_names) - 1:
307 raise ValueError(
308 "'interleaved_parallel' should be at the last dim of mesh_dim_names, means virtual sharding."
309 )
310 self._dev_name_to_dev_id = {
311 name: self._dev_rank - i - 1 for i, name in enumerate(self.mesh_dim_names)
312 }
313 self._dev_name_to_index = {name: i for i, name in enumerate(self.mesh_dim_names)}
315 def _initialize_runtime_state(self, init_backend: bool) -> None:
316 """Initialize caches and optional process-group state for the mesh view."""
317 self._cache_rank_list_along_axis = {}
318 self._global_shape_map = {}
319 self._sub_mesh_cache = {}
320 self._flatten_mapping: dict[str, 'DeviceMesh'] = {}
321 self._ndim = len(self._mesh_shape)
322 self._dim_group_backends = (None,) * self._ndim
323 self._dim_group_sources = tuple((self, dim) for dim in range(self._ndim))
324 self._sub_mesh: List['DeviceMesh'] = []
325 if not init_backend:
326 return
327 self._dim_group_names = self._init_process_groups(
328 self._mesh_shape,
329 self.mesh_dim_names,
330 self._rank_list,
331 )
333 @staticmethod
334 def _build_layout_from_mesh(mesh: Tensor) -> _MeshLayout:
335 mesh_shape = tuple(mesh.shape)
336 return _MeshLayout(mesh_shape, _contiguous_strides(mesh_shape))
338 @staticmethod
339 def _build_rank_map_from_mesh(mesh: Tensor) -> Tensor:
340 return _host_tensor_from_numpy(platform.tensor_to_numpy(mesh).reshape(-1).astype(np.int32))
342 @staticmethod
343 def _convert_rank_map_to_tensor(rank_map: Tensor) -> Tensor:
344 """Normalize a rank-map input into the flat int32 Tensor stored on the mesh.
346 Tensor input is returned as-is to preserve its original device; list /
347 tuple / numpy input is built into a fresh flat int32 Tensor.
348 """
349 if isinstance(rank_map, Tensor):
350 # Reuse the existing tensor as-is so we preserve its real device.
351 # Going through ``Tensor(np_array)`` would re-create on whatever
352 # device context is active (e.g. ``ms.DeviceCtx("meta")`` while
353 # ``DeviceMesh.concatenate`` runs under ``fully_shard``), which then
354 # breaks the immediate ``asnumpy()`` in ``_refresh_mesh_view``.
355 # All in-tree callers that pass a Tensor pass an existing
356 # ``DeviceMesh._rank_map`` — already a flat int32 tensor, so no
357 # reshape/cast is needed.
358 return rank_map
359 rank_map_np = np.array(rank_map)
360 return _host_tensor_from_numpy(rank_map_np.reshape(-1).astype(np.int32))
362 @staticmethod
363 def _get_mesh_tensor_from_full_mesh(full_mesh: Tensor, current_rank: Optional[int] = None) -> Tensor:
364 """Select the per-rank mesh view from a fully materialized layout remap."""
365 if full_mesh.shape[0] == 1:
366 return full_mesh[0]
368 if current_rank is None:
369 current_rank = platform.get_rank()
371 rank_coords = (full_mesh == current_rank).nonzero()
372 if rank_coords.shape[0] > 0:
373 return full_mesh[rank_coords[0, 0]]
374 raise RuntimeError(
375 "In order to get the mesh tensor of a DeviceMesh it needs to "
376 "either have all its original dimensions or contain the local rank."
377 )
379 def _compute_coordinate_on_dim(self):
380 """Compute the current rank coordinates inside this mesh view."""
381 # Use the cached numpy view rather than ``self.mesh`` so this works
382 # even when the mesh tensor lives on the meta device (DeviceMesh
383 # constructed under ``ms.DeviceCtx("meta")`` via ``fully_shard``).
384 per_rank_mesh_np = getattr(self, "_per_rank_mesh_np", None)
385 if per_rank_mesh_np is not None:
386 rank_coords = np.argwhere(per_rank_mesh_np == self._rank)
387 if rank_coords.shape[0] not in (0, 1):
388 raise AssertionError(
389 f"rank_coords.shape[0] must be 0 or 1, got {rank_coords.shape[0]}"
390 )
391 if rank_coords.shape[0] == 0:
392 return None
393 return tuple(int(x) for x in rank_coords[0])
394 return self._compute_coordinates_from_mesh(self.mesh, self._rank)
396 @staticmethod
397 def _compute_coordinates_from_mesh(
398 mesh_tensor: Tensor,
399 rank: int,
400 ):
401 """Locate one rank inside a mesh tensor and return its coordinates."""
402 rank_coords = (mesh_tensor == rank).nonzero()
403 if rank_coords.shape[0] not in (0, 1):
404 raise AssertionError(
405 f"rank_coords.shape[0] must be 0 or 1, got {rank_coords.shape[0]}"
406 )
408 if rank_coords.shape[0] == 0:
409 return None
411 coords = rank_coords[0].tolist()
412 return tuple(coords)
414 def size(self, mesh_dim=None) -> int:
415 """Return the size along a specific mesh dimension, or total number of devices if mesh_dim is None."""
416 if mesh_dim is not None:
417 return self.mesh.shape[mesh_dim]
418 return self.mesh.numel()
420 def get_coordinate(self):
421 """Return the multi-dimensional coordinate of the current rank in the mesh, or None."""
422 return self._coordinate_on_dim if self._coordinate_on_dim else None
424 def __enter__(self) -> "DeviceMesh":
425 _mesh_resources.mesh_stack.append(self)
426 return self
428 def __exit__(
429 self,
430 exc_type: Optional[Type[BaseException]],
431 exc_val: Optional[BaseException],
432 exc_tb: Optional[TracebackType],
433 ) -> None:
434 del self
435 _mesh_resources.mesh_stack.pop()
437 @staticmethod
438 def _convert_mesh_to_tensor(mesh: Union[Tensor, list, tuple, np.ndarray]) -> Tensor:
439 """Convert a public mesh input into an int32 platform tensor."""
440 if isinstance(mesh, Tensor):
441 mesh = platform.tensor_to_numpy(mesh)
442 elif isinstance(mesh, (list, tuple)):
443 mesh = np.array(mesh)
444 elif not isinstance(mesh, np.ndarray):
445 raise TypeError(
446 f"mesh must be Tensor, list, tuple or numpy array, but got {type(mesh)}"
447 )
449 mesh = mesh.astype(np.int32)
450 return _host_tensor_from_numpy(mesh)
452 @staticmethod
453 def _init_one_process_group(mesh_shape: tuple[int, ...], mesh_dim_names: tuple[str, ...],
454 dim_name: str, rank_list: tuple[int, ...]) -> str:
455 """Create one process-group family for the named mesh dimension."""
456 group_key = None
457 split_ranks = set()
458 if not isinstance(dim_name, tuple):
459 dim_name = (dim_name,)
460 for rank in rank_list:
461 split_rank = _get_sub_rank_list(mesh_shape, mesh_dim_names, rank_list, dim_name, rank)
462 sorted_rank = tuple(sorted(split_rank))
463 split_ranks.add(sorted_rank)
464 if rank == platform.get_rank():
465 group_key = str(sorted_rank)
466 split_ranks = sorted([list(item) for item in split_ranks])
467 platform.split_group(split_ranks=split_ranks)
468 return group_key
470 @staticmethod
471 def _build_dim_split_ranks(
472 sub_layout: _MeshLayout,
473 rank_map: Tensor,
474 ) -> tuple[list[list[int]], Optional[str]]:
475 """Build rank lists and the local cache key for one logical mesh axis."""
476 pg_ranks_by_dim = sub_layout.remap_to_numpy(platform.tensor_to_numpy(rank_map))
477 current_rank = platform.get_rank()
478 split_ranks = []
479 split_ranks_set = set()
480 group_key = None
481 for dim_mesh in np.array(pg_ranks_by_dim):
482 subgroup_ranks = tuple(int(rank) for rank in np.array(dim_mesh).reshape(-1).tolist())
483 subgroup_ranks_sorted = tuple(sorted(subgroup_ranks))
484 if subgroup_ranks_sorted not in split_ranks_set:
485 split_ranks_set.add(subgroup_ranks_sorted)
486 split_ranks.append(list(subgroup_ranks_sorted))
487 if current_rank in subgroup_ranks:
488 if group_key is not None:
489 raise RuntimeError(
490 "Each device mesh dimension should get only one process group per rank."
491 )
492 group_key = str(subgroup_ranks_sorted)
493 split_ranks = sorted(split_ranks)
494 return split_ranks, group_key
496 @staticmethod
497 def _cache_group_if_needed(group_key: Optional[str], group: Any) -> None:
498 if group_key is not None and group is not None and group_key not in EXISTING_COMM_GROUPS:
499 EXISTING_COMM_GROUPS[group_key] = group
501 @staticmethod
502 def _init_process_groups_for_layout(
503 layout: _MeshLayout,
504 rank_map: Tensor,
505 mesh_dim_names: Union[tuple[str, ...], None],
506 backend_override: Optional[tuple[BackendConfig, ...]] = None,
507 ) -> list:
508 """Initialize process groups for each top-level axis in the given layout."""
509 if mesh_dim_names is None:
510 mesh_dim_names = tuple(f"dim_{dim}" for dim in range(len(layout)))
511 if backend_override is None:
512 backend_override = (None,) * len(layout)
513 if len(backend_override) != len(layout):
514 raise ValueError(
515 f"backend_override length {len(backend_override)} must match layout rank {len(layout)}"
516 )
518 dim_group_names = []
519 for dim, sub_layout in enumerate(layout):
520 split_ranks, group_key = DeviceMesh._build_dim_split_ranks(sub_layout, rank_map)
521 if _should_defer_group_init(sub_layout, backend_override[dim]):
522 dim_group_names.append(None)
523 continue
524 group = platform.split_group(split_ranks=split_ranks, pg_options=_get_cp_pg_options(mesh_dim_names, dim))
525 DeviceMesh._cache_group_if_needed(group_key, group)
526 dim_group_names.append(group_key)
527 return dim_group_names
529 @staticmethod
530 def _init_process_groups(mesh_shape: tuple[int, ...], mesh_dim_names: Union[tuple[str, ...], None],
531 rank_list: tuple[int, ...],
532 backend_override: Optional[tuple[BackendConfig, ...]] = None) -> list:
533 layout = _MeshLayout(mesh_shape, _contiguous_strides(mesh_shape))
534 rank_map = DeviceMesh._convert_rank_map_to_tensor(rank_list)
535 return DeviceMesh._init_process_groups_for_layout(
536 layout,
537 rank_map,
538 mesh_dim_names,
539 backend_override=backend_override,
540 )
542 @property
543 def rank(self):
544 """Return the global rank of the current process within this device mesh."""
545 return self._rank
547 @property
548 def mesh_shape(self):
549 """Return the shape of the device mesh as a tuple of integers."""
550 return self._mesh_shape
552 @property
553 def rank_list(self):
554 """Return the tuple of ranks participating in this device mesh."""
555 return self._rank_list
557 @property
558 def ndim(self) -> int:
559 """Return the number of dimensions in the device mesh."""
560 return self._ndim
562 @property
563 def shape(self) -> tuple:
564 """Return the shape of the device mesh as a tuple."""
565 return self._mesh_shape
567 @property
568 def root_mesh(self) -> Optional['DeviceMesh']:
569 """Return the root DeviceMesh from which this mesh was sliced, or None."""
570 return self._root_mesh
572 @root_mesh.setter
573 def root_mesh(self, value: Optional['DeviceMesh']):
574 """Set the root DeviceMesh from which this mesh was sliced."""
575 self._root_mesh = value
577 @property
578 def sub_mesh(self) -> List['DeviceMesh']:
579 """Return the list of sub-meshes derived from this device mesh."""
580 return self._sub_mesh
582 def get_flatten_mapping(self) -> dict:
583 """Return the mapping of sub-mesh names to their flattened DeviceMesh instances."""
584 return self._flatten_mapping
586 def add_flatten_mapping(self, name: str, mesh: 'DeviceMesh') -> None:
587 """Register a named DeviceMesh in the flatten mapping cache."""
588 self._flatten_mapping[name] = mesh
590 def __getitem__(self, sub_mesh_dim_names: Union[str, tuple[str, ...]]) -> 'DeviceMesh':
591 if not self.mesh_dim_names:
592 raise RuntimeError("Cannot slice a DeviceMesh without mesh_dim_names!")
594 sub_mesh_dim_names = DeviceMesh._normalize_sub_mesh_dim_names(sub_mesh_dim_names)
595 flatten_mapping = self._get_root_mesh().get_flatten_mapping()
597 flattened_result = self._try_get_from_flatten_mapping(sub_mesh_dim_names, flatten_mapping)
598 if flattened_result is not None:
599 return flattened_result
601 layout = self._get_slice_mesh_layout(sub_mesh_dim_names)
602 if sub_mesh_dim_names in self._sub_mesh_cache:
603 return self._sub_mesh_cache[sub_mesh_dim_names]
604 if layout == self._layout:
605 return self
606 return self._create_and_cache_sub_mesh(sub_mesh_dim_names, layout)
608 @staticmethod
609 def _normalize_sub_mesh_dim_names(sub_mesh_dim_names: Union[str, tuple[str, ...]]) -> tuple[str, ...]:
610 """Normalize a slice selector into a non-empty tuple of mesh dim names."""
611 if isinstance(sub_mesh_dim_names, str):
612 sub_mesh_dim_names = (sub_mesh_dim_names,)
614 if not isinstance(sub_mesh_dim_names, tuple):
615 raise TypeError(
616 f"sub_mesh_dim_names must be str or tuple, but got {type(sub_mesh_dim_names)}"
617 )
619 if len(sub_mesh_dim_names) == 0:
620 raise ValueError("sub_mesh_dim_names cannot be empty")
622 return sub_mesh_dim_names
624 @staticmethod
625 def _try_get_from_flatten_mapping(sub_mesh_dim_names: tuple[str, ...],
626 flatten_mapping: dict) -> Optional['DeviceMesh']:
627 if len(sub_mesh_dim_names) == 1 and sub_mesh_dim_names[0] in flatten_mapping:
628 return flatten_mapping[sub_mesh_dim_names[0]]
629 return None
631 def _get_mesh_dim_by_name(self, mesh_dim_name: str) -> int:
632 """Resolve a named mesh axis to its integer position."""
633 mesh_dim_names = self.mesh_dim_names or ()
634 if len(mesh_dim_names) == 0:
635 raise KeyError("No mesh_dim_names found.")
636 if mesh_dim_name not in mesh_dim_names:
637 raise KeyError(
638 f"Mesh dimension '{mesh_dim_name}' does not exist. "
639 f"Available mesh dimensions are: {mesh_dim_names}"
640 )
641 return mesh_dim_names.index(mesh_dim_name)
643 def _get_slice_mesh_layout(self, sub_mesh_dim_names: tuple[str, ...]) -> _MeshLayout:
644 """Construct the layout corresponding to one named sub-mesh slice request."""
645 root_mesh = self._get_root_mesh()
646 slice_from_root = self == root_mesh
647 flatten_name_to_layout = (
648 {key: mesh._layout for key, mesh in root_mesh.get_flatten_mapping().items()}
649 if slice_from_root else {}
650 )
651 valid_dim_names = [*(self.mesh_dim_names or ()), *flatten_name_to_layout]
652 if not all(name in valid_dim_names for name in sub_mesh_dim_names):
653 raise KeyError(
654 f"Invalid mesh_dim_names {sub_mesh_dim_names} specified. "
655 f"Valid mesh_dim_names are {valid_dim_names}."
656 )
658 if all(name in (self.mesh_dim_names or ()) for name in sub_mesh_dim_names):
659 indices = [self.mesh_dim_names.index(name) for name in sub_mesh_dim_names]
660 if indices != sorted(indices):
661 raise ValueError(
662 f"sub_mesh_dim_names {sub_mesh_dim_names} must follow the order of "
663 f"original mesh_dim_names {self.mesh_dim_names}"
664 )
666 sliced_sizes: list[IntTuple] = []
667 sliced_strides: list[IntTuple] = []
668 for name in sub_mesh_dim_names:
669 if name in (self.mesh_dim_names or ()):
670 layout = self._layout[self.mesh_dim_names.index(name)]
671 else:
672 layout = flatten_name_to_layout[name]
673 sliced_sizes.append(layout.sizes)
674 sliced_strides.append(layout.strides)
676 pre_stride = -1
677 for stride in reversed(sliced_strides):
678 if not _is_int(stride):
679 raise NotImplementedError(
680 "Currently, this only allows slicing out a contiguous flattened dim."
681 )
682 if stride < pre_stride:
683 raise ValueError(
684 f"Invalid mesh_dim_names {sub_mesh_dim_names} specified. "
685 "Mesh dim indices should be in ascending order."
686 )
687 pre_stride = stride
689 if len(sliced_sizes) == 1:
690 layout = _MeshLayout(sliced_sizes[0], sliced_strides[0])
691 else:
692 layout = _MeshLayout(tuple(sliced_sizes), tuple(sliced_strides))
693 if not layout.check_non_overlap():
694 raise RuntimeError(f"Slicing overlapping dim_names {sub_mesh_dim_names} is not allowed.")
695 return layout
697 def _create_and_cache_sub_mesh(self, sub_mesh_dim_names: tuple[str, ...], layout: _MeshLayout) -> 'DeviceMesh':
698 """Create a sub-mesh view, copy group metadata, and cache the result."""
699 root_mesh = self._get_root_mesh()
700 sub_mesh = DeviceMesh(
701 device_type=self.device_type,
702 mesh_dim_names=sub_mesh_dim_names,
703 _init_backend=False,
704 _layout=layout,
705 _rank_map=root_mesh._rank_map,
706 _root_mesh=root_mesh,
707 )
709 slice_dim_group_name = []
710 slice_dim_group_backends: list[BackendConfig] = []
711 slice_dim_group_sources: list[tuple['DeviceMesh', int]] = []
712 for name in sub_mesh_dim_names:
713 if name in (self.mesh_dim_names or ()):
714 dim_index = self.mesh_dim_names.index(name)
715 if hasattr(self, "_dim_group_names"):
716 slice_dim_group_name.append(self._dim_group_names[dim_index])
717 slice_dim_group_backends.append(self._dim_group_backends[dim_index])
718 if hasattr(self, "_dim_group_sources"):
719 slice_dim_group_sources.append(self._dim_group_sources[dim_index]) # pylint: disable=W0212
720 else:
721 slice_dim_group_sources.append((self, dim_index))
722 elif name in root_mesh.get_flatten_mapping():
723 flatten_mesh = root_mesh.get_flatten_mapping()[name]
724 if hasattr(flatten_mesh, "_dim_group_names"):
725 slice_dim_group_name.append(flatten_mesh._dim_group_names[0])
726 slice_dim_group_backends.append(flatten_mesh._dim_group_backends[0])
727 if hasattr(flatten_mesh, "_dim_group_sources"):
728 slice_dim_group_sources.append(flatten_mesh._dim_group_sources[0]) # pylint: disable=W0212
729 else:
730 slice_dim_group_sources.append((flatten_mesh, 0))
731 if slice_dim_group_name:
732 sub_mesh._dim_group_names = slice_dim_group_name # pylint: disable=W0212
733 if slice_dim_group_backends:
734 sub_mesh._dim_group_backends = tuple(slice_dim_group_backends) # pylint: disable=W0212
735 if slice_dim_group_sources:
736 sub_mesh._dim_group_sources = tuple(slice_dim_group_sources) # pylint: disable=W0212
738 self._sub_mesh_cache[sub_mesh_dim_names] = sub_mesh
739 self.sub_mesh.append(sub_mesh)
740 _register_device_mesh(sub_mesh)
741 return sub_mesh
743 def get_group(self, mesh_dim: Optional[Union[int, str]] = None):
744 """Return the communication group for one mesh axis."""
745 if not hasattr(self, "_dim_group_names"):
746 raise RuntimeError("DeviceMesh process groups not initialized!")
748 if self.ndim > 1 and mesh_dim is None:
749 raise RuntimeError(
750 f"Found the DeviceMesh have {self.ndim} dimensions. "
751 "Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1."
752 )
754 root_mesh = self._get_root_mesh()
755 if isinstance(mesh_dim, str) and mesh_dim in root_mesh.get_flatten_mapping():
756 flattened_mesh = root_mesh.get_flatten_mapping()[mesh_dim]
757 return flattened_mesh.get_comm_group_by_axis(mesh_dim)
759 return self.get_comm_group_by_axis(mesh_dim)
761 def get_all_groups(self) -> list:
762 """Return the communication groups for all mesh dimensions."""
763 if not hasattr(self, "_dim_group_names"):
764 raise RuntimeError("DeviceMesh process groups not initialized!")
766 return [self.get_group(i) for i in range(self.ndim)]
768 @staticmethod
769 def from_group(group: Union[Any, list[Any]],
770 device_type: str,
771 mesh: Union[Tensor, list, tuple, np.ndarray] = None,
772 mesh_dim_names: Union[tuple[str, ...], list[str]] = None
773 ) -> 'DeviceMesh':
774 """Build a DeviceMesh from an existing process group or a list of groups."""
775 if not isinstance(group, list):
776 group_ranks = platform.get_process_group_ranks(group)
777 group_key = str(tuple(sorted(group_ranks)))
778 if not platform.get_created_group(group_ranks):
779 EXISTING_COMM_GROUPS[group_key] = group
780 tensor_type_mesh_invalid = isinstance(mesh, Tensor) and mesh.tolist() != group_ranks
781 not_tensor_type_mesh_invalid = mesh is not None and not isinstance(mesh, Tensor) and mesh != group_ranks
782 if tensor_type_mesh_invalid or not_tensor_type_mesh_invalid:
783 raise ValueError(
784 f"Invalid mesh_shape {str(mesh)} for 1D group with ranks {group_ranks}"
785 )
786 device_mesh = DeviceMesh(device_type, group_ranks, mesh_dim_names=mesh_dim_names, _init_backend=False)
787 device_mesh._dim_group_names = [group_key] # pylint: disable=W0212
788 return device_mesh
790 groups = list(group)
791 if len(groups) == 0:
792 raise ValueError("Expect at least one group be specified.")
793 if mesh is None:
794 raise ValueError("mesh_shape is must specified when group is a list.")
795 mesh = DeviceMesh._convert_mesh_to_tensor(mesh)
796 if mesh.ndim != len(groups):
797 raise ValueError("mesh dimensions must match group dimensions.")
798 device_mesh = DeviceMesh(device_type, mesh, mesh_dim_names=mesh_dim_names, _init_backend=False)
799 device_mesh._dim_group_names = [] # pylint: disable=W0212
800 for dim_group in groups:
801 group_ranks = platform.get_process_group_ranks(dim_group)
802 group_key = str(tuple(sorted(group_ranks)))
803 if not platform.get_created_group(group_ranks):
804 EXISTING_COMM_GROUPS[group_key] = dim_group
805 device_mesh._dim_group_names.append(group_key) # pylint: disable=W0212
806 return device_mesh
808 def get_local_rank(self, mesh_dim: Optional[Union[int, str]] = None) -> int:
809 """Return the local coordinate of the current rank along one mesh dimension."""
810 if self.ndim > 1 and mesh_dim is None:
811 raise RuntimeError(
812 f"Found the DeviceMesh have {self.ndim} dimensions. "
813 "Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1."
814 )
816 if mesh_dim is None:
817 mesh_dim = 0
819 if isinstance(mesh_dim, str):
820 if mesh_dim not in self.mesh_dim_names: # pylint: disable=E1135
821 raise ValueError(
822 f"mesh_dim '{mesh_dim}' not found in mesh_dim_names {self.mesh_dim_names}"
823 )
824 dim_index = self.mesh_dim_names.index(mesh_dim)
825 else:
826 if not isinstance(mesh_dim, int) or mesh_dim < 0 or mesh_dim >= self.ndim:
827 raise ValueError(
828 f"mesh_dim must be an integer in range [0, {self.ndim}), "
829 f"but got {mesh_dim}"
830 )
831 dim_index = mesh_dim
833 if self._rank not in self._rank_list:
834 raise ValueError(
835 f"Current rank {self._rank} not found in rank_list {self._rank_list}"
836 )
838 idx = self._rank_list.index(self._rank)
839 coord = [0] * len(self._mesh_shape)
840 temp = idx
841 for i in range(len(self._mesh_shape) - 1, -1, -1):
842 coord[i] = temp % self._mesh_shape[i]
843 temp //= self._mesh_shape[i]
845 return coord[dim_index]
847 def flatten(self, mesh_dim_name: Optional[str] = None) -> 'DeviceMesh':
848 """Flatten the device mesh into a 1-D mesh and return the new DeviceMesh."""
849 return self._create_flatten_mesh(mesh_dim_name)
851 def _get_root_mesh(self) -> 'DeviceMesh':
852 """Return the canonical root mesh for this view."""
853 if self._root_mesh is None:
854 return self
855 return self._root_mesh._get_root_mesh() # pylint: disable=protected-access
857 @staticmethod
858 def _validate_concatenate_inputs(
859 meshes: Sequence['DeviceMesh'],
860 ) -> tuple['DeviceMesh', tuple['DeviceMesh', ...], tuple[str, ...], tuple[int, ...]]:
861 """Validate concatenate inputs and return root metadata plus canonical mesh views."""
862 if len(meshes) == 0:
863 raise ValueError("DeviceMesh.concatenate expects at least one mesh.")
864 if len(meshes) == 1:
865 return (
866 meshes[0]._get_root_mesh(),
867 tuple(meshes),
868 tuple(meshes[0].mesh_dim_names or ()),
869 meshes[0]._flatten_rank_map,
870 )
872 # Torch treats the flattened rank map as the common root tensor identity.
873 # If a peer view lost root metadata, recover canonical views from any input that still has it.
874 root_mesh = next(
875 (mesh._get_root_mesh() for mesh in meshes if mesh.root_mesh is not None), # pylint: disable=protected-access
876 meshes[0]._get_root_mesh(), # pylint: disable=protected-access
877 )
878 requested_dim_names: list[str] = []
879 canonical_meshes: list['DeviceMesh'] = []
880 flatten_rank_map = root_mesh._flatten_rank_map # pylint: disable=protected-access
881 anchor_meshes = DeviceMesh._collect_concatenate_anchor_meshes(meshes, root_mesh)
882 for mesh in meshes:
883 if not mesh.mesh_dim_names:
884 raise ValueError("DeviceMesh.concatenate requires mesh_dim_names on every input mesh.")
885 if mesh._flatten_rank_map == flatten_rank_map: # pylint: disable=protected-access
886 canonical_mesh = mesh
887 else:
888 canonical_mesh = DeviceMesh._recover_concatenate_mesh_from_anchors(
889 mesh,
890 anchor_meshes,
891 flatten_rank_map,
892 )
893 if canonical_mesh is None:
894 raise ValueError("DeviceMesh.concatenate expects all meshes to share the same root mesh.")
895 canonical_meshes.append(canonical_mesh)
896 requested_dim_names.extend(canonical_mesh.mesh_dim_names)
897 return root_mesh, tuple(canonical_meshes), tuple(requested_dim_names), flatten_rank_map
899 @staticmethod
900 def _collect_concatenate_anchor_meshes(
901 meshes: Sequence['DeviceMesh'],
902 root_mesh: 'DeviceMesh',
903 ) -> list['DeviceMesh']:
904 """Collect mesh views that can recover orphaned concatenate inputs by dim name."""
905 anchor_meshes: list['DeviceMesh'] = []
906 seen_ids: set[int] = set()
908 def add_anchor(mesh: Optional['DeviceMesh']) -> None:
909 if mesh is None or id(mesh) in seen_ids:
910 return
911 seen_ids.add(id(mesh))
912 anchor_meshes.append(mesh)
914 add_anchor(root_mesh)
915 for flatten_mesh in root_mesh.get_flatten_mapping().values():
916 add_anchor(flatten_mesh)
918 for mesh in meshes:
919 if mesh.root_mesh is None:
920 continue
921 add_anchor(mesh)
922 add_anchor(mesh._get_root_mesh()) # pylint: disable=protected-access
923 for source_mesh, _ in getattr(mesh, "_dim_group_sources", ()):
924 if isinstance(source_mesh, DeviceMesh):
925 add_anchor(source_mesh)
926 add_anchor(source_mesh._get_root_mesh()) # pylint: disable=protected-access
928 return anchor_meshes
930 @staticmethod
931 def _recover_concatenate_mesh_from_anchors(
932 mesh: 'DeviceMesh',
933 anchor_meshes: Sequence['DeviceMesh'],
934 flatten_rank_map: tuple[int, ...],
935 ) -> Optional['DeviceMesh']:
936 """Recover an orphan mesh as a view in the shared root coordinate system."""
937 mesh_dim_names = tuple(mesh.mesh_dim_names or ())
938 for anchor_mesh in anchor_meshes:
939 try:
940 candidate = anchor_mesh[mesh_dim_names]
941 except (KeyError, ValueError, RuntimeError, NotImplementedError):
942 continue
943 candidate_attrs = (
944 candidate.device_type == mesh.device_type,
945 candidate.mesh_shape == mesh.mesh_shape,
946 candidate.rank_list == mesh.rank_list,
947 candidate._flatten_rank_map == flatten_rank_map, # pylint: disable=protected-access
948 )
949 if all(candidate_attrs):
950 return candidate
951 return None
953 @staticmethod
954 def _validate_concatenate_root_order(root_mesh: 'DeviceMesh', requested_dim_names: tuple[str, ...]) -> None:
955 """Require original root dims to stay in root order when concatenating by name."""
956 root_dim_names = tuple(root_mesh.mesh_dim_names) if root_mesh.mesh_dim_names else ()
957 if not root_dim_names or not all(dim_name in root_dim_names for dim_name in requested_dim_names):
958 return
960 requested_indices = [root_dim_names.index(dim_name) for dim_name in requested_dim_names]
961 if requested_indices != sorted(requested_indices):
962 raise ValueError(
963 "DeviceMesh.concatenate expects meshes to follow the root mesh order. "
964 f"Got root mesh dims {root_dim_names} and requested dims {requested_dim_names}."
965 )
967 @staticmethod
968 def _collect_concatenate_metadata(
969 meshes: Sequence['DeviceMesh'],
970 ) -> tuple[
971 list[str],
972 list[IntTuple],
973 list[IntTuple],
974 list[Optional[str]],
975 list[BackendConfig],
976 list[tuple['DeviceMesh', int]],
977 ]:
978 """Collect layout and process-group metadata from all concatenate inputs."""
979 concat_dim_names: list[str] = []
980 concat_sizes: list[IntTuple] = []
981 concat_strides: list[IntTuple] = []
982 concat_dim_group_names: list[Optional[str]] = []
983 concat_dim_group_backends: list[BackendConfig] = []
984 concat_dim_group_sources: list[tuple['DeviceMesh', int]] = []
986 for mesh in meshes:
987 for dim, sub_layout in enumerate(mesh._layout): # pylint: disable=protected-access
988 concat_sizes.append(sub_layout.sizes)
989 concat_strides.append(sub_layout.strides)
990 if hasattr(mesh, "_dim_group_names"):
991 concat_dim_group_names.append(mesh._dim_group_names[dim]) # pylint: disable=protected-access
992 concat_dim_group_backends.append(mesh._dim_group_backends[dim]) # pylint: disable=protected-access
993 if hasattr(mesh, "_dim_group_sources"):
994 concat_dim_group_sources.append(mesh._dim_group_sources[dim]) # pylint: disable=protected-access
995 else:
996 concat_dim_group_sources.append((mesh, dim))
997 concat_dim_names.extend(mesh.mesh_dim_names)
999 if len(set(concat_dim_names)) != len(concat_dim_names):
1000 raise ValueError(
1001 f"DeviceMesh.concatenate expects disjoint mesh dims, but got {tuple(concat_dim_names)}."
1002 )
1003 return (
1004 concat_dim_names,
1005 concat_sizes,
1006 concat_strides,
1007 concat_dim_group_names,
1008 concat_dim_group_backends,
1009 concat_dim_group_sources,
1010 )
1012 @staticmethod
1013 def _build_concatenate_layout(concat_sizes: list[IntTuple], concat_strides: list[IntTuple]) -> _MeshLayout:
1014 """Build the layout represented by concatenated top-level mesh axes."""
1015 if len(concat_sizes) == 1:
1016 return _MeshLayout(concat_sizes[0], concat_strides[0])
1017 return _MeshLayout(tuple(concat_sizes), tuple(concat_strides))
1019 @staticmethod
1020 def _set_concatenated_group_state(
1021 mesh: 'DeviceMesh',
1022 dim_group_names: list[Optional[str]],
1023 dim_group_backends: list[BackendConfig],
1024 dim_group_sources: list[tuple['DeviceMesh', int]],
1025 ) -> None:
1026 """Attach inherited process-group metadata to a concatenated mesh view."""
1027 if dim_group_names:
1028 mesh._dim_group_names = dim_group_names # pylint: disable=W0212
1029 if dim_group_backends:
1030 mesh._dim_group_backends = tuple(dim_group_backends) # pylint: disable=W0212
1031 if dim_group_sources:
1032 mesh._dim_group_sources = tuple(dim_group_sources) # pylint: disable=W0212
1034 @staticmethod
1035 def concatenate(meshes: Sequence['DeviceMesh']) -> 'DeviceMesh':
1036 """Concatenate multiple sub-mesh views into one wider layout-backed mesh."""
1037 if len(meshes) == 1:
1038 return meshes[0]
1039 root_mesh, canonical_meshes, requested_dim_names, _ = DeviceMesh._validate_concatenate_inputs(meshes)
1040 DeviceMesh._validate_concatenate_root_order(root_mesh, requested_dim_names)
1041 (
1042 concat_dim_names,
1043 concat_sizes,
1044 concat_strides,
1045 concat_dim_group_names,
1046 concat_dim_group_backends,
1047 concat_dim_group_sources,
1048 ) = DeviceMesh._collect_concatenate_metadata(canonical_meshes)
1049 concat_layout = DeviceMesh._build_concatenate_layout(concat_sizes, concat_strides)
1050 if not concat_layout.check_non_overlap():
1051 raise ValueError(f"Cannot concatenate overlapping meshes: {meshes}")
1053 res_mesh = DeviceMesh(
1054 root_mesh.device_type,
1055 mesh_dim_names=tuple(concat_dim_names),
1056 _init_backend=False,
1057 _layout=concat_layout,
1058 _rank_map=root_mesh._rank_map, # pylint: disable=protected-access
1059 _root_mesh=root_mesh,
1060 )
1061 DeviceMesh._set_concatenated_group_state(
1062 res_mesh,
1063 concat_dim_group_names,
1064 concat_dim_group_backends,
1065 concat_dim_group_sources,
1066 )
1067 _register_device_mesh(res_mesh)
1068 return res_mesh
1070 _concatenate = concatenate
1072 def _create_flatten_mesh(
1073 self,
1074 mesh_dim_name: Optional[str] = None,
1075 backend_override: BackendConfig = None,
1076 ) -> 'DeviceMesh':
1077 """Create or reuse a flattened one-dimensional mesh view."""
1078 root_mesh = self._get_root_mesh()
1080 if mesh_dim_name is None:
1081 mesh_dim_name = "_".join(self.mesh_dim_names)
1083 if self.ndim == 1 and mesh_dim_name in self.mesh_dim_names: # pylint: disable=E1135
1084 return self
1086 invalid_dim_names = root_mesh.mesh_dim_names
1087 if mesh_dim_name in invalid_dim_names:
1088 raise ValueError(
1089 f"'{mesh_dim_name}' already exists in the root mesh mesh_dim_names "
1090 f"{invalid_dim_names}. Please specify another valid mesh_dim_name."
1091 )
1093 flattened_mesh_layout = self._layout.coalesce()
1094 if len(flattened_mesh_layout) > 1:
1095 flattened_mesh_layout = flattened_mesh_layout.nest()
1097 flatten_mapping = root_mesh.get_flatten_mapping()
1098 if mesh_dim_name in flatten_mapping:
1099 cached_mesh = flatten_mapping[mesh_dim_name]
1100 if cached_mesh._layout == flattened_mesh_layout: # pylint: disable=protected-access
1101 return cached_mesh
1102 raise ValueError(
1103 f"Flatten mesh with mesh_dim_name '{mesh_dim_name}' has been created "
1104 f"before with different layout. Please specify another valid mesh_dim_name."
1105 )
1107 res_flattened_mesh = DeviceMesh(
1108 device_type=root_mesh.device_type,
1109 mesh_dim_names=(mesh_dim_name,),
1110 _init_backend=False,
1111 _layout=flattened_mesh_layout,
1112 _rank_map=root_mesh._rank_map,
1113 _root_mesh=root_mesh,
1114 )
1115 res_flattened_mesh._dim_group_backends = (backend_override,) # pylint: disable=W0212
1116 if hasattr(self, "_dim_group_names"):
1117 res_flattened_mesh._dim_group_names = DeviceMesh._init_process_groups_for_layout( # pylint: disable=W0212
1118 res_flattened_mesh._layout,
1119 root_mesh._rank_map,
1120 res_flattened_mesh.mesh_dim_names,
1121 backend_override=(backend_override,),
1122 )
1124 root_mesh.add_flatten_mapping(mesh_dim_name, res_flattened_mesh)
1125 root_mesh._sub_mesh_cache[(mesh_dim_name,)] = res_flattened_mesh # pylint: disable=W0212
1126 root_mesh.sub_mesh.append(res_flattened_mesh)
1127 _register_device_mesh(res_flattened_mesh)
1129 return res_flattened_mesh
1131 def _create_unflatten_mesh(
1132 self,
1133 dim: int,
1134 mesh_sizes: tuple[int, ...],
1135 mesh_dim_names: tuple[str, ...],
1136 backend_override: tuple[BackendConfig, ...],
1137 ) -> 'DeviceMesh':
1138 """Split one logical mesh axis into multiple named axes."""
1139 inner_layout = _MeshLayout(mesh_sizes, _contiguous_strides(mesh_sizes))
1140 original_layout = self._layout[dim]
1141 if inner_layout.numel() != original_layout.numel():
1142 raise ValueError(
1143 f"The product of mesh_sizes={mesh_sizes} is {inner_layout.numel()}, "
1144 f"but the original dimension at dim={dim} has size {original_layout.numel()}."
1145 )
1147 partial_layout = original_layout.composition(inner_layout)
1148 unflattened_layout = self._layout.splice(dim, dim + 1, partial_layout)
1149 unflattened_mesh_dim_names = list(self.mesh_dim_names or ())
1150 unflattened_mesh_dim_names[dim: dim + 1] = list(mesh_dim_names)
1152 root_mesh = self._get_root_mesh()
1153 res_mesh = DeviceMesh(
1154 self.device_type,
1155 mesh_dim_names=tuple(unflattened_mesh_dim_names),
1156 _init_backend=False,
1157 _layout=unflattened_layout,
1158 _rank_map=root_mesh._rank_map,
1159 _root_mesh=root_mesh,
1160 )
1162 dim_group_backends = list(self._dim_group_backends)
1163 dim_group_backends[dim: dim + 1] = list(backend_override)
1164 res_mesh._dim_group_backends = tuple(dim_group_backends) # pylint: disable=W0212
1166 if hasattr(self, "_dim_group_names"):
1167 dim_group_names = list(self._dim_group_names)
1168 dim_group_names[dim: dim + 1] = DeviceMesh._init_process_groups_for_layout(
1169 partial_layout,
1170 root_mesh._rank_map,
1171 mesh_dim_names,
1172 backend_override=backend_override,
1173 )
1174 res_mesh._dim_group_names = dim_group_names # pylint: disable=W0212
1176 _register_device_mesh(res_mesh)
1177 return res_mesh
1179 def _flatten(self, mesh_dim_name: Optional[str] = None, backend_override: Any = None) -> 'DeviceMesh':
1180 return self._create_flatten_mesh(
1181 mesh_dim_name,
1182 backend_override=_normalize_backend_value(backend_override),
1183 )
1185 def _unflatten(
1186 self,
1187 dim: Union[int, str],
1188 mesh_sizes: tuple[int, ...],
1189 mesh_dim_names: tuple[str, ...],
1190 backend_override: Optional[dict[Union[int, str], Any]] = None,
1191 ) -> 'DeviceMesh':
1192 """Torch-compatible helper that expands one mesh axis into a nested layout."""
1193 if isinstance(dim, int):
1194 if dim < 0 or dim >= self.ndim:
1195 raise ValueError(f"dim {dim} specified in `_unflatten` is out of range {self.ndim}")
1196 else:
1197 mesh_dim_names_tuple = self.mesh_dim_names or ()
1198 if dim not in mesh_dim_names_tuple:
1199 raise ValueError(f"dim {dim} specified in `_unflatten` is not in {mesh_dim_names_tuple}")
1200 dim = mesh_dim_names_tuple.index(dim)
1202 if len(mesh_sizes) != len(mesh_dim_names):
1203 raise RuntimeError("mesh_dim_names must have same length as mesh_sizes in _unflatten!")
1205 backend_override_tuple = (
1206 _normalize_backend_override(backend_override, len(mesh_sizes), mesh_dim_names)
1207 if backend_override is not None
1208 else (None,) * len(mesh_dim_names)
1209 )
1210 return self._create_unflatten_mesh(dim, mesh_sizes, mesh_dim_names, backend_override_tuple)
1212 def assert_axis(self, axis, operate_name):
1213 """Validate that the given axis name exists in mesh_dim_names.
1215 Raises RuntimeError if mesh_dim_names is not set, or ValueError
1216 if the axis is not among the declared mesh dimension names.
1217 """
1218 if not self.mesh_dim_names:
1219 raise RuntimeError(f"mesh_dim_names not specified, {operate_name} is not supported.")
1220 if axis not in self.mesh_dim_names: # pylint: disable=E1135
1221 raise ValueError(
1222 f"The axis name must be one of mesh dim name {self.mesh_dim_names}, but got {axis}"
1223 )
1225 def axis_id(self, axis):
1226 """Return the reverse-indexed device dimension id for the named axis, or -1 for 'None'."""
1227 if axis == "None":
1228 return -1
1229 self.assert_axis(axis, "axis_id")
1230 return self._dev_name_to_dev_id[axis]
1232 def axis_index(self, axis):
1233 """Return the positional index of the named axis within the mesh dimensions."""
1234 self.assert_axis(axis, "axis_index")
1235 return self._dev_name_to_index[axis]
1237 def get_device_num_along_axis(self, axis):
1238 """Return the number of devices along the named mesh axis."""
1239 self.assert_axis(axis, "get_device_num_along_axis")
1240 return self.mesh_shape[self.mesh_dim_names.index(axis)]
1242 def get_rank_list_along_axis(self, mesh_dim):
1243 """Return the ranks that share every other coordinate with the current rank."""
1244 if mesh_dim in self._cache_rank_list_along_axis:
1245 return self._cache_rank_list_along_axis[mesh_dim]
1246 self.assert_axis(mesh_dim, "get_rank_list_along_axis")
1248 mesh_shape = self.mesh_shape
1249 mesh_dim_names = self.mesh_dim_names
1250 rank_list = self.rank_list
1251 rank = self.rank
1253 if rank not in rank_list:
1254 raise ValueError(f"Rank {rank} not found in rank_list")
1256 idx = rank_list.index(rank)
1257 coord = [0] * len(mesh_shape)
1258 temp = idx
1259 for i in range(len(mesh_shape) - 1, -1, -1):
1260 coord[i] = temp % mesh_shape[i]
1261 temp //= mesh_shape[i]
1263 dim_index = mesh_dim_names.index(mesh_dim)
1264 strides = [1] * len(mesh_shape)
1265 for i in range(len(mesh_shape) - 2, -1, -1):
1266 strides[i] = strides[i + 1] * mesh_shape[i + 1]
1268 result_ranks = []
1269 for v in range(mesh_shape[dim_index]):
1270 new_coord = coord.copy()
1271 new_coord[dim_index] = v
1272 new_idx = 0
1273 for i in range(len(mesh_shape)):
1274 new_idx += new_coord[i] * strides[i]
1276 result_ranks.append(rank_list[new_idx])
1278 self._cache_rank_list_along_axis[mesh_dim] = result_ranks
1279 return result_ranks
1281 def get_global_shape(self, slice_shape, tensor_map):
1282 """Infer the global tensor shape from a shard shape and tensor-map metadata."""
1283 map_key = hash((slice_shape, tensor_map))
1284 if map_key in self._global_shape_map:
1285 return self._global_shape_map[map_key]
1286 if tensor_map is None:
1287 raise ValueError(
1288 "tensor_map is not set. Please configure the tensor map by calling the layout."
1289 )
1290 if len(slice_shape) != len(tensor_map):
1291 raise ValueError(
1292 f"Length of slice_shape ({len(slice_shape)}) must match "
1293 f"the length of tensor_map ({len(tensor_map)})."
1294 )
1296 n_dims = len(self._mesh_shape)
1297 factors = [1] * len(slice_shape)
1299 for dev_idx, size in enumerate(self._mesh_shape):
1300 reverse_idx = n_dims - 1 - dev_idx
1301 for axis_idx, mapping in enumerate(tensor_map):
1302 if isinstance(mapping, int):
1303 if mapping == -1:
1304 continue
1305 if mapping == reverse_idx:
1306 factors[axis_idx] *= size
1307 break
1308 elif isinstance(mapping, tuple):
1309 if reverse_idx in mapping:
1310 factors[axis_idx] *= size
1311 break
1313 global_shape = []
1314 for i, dim in enumerate(slice_shape):
1315 global_shape.append(dim * factors[i])
1316 self._global_shape_map[map_key] = tuple(global_shape)
1317 return tuple(global_shape)
1319 def _materialize_dim_group(self, mesh_dim: int) -> Optional[str]:
1320 """Create a deferred process group for one mesh dimension on first use."""
1321 if not hasattr(self, "_dim_group_names"):
1322 self._dim_group_names = [None] * self.ndim # pylint: disable=W0201
1324 if hasattr(self, "_dim_group_sources"):
1325 source_mesh, source_dim = self._dim_group_sources[mesh_dim] # pylint: disable=W0212
1326 if source_mesh is not self or source_dim != mesh_dim:
1327 source_group_key = source_mesh._materialize_dim_group(source_dim) # pylint: disable=W0212
1328 self._dim_group_names[mesh_dim] = source_group_key
1329 return source_group_key
1331 group_key = self._dim_group_names[mesh_dim]
1332 if group_key is not None and group_key in EXISTING_COMM_GROUPS:
1333 return group_key
1335 split_ranks, group_key = DeviceMesh._build_dim_split_ranks(self._layout[mesh_dim], self._rank_map)
1336 group = platform.split_group(
1337 split_ranks=split_ranks,
1338 pg_options=_get_cp_pg_options(self.mesh_dim_names, mesh_dim),
1339 )
1340 DeviceMesh._cache_group_if_needed(group_key, group)
1341 self._dim_group_names[mesh_dim] = group_key
1342 return group_key
1344 def get_comm_group_by_axis(self, mesh_dim: Union[str, int]):
1345 """Return the cached or lazily materialized process group for one mesh axis."""
1346 if self.ndim == 1 and mesh_dim is None:
1347 mesh_dim = 0
1349 if isinstance(mesh_dim, str):
1350 if self.mesh_dim_names is None or len(self.mesh_dim_names) == 0:
1351 raise ValueError(f"DeviceMesh mesh_dim_names is not set, string mesh_dim {mesh_dim}, is not support.")
1352 if mesh_dim not in self.mesh_dim_names: # pylint: disable=E1135
1353 raise ValueError(
1354 f"mesh_dim can pass a string or integer, but string mesh_dim '{mesh_dim}' not found in "
1355 f"mesh_dim_names {self.mesh_dim_names}"
1356 )
1357 mesh_dim = self.mesh_dim_names.index(mesh_dim)
1358 else:
1359 if not isinstance(mesh_dim, int) or mesh_dim < 0 or mesh_dim >= self.ndim:
1360 raise ValueError(
1361 f"mesh_dim can pass a string or integer, if not string, mesh_dim should be a integer in range "
1362 f"[0, {self.ndim}), but got {mesh_dim}"
1363 )
1365 if not hasattr(self, "_dim_group_names"):
1366 raise RuntimeError("DeviceMesh process groups not initialized!")
1368 group_key = self._dim_group_names[mesh_dim]
1369 if group_key is None or group_key not in EXISTING_COMM_GROUPS:
1370 group_key = self._materialize_dim_group(mesh_dim)
1371 if group_key not in EXISTING_COMM_GROUPS:
1372 raise ValueError(f"{group_key} not in group cache {EXISTING_COMM_GROUPS.keys()}")
1373 return EXISTING_COMM_GROUPS[group_key]
1375 def get_devices_for_axis(self, mesh_dim: Union[str, int], rank: int):
1376 """List peer ranks that share all coordinates except the requested axis."""
1377 if isinstance(mesh_dim, str):
1378 if not self.mesh_dim_names:
1379 raise ValueError("_mesh_dim_names is not set, string mesh_dim is not supported, please pass a integer.")
1380 mesh_dim_names = self.mesh_dim_names
1381 if mesh_dim not in mesh_dim_names: # pylint: disable=E1135
1382 raise ValueError(f"mesh_dim '{mesh_dim}' not found in mesh_dim_names {mesh_dim_names}")
1383 mesh_dim = mesh_dim_names.index(mesh_dim)
1385 mesh_shape = self._mesh_shape
1386 if mesh_dim < 0 or mesh_dim >= self.ndim:
1387 raise ValueError(f"mesh_dim {mesh_dim} can not out of range [0, {self.ndim})")
1388 rank_list = self._rank_list
1389 if rank not in rank_list:
1390 raise ValueError(f"Rank {rank} not found in rank_list")
1392 idx = rank_list.index(rank)
1393 coord = [0] * len(mesh_shape)
1394 temp = idx
1395 for i in range(len(mesh_shape) - 1, -1, -1):
1396 coord[i] = temp % mesh_shape[i]
1397 temp //= mesh_shape[i]
1399 strides = [1] * len(mesh_shape)
1400 for i in range(len(mesh_shape) - 2, -1, -1):
1401 strides[i] = strides[i + 1] * mesh_shape[i + 1]
1403 result_ranks = []
1404 for v in range(mesh_shape[mesh_dim]):
1405 new_coord = coord.copy()
1406 new_coord[mesh_dim] = v
1407 new_idx = 0
1408 for i in range(len(mesh_shape)):
1409 new_idx += new_coord[i] * strides[i]
1411 result_ranks.append(rank_list[new_idx])
1413 return result_ranks
1415 def to_hash(self):
1416 """Return a hashable tuple that uniquely identifies this device mesh configuration."""
1417 map_key = (self.mesh_shape, self.mesh_dim_names, self.rank_list)
1418 return map_key
1420 def __repr__(self):
1421 return (
1422 f"DeviceMesh(device_type='{self.device_type}', mesh_shape={self._mesh_shape}, "
1423 f"mesh_dim_names={self.mesh_dim_names}, rank_list={self._rank_list})"
1424 )
1426 def __str__(self):
1427 return self.__repr__()
1429 def __deepcopy__(self, memo):
1430 cls = self.__class__
1431 result = cls.__new__(cls)
1432 memo[id(self)] = result
1433 for k, v in self.__dict__.items():
1434 if k in ("_root_mesh", "_dim_group_sources"):
1435 setattr(result, k, v)
1436 else:
1437 setattr(result, k, copy.deepcopy(v, memo))
1438 return result
1441_DEVICE_MESH_MAP = {}
1444def _device_mesh_map_key(
1445 mesh_shape: tuple[int, ...],
1446 mesh_dim_names: Union[tuple[str, ...], list[str], None],
1447 rank_list: tuple[int, ...],
1448) -> int:
1449 mesh_dim_names = tuple(mesh_dim_names) if mesh_dim_names else None
1450 return hash((tuple(mesh_shape), mesh_dim_names, tuple(rank_list)))
1453def _register_device_mesh(mesh: DeviceMesh) -> DeviceMesh:
1454 """Register a DeviceMesh in the global cache without dropping root metadata."""
1455 map_key = _device_mesh_map_key(mesh.mesh_shape, mesh.mesh_dim_names, mesh.rank_list)
1456 existing = _DEVICE_MESH_MAP.get(map_key)
1457 if existing is None:
1458 _DEVICE_MESH_MAP[map_key] = mesh
1459 return mesh
1461 if existing.root_mesh is None and mesh.root_mesh is not None:
1462 if existing is mesh._get_root_mesh(): # pylint: disable=protected-access
1463 return existing
1464 _DEVICE_MESH_MAP[map_key] = mesh
1465 return mesh
1467 return existing
1470def _create_device_mesh(device_type: str,
1471 mesh_shape: tuple[int, ...],
1472 *,
1473 mesh_dim_names: Union[tuple[str, ...], list[str], None] = None,
1474 rank_list: tuple[int, ...],
1475 init_backend: bool = True, ):
1476 """Create or reuse a cached DeviceMesh with the requested topology."""
1477 mesh = np.array(rank_list).reshape(mesh_shape)
1478 mesh_dim_names = tuple(mesh_dim_names) if mesh_dim_names else None
1479 map_key = _device_mesh_map_key(mesh_shape, mesh_dim_names, rank_list)
1480 if map_key not in _DEVICE_MESH_MAP:
1481 _register_device_mesh(
1482 DeviceMesh(device_type, mesh, mesh_dim_names=mesh_dim_names, _init_backend=init_backend)
1483 )
1484 return _DEVICE_MESH_MAP.get(map_key, None)
1487def init_device_mesh(
1488 device_type: str,
1489 mesh_shape: tuple[int, ...],
1490 *,
1491 mesh_dim_names: Union[tuple[str, ...], list[str], None] = None,
1492 rank_list: Optional[tuple[int, ...]] = None,
1493 init_backend: bool = True,
1494) -> DeviceMesh:
1495 """Initialize a cached DeviceMesh from the provided shape, names, and ranks."""
1496 total_devices = int(np.prod(np.array(mesh_shape)))
1497 if rank_list is not None:
1498 if len(rank_list) != total_devices:
1499 raise ValueError(
1500 f"rank_list length ({len(rank_list)}) must equal mesh size ({total_devices})"
1501 )
1502 else:
1503 if init_backend:
1504 platform.init_process_group()
1505 try:
1506 current_rank = platform.get_rank()
1507 except Exception as exc:
1508 raise RuntimeError(
1509 "init_device_mesh: failed to get current rank for automatic rank_list generation. "
1510 "Either pass rank_list explicitly, or ensure the process group is initialized before calling "
1511 "init_device_mesh (or set init_backend=True to let init_device_mesh initialize it)."
1512 ) from exc
1513 base = current_rank - (current_rank % total_devices)
1514 rank_list = tuple(range(base, base + total_devices))
1516 if not isinstance(mesh_shape, tuple):
1517 raise TypeError(f'mesh_shape must be a tuple, but got {type(mesh_shape)}')
1519 for size in mesh_shape:
1520 if not isinstance(size, int) or size <= 0:
1521 raise ValueError(
1522 f"Each element of mesh_shape must be a positive integer, but got {mesh_shape}"
1523 )
1525 if mesh_dim_names is not None:
1526 if not isinstance(mesh_dim_names, (tuple, list)):
1527 raise TypeError(
1528 f'mesh_dim_names must be a tuple or list, but got {type(mesh_dim_names)}'
1529 )
1530 mesh_dim_names = tuple(mesh_dim_names)
1531 if len(mesh_shape) != len(mesh_dim_names):
1532 raise ValueError(
1533 f'mesh_shape ({len(mesh_shape)}) and mesh_dim_names '
1534 f'({len(mesh_dim_names)}) should have same length'
1535 )
1536 if len(set(mesh_dim_names)) != len(mesh_dim_names):
1537 raise ValueError(f'Each element of mesh_dim_names {mesh_dim_names} should be different')
1538 if any(not isinstance(name, str) or name == "" for name in mesh_dim_names):
1539 raise ValueError(f'Each element of mesh_dim_names {mesh_dim_names} should be a non-empty string')
1541 return _create_device_mesh(
1542 device_type,
1543 mesh_shape,
1544 mesh_dim_names=mesh_dim_names,
1545 rank_list=rank_list,
1546 init_backend=init_backend,
1547 )