如何解决CPU/GPU 设备之间的 Torch 分布式广播和减少
使用 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 Inds = [0,1] GPU0 数据 = [1,2] GPU1 Inds = [2,3] GPU1 数据 = [2,3]
-
在 GPU0 和 GPU1 上独立运行一个进程。为此,可以使用简单的 Add() 函数。
-
在必要时使用 GPU 数据更新 CPU 版本(对于 GPU 抓取的指标)。这是我可能会使用 reduce 从设备获取两个张量的地方。我可能会将它存储在一个键值字典中,其中 key 是设备 ID(GPU 0 为 0,GPU 1 为 1)并将 inds 和数据存储在一个元组中。然后我需要更新CPU版本并再次运行整个过程。
我计划使用 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 举报,一经查实,本站将立刻删除。