Torch distributed broadcast and reduce between CPU/ GPU devices

Using the torch.distributed package. I am trying to move tensors from CPU -> GPU0, GPU1 in two separate processes and update the master version (on CPU).

Assume I have two GPU's connected. One on Device0, the other on Device1.

1. Store a very large array on CPU (something that can't fit onto a single device/gpu) X = [1,2,3,4,5,6] for example.

2. Broadcast part of the array to GPU device 0 and GPU device 1. 0,1 have different chunks of that array. GPU0 Inds = [0,1] GPU0 data = [1,2] GPU1 Inds = [2,3] GPU1 data = [2,3]

3. Run a process on GPU0 and GPU1 independently. For this purpose a simple Add() function will do.

4. Update the CPU version with GPU data where necessary (for the inds that GPU grabbed). This is where I would probably use reduce to get both tensors from the devices. I would probably store it in a key-value dict where key is the device ID (0 for GPU 0, 1 for GPU 1) and have the inds and data stored in a tuple. Then I need to update the CPU version and run the whole process again.

What I am trying to do can be seen in this diagram.

I plan on using an NCCL backend which apparently supports broadcast and reduce.

My code should look something like this:

Main() function spawns two processes and holds the cpu tensor Foo() starts the two processes and allows broadcasting, and updating between them (what I want to do in diagram)

def main():
  args.world_size=2
  args.backend='nccl'
  os.environ['MASTER_ADDR'] = '127.0.0.2'
  os.environ['MASTER_PORT'] = '29500'
  args.data = t.Tensor([1,2,3,4,5,6])
  mp.spawn(foo, nprocs=2, args=(args,), join=True)

def foo(rank):
  data = args.data # cpu data
  dist.init_process_group(backend=args.backend, init_method='env://', world_size=args.world_size, rank=rank)
  if rank==0:
    inds=[0,1]
  elif rank == 1:
    inds=[2,3]
  
  gpu_data = data[inds].cuda(rank) # will send to GPU0 or GPU1. # probably need to use the torch.dist.broadcast operation here but idk how.
  data[inds].data=gpu.data # this would be the update step. 
  
  gpu_data +=1
  dist.destroy_process_group()
  print('data: ', data) >> is not [2,3,4,5,5,6] which is what it should be