mirror of https://github.com/hpcaitech/ColossalAI
204 lines
8.5 KiB
Python
204 lines
8.5 KiB
Python
import itertools
|
|
from functools import reduce
|
|
from operator import mul
|
|
from typing import Dict, List, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
import torch.distributed as dist
|
|
from torch.distributed import ProcessGroup
|
|
|
|
|
|
def prod(nums: List[int]) -> int:
|
|
"""Product of a list of numbers.
|
|
|
|
Args:
|
|
nums (List[int]): A list of numbers.
|
|
|
|
Returns:
|
|
int: The product of the numbers.
|
|
"""
|
|
return reduce(mul, nums)
|
|
|
|
|
|
class ProcessGroupMesh:
|
|
"""A helper class to manage the process group mesh. It only describes how to organize process groups, and it's decoupled with parallel method.
|
|
It just initialize process groups and cache them. The parallel method should manage them and use them to do the parallel computation.
|
|
|
|
We use a ND-tuple to represent the process group mesh. And a ND-coordinate is to represent each process.
|
|
For example, ``(0, 1, 0)`` represents the process whose rank is 2 in a 3D process group mesh with size ``(2, 2, 2)``.
|
|
|
|
Args:
|
|
*size (int): The size of each dimension of the process group mesh. The product of the size must be equal to the world size.
|
|
|
|
Attributes:
|
|
shape (Tuple[int, ...]): The shape of the process group mesh.
|
|
rank (int): The rank of the current process.
|
|
"""
|
|
|
|
def __init__(self, *size: int) -> None:
|
|
assert dist.is_initialized(), "Please initialize torch.distributed first."
|
|
assert prod(size) == dist.get_world_size(), "The product of the size must be equal to the world size."
|
|
self._shape = size
|
|
self._rank = dist.get_rank()
|
|
self._coord = ProcessGroupMesh.unravel(self._rank, self._shape)
|
|
self._ranks_to_group: Dict[Tuple[int, ...], ProcessGroup] = {}
|
|
self._group_to_ranks: Dict[ProcessGroup, Tuple[int, ...]] = {}
|
|
|
|
@property
|
|
def shape(self) -> Tuple[int, ...]:
|
|
return self._shape
|
|
|
|
@property
|
|
def rank(self) -> int:
|
|
return self._rank
|
|
|
|
def size(self, dim: Optional[int] = None) -> Union[int, Tuple[int, ...]]:
|
|
"""Get the size of the process group mesh.
|
|
|
|
Args:
|
|
dim (Optional[int], optional): Dimension of the process group mesh. `None` means all dimensions. Defaults to None.
|
|
|
|
Returns:
|
|
Union[int, Tuple[int, ...]]: Size of the target dimension or the whole process group mesh.
|
|
"""
|
|
if dim is None:
|
|
return self._shape
|
|
else:
|
|
return self._shape[dim]
|
|
|
|
def coordinate(self, dim: Optional[int] = None) -> Union[int, Tuple[int, ...]]:
|
|
"""Get the coordinate of the process group mesh.
|
|
|
|
Args:
|
|
dim (Optional[int], optional): Dimension of the process group mesh. `None` means all dimensions. Defaults to None.
|
|
|
|
Returns:
|
|
Union[int, Tuple[int, ...]]: Coordinate of the target dimension or the whole process group mesh.
|
|
"""
|
|
if dim is None:
|
|
return self._coord
|
|
else:
|
|
return self._coord[dim]
|
|
|
|
@staticmethod
|
|
def unravel(rank: int, shape: Tuple[int, ...]) -> Tuple[int, ...]:
|
|
"""Convert a rank to a coordinate.
|
|
|
|
Args:
|
|
rank (int): Rank to be converted.
|
|
shape (Tuple[int, ...]): Shape of the process group mesh.
|
|
|
|
Returns:
|
|
Tuple[int, ...]: Coordinate of the rank.
|
|
"""
|
|
return np.unravel_index(rank, shape)
|
|
|
|
@staticmethod
|
|
def ravel(coord: Tuple[int, ...], shape: Tuple[int, ...]) -> int:
|
|
"""Convert a coordinate to a rank.
|
|
|
|
Args:
|
|
coords (Tuple[int, ...]): Coordinate to be converted.
|
|
shape (Tuple[int, ...]): Shape of the process group mesh.
|
|
|
|
Returns:
|
|
int: Rank of the coordinate.
|
|
"""
|
|
return np.ravel_multi_index(coord, shape)
|
|
|
|
def get_group(self, ranks_in_group: List[int], backend: Optional[str] = None) -> ProcessGroup:
|
|
"""Get the process group with the given ranks. It the process group doesn't exist, it will be created.
|
|
|
|
Args:
|
|
ranks_in_group (List[int]): Ranks in the process group.
|
|
backend (Optional[str], optional): Backend of the process group. Defaults to None.
|
|
|
|
Returns:
|
|
ProcessGroup: The process group with the given ranks.
|
|
"""
|
|
ranks_in_group = sorted(ranks_in_group)
|
|
if tuple(ranks_in_group) not in self._group_to_ranks:
|
|
group = dist.new_group(ranks_in_group, backend=backend)
|
|
self._ranks_to_group[tuple(ranks_in_group)] = group
|
|
self._group_to_ranks[group] = tuple(ranks_in_group)
|
|
return self._ranks_to_group[tuple(ranks_in_group)]
|
|
|
|
def get_ranks_in_group(self, group: ProcessGroup) -> List[int]:
|
|
"""Get the ranks in the given process group. The process group must be created by this class.
|
|
|
|
Args:
|
|
group (ProcessGroup): The process group.
|
|
|
|
Returns:
|
|
List[int]: Ranks in the process group.
|
|
"""
|
|
return list(self._group_to_ranks[group])
|
|
|
|
@staticmethod
|
|
def get_coords_along_axis(base_coord: Tuple[int, ...], axis: int,
|
|
indices_at_axis: List[int]) -> List[Tuple[int, ...]]:
|
|
"""Get coordinates along the given axis.
|
|
|
|
Args:
|
|
base_coord (Tuple[int, ...]): Base coordinate which the coordinates along the axis are based on.
|
|
axis (int): Axis along which the coordinates are generated.
|
|
indices_at_axis (List[int]): Indices at the axis.
|
|
|
|
Returns:
|
|
List[Tuple[int, ...]]: Coordinates along the axis.
|
|
"""
|
|
coords_in_group = []
|
|
for idx in indices_at_axis:
|
|
coords_in_group.append(base_coord[:axis] + (idx,) + base_coord[axis + 1:])
|
|
return coords_in_group
|
|
|
|
def create_group_along_axis(self,
|
|
axis: int,
|
|
indices_at_axis: Optional[List[int]] = None,
|
|
backend: Optional[str] = None) -> ProcessGroup:
|
|
"""Create all process groups along the given axis, and return the one which the current process belongs to.
|
|
|
|
Args:
|
|
axis (int): Axis along which the process groups are created.
|
|
indices_at_axis (Optional[List[int]], optional): Indices at the axis. Defaults to None.
|
|
backend (Optional[str], optional): Backend of the process group. Defaults to None.
|
|
|
|
Returns:
|
|
ProcessGroup: The process group along the given axis which the current process belongs to.
|
|
"""
|
|
indices_at_axis = indices_at_axis or list(range(self._shape[axis]))
|
|
reduced_shape = list(self._shape)
|
|
# the choices on the axis are reduced to 1, since it's determined by `indices_at_axis`
|
|
reduced_shape[axis] = 1
|
|
target_group = None
|
|
# use Cartesian product to generate all combinations of coordinates
|
|
for base_coord in itertools.product(*[range(s) for s in reduced_shape]):
|
|
coords_in_group = ProcessGroupMesh.get_coords_along_axis(base_coord, axis, indices_at_axis)
|
|
ranks_in_group = tuple([ProcessGroupMesh.ravel(coord, self._shape) for coord in coords_in_group])
|
|
group = self.get_group(ranks_in_group, backend=backend)
|
|
if self._rank in ranks_in_group:
|
|
target_group = group
|
|
return target_group
|
|
|
|
def get_group_along_axis(self,
|
|
axis: int,
|
|
indices_at_axis: Optional[List[int]] = None,
|
|
backend: Optional[str] = None) -> ProcessGroup:
|
|
"""Get the process group along the given axis which the current process belongs to. If the process group doesn't exist, it will be created.
|
|
|
|
Args:
|
|
axis (int): Axis along which the process groups are created.
|
|
indices_at_axis (Optional[List[int]], optional): Indices at the axis. Defaults to None.
|
|
backend (Optional[str], optional): Backend of the process group. Defaults to None.
|
|
|
|
Returns:
|
|
ProcessGroup: The process group along the given axis which the current process belongs to.
|
|
"""
|
|
indices_at_axis = indices_at_axis or list(range(self._shape[axis]))
|
|
coords_in_group = ProcessGroupMesh.get_coords_along_axis(self._coord, axis, indices_at_axis)
|
|
ranks_in_group = tuple([ProcessGroupMesh.ravel(coord, self._shape) for coord in coords_in_group])
|
|
if ranks_in_group not in self._ranks_to_group:
|
|
# no need to cache it explicitly, since it will be cached in `create_group_along_axis`
|
|
return self.create_group_along_axis(axis, indices_at_axis, backend=backend)
|
|
return self._ranks_to_group[ranks_in_group]
|