将 torch.distributed.rpc.rpc_aync 放在不同的 .py 文件时，结果会有所不同

Posted 2023-02-16

【中文标题】将 torch.distributed.rpc.rpc_aync 放在不同的 .py 文件时，结果会有所不同

【英文标题】：The results is different when placing torch.distributed.rpc.rpc_aync at a different .py file

【发布时间】：2021-12-06 03:38:43

【问题描述】：

我想在worker端执行一个函数并将结果返回给master。但是，我发现将 rpc_async 放在不同的 .py 文件时结果不同

方法一

master.py：

import os
import torch
import torch.distributed.rpc as rpc
from torch.distributed.rpc import RRef
from test import sub_fun

os.environ['MASTER_ADDR'] = '10.5.26.19'
os.environ['MASTER_PORT'] = '5677'

rpc.init_rpc("master", rank=0, world_size=2)
rref = torch.Tensor([0])
sub_fun(rref)
rpc.shutdown()

test.py

def f(rref):
print("function is executed on master")

def sub_fun(rref):
x = rpc.rpc_async("worker", f, args=(rref,))

worker.py：

import os
import torch
import torch.distributed.rpc as rpc
from torch.distributed.rpc import RRef

os.environ['MASTER_ADDR'] = '10.5.26.19'
os.environ['MASTER_PORT'] = '5677'

def f(rref):
print("function is executed on worker")
rpc.init_rpc("worker", rank=1, world_size=2)
rpc.shutdown()

我发现worker端的输出是“function is executed on master”。

方法二

当我把 sub_fun 和 f 这两个函数放在 master.py 而不是 test.py 时，结果是“function is executed on worker”。

为什么这两种方式输出不同的结果。以及如何使用方法 1 得到结果 2。

【参考方案1】：

RPC 假设模块和功能在所有进程中都是一致的。在这种情况下，您拥有相同的函数f，但在两个进程中具有不同的实现。请注意，RPC 不会将整个函数代码发送到远程端执行，而只会在远程端找到函数f 并执行那里存在的任何内容。

我建议使用两个不同的函数来实现您在此处尝试执行的操作，因为根据您导入文件的方式，一个函数实现可能会覆盖另一个。