Dear @mrshenli, thanks very much for your response. Indeed that works.
Let me post my code adapted to your nice illustration that I actually need to work with as a part of a much longer implementation:
import torch
import torch.multiprocessing as mp
import os


def run(rank, groups, world_size):
    print("Rank ", rank, " Group ", groups[str(rank)])
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '29500'
    torch.cuda.set_device(groups[str(rank)]['gpu'])
    # global group
    torch.distributed.init_process_group(backend='nccl',
                                         rank=rank,
                                         world_size=world_size)
    my_group = torch.distributed.new_group(ranks=groups[str(rank)]['grp'])
    # tensor to bcast over group
    t = torch.tensor([1]).float().cuda().fill_(rank)
    torch.distributed.all_reduce(t, group=my_group)
    print('rank: {} - val: {}'.format(rank, t.item()))


def assign_groups(num_masters, workers_per_master, available_gpus):
    groups = {}
    distranks = num_masters
    gpu_allocated = 1
    for i in range(num_masters):
        my_group = [i]
        gpus = [available_gpus[0]]
        for j in range(workers_per_master):
            my_group.append(distranks)
            gpus.append(available_gpus[gpu_allocated])
            distranks += 1
            gpu_allocated += 1
        for r, g in zip(my_group, gpus):
            groups.update({
                str(r): {
                    'grp': my_group,
                    'gpu': g
                }
            })
    return groups


def main():
    num_masters = 3
    workers_per_master = 1
    available_gpus = [0, 1, 2, 3]
    groups = assign_groups(num_masters, workers_per_master, available_gpus)
    world_size = 6
    mp.spawn(run, args=(groups, world_size), nprocs=world_size, join=True)


if __name__ == "__main__":
    main()


Essentially, I have 3 master processes here and each of the masters has a worker. The masters with their respective workers make a group. Thus there are three groups altogether. Although in the above code the masters are assigned to cuda:0 and each of the workers to the next cuda devices, it may change depending on the setting. Thus it is immaterial in this illustration whether I comment out the line torch.cuda.set_device(groups[str(rank)]['gpu']).
Now, for sure for each of the masters and workers, I have its assigned group as groups[str(rank)]['grp']. Thus, when I run this code it should behave like those distributed communications being called concurrently as also illustrated by your example. However, it results in NCCL INFO Call to connect returned Connection refused.
Definitely I am doing something wrong here, not sure what.
Thanks again.