1

使用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 数据 = [0,1] GPU0 数据 = [1,2] GPU1 数据 = [2,3] GPU1 数据 = [2,3]

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

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

我正在尝试做的事情可以在这张图中看到。在此处输入图像描述

我计划使用显然支持广播和减少的 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
4

1 回答 1

0

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

于 2021-11-01T18:54:28.260 回答