954 lines
44 KiB
Python
954 lines
44 KiB
Python
# mypy: allow-untyped-defs
|
|
# Copyright (c) Meta Platforms, Inc. and affiliates
|
|
import logging
|
|
import math
|
|
import threading
|
|
from functools import reduce
|
|
from itertools import chain
|
|
from typing import Dict, List, Optional, Tuple, TYPE_CHECKING, Union
|
|
|
|
import torch
|
|
from torch.distributed import is_available
|
|
from torch.utils._typing_utils import not_none
|
|
|
|
|
|
__all__ = ["init_device_mesh", "DeviceMesh"]
|
|
|
|
|
|
if not is_available():
|
|
import sys
|
|
|
|
# We need to create the stubs when distributed is not available.
|
|
# Otherwise, we would fail the doc tests (```./.ci/pytorch/docs-test.sh```),
|
|
# since it would try to import ``torch.distributed.device_mesh`` or
|
|
# ``torch.distributed.init_device_mesh`` but cannot find them.
|
|
|
|
class _DeviceMeshStub:
|
|
pass
|
|
|
|
def _init_device_mesh_stub():
|
|
pass
|
|
|
|
sys.modules["torch.distributed.device_mesh"].DeviceMesh = _DeviceMeshStub # type: ignore[attr-defined]
|
|
sys.modules[
|
|
"torch.distributed.device_mesh"
|
|
].init_device_mesh = _init_device_mesh_stub # type: ignore[attr-defined]
|
|
|
|
|
|
else:
|
|
from torch.distributed.distributed_c10d import (
|
|
_find_pg_by_ranks_and_tag,
|
|
_get_default_group,
|
|
_get_group_tag,
|
|
get_backend,
|
|
get_process_group_ranks,
|
|
get_rank,
|
|
get_world_size,
|
|
init_process_group,
|
|
is_initialized,
|
|
new_group,
|
|
ProcessGroup,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# only import numpy typing when type checking
|
|
if TYPE_CHECKING:
|
|
try:
|
|
from numpy.typing import ArrayLike
|
|
except ImportError:
|
|
logger.warning(
|
|
"DeviceMesh requires numpy >= 1.21 to be installed for type checking"
|
|
)
|
|
|
|
class _MeshEnv(threading.local):
|
|
def __init__(self) -> None:
|
|
self.mesh_stack: List[DeviceMesh] = []
|
|
self.child_to_root_mapping: Dict[DeviceMesh, DeviceMesh] = {}
|
|
self.mesh_dim_group_options: Dict[
|
|
int, Tuple[str, Optional[ProcessGroup.Options]]
|
|
] = {}
|
|
self.root_to_flatten_mapping: Dict[DeviceMesh, Dict[str, DeviceMesh]] = {}
|
|
# Record flatten mesh name to its mesh dim index in root mesh.
|
|
self.flatten_name_to_root_dims: Dict[
|
|
DeviceMesh, Dict[str, Tuple[int, ...]]
|
|
] = {}
|
|
|
|
def get_current_mesh(self) -> "DeviceMesh":
|
|
if len(self.mesh_stack) == 0:
|
|
raise RuntimeError("No device mesh is currently active!")
|
|
return self.mesh_stack[-1]
|
|
|
|
def create_sub_mesh(
|
|
self,
|
|
device_mesh: "DeviceMesh",
|
|
submesh_dim_names: Tuple[str, ...],
|
|
submesh_dims: List[Tuple[int, ...]],
|
|
) -> "DeviceMesh":
|
|
# Get the submesh dim size from the submesh_dims.
|
|
# For example, if we have a 3D mesh with mesh_shape (2, 2, 2) mesh_dim_names ("dp", "cp", "tp") and we want
|
|
# to slice out mesh["dp_cp"], then submesh_dims = [(0, 1), (2,)] and submesh_dim_size = [2 * 2, 2] = [4, 2].
|
|
# If we want to slice out mesh["dp", "cp"], then submesh_dims = [(0,), (1,)] and submesh_dim_size = [2, 2].
|
|
slice_dim_size = [
|
|
reduce(
|
|
lambda x, y: device_mesh.mesh.size(x) * device_mesh.mesh.size(y),
|
|
mesh_dim,
|
|
)
|
|
if len(mesh_dim) > 1
|
|
else device_mesh.mesh.size(mesh_dim[0])
|
|
for mesh_dim in submesh_dims
|
|
]
|
|
|
|
mesh_tensor = device_mesh.mesh
|
|
# slice_dim_idx could be differnt from submesh_dims, as we may need to flatten out some dims.
|
|
slice_dim_idx = []
|
|
slice_dim_group_info = []
|
|
# keep track of the number of dims that have been flattened so we can get the correct slice_dim_idx in the
|
|
# flattened mesh tensor.
|
|
num_dims_flatten = 0
|
|
for mesh_dim_indices, mesh_dim_name in zip(submesh_dims, submesh_dim_names):
|
|
# Currently, this only allows slicing out a contiguous flattened dim.
|
|
# TODO: we need to handle reconstructing a non-contiguous flattened dim.
|
|
if len(mesh_dim_indices) > 1:
|
|
# We need to move the start_dim and end_dim to the left if some dims are already flattened.
|
|
mesh_tensor = mesh_tensor.flatten(
|
|
start_dim=mesh_dim_indices[0] - num_dims_flatten,
|
|
end_dim=mesh_dim_indices[-1] - num_dims_flatten,
|
|
)
|
|
# If some dims are already flattened, we need to adjust the slice_dim_idx accordingly.
|
|
# For example, if the submesh_dims = [(0, 1), (2,), (3, 4)] with 0-1 flattened and 3-4 flattened,
|
|
# then the final slice_dim_idx should be [0, 1, 2].
|
|
slice_dim_idx.append(mesh_dim_indices[0] - num_dims_flatten)
|
|
num_dims_flatten += len(mesh_dim_indices) - 1
|
|
slice_dim_group_info.append(
|
|
self.root_to_flatten_mapping[device_mesh][
|
|
mesh_dim_name
|
|
]._dim_group_infos[0]
|
|
)
|
|
else:
|
|
slice_dim_idx.append(mesh_dim_indices[0] - num_dims_flatten)
|
|
slice_dim_group_info.append(
|
|
device_mesh._dim_group_infos[mesh_dim_indices[0]]
|
|
)
|
|
|
|
# mesh_tensor has already been flattened if needed. So mesh_tensor.ndim <= device_mesh.mesh.ndim now.
|
|
mesh_dims_remained_idx = list(range(mesh_tensor.ndim))
|
|
for idx in slice_dim_idx:
|
|
mesh_dims_remained_idx.remove(idx)
|
|
|
|
# pg_ranks_by_dim is the size of [number of local ranks of the outermost submesh dimension, *slice_dim_idx]
|
|
# This means on each local rank of the outermost slice mesh dim, we have a tensor of submesh size with
|
|
# the pg ranks of the submesh. From this, we can extract the submesh mesh tensor contains the current rank.
|
|
pg_ranks_by_dim = mesh_tensor.permute(
|
|
*mesh_dims_remained_idx, *slice_dim_idx
|
|
).reshape(-1, *slice_dim_size)
|
|
|
|
cur_rank = device_mesh.get_rank()
|
|
for mesh_nd in pg_ranks_by_dim:
|
|
submesh = DeviceMesh(
|
|
device_mesh.device_type,
|
|
mesh_nd,
|
|
mesh_dim_names=submesh_dim_names,
|
|
_init_backend=False,
|
|
)
|
|
if cur_rank in mesh_nd:
|
|
res_submesh = submesh
|
|
|
|
res_submesh._dim_group_infos = slice_dim_group_info # type: ignore[possibly-undefined]
|
|
self.child_to_root_mapping[res_submesh] = device_mesh
|
|
|
|
return res_submesh
|
|
|
|
def create_flatten_mesh(
|
|
self, device_mesh: "DeviceMesh", mesh_dim_name: Optional[str] = None
|
|
) -> "DeviceMesh":
|
|
root_mesh = _mesh_resources.get_root_mesh(device_mesh)
|
|
|
|
flatten_dims_in_root = [
|
|
not_none(root_mesh.mesh_dim_names).index(flattened_mesh_dim_name)
|
|
for flattened_mesh_dim_name in not_none(device_mesh.mesh_dim_names)
|
|
]
|
|
|
|
if not mesh_dim_name:
|
|
mesh_dim_name = "_".join(
|
|
[
|
|
not_none(root_mesh.mesh_dim_names)[dim]
|
|
for dim in flatten_dims_in_root
|
|
]
|
|
)
|
|
|
|
# Check whether the mesh_dim_name for flattened mesh is valid.
|
|
self.flatten_name_to_root_dims.setdefault(root_mesh, {})
|
|
invalid_dim_names = chain(
|
|
*list(not_none(root_mesh.mesh_dim_names)),
|
|
*self.flatten_name_to_root_dims[root_mesh].keys(),
|
|
)
|
|
if mesh_dim_name in invalid_dim_names:
|
|
raise RuntimeError(
|
|
f"{mesh_dim_name} already exists for submesh of the {root_mesh}. ",
|
|
f"The mesh_dim_names of submesh and flattened mesh are {invalid_dim_names}. "
|
|
f"Please specify another valid mesh_dim_name.",
|
|
)
|
|
|
|
# Quick return if the flatten mesh has been created before.
|
|
# TODO: If we decide to restrict flatten initialization once, we should remove
|
|
# this check and throw an error if the flatten mesh is already created before.
|
|
if (
|
|
root_mesh in self.root_to_flatten_mapping
|
|
and mesh_dim_name in self.root_to_flatten_mapping[root_mesh]
|
|
):
|
|
return self.root_to_flatten_mapping[root_mesh][mesh_dim_name]
|
|
|
|
flattened_mesh_dim_size = math.prod(device_mesh.mesh.size())
|
|
|
|
remained_dims_in_root = list(range(root_mesh.mesh.ndim))
|
|
for flatten_dim_in_root in flatten_dims_in_root:
|
|
remained_dims_in_root.remove(flatten_dim_in_root)
|
|
|
|
pg_ranks_by_dim = root_mesh.mesh.permute(
|
|
*remained_dims_in_root, *flatten_dims_in_root
|
|
).reshape(-1, flattened_mesh_dim_size)
|
|
|
|
cur_rank = root_mesh.get_rank()
|
|
for mesh_nd in pg_ranks_by_dim:
|
|
# need to init backend here since the flattened pg doesn't exist in root mesh.
|
|
flattened_mesh = DeviceMesh(
|
|
root_mesh.device_type,
|
|
mesh_nd,
|
|
mesh_dim_names=(mesh_dim_name,),
|
|
)
|
|
if cur_rank in mesh_nd:
|
|
res_flattened_mesh = flattened_mesh
|
|
self.child_to_root_mapping[res_flattened_mesh] = root_mesh # type: ignore[possibly-undefined]
|
|
self.root_to_flatten_mapping.setdefault(root_mesh, {})[mesh_dim_name] = res_flattened_mesh # type: ignore[possibly-undefined]
|
|
self.flatten_name_to_root_dims[root_mesh][mesh_dim_name] = tuple(flatten_dims_in_root) # type: ignore[possibly-undefined]
|
|
|
|
return res_flattened_mesh
|
|
|
|
def get_root_mesh(self, device_mesh: "DeviceMesh") -> "DeviceMesh":
|
|
# If a mesh could not be found in the child_to_root_mapping, it is a root mesh itself.
|
|
# A root mesh is not created through slicing.
|
|
# We considers the root mesh of a root mesh is itself.
|
|
root_mesh = self.child_to_root_mapping.get(device_mesh, None)
|
|
return device_mesh if not root_mesh else root_mesh
|
|
|
|
def get_root_mesh_dim(self, device_mesh: "DeviceMesh") -> Optional[int]:
|
|
"""
|
|
Returns the index of the mesh dim in the root mesh.
|
|
The device_mesh passed in needs to be sliced out from the root mesh
|
|
or submesh of the root mesh.
|
|
"""
|
|
root_mesh = self.get_root_mesh(device_mesh)
|
|
child_mesh_dim_names = device_mesh.mesh_dim_names
|
|
if root_mesh and child_mesh_dim_names:
|
|
assert (
|
|
len(child_mesh_dim_names) == 1
|
|
), "The submesh can only be a 1D mesh."
|
|
child_mesh_dim_name = child_mesh_dim_names[0]
|
|
return self.get_mesh_dim_by_name(root_mesh, child_mesh_dim_name)
|
|
return None
|
|
|
|
@staticmethod
|
|
def num_devices_per_host(device_type: str) -> int:
|
|
return _get_device_handle(device_type).device_count()
|
|
|
|
@staticmethod
|
|
def num_hosts(device_type: str) -> int:
|
|
# ProcessGroup can't tell us this info so we have to infer it, assume
|
|
# homogeneous hardware for now
|
|
return get_world_size() // _MeshEnv.num_devices_per_host(device_type)
|
|
|
|
def get_mesh_dim_by_name(
|
|
self, device_mesh: "DeviceMesh", mesh_dim_name: str
|
|
) -> int:
|
|
if (
|
|
device_mesh.mesh_dim_names is None
|
|
or len(device_mesh.mesh_dim_names) == 0
|
|
):
|
|
raise KeyError(
|
|
"No `mesh_dim_names` found.",
|
|
)
|
|
if mesh_dim_name not in device_mesh.mesh_dim_names:
|
|
raise KeyError(
|
|
f"Mesh dimension '{mesh_dim_name}' does not exist.",
|
|
f"Available mesh dimensions are: mesh_dim_names={device_mesh.mesh_dim_names}",
|
|
)
|
|
return not_none(device_mesh.mesh_dim_names.index(mesh_dim_name))
|
|
|
|
def _set_mesh_dim_group_options(
|
|
self,
|
|
dim: int,
|
|
backend: str,
|
|
pg_options: Optional[ProcessGroup.Options] = None,
|
|
) -> None:
|
|
self.mesh_dim_group_options[dim] = (backend, pg_options)
|
|
|
|
def _get_slice_mesh_dims(
|
|
self, device_mesh, mesh_dim_names
|
|
) -> List[Tuple[int, ...]]:
|
|
"""
|
|
Validate whether the mesh_dim_names is valid for slicing the given device_mesh.
|
|
If valid, return dim indexes of the slice mesh in the device mesh.
|
|
"""
|
|
if device_mesh != self.get_root_mesh(device_mesh):
|
|
raise RuntimeError("Cannot create a submesh from a submesh.")
|
|
|
|
# The slice mesh_dim_names should consist either the device_mesh's mesh_dim_names
|
|
# or its flattened mesh's mesh_dim_names.
|
|
self.flatten_name_to_root_dims.setdefault(device_mesh, {})
|
|
flatten_name_to_root_dims = self.flatten_name_to_root_dims[device_mesh]
|
|
valid_mesh_dim_names = [
|
|
*device_mesh.mesh_dim_names,
|
|
*flatten_name_to_root_dims,
|
|
]
|
|
|
|
if not all(
|
|
mesh_dim_name in valid_mesh_dim_names
|
|
for mesh_dim_name in mesh_dim_names
|
|
):
|
|
raise KeyError(
|
|
f"Invalid mesh_dim_names {mesh_dim_names} specified. "
|
|
f"Valid mesh_dim_names are {valid_mesh_dim_names}."
|
|
)
|
|
|
|
# Validate the order of the slice mesh dim indices.
|
|
# This needs to be in ascending order.
|
|
curr_idx = -1
|
|
slice_mesh_dims = []
|
|
for mesh_dim_name in mesh_dim_names:
|
|
if mesh_dim_name in flatten_name_to_root_dims:
|
|
mesh_indices = flatten_name_to_root_dims[mesh_dim_name]
|
|
# TODO: this doesn't allow non-contiguous slicing with flatten dim yet. next_idx
|
|
# should be mesh_indices[0] once we support non-contiguous slicing with flatten dim.
|
|
next_idx = mesh_indices[-1]
|
|
slice_mesh_dims.append(mesh_indices)
|
|
else:
|
|
next_idx = device_mesh.mesh_dim_names.index(mesh_dim_name)
|
|
slice_mesh_dims.append((next_idx,))
|
|
if next_idx <= curr_idx:
|
|
raise KeyError(
|
|
f"Invalid mesh_dim_names {mesh_dim_names} specified. ",
|
|
f"Found mesh dim indices to slice: {slice_mesh_dims}. ",
|
|
"Mesh dim indices should be in ascending order.",
|
|
)
|
|
curr_idx = next_idx
|
|
|
|
return slice_mesh_dims
|
|
|
|
def _get_all_submeshes(
|
|
self, device_mesh: "DeviceMesh", mesh_dim_name: str
|
|
) -> List["DeviceMesh"]:
|
|
"""
|
|
Return all the submeshes of a given mesh dimension of the device mesh.
|
|
"""
|
|
mesh_dim = self.get_mesh_dim_by_name(device_mesh, mesh_dim_name)
|
|
pg_ranks_by_dim = device_mesh.mesh.swapdims(-1, mesh_dim).reshape(
|
|
-1, device_mesh.mesh.size(mesh_dim)
|
|
)
|
|
|
|
cur_rank = device_mesh.get_rank()
|
|
res_submeshes = []
|
|
for mesh_1d in pg_ranks_by_dim:
|
|
submesh = DeviceMesh(
|
|
device_mesh.device_type,
|
|
mesh_1d,
|
|
mesh_dim_names=(mesh_dim_name,),
|
|
_init_backend=False,
|
|
)
|
|
submesh._dim_group_infos = (
|
|
[device_mesh._dim_group_infos[mesh_dim]]
|
|
if cur_rank in mesh_1d
|
|
else []
|
|
)
|
|
res_submeshes.append(submesh)
|
|
|
|
return res_submeshes
|
|
|
|
_mesh_resources: _MeshEnv = _MeshEnv()
|
|
|
|
def _get_device_handle(device_type: str = "cuda"):
|
|
"""
|
|
Get the module corresponding to the device_type which is cuda or cuda-like device.
|
|
For example, when the device_type is cuda, the module `torch.cuda` is returned.
|
|
Return None when there is no corresponding module for device_type, otherwise
|
|
return the corresponding module.
|
|
"""
|
|
return getattr(torch, device_type, None)
|
|
|
|
class DeviceMesh:
|
|
"""
|
|
DeviceMesh represents a mesh of devices, where layout of devices could be
|
|
represented as a n-d dimension array, and each value of the n-d dimensional
|
|
array is the global id of the default process group ranks.
|
|
|
|
DeviceMesh could be used to describe the layout of devices across the cluster,
|
|
and serves as a proxy for communication among the device lists within the cluster.
|
|
|
|
DeviceMesh can be used as a context manager.
|
|
|
|
.. note::
|
|
DeviceMesh follows SPMD programming model, which means the same PyTorch Python program
|
|
is running on all processes/ranks in the cluster. Therefore, users need to make sure the
|
|
`mesh` array (which describes the layout of devices) should be identical across all ranks.
|
|
Inconsistent `mesh` will lead to silent hang.
|
|
|
|
Args:
|
|
device_type (str): The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like".
|
|
mesh (ndarray): A multi-dimensional array or an integer tensor describing the layout
|
|
of devices, where the IDs are global IDs of the default process group.
|
|
|
|
Returns:
|
|
DeviceMesh: A :class:`DeviceMesh` object representing the device layout.
|
|
|
|
The following program runs on each process/rank in an SPMD manner. In this example, we have 2
|
|
hosts with 4 GPUs each.
|
|
A reduction over the first dimension of mesh will reduce across
|
|
columns (0, 4), .. and (3, 7), a reduction over the second dimension
|
|
of mesh reduces across rows (0, 1, 2, 3) and (4, 5, 6, 7).
|
|
|
|
Example::
|
|
>>> # xdoctest: +SKIP("no rank")
|
|
>>> from torch.distributed.device_mesh import DeviceMesh
|
|
>>>
|
|
>>> # Initialize device mesh as (2, 4) to represent the topology
|
|
>>> # of cross-host(dim 0), and within-host (dim 1).
|
|
>>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])
|
|
"""
|
|
|
|
device_type: str
|
|
mesh: torch.Tensor
|
|
mesh_dim_names: Optional[Tuple[str, ...]]
|
|
|
|
def __init__(
|
|
self,
|
|
device_type: str,
|
|
mesh: Union[torch.Tensor, "ArrayLike"],
|
|
*,
|
|
mesh_dim_names: Optional[Tuple[str, ...]] = None,
|
|
_init_backend: bool = True,
|
|
) -> None:
|
|
self.device_type = device_type
|
|
if isinstance(mesh, torch.Tensor) and mesh.device.type != "cpu":
|
|
raise ValueError(f"`mesh` must be a CPU tensor, got {mesh}")
|
|
self.mesh = (
|
|
mesh.detach().to(dtype=torch.int)
|
|
if isinstance(mesh, torch.Tensor)
|
|
else torch.tensor(mesh, device="cpu", dtype=torch.int)
|
|
)
|
|
self.mesh_dim_names = tuple(mesh_dim_names) if mesh_dim_names else None
|
|
|
|
# private field to pre-generate DeviceMesh's hash
|
|
self._flatten_mesh_list = tuple(self.mesh.flatten().tolist())
|
|
self._thread_id = None
|
|
|
|
# Skip process group initialization if xla device or init backend is False
|
|
# TODO(yeounoh) implement DeviceMesh backend and register XLA backend.
|
|
if device_type != "xla":
|
|
# always try to create default (world) pg, even if it is not initialized
|
|
# already. The world pg is used for device mesh identity (rank) on each
|
|
# process (we need to know if the current global rank is in the mesh or not).
|
|
if _init_backend:
|
|
self._get_or_create_default_group()
|
|
self._init_process_groups()
|
|
|
|
if is_initialized() and get_backend() == "threaded":
|
|
self._thread_id = threading.get_ident()
|
|
|
|
# calculate the coordinates of the current global rank on the mesh
|
|
rank_coords = (self.mesh == get_rank()).nonzero()
|
|
assert rank_coords.size(0) in (0, 1)
|
|
self._coordinate_on_dim: Optional[List[int]] = (
|
|
rank_coords[0].tolist() if rank_coords.size(0) > 0 else None
|
|
)
|
|
|
|
def _get_or_create_default_group(self):
|
|
default_initialized = is_initialized()
|
|
if not default_initialized:
|
|
init_process_group()
|
|
|
|
world_size = get_world_size()
|
|
if self.mesh.numel() > world_size:
|
|
raise RuntimeError(
|
|
f"Mesh should not be bigger than default world size {world_size}, but found {self.mesh.numel()} ranks!"
|
|
)
|
|
|
|
device_handle = _get_device_handle(self.device_type)
|
|
# TODO: if user want to pass pg_options, offer a way to do it
|
|
if not default_initialized and device_handle:
|
|
# automatically set the current cuda/cuda-like device base on num of gpu devices available in each host
|
|
# NOTE: This device selection would only work for homogeneous hardware.
|
|
num_devices_per_host = device_handle.device_count()
|
|
if (
|
|
world_size > num_devices_per_host
|
|
and world_size % num_devices_per_host != 0
|
|
):
|
|
raise RuntimeError(
|
|
f"DeviceMesh only support homogeneous hardware, but found "
|
|
f"{world_size} ranks and {num_devices_per_host} {self.device_type} devices!"
|
|
)
|
|
device_handle.set_device(get_rank() % num_devices_per_host)
|
|
|
|
return _get_default_group()
|
|
|
|
def _init_process_groups(self):
|
|
# tag/ranks/group_name associated with each mesh dimension, each
|
|
# mesh dimension should have one sub-group per rank
|
|
#
|
|
# TODO(yifu): remove tag and ranks once we fully migrate to native
|
|
# functional collectives. See details in:
|
|
# https://github.com/pytorch/pytorch/issues/93173#issuecomment-1907095208
|
|
dim_group_infos: List[Tuple[str, List[int], str]] = []
|
|
|
|
if self.mesh.ndim == 1 and self.mesh.numel() == get_world_size():
|
|
# Append the default pg to the first dim groups only if the default pg is compatible with `self.device_type`.
|
|
# Otherwise, create new pg.
|
|
default_group = _get_default_group()
|
|
ranks = list(range(get_world_size()))
|
|
dim_group = (
|
|
new_group(backend="cpu:gloo,cuda:nccl", ranks=ranks)
|
|
if torch.cuda.is_available()
|
|
and get_backend(default_group) == "gloo"
|
|
else default_group
|
|
)
|
|
dim_group_infos.append(
|
|
(
|
|
_get_group_tag(dim_group),
|
|
ranks,
|
|
dim_group.group_name,
|
|
)
|
|
)
|
|
else:
|
|
# create sub pgs base on the mesh argument specified
|
|
for dim in range(self.mesh.ndim):
|
|
# swap the current dim to the last dim
|
|
# then reshape to flatten out other dims
|
|
pg_ranks_by_dim = self.mesh.swapdims(-1, dim).reshape(
|
|
-1, self.mesh.size(dim)
|
|
)
|
|
# multi-dim mesh, create subgroups by looping over the pg_ranks
|
|
# for each dim and append the groups
|
|
for dim_mesh in pg_ranks_by_dim:
|
|
subgroup_ranks = dim_mesh.tolist()
|
|
|
|
# Respect dim group options specified via _MeshEnv.set_dim_group_options().
|
|
# Inherit from the parent group if no options are specified for the group.
|
|
if dim in _mesh_resources.mesh_dim_group_options:
|
|
(
|
|
backend,
|
|
pg_options,
|
|
) = _mesh_resources.mesh_dim_group_options[dim]
|
|
else:
|
|
backend, pg_options = None, None
|
|
|
|
# We temporarily revert the re-use subgroup, since it breaks two internal tests.
|
|
# Temporarily reverting to resolve test timeout while root-causing.
|
|
# TODO: Add two tests to cover internal tests scenarios and re-enable reuse subgroup if exists.
|
|
dim_group = new_group(
|
|
ranks=subgroup_ranks,
|
|
backend=backend,
|
|
pg_options=pg_options,
|
|
)
|
|
|
|
# only add to dim_groups if the current rank in the subgroup
|
|
if self.get_rank() in subgroup_ranks:
|
|
if len(dim_group_infos) > dim:
|
|
raise RuntimeError(
|
|
f"Each device mesh dimension should get only one process group, but got {self.get_rank()} "
|
|
f"in {subgroup_ranks}!"
|
|
)
|
|
dim_group_infos.append(
|
|
(
|
|
_get_group_tag(not_none(dim_group)),
|
|
subgroup_ranks,
|
|
dim_group.group_name,
|
|
)
|
|
)
|
|
self._dim_group_infos = dim_group_infos
|
|
|
|
def __enter__(self) -> "DeviceMesh":
|
|
# set this mesh as the current mesh in mesh env
|
|
_mesh_resources.mesh_stack.append(self)
|
|
return self
|
|
|
|
# pyre-fixme[2]: Parameter must be annotated.
|
|
def __exit__(self, exc_type, exc_value, exc_traceback) -> None:
|
|
# pop this mesh from mesh env
|
|
_mesh_resources.mesh_stack.pop()
|
|
|
|
def __repr__(self) -> str:
|
|
device_mesh_repr = (
|
|
f"DeviceMesh('{self.device_type}', {self.mesh.tolist()})"
|
|
if not self.mesh_dim_names
|
|
else f"DeviceMesh('{self.device_type}', {self.mesh.tolist()}, mesh_dim_names={self.mesh_dim_names})"
|
|
)
|
|
return device_mesh_repr
|
|
|
|
def __hash__(self):
|
|
# lazily compute hash
|
|
self._hash = getattr(self, "_hash", None)
|
|
if not self._hash:
|
|
self._hash = hash(
|
|
(
|
|
self._flatten_mesh_list,
|
|
self.mesh.shape,
|
|
self.device_type,
|
|
self.mesh_dim_names,
|
|
self._thread_id,
|
|
)
|
|
)
|
|
return self._hash
|
|
|
|
def __eq__(self, other: object) -> bool:
|
|
if not isinstance(other, DeviceMesh):
|
|
return False
|
|
if id(self) == id(other):
|
|
return True
|
|
else:
|
|
return (
|
|
self._flatten_mesh_list == other._flatten_mesh_list
|
|
and self.mesh.shape == other.mesh.shape
|
|
and self.device_type == other.device_type
|
|
and self.mesh_dim_names == other.mesh_dim_names
|
|
and self._thread_id == other._thread_id
|
|
)
|
|
|
|
def __getitem__(
|
|
self, mesh_dim_names: Union[str, Tuple[str, ...]]
|
|
) -> "DeviceMesh":
|
|
"""
|
|
Slice the current DeviceMesh based on the mesh_dim_names given to create a submesh.
|
|
The submesh created consists of the dimensions and the communicators indicated by
|
|
``mesh_dim_names``
|
|
|
|
Args:
|
|
mesh_dim_names (Union[str, Tuple[str]]): the name or the tuple of names of the
|
|
mesh dimension of the DeviceMesh to create the submesh for.
|
|
Returns:
|
|
A :class:`DeviceMesh` object
|
|
|
|
The following program runs on each process/rank in an SPMD manner in a world size of 8.
|
|
In the first example:
|
|
Calling mesh_2d["tp"] on rank 0, 1, 2, 3 returns a 1D submesh of DeviceMesh:([0, 1, 2, 3]).
|
|
Calling mesh_2d["tp"] on rank 4, 5, 6, 7 returns a 1D submesh of DeviceMesh:([4, 5, 6, 7]).
|
|
Calling mesh_2d["dp"] on rank 0, 4 returns a 1D submesh of DeviceMesh:([0, 4]).
|
|
Calling mesh_2d["dp"] on rank 1, 5 returns a 1D submesh of DeviceMesh:([1, 5]).
|
|
Calling mesh_2d["dp"] on rank 2, 6 returns a 1D submesh of DeviceMesh:([2, 6]).
|
|
Calling mesh_2d["dp"] on rank 3, 7 returns a 1D submesh of DeviceMesh:([3, 7]).
|
|
|
|
In the second example:
|
|
Calling mesh_3d["dp", "cp"] on rank 0, 1, 4, 5 returns a 2D submesh of DeviceMesh:([[0, 1], [4, 5]]).
|
|
Calling mesh_3d["dp", "cp"] on rank 2, 3, 6, 7 returns a 2D submesh of DeviceMesh:([[2, 3], [6, 7]]).
|
|
Calling mesh_3d["cp", "dp"] on rank 0, 1, 4, 5 returns a 2D submesh of DeviceMesh:([[0, 4], [1, 5]]).
|
|
Calling mesh_3d["cp", "dp"] on rank 2, 3, 6, 7 returns a 2D submesh of DeviceMesh:([[2, 6], [3, 7]]).
|
|
|
|
Example::
|
|
>>> # xdoctest: +SKIP("no rank")
|
|
>>> from torch.distributed.device_mesh import DeviceMesh
|
|
>>>
|
|
>>> # Initialize a 2D device mesh as (2, 4) to represent the topology
|
|
>>> # of cross-host(dim 0), and within-host (dim 1).
|
|
>>> mesh_2d = init_device_mesh(device_type="cuda", (2,4), mesh_dim_names=("dp", "tp"))
|
|
>>> tp_mesh = mesh_2d["tp"]
|
|
>>> dp_mesh = mesh_2d["dp"]
|
|
>>>
|
|
>>> # Initialize a 3D mesh.
|
|
>>> mesh_3d = init_device_mesh(device_type="cuda", (2,2,2), mesh_dim_names=("dp", "pp", "cp"))
|
|
>>> # The order of the mesh_dim_names provided deteremines the order of dimensions in the submesh.
|
|
>>> dp_cp_mesh = mesh_3d["dp", "cp"]
|
|
>>> cp_dp_mesh = mesh_3d["cp", "dp"]
|
|
"""
|
|
if not self.mesh_dim_names:
|
|
raise RuntimeError("Cannot slice a DeviceMesh without mesh_dim_names!")
|
|
|
|
mesh_dim_names = (
|
|
(mesh_dim_names,) if isinstance(mesh_dim_names, str) else mesh_dim_names
|
|
)
|
|
|
|
if mesh_dim_names == self.mesh_dim_names:
|
|
return self
|
|
else:
|
|
slice_mesh_dims = _mesh_resources._get_slice_mesh_dims(
|
|
self, mesh_dim_names
|
|
)
|
|
submesh = _mesh_resources.create_sub_mesh(
|
|
self, mesh_dim_names, slice_mesh_dims
|
|
)
|
|
return submesh
|
|
|
|
def get_group(self, mesh_dim: Optional[Union[int, str]] = None) -> ProcessGroup:
|
|
"""
|
|
Returns the single ProcessGroup specified by mesh_dim, or, if mesh_dim is not specified and the
|
|
DeviceMesh is 1-dimensional, returns the only ProcessGroup in the mesh.
|
|
|
|
Args:
|
|
mesh_dim (str/int, optional): it can be the name of the mesh dimension or the index
|
|
of the mesh dimension. Default is None.
|
|
|
|
Returns:
|
|
A :class:`ProcessGroup` object.
|
|
"""
|
|
if not hasattr(self, "_dim_group_infos"):
|
|
raise RuntimeError("DeviceMesh process groups not initialized!")
|
|
|
|
if self.mesh.ndim > 1 and mesh_dim is None:
|
|
raise RuntimeError(
|
|
f"Found the DeviceMesh have {self.mesh.ndim} dimensions",
|
|
"Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.",
|
|
"If you want to get the list of all the ProcessGroups in the DeviceMesh,"
|
|
"please use `get_all_groups()` instead.",
|
|
)
|
|
|
|
# Quick return if the current device_mesh is a 1D mesh.
|
|
if self.mesh.ndim == 1 and mesh_dim is None:
|
|
return not_none(
|
|
_find_pg_by_ranks_and_tag(*self._dim_group_infos[0][:2]) # type: ignore[index]
|
|
)
|
|
|
|
root_mesh = _mesh_resources.get_root_mesh(self)
|
|
root_to_flatten_mapping = _mesh_resources.root_to_flatten_mapping.get(
|
|
root_mesh, None
|
|
)
|
|
if root_to_flatten_mapping and mesh_dim in root_to_flatten_mapping.keys():
|
|
dim_group_infos = root_to_flatten_mapping[mesh_dim]._dim_group_infos[0][:2] # type: ignore[index]
|
|
return not_none(_find_pg_by_ranks_and_tag(*dim_group_infos))
|
|
else:
|
|
mesh_dim = (
|
|
_mesh_resources.get_mesh_dim_by_name(self, mesh_dim)
|
|
if isinstance(mesh_dim, str)
|
|
else mesh_dim
|
|
)
|
|
return not_none(
|
|
_find_pg_by_ranks_and_tag(*self._dim_group_infos[mesh_dim][:2]) # type: ignore[index]
|
|
)
|
|
|
|
def get_all_groups(self) -> List[ProcessGroup]:
|
|
"""
|
|
Returns a list of ProcessGroups for all mesh dimensions.
|
|
|
|
Returns:
|
|
A list of :class:`ProcessGroup` object.
|
|
"""
|
|
return [self.get_group(i) for i in range(self.mesh.ndim)]
|
|
|
|
@staticmethod
|
|
def from_group(
|
|
group: Union[ProcessGroup, List[ProcessGroup]],
|
|
device_type: str,
|
|
mesh: Optional[Union[torch.Tensor, "ArrayLike"]] = None,
|
|
*,
|
|
mesh_dim_names: Optional[Tuple[str, ...]] = None,
|
|
) -> "DeviceMesh":
|
|
"""
|
|
Constructs a :class:`DeviceMesh` with ``device_type`` from an
|
|
existing :class:`ProcessGroup`.
|
|
|
|
The constructed device mesh has number of dimensions equal to the
|
|
number of groups passed. If more than one group is passed, then the
|
|
``mesh`` argument is required.
|
|
"""
|
|
if isinstance(group, ProcessGroup):
|
|
group_ranks = get_process_group_ranks(group)
|
|
if (
|
|
isinstance(mesh, torch.Tensor) and mesh.tolist() != group_ranks
|
|
) or (mesh is not None and mesh != group_ranks):
|
|
raise ValueError(
|
|
f"Invalid mesh {str(mesh)} for ProcessGroup with ranks {group_ranks}"
|
|
)
|
|
mesh = torch.tensor(group_ranks, device="cpu", dtype=torch.int)
|
|
device_mesh = DeviceMesh(
|
|
device_type,
|
|
mesh,
|
|
mesh_dim_names=mesh_dim_names,
|
|
_init_backend=False,
|
|
)
|
|
device_mesh._dim_group_infos = [
|
|
(_get_group_tag(group), group_ranks, group.group_name)
|
|
]
|
|
return device_mesh
|
|
groups = list(group)
|
|
if len(groups) == 0:
|
|
raise ValueError("Expects at least one ProcessGroup to be passed")
|
|
if mesh is None:
|
|
raise ValueError("Must pass mesh if passing multiple ProcessGroups")
|
|
mesh = (
|
|
mesh.detach().to(dtype=torch.int, device="cpu")
|
|
if isinstance(mesh, torch.Tensor)
|
|
else torch.tensor(mesh, device="cpu", dtype=torch.int)
|
|
)
|
|
if mesh.ndim != len(groups):
|
|
raise ValueError(
|
|
"Expects mesh with ndim equal to number of ProcessGroups but got "
|
|
f"mesh {mesh.tolist()} and {len(groups)} ProcessGroups"
|
|
)
|
|
device_mesh = DeviceMesh(
|
|
device_type, mesh, mesh_dim_names=mesh_dim_names, _init_backend=False
|
|
)
|
|
device_mesh._dim_group_infos = [
|
|
(
|
|
_get_group_tag(group),
|
|
get_process_group_ranks(group),
|
|
group.group_name,
|
|
)
|
|
for group in groups
|
|
]
|
|
return device_mesh
|
|
|
|
def size(self, mesh_dim: Optional[int] = None) -> int:
|
|
return self.mesh.numel() if mesh_dim is None else self.mesh.size(mesh_dim)
|
|
|
|
@property
|
|
def ndim(self) -> int:
|
|
return self.mesh.ndim
|
|
|
|
@property
|
|
def shape(self) -> Tuple[int, ...]:
|
|
return tuple(self.mesh.shape)
|
|
|
|
def get_rank(self) -> int:
|
|
"""
|
|
Returns the current global rank.
|
|
"""
|
|
return get_rank()
|
|
|
|
def get_local_rank(self, mesh_dim: Optional[Union[int, str]] = None) -> int:
|
|
"""
|
|
Returns the local rank of the given mesh_dim of the DeviceMesh.
|
|
|
|
Args:
|
|
mesh_dim (str/int, optional): it can be the name of the mesh dimension or the index
|
|
of the mesh dimension. Default is None.
|
|
|
|
Returns:
|
|
An integer denotes the local rank.
|
|
|
|
The following program runs on each process/rank in an SPMD manner. In this example, we have 2
|
|
hosts with 4 GPUs each.
|
|
Calling mesh_2d.get_local_rank(mesh_dim=0) on rank 0, 1, 2, 3 would return 0.
|
|
Calling mesh_2d.get_local_rank(mesh_dim=0) on rank 4, 5, 6, 7 would return 1.
|
|
Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 0, 4 would return 0.
|
|
Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 1, 5 would return 1.
|
|
Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 2, 6 would return 2.
|
|
Calling mesh_2d.get_local_rank(mesh_dim=1) on rank 3, 7 would return 3.
|
|
|
|
Example::
|
|
>>> # xdoctest: +SKIP("no rank")
|
|
>>> from torch.distributed.device_mesh import DeviceMesh
|
|
>>>
|
|
>>> # Initialize device mesh as (2, 4) to represent the topology
|
|
>>> # of cross-host(dim 0), and within-host (dim 1).
|
|
>>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])
|
|
"""
|
|
if self.ndim > 1 and mesh_dim is None:
|
|
raise RuntimeError(
|
|
f"Found the DeviceMesh have {self.mesh.ndim} dimensions",
|
|
"Optional kwarg `mesh_dim` needs to be specified when device_mesh.ndim > 1.",
|
|
)
|
|
elif mesh_dim is None:
|
|
mesh_dim = 0
|
|
|
|
mesh_dim_group = not_none(self.get_group(mesh_dim))
|
|
assert isinstance(
|
|
mesh_dim_group, ProcessGroup
|
|
), "We expect ProcessGroup before calling `get_rank`!"
|
|
return not_none(get_rank(mesh_dim_group))
|
|
|
|
def get_coordinate(self) -> Optional[List[int]]:
|
|
"""
|
|
Return the relative indices of this rank relative to all
|
|
dimensions of the mesh. If this rank is not part of the mesh, return None.
|
|
"""
|
|
return self._coordinate_on_dim if self._coordinate_on_dim else None
|
|
|
|
def _flatten(self, mesh_dim_name: Optional[str] = None) -> "DeviceMesh":
|
|
"""
|
|
Returns a 1D DeviceMesh by flattening the current DeviceMesh.
|
|
|
|
If no mesh_dim_name is provided, the default is a string concatentaing the mesh_dim_names of the
|
|
given submesh with each mesh_dim_name separated by "_". For example, if we have a 3D mesh
|
|
DeviceMesh([[[0, 1], [2, 3]], [[4, 5], [6, 7]]], mesh_dim_names=("dp", "cp", "tp")), calling
|
|
mesh_3d["dp", "cp"]._flatten() will create a 1D submesh DeviceMesh([0, 1, 2, 3], mesh_dim_names=("dp_cp",))
|
|
on rank 0, 1, 2, 3 and a 1D submesh DeviceMesh([4, 5, 6, 7], mesh_dim_names=("dp_cp",)) on rank 4, 5, 6, 7.
|
|
|
|
After the flattened dimension is created, to access the flattened dimesnion in mesh_3d, one can use the
|
|
existing slicing method to obtain the flattened mesh through calling mesh_3d["dp_cp"].
|
|
"""
|
|
if not self.mesh_dim_names:
|
|
raise RuntimeError(
|
|
"Cannot flatten a DeviceMesh without mesh_dim_names!"
|
|
)
|
|
|
|
return _mesh_resources.create_flatten_mesh(self, mesh_dim_name)
|
|
|
|
def init_device_mesh(
|
|
device_type: str,
|
|
mesh_shape: Tuple[int, ...],
|
|
*,
|
|
mesh_dim_names: Optional[Tuple[str, ...]] = None,
|
|
) -> DeviceMesh:
|
|
"""
|
|
Initializes a `DeviceMesh` based on `device_type`, `mesh_shape`, and `mesh_dim_names` parameters.
|
|
|
|
This creates a DeviceMesh with an n-dimensional array layout, where `n` is the length of `mesh_shape`.
|
|
If `mesh_dim_names` is provided, each dimension is labeled as `mesh_dim_names[i]`.
|
|
|
|
.. note::
|
|
`init_device_mesh` follows SPMD programming model, meaning the same PyTorch Python program
|
|
runs on all processes/ranks in the cluster. Ensure `mesh_shape` (the dimensions of the nD array
|
|
describing device layout) is identical across all ranks. Inconsistent `mesh_shape` may lead to hanging.
|
|
|
|
.. note::
|
|
If no process group is found, init_device_mesh will initialize distributed process group/groups
|
|
required for distributed communications behind the scene.
|
|
|
|
Args:
|
|
device_type (str): The device type of the mesh. Currently supports: "cpu", "cuda/cuda-like".
|
|
Passing in a device type with a GPU index, such as "cuda:0", is not allowed.
|
|
mesh_shape (Tuple[int]): A tuple defining the dimensions of the multi-dimensional array
|
|
describing the layout of devices.
|
|
mesh_dim_names (Tuple[str], optional): A tuple of mesh dimension names to assign to each dimension
|
|
of the multi-dimensional array describing the layout of devices. Its length must match the length
|
|
of `mesh_shape`. Each string in `mesh_dim_names` must be unique.
|
|
|
|
Returns:
|
|
DeviceMesh: A :class:`DeviceMesh` object representing the device layout.
|
|
|
|
Example::
|
|
>>> # xdoctest: +SKIP("no rank")
|
|
>>> from torch.distributed.device_mesh import init_device_mesh
|
|
>>>
|
|
>>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,))
|
|
>>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp"))
|
|
|
|
"""
|
|
if mesh_dim_names is not None:
|
|
if len(set(mesh_dim_names)) != len(mesh_dim_names):
|
|
raise RuntimeError(
|
|
"Each mesh_dim_name must be unique.",
|
|
f"Found repeated mesh_dim_name in mesh_dim_names {mesh_dim_names}",
|
|
)
|
|
|
|
if len(mesh_shape) != len(mesh_dim_names):
|
|
raise RuntimeError(
|
|
"mesh_shape and mesh_dim_names should have same length!",
|
|
f"Found len(mesh_dim_names): {len(mesh_dim_names)} and len(mesh_shape):{len(mesh_shape)}.",
|
|
)
|
|
|
|
# assume valid device types are all letters
|
|
if device_type and not device_type.isalpha():
|
|
raise RuntimeError(
|
|
f"Device type with GPU index is not supported but got {device_type}. ",
|
|
"If you maintained a 'torch.device' object, it's recommended to pass in 'device.type'.",
|
|
)
|
|
|
|
# Always initialize the mesh's tensor on CPU, regardless of what the
|
|
# external device type has been set to be (e.g. meta)
|
|
with torch.device("cpu"):
|
|
mesh = torch.arange(math.prod(mesh_shape), dtype=torch.int).view(mesh_shape)
|
|
device_mesh = DeviceMesh(
|
|
device_type=device_type,
|
|
mesh=mesh,
|
|
mesh_dim_names=mesh_dim_names,
|
|
)
|
|
|
|
return device_mesh
|