title: Torch distributed broadcast and reduce between CPU/ GPU devices
body: <p>Using the <a href="https://pytorch.org/docs/stable/distributed.html#module-torch.distributed" rel="nofollow noreferrer">torch.distributed</a> package. I am trying to move tensors from CPU -&gt; GPU0, GPU1 in two separate processes and update the master version (on CPU).</p>
  <p>Assume I have two GPU's connected. One on Device0, the other on Device1.</p>
  <ol>
  <li><p>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.</p>
  </li>
  <li><p><strong>Broadcast</strong> 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]</p>
  </li>
  <li><p>Run a process on GPU0 and GPU1 independently. For this purpose a simple Add() function will do.</p>
  </li>
  <li><p>Update the CPU version with GPU data where necessary (for the inds that GPU grabbed). This is where I would probably use <a href="https://pytorch.org/docs/stable/distributed.html#module-torch.distributed" rel="nofollow noreferrer">reduce</a> 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.</p>
  </li>
  </ol>
  <p>What I am trying to do can be seen in this diagram.<a href="https://i.stack.imgur.com/iyqKw.jpg" rel="nofollow noreferrer"><img src="https://i.stack.imgur.com/iyqKw.jpg" alt="enter image description here" /></a></p>
  <p>I plan on using an NCCL backend which apparently <a href="https://pytorch.org/docs/stable/distributed.html#module-torch.distributed" rel="nofollow noreferrer">supports</a> broadcast and reduce.</p>
  <p>My code should look something like this:</p>
  <p>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)</p>
  <pre><code>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) &gt;&gt; is not [2,3,4,5,5,6] which is what it should be
  </code></pre>