Repository: hpcaitech/ColossalAI
File: tests/test_shardformer/test_layer/test_ring_attn.py (Python)

Code before:
    # Check grads
    labels = torch.ones(out.shape[0], dtype=dtype, device=device)
    F.mse_loss(out.sum((-2, -1)), labels).backward()
    F.mse_loss(ring_out.sum((-2, -1)), labels[: ring_out.shape[0]]).backward()
    dq, dk, dv = [
        split_varlen_zigzag(
            qkv.grad[:, i], mask_info["cu_seqlens"] * sp_size, sp_group, mask_info["max_seqlen"] * sp_size
        )
        for i in range(3)
    ]
    dq_ring, dk_ring, dv_ring = [
        x.transpose(1, 2).reshape(-1, nheads, d)[mask_info["valid_indices"]]
        for x in (q_ring.grad, k_ring.grad, v_ring.grad)
    ]

    assert_close(dq, dq_ring, atol=atol, rtol=rtol)
    assert_close(dk, dk_ring, atol=atol, rtol=rtol)
    assert_close(dv, dv_ring, atol=atol, rtol=rtol)


def launch_single_ring(rank, world_size, port):
    colossalai.launch(rank, world_size, "localhost", port)
    check_packed_seq()
    check_ring_attn(sp_size=world_size)


def launch_double_ring(rank, world_size, port):
    colossalai.launch(rank, world_size, "localhost", port)
    check_ring_attn(sp_size=4, tp_size=2)


@rerun_if_address_is_in_use()
@parameterize("world_size", [2])
def test_ring_attn(world_size):
    spawn(launch_single_ring, nprocs=world_size)


@rerun_if_address_is_in_use()
@parameterize("world_size", [8])
def test_double_ring(world_size):
    spawn(launch_double_ring, nprocs=world_size)


if __name__ == "__main__":
    test_ring_attn()
    test_double_ring()


Review comment (suggestion):
Perhaps add a single test for the complext rank generation logic(https://github.com/hpcaitech/ColossalAI/blob/703bb5c18dd3a701209017de828e9cdc8ca58356/colossalai/shardformer/layer/attn.py#L483) for `world_size = 16`?
e.g. by putting rank generation into a single function `def get_double_ring_ranks(world_size, tp_size, sp_size, inner_ring_size, dp_pp_size)` and group creation after it.
Up to you🙂

Code after:
    assert_close(lse, ring_lse, atol=atol, rtol=rtol)
    assert_close(out, ring_out, atol=atol, rtol=rtol)

    # Check grads
    labels = torch.ones(out.shape[0], dtype=dtype, device=device)
    F.mse_loss(out.sum((-2, -1)), labels).backward()
    F.mse_loss(ring_out.sum((-2, -1)), labels[: ring_out.shape[0]]).backward()
    dq, dk, dv = [
        split_varlen_zigzag(
            qkv.grad[:, i], mask_info["cu_seqlens"] * sp_size, sp_group, mask_info["max_seqlen"] * sp_size
        )
        for i in range(3)
    ]
    dq_ring, dk_ring, dv_ring = [
        x.transpose(1, 2).reshape(-1, nheads, d)[mask_info["valid_indices"]]
        for x in (q_ring.grad, k_ring.grad, v_ring.grad)
    ]

    assert_close(dq, dq_ring, atol=atol, rtol=rtol)
    assert_close(dk, dk_ring, atol=atol, rtol=rtol)
    assert_close(dv, dv_ring, atol=atol, rtol=rtol)


def launch_single_ring(rank, world_size, port):
    colossalai.launch(rank, world_size, "localhost", port)
    check_packed_seq()
    check_ring_attn(inner_ring_size=None)


def launch_double_ring(rank, world_size, port):
    colossalai.launch(rank, world_size, "localhost", port)
    check_ring_attn(inner_ring_size=2)


@rerun_if_address_is_in_use()
@parameterize("world_size", [2])
def test_ring_attn(world_size):
    spawn(launch_single_ring, nprocs=world_size)


@rerun_if_address_is_in_use()
@parameterize("world_size", [4])
def test_double_ring(world_size):
    spawn(launch_double_ring, nprocs=world_size)


if __name__ == "__main__":
    test_ring_attn()
    test_double_ring()