Hi @mrshenli, yes, I did that for each process. Don’t understand where is the mistake. Here is the sample code:
import os
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import time
import datetime

def run(rank, a, q):
   os.environ["NCCL_BLOCKING_WAIT"] = "1"
   print('Rank ', rank, 'NCCL_BLOCKING_WAIT is: ', os.environ["NCCL_BLOCKING_WAIT"])
   dist_init_method = 'tcp://{master_ip}:{master_port}'.format(
               master_ip='127.0.0.1', master_port='12346')
   torch.distributed.init_process_group(backend="nccl",
                                       init_method=dist_init_method,
                                       timeout=datetime.timedelta(seconds=5),
                                       world_size=4,
                                       rank=rank)
   group1 = dist.new_group([0, 1])
   group2 = dist.new_group([2, 3])
   tensor = torch.ones(1)

   device = torch.device('cuda', rank)
   tensor = tensor.to(device)

   while not q.empty():
       print('Rank ', rank, ' in the loop ')
       current_index = q.get()
       print(f'Process {rank} current index is: {current_index}')
       try:
           if rank == 0 or rank == 1:
               dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group1)
               print(f'Process {rank} all_reduce tensor is: {tensor}')
           else:
               dist.all_reduce(tensor, op=dist.ReduceOp.SUM, group=group2)
               print(f'Process {rank} all_reduce tensor is: {tensor}')
       except Exception:
           pass
   print('Rank ', rank, ' has data ', tensor[0])

if __name__ == "__main__":
   a = 1
   ctx = mp.get_context('spawn')
   q = ctx.Queue()
   flag = ctx.Queue()
   for index in range(5):
       q.put(index)
   
   mp.spawn(run, args=(a, q), nprocs=4)

An update, I tried “gloo” without setting timeout, and it can terminate properly. I’m wondering maybe “gloo” takes care of this situation by itself? It doesn’t have something to do with the “timeout”?