Profiling `init_process_group('nccl')`

benchmark_dist_init.py

# run as:
# python -u -m torch.distributed.run --nproc_per_node=8 --rdzv_endpoint localhost:6000  --rdzv_backend c10d benchmark_dist_init.py

import torch
import os
import cProfile

import torch.distributed as dist
import timeit

def init():
    dist.init_process_group('nccl')

local_rank = int(os.environ["LOCAL_RANK"])

if local_rank == 0:
    cProfile.run("init()", sort="tottime")
else:
    init()