CPU/GPU 设备之间的 Torch 分布式广播和减少

编程问答 2022-05-30

问题描述

使用 torch.distributed 包。我正在尝试在两个单独的进程中从 cpu -> GPU0、GPU1 移动张量并更新主版本（在 cpu 上）。

假设我连接了两个 GPU。一个在 Device0 上，另一个在 Device1 上。

  1. 在 cpu 上存储一个非常大的数组（无法放入单个设备/gpu 的东西） 例如，X = [1,2,3,4,5,6]。

  2. 广播部分阵列到 GPU 设备 0 和 GPU 设备 1。0,1 具有该阵列的不同块。 GPU0 Inds = [0,1] GPU0 数据 = [1,2] GPU1 Inds = [2,3] GPU1 数据 = [2,3]

  3. 在 GPU0 和 GPU1 上独立运行一个进程。为此，可以使用简单的 Add() 函数。

  4. 在必要时使用 GPU 数据更新 cpu 版本（对于 GPU 抓取的指标）。这是我可能会使用 reduce 从设备获取两个张量的地方。我可能会将它存储在一个键值字典中，其中 key 是设备 ID（GPU 0 为 0，GPU 1 为 1）并将 inds 和数据存储在一个元组中。然后我需要更新cpu版本并再次运行整个过程。

我正在尝试做的事情可以在这张图表中看到。

enter image description here

我计划使用 Nccl 后端，它显然 supports 广播和减少。

我的代码应该是这样的：

Main() 函数产生两个进程并保存 cpu 张量 Foo() 启动两个进程并允许在它们之间进行广播和更新（我想在图中做什么）

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,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,6] which is what it should be

解决方法

暂无找到可以解决该程序问题的有效方法，小编努力寻找整理中！

如果你已经找到好的解决方法，欢迎将解决方案带上本链接一起发送给小编。

小编邮箱:dio#foxmail.com (将#修改为@）