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

Posted 2023-03-29

【中文标题】CPU/GPU 设备之间的 Torch 分布式广播和减少

【英文标题】：Torch distributed broadcast and reduce between CPU/ GPU devices

【发布时间】：2021-04-15 21:50:16

【问题描述】：

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

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

在 CPU 上存储一个非常大的数组（不适合单个设备/gpu 的东西） 例如 X = [1,2,3,4,5,6]。

将数组的一部分广播到 GPU 设备 0 和 GPU 设备 1。0,1 具有该数组的不同块。 GPU0 数据 = [0,1] GPU0 数据 = [1,2] GPU1 指标 = [2,3] GPU1 数据 = [2,3]

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

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

我想要做的可以在这张图中看到。

我计划使用显然supports 广播和减少的 NCCL 后端。

我的代码应该是这样的：

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,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

【参考方案1】：

从外观上看，您想使用三个副本：一个在 CPU 上，两个在 GPU 上。如果是这种情况，您可能需要与nccl 不同的后端。请看https://pytorch.org/docs/stable/distributed.html