Pytorch Collective Communications

import os
from typing import Callable

import torch
import torch.distributed as dist

def init_process(rank: int, world_size: int, fn: Callable([int, int], None], backend="gloo")):
    """Initialize the distributed environment."""
    os.environ["MASTER ADDR"] = "127.0.0.1"
    os.environ["MASTER_PORT"] = "12345"
    dist.init_process_group(backend, rank=rank, world_size=world_size, fn(rnk, size))

import torch.multiprocessing as mp

def func(rank: int, size: int):
    # each process will call this function
    continue

if __name__ == "__main__":
    size = 4
    processes = []
    mp.set_start_method("spawn")
    for rank in range(size):
        p = mp.Process(target=init_process, args=(rank, size, func))
        p.start()
        processes.append(p)

    for p in processes:
        p.join()

python -m torch.distributed.launch --nnodes 1 --proc_per_node 4 main.py

if __name__ == "__main__":
    init_process(
        rank = torch.distributed.get_rank()
        world_size = torch.distributed.get_world_size()
        fn = func
    )