reduced test case for all_reduce failure

all_reduce_bug.py

import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch_ucc

def worker(rank, world_size):
    os.environ["RANK"] = str(rank)
    os.environ["WORLD_SIZE"] = str(world_size)
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "10638"

    dist.init_process_group("ucc", rank=rank, world_size=world_size)

    device = torch.device("cuda", rank)

    t = torch.ones(2, device=device)

    dist.all_reduce(t)

    torch.cuda.synchronize()


if "OMPI_COMM_WORLD_RANK" in os.environ:
    print(f"hello test yay")
    worker(int(os.environ["OMPI_COMM_WORLD_RANK"]), int(os.environ["OMPI_COMM_WORLD_SIZE"]))

elif __name__ == "__main__":
    world_size = 8
    mp.spawn(worker, args=(world_size,), nprocs=world_size, join=True)