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