def METHOD_NAME(rank, world_size):
    # params
    batch_size = 4
    num_heads = 4
    seq_length = 32
    attention_head_size = 32
    sub_seq_length = seq_length // world_size
    # create master tensors
    q = torch.rand(batch_size * num_heads, seq_length, attention_head_size).cuda()
    k = torch.rand(batch_size * num_heads, seq_length, attention_head_size).cuda()
    dist.broadcast(q, src=0, group=gpc.get_group(ParallelMode.SEQUENCE))
    dist.broadcast(k, src=0, group=gpc.get_group(ParallelMode.SEQUENCE))
    # create distributed tensors
    sub_q = q.clone()[:, rank * sub_seq_length:(rank + 1) * sub_seq_length].contiguous()
    sub_k = k.clone()[:, rank * sub_seq_length:(rank + 1) * sub_seq_length].contiguous()
    # set autograd attributes
    q.requires_grad = True
    k.requires_grad = True
    q.retain_grad()
    k.retain_grad()
    sub_q.requires_grad = True
    sub_k.requires_grad = True
    sub_q.retain_grad()
    sub_k.retain_grad()
    # compute master attention scores
    a = torch.matmul(q, k.transpose(2, 1))
    # compute distributed attention scores
    ring_qk = RingQK.apply
    sub_a = ring_qk(sub_q, sub_k, batch_size, num_heads, sub_seq_length)
    # check master and distributed attention scores
    sub_master_a = a[:, rank * sub_seq_length:(rank + 1) * sub_seq_length]
    assert torch.allclose(sub_a, sub_master_a, rtol=1e-5, atol=1e-2)
    # run master backward
    a.retain_grad()
    a.mean().backward()
    # run distributed backward
    partial_master_a_grad = a.grad[:, rank * sub_seq_length:(rank + 1) * sub_seq_length]
    torch.autograd.backward(sub_a, partial_master_a_grad)
    # check master and distributed grads
    partial_master_q_grad = q.grad[:, rank * sub_seq_length:(rank + 1) * sub_seq_length]
    assert torch.allclose(sub_q.grad, partial_master_q_grad, rtol=1e-5, atol=1e-2), \
        'attention score cannot match'