Class ProcessGroupMesh

colossalai/cluster/process_group_mesh.py:25–276

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

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."
        world_size = dist.get_world_size()
        prod_size = prod(size)
        assert (
            prod_size == world_size
        ), f"The product of the size({prod_size}) must be equal to the world size({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, ...], Union[ProcessGroup, GroupMember.NON_GROUP_MEMBER]] = {}
        self._group_to_ranks: Dict[ProcessGroup, Tuple[int, ...]] = {}

    def destroy_mesh_process_groups(self):
        r"""
        Destructor method for the ProcessGroupMesh class.

        When the ProcessGroupMesh object is deleted or goes out of scope, this method is called. It is responsible for
        cleaning up any process groups that were created during the lifetime of the object.

        Note:
            All process groups in PyTorch are represented as global variables, and they may not be automatically destroyed
            when the ProcessGroupMesh's lifetime ends. This method manually destroys the process groups to release
            system resources.
        """
        for group in self._ranks_to_group.values():
            try:
                dist.destroy_process_group(group)
            except ValueError:
                pass

        # Manually clear all process groups to save memory
        gc.collect()

    @property
    def shape(self) -> Tuple[int, ...]:
        return self._shape

    @property
    def rank(self) -> int:
        return self._rank