Repository: hpcaitech/ColossalAI
File: colossalai/shardformer/layer/attn.py (Python)

Code before:
    INNER_RING_GROUP: dist.ProcessGroup = None
    INNER_RING_GROUP_COPY: dist.ProcessGroup = None
    INTER_RING_GROUP: dist.ProcessGroup = None
    INTER_RING_GROUP_COPY: dist.ProcessGroup = None

    @staticmethod
    def get_double_ring_groups(sp_group, inner_ring_size=None):
        """
        Get 2D ring groups for the given process group. Generally, to avoid congestion, the inner ring size
        shouldn't be larger than the number of NICs on each node.
        Args:
            sp_group (dist.ProcessGroup): Process group for sequence parallelism
            inner_ring_size (Optional[int], optional): Inner ring size. Defaults to None.
        Returns:
            Tuple[dist.ProcessGroup, dist.ProcessGroup]: Inner-ring process group and inter-ring process group.
        """
        sp_size = dist.get_world_size(sp_group)
        sp_rank = dist.get_rank(sp_group)

        assert inner_ring_size is not None

        assert (
            inner_ring_size <= sp_size and sp_size % inner_ring_size == 0
        ), f"Error: sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"

        if inner_ring_size == sp_size:
            return sp_group, sp_group
        assert (
            sp_size % inner_ring_size == 0
        ), f"sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"
        logger = get_dist_logger()
        logger.info(
            f"Using 2D Ring Attention with inner ring size {inner_ring_size} to maximze NIC util for inter-node comm. Cross your fingers for speed-ups!",
            ranks=[0],
        )
        num_rings = sp_size // inner_ring_size
        inner_ring_group = None
        inter_ring_group = None

        inter_axis, inner_axis = 0, 1
        pg_mesh = ProcessGroupMesh(num_rings, inner_ring_size)

        # Create inner ring groups
        for i in range(inner_ring_size):
            ranks = list(range(i * inner_ring_size, (i + 1) * inner_ring_size))
            group = pg_mesh.get_group_along_axis(inner_axis)
            if sp_rank in ranks:
                inner_ring_group = group

        # Create inter ring groups
        for i in range(num_rings):

Review comment (suggestion):
We can keep this rule of thumb logic for easy topology construction? Users most likely have no idea what "double ring" means and won't touch `inner_ring_size` at all. Thanks.

Code after:
    INNER_RING_GROUP_COPY: dist.ProcessGroup = None
    INTER_RING_GROUP: dist.ProcessGroup = None
    INTER_RING_GROUP_COPY: dist.ProcessGroup = None

    @staticmethod
    def get_double_ring_groups(sp_axis, pg_mesh, inner_ring_size=None):
        """
        Get 2D ring groups for the given process group. Generally, to avoid congestion, the inner ring size
        shouldn't be larger than the number of NICs on each node.
        Args:
            sp_group (dist.ProcessGroup): Process group for sequence parallelism
            inner_ring_size (Optional[int], optional): Inner ring size. Defaults to None.
        Returns:
            Tuple[dist.ProcessGroup, dist.ProcessGroup]: Inner-ring process group and inter-ring process group.
        """
        assert pg_mesh is not None, f"Error: The pg mesh is None! please check the process group initialization."

        sp_group = pg_mesh.get_group_along_axis(sp_axis)
        sp_size = dist.get_world_size(sp_group)
        sp_rank = dist.get_rank(sp_group)

        assert inner_ring_size is not None

        assert (
            inner_ring_size <= sp_size and sp_size % inner_ring_size == 0
        ), f"Error: sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"

        if inner_ring_size == sp_size:
            return sp_group, sp_group
        assert (
            sp_size % inner_ring_size == 0
        ), f"sp_size {sp_size} should be divisible by inner_ring_size {inner_ring_size}"
        logger = get_dist_logger()
        logger.info(
            f"Using 2D Ring Attention with inner ring size {inner_ring_size} to maximze NIC util for inter-node comm. Cross your fingers for speed-ups!",
            ranks=[0],
        )
        num_rings = sp_size // inner_ring_size
        inner_ring_group = None
        inter_ring_group = None

        # Create inner ring groups
        for i in range(inner_ring_size):
            ranks = list(range(i * inner_ring_size, (i + 1) * inner_ring_size))
            group = pg_mesh.get_group_along_axis(sp_axis, ranks)
            if sp_rank in ranks:
                inner_ring_group = group

        # Create inter ring groups
        for i in range(num_rings):
            ranks = list(range(i, sp_size, num_rings))