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

Code before:
                inner_ring_size = min(4, sp_size)
        else:
            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

        world_size = dist.get_world_size()
        rank = dist.get_rank()
        groups = int(world_size / sp_size)

        if tp_size > 1:
            for group_id in range(groups):
                for i in range(inner_ring_size):
                    ranks = list(range(i + (group_id * sp_size), (1 + group_id) * sp_size, inner_ring_size))
                    group = dist.new_group(ranks)
                    if rank in ranks:
                        inner_ring_group = group
            for group_id in range(groups):
                for i in range(num_rings):
                    ranks = list(range(i + group_id * num_rings, world_size, sp_size))
                    group = dist.new_group(ranks)
                    if rank in ranks:
                        inter_ring_group = group
        else:
            for i in range(sp_size // 2):
                ranks = list(range((i) * num_rings, (i + 1) * num_rings, 1))
                if rank in ranks:
                    print(
                        "rank:",
                        rank,
                        "inner ranks:",
                        ranks,
                    )
                    group = dist.new_group(ranks)
                    inner_ring_group = group
            for group_id in range(num_rings):
                for i in range(num_rings):

Review comment (refactor):
Maybe it should look like this?
```
world_size = 16
sp_size = 4
outer_ring_size = 2
inner_ring_size = sp_size // outer_ring_size
tp_size = 2
total_inner_rings = world_size // (outer_ring_size * tp_size) # loop through groups of size (inner_ring * tp) 
total_inner_size = tp_size * inner_ring_size
for j in range(total_inner_rings):
    # inside each group, duplicate tp group inner_sp_size times
    for k in range(tp_size):
        print(f"inner ring ranks: {list(range(k + j * total_inner_size, (j + 1) * total_inner_size, tp_size))}")
print("---------------------------------")


sp_tp_size = total_inner_size * outer_ring_size
n_groups = world_size // sp_tp_size # dp * pp
print(f"n_groups: {n_groups}")
for i in range(0, n_groups):
    start = i * sp_tp_size
    end = (i + 1) * sp_tp_size
    for j in range(outer_ring_size):    
        for k in range(tp_size):
            print(f"inter ring ranks: {list(range(start + k + j * tp_size, end, total_inner_size))}")
```
<img width="313" alt="image" src="https://github.com/user-attachments/assets/6eeacb11-3494-40e0-976c-7e22989fa31a">
<img width="579" alt="image" src="https://github.com/user-attachments/assets/54aa596f-4f60-4f60-8ffc-e0a5f3025c4b">


Code after:
        ), 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))
            group = pg_mesh.get_group_along_axis(sp_axis, ranks)
            if sp_rank in ranks:
                inter_ring_group = group

        return inner_ring_group, inter_ring_group

    @staticmethod
    def attention(
        q,  # (B, H, Sq, D)
        k,
        v,
        sp_axis,
        attention_mask_type,
        cu_seqlens=None,
        max_seqlen=None,
        valid_indices=None,
        dropout_p=0.0,
        softmax_scale=None,
        deterministic=False,
        return_softmax=False,
        inner_ring_size=None,
        pg_mesh=None,
        **kwargs,
    ):
        """