[Bugfix] Pad hidden_states to avoid cross-ring AllGatherV

[ApsarasX]

What this PR does / why we need it?

On the 910B2C hardware platform, each machine is equipped with 16 NPU cards. When using Tensor Parallelism (TP)=16, if the num_tokens (input sequence length) is not divisible by tp_size, it triggers cross-ring AllGatherV communication operations, which may result in unexpected errors(see figure below).

Does this PR introduce any user-facing change?

No

How was this patch tested?

[ganyi1996ppo]

Have you tried torch.distributed.all_gather and cat the data collected from all rank?

[ApsarasX]

Have you tried torch.distributed.all_gather and cat the data collected from all rank?

import torch
import torch_npu
import torch.distributed as dist
import torch.multiprocessing as mp

def verify_tp_o_method(rank, world_size: int, hidden_statses: list[torch.Tensor]):
    torch.npu.set_device(rank)

    dist.init_process_group(
        backend='hccl',
        init_method='tcp://127.0.0.1:55223',
        world_size=world_size,
        rank=rank
    )

    hidden_statses = hidden_statses.to(f"npu:{rank}")

    chunk_hidden_states = torch.tensor_split(hidden_statses, world_size, dim=0)

    router_hidden_states = chunk_hidden_states[rank]

    print(f"[rank={rank}] router_hidden_states.shape = {router_hidden_states.shape}")

    dist.all_gather(list(chunk_hidden_states), router_hidden_states)

    dist.barrier()

    dist.destroy_process_group()

def main():
    world_size = 16
    inter_dim = 16369
    # inter_dim = 16384
    output_dim = 7168
    hidden_statses = torch.randn((inter_dim, output_dim), dtype=torch.half)

    mp.spawn(verify_tp_o_method, args=(world_size, hidden_statses,), nprocs=world_size, join=True)

if __name__ == "__main__":
    main()

The above code will cause the following error:

[ganyi1996ppo]

import torch
import torch_npu
import torch.distributed as dist
import torch.multiprocessing as mp

def verify_tp_o_method(rank, world_size: int, hidden_statses: list[torch.Tensor]):
    torch.npu.set_device(rank)

    dist.init_process_group(
        backend='hccl',
        init_method='tcp://127.0.0.1:55223',
        world_size=world_size,
        rank=rank
    )

    hidden_statses = hidden_statses.to(f"npu:{rank}")

    chunk_hidden_states = torch.tensor_split(hidden_statses, world_size, dim=0)

    router_hidden_states = chunk_hidden_states[rank]

    print(f"[rank={rank}] router_hidden_states.shape = {router_hidden_states.shape}")

    dist.all_gather(list(chunk_hidden_states), router_hidden_states)

    dist.barrier()

    dist.destroy_process_group()

def main():
    world_size = 16
    inter_dim = 16369
    # inter_dim = 16384
    output_dim = 7168
    hidden_statses = torch.randn((inter_dim, output_dim), dtype=torch.half)

    mp.spawn(verify_tp_o_method, args=(world_size, hidden_statses,), nprocs=world_size, join=True)

if __name__ == "__main__":
    main()

The above code will cause the following error:

Strange, this works fine on my machine.....

[ApsarasX]

@ganyi1996ppo What progress about this PR?

[github-actions]

This pull request has conflicts, please resolve those before we can evaluate the pull request.

[jianzs]

import torch
import torch_npu
import torch.distributed as dist
import torch.multiprocessing as mp

def verify_tp_o_method(rank, world_size: int, hidden_statses: list[torch.Tensor]):
    torch.npu.set_device(rank)

    dist.init_process_group(
        backend='hccl',
        init_method='tcp://127.0.0.1:55223',
        world_size=world_size,
        rank=rank
    )

    hidden_statses = hidden_statses.to(f"npu:{rank}")

    chunk_hidden_states = torch.tensor_split(hidden_statses, world_size, dim=0)

    router_hidden_states = chunk_hidden_states[rank]

    print(f"[rank={rank}] router_hidden_states.shape = {router_hidden_states.shape}")

    dist.all_gather(list(chunk_hidden_states), router_hidden_states)

    dist.barrier()

    dist.destroy_process_group()

def main():
    world_size = 16
    inter_dim = 16369
    # inter_dim = 16384
    output_dim = 7168
    hidden_statses = torch.randn((inter_dim, output_dim), dtype=torch.half)

    mp.spawn(verify_tp_o_method, args=(world_size, hidden_statses,), nprocs=world_size, join=True)

if __name__ == "__main__":
    main()

The above code will cause the following error:

Strange, this works fine on my machine.....

Is it because your device have only 8 cards instead of 16?