fix dist attn reshape error

[tkdcjf159]

By default, DeepSpeed's DistributedAttention is set with scatter_idx = 2 and gather_idx = 0. However, if I set gather_idx to 1 and have a batch size greater than 1, an error will occur during the output all to all operation, as illustrated below. To fix this, modify the seq_world_size to -1.

def single_all_to_all(input, scatter_idx, gather_idx, group):
    # Assume input shape [2, 1024, 8, 16], scatter_idx = 1, gather_idx=2, seq_world_size=8
    seq_world_size = dist.get_world_size(group)
    inp_shape = list(input.shape) # inp_shape = [2, 1024, 8, 16]
    inp_shape[scatter_idx] = inp_shape[scatter_idx] // seq_world_size # inp_shape = [2, 128, 8, 16]
    if scatter_idx < 2:
        # Reshaping from [2, 1024, 8, 16] to [8, 128, 8, 16]: ERROR! (2 * 1024 * 8 * 16) != (8 * 128 * 8 * 16)
        # Use -1 to fix issue
        input_t = input.reshape(
            [-1, inp_shape[scatter_idx]] + \
            # [seq_world_size, inp_shape[scatter_idx]] + \
            inp_shape[scatter_idx + 1:]
        ).contiguous()
    else:
        # Transpose groups of heads with the seq-len parallel dimension to scatter them
        input_t = input.reshape(
            [-1, seq_world_size, inp_shape[scatter_idx]] + \
            inp_shape[scatter_idx + 1:]
        ).transpose(0, 1).contiguous()

    output = torch.empty_like(input_t)
    dist.all_to_all_single(output, input_t, group=group)

    # If scattering the seq-dim, transpose the heads back to the original dimension
    if scatter_idx < 2:
        output = output.transpose(0, 1).contiguous()

    return output.reshape(
        inp_shape[: gather_idx] + \
        [inp_shape[gather_idx] * seq_world_size,] + \
        inp_shape[gather_idx + 1:]).contiguous()

[tkdcjf159]

@microsoft-github-policy-service agree company="Upstage"

[loadams]

@tkdcjf159 - if you're still interested in this PR could you resolve the conflicts and we will get it reviewed?

[loadams]

@tkdcjf159 - closing this PR as this code has been refactored, if you believe this bug still remains, please comment and we can re-open this PR.