评价此页

★ ★ ★ ★ ★

分布式 RPC 框架#

创建于：2019 年 11 月 14 日 | 最后更新于：2025 年 07 月 09 日

分布式 RPC 框架提供了一种通过一系列原始接口实现多机模型训练的机制，这些接口允许远程通信，并提供了一个更高级别的 API，可以自动区分分布在多台机器上的模型。

警告

RPC 包中的 API 是稳定的，并且处于维护模式。

警告

CUDA 支持是一个Beta功能。RPC 包并非所有功能都与 CUDA 支持兼容，因此不建议使用它们。这些不支持的功能包括：RRefs、JIT 兼容性、分布式自动微分和分布式优化器以及性能分析。有关分布式训练所有功能的简要介绍，请参阅PyTorch 分布式概述 <https://pytorch.ac.cn/tutorials/beginner/dist_overview.html>__。

注意

请参阅 PyTorch 分布式概述 <https://pytorch.ac.cn/tutorials/beginner/dist_overview.html>__ 以获取有关分布式训练所有功能的简要介绍。

基础#

分布式 RPC 框架可以轻松地远程运行函数，支持在不复制实际数据的情况下引用远程对象，并提供 autograd 和 optimizer API，以便在 RPC 边界之间透明地运行反向传播和更新参数。这些功能可以分为四组 API。

远程过程调用 (RPC) 支持在指定的远程工作节点上使用给定参数运行函数，并获取返回值或创建返回值的引用。有三个主要的 RPC API：rpc_sync()（同步）、rpc_async()（异步）和remote()（异步并返回远程返回值引用）。如果用户代码在没有返回值的情况下无法继续执行，请使用同步 API。否则，请使用异步 API 获取一个 future，并在需要返回值时等待该 future。当需要远程创建某物但永远不需要将其获取到调用方时，remote() API 非常有用。例如，假设一个驱动进程正在设置一个参数服务器和一个训练器。驱动程序可以在参数服务器上创建一个嵌入表，然后将该表引用共享给训练器，但驱动程序本身永远不会在本地使用该嵌入表。在这种情况下，rpc_sync() 和 rpc_async() 不再适用，因为它们始终意味着返回值将立即或在未来返回给调用方。
远程引用 (RRef) 用作本地或远程对象的分布式共享指针。它可以与其他工作节点共享，并且引用计数将得到透明处理。每个 RRef 只有一个所有者，对象仅存在于该所有者处。持有 RRef 的非所有者工作节点可以通过显式请求来获取该对象的副本。当某个工作节点需要访问某个数据对象，但该节点本身既不是创建者（remote() 的调用者）也不是该对象的所有者时，这非常有用。如下文将讨论的分布式优化器就是此类用例之一。
分布式自动微分 (Distributed Autograd) 将所有参与前向传播的工作节点上的本地自动微分引擎缝合在一起，并在反向传播过程中自动访问它们以计算梯度。如果前向传播需要跨多台机器进行，例如分布式模型并行训练、参数服务器训练等，这尤其有帮助。有了此功能，用户代码不再需要担心如何跨 RPC 边界发送梯度以及应以何种顺序启动本地自动微分引擎，这在存在嵌套和相互依赖的 RPC 调用时可能会非常复杂。
分布式优化器 (Distributed Optimizer) 的构造函数接受一个 Optimizer()（例如 SGD()、Adagrad() 等）和一个参数 RRef 列表，在每个不同的 RRef 所有者上创建一个 Optimizer() 实例，并在运行 step() 时相应地更新参数。当您拥有分布式前向和后向传播时，参数和梯度将分散在多个工作节点上，因此需要每个参与工作节点上的优化器。分布式优化器将所有这些本地优化器包装在一个之中，并提供简洁的构造函数和 step() API。

RPC#

在使用 RPC 和分布式自动微分原语之前，必须进行初始化。要初始化 RPC 框架，我们需要使用 init_rpc()，它将初始化 RPC 框架、RRef 框架和分布式自动微分。

torch.distributed.rpc.init_rpc(name, backend=None, rank=-1, world_size=None, rpc_backend_options=None)[source]#

初始化 RPC 原语，例如本地 RPC 代理和分布式自动微分，这会立即使当前进程准备好发送和接收 RPC。

参数

name (str) – 此节点的全局唯一名称。（例如，Trainer3、ParameterServer2、Master、Worker1）名称只能包含数字、字母、下划线、冒号和/或连字符，且长度必须小于 128 个字符。
backend (BackendType, optional) – RPC 后端实现的类型。支持的值为 BackendType.TENSORPIPE（默认值）。有关更多信息，请参阅后。
rank (int) – 此节点的全局唯一 ID/rank。
world_size (int) – 组中的工作节点数量。
rpc_backend_options (RpcBackendOptions, optional) – 传递给 RpcAgent 构造函数的选项。它必须是 RpcBackendOptions 的代理特定子类，并包含代理特定的初始化配置。默认情况下，对于所有代理，它将默认超时设置为 60 秒，并使用使用 init_method = "env://"初始化的底层进程组进行协调。这意味着需要正确设置环境变量 MASTER_ADDR 和 MASTER_PORT。有关更多信息，请参阅后端，并查找可用的选项。

以下 API 允许用户远程执行函数以及创建远程数据对象的引用（RRefs）。在这些 API 中，当将 Tensor 作为参数或返回值传递时，目标工作节点将尝试创建一个具有相同元数据（即形状、步幅等）的 Tensor。我们故意禁止传输 CUDA 张量，因为如果源工作节点和目标工作节点上的设备列表不匹配，可能会导致崩溃。在这种情况下，应用程序始终可以显式地将输入张量移至调用方的 CPU，并在必要时将它们移至被调用方的所需设备。

torch.distributed.rpc.rpc_sync(to, func, args=None, kwargs=None, timeout=-1.0)[source]#

执行一个阻塞式 RPC 调用，在工作节点 to 上运行函数 func。RPC 消息与 Python 代码的执行并行发送和接收。此方法是线程安全的。

参数

to (str 或 WorkerInfo 或 int) – 目标工作节点的名称/rank/WorkerInfo。
func (Callable) – 一个可调用函数，例如 Python 可调用对象、内置运算符（例如 add()）和注解的 TorchScript 函数。
args (tuple) – 调用 func 的参数元组。
kwargs (dict) – 调用 func 的关键字参数字典。
timeout (float, optional) – 此 RPC 的超时时间（秒）。如果 RPC 在此时间内未完成，将引发指示超时的异常。值为 0 表示无限超时，即永远不会引发超时错误。如果未提供，将使用初始化期间或使用 _set_rpc_timeout 设置的默认值。

返回

返回运行 func 并带参数 args 和 kwargs 的结果。

示例：

请确保在两个工作节点上都正确设置了 MASTER_ADDR 和 MASTER_PORT。有关更多详细信息，请参阅 init_process_group() API。例如，

export MASTER_ADDR=localhost export MASTER_PORT=5678

然后，在两个不同的进程中运行以下代码

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> ret = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(2), 3))
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

下面是一个使用 RPC 运行 TorchScript 函数的示例。

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(tensor: torch.Tensor, scalar: int):
>>>    return torch.add(tensor, scalar)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> ret = rpc.rpc_sync("worker1", my_script_add, args=(torch.ones(2), 3))
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

torch.distributed.rpc.rpc_async(to, func, args=None, kwargs=None, timeout=-1.0)[source]#

执行一个非阻塞 RPC 调用，在工作节点 to 上运行函数 func。RPC 消息与 Python 代码的执行并行发送和接收。此方法是线程安全的。此方法将立即返回一个可以等待的 Future。

参数

to (str 或 WorkerInfo 或 int) – 目标工作节点的名称/rank/WorkerInfo。
func (Callable) – 一个可调用函数，例如 Python 可调用对象、内置运算符（例如 add()）和注解的 TorchScript 函数。
args (tuple) – 调用 func 的参数元组。
kwargs (dict) – 调用 func 的关键字参数字典。
timeout (float, optional) – 此 RPC 的超时时间（秒）。如果 RPC 在此时间内未完成，将引发指示超时的异常。值为 0 表示无限超时，即永远不会引发超时错误。如果未提供，将使用初始化期间或使用 _set_rpc_timeout 设置的默认值。

返回

返回一个可以等待的 Future 对象。当完成时，可以在 Future 对象中检索在 args 和 kwargs 上运行 func 的返回值。

警告

不支持使用 GPU 张量作为 func 的参数或返回值，因为我们不支持通过网络传输 GPU 张量。在使用它们作为 func 的参数或返回值之前，您需要显式地将 GPU 张量复制到 CPU。

警告

rpc_async API 在发送张量参数到线路上之前不会复制它们的存储，这可能由不同的线程完成，具体取决于 RPC 后端类型。调用方应确保在返回的 Future 完成之前，这些张量的内容保持不变。

示例：

请确保在两个工作节点上都正确设置了 MASTER_ADDR 和 MASTER_PORT。有关更多详细信息，请参阅 init_process_group() API。例如，

export MASTER_ADDR=localhost export MASTER_PORT=5678

然后，在两个不同的进程中运行以下代码

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> fut1 = rpc.rpc_async("worker1", torch.add, args=(torch.ones(2), 3))
>>> fut2 = rpc.rpc_async("worker1", min, args=(1, 2))
>>> result = fut1.wait() + fut2.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

下面是一个使用 RPC 运行 TorchScript 函数的示例。

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(tensor: torch.Tensor, scalar: int):
>>>    return torch.add(tensor, scalar)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> fut = rpc.rpc_async("worker1", my_script_add, args=(torch.ones(2), 3))
>>> ret = fut.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

torch.distributed.rpc.remote(to, func, args=None, kwargs=None, timeout=-1.0)[source]#

执行远程调用，在工作节点 to 上运行 func，并立即返回结果值的 RRef。工作节点 to 将是返回的 RRef 的所有者，而调用 remote 的工作节点是用户。所有者管理其 RRef 的全局引用计数，并且所有者 RRef 仅在其不再有任何全局引用时才会被销毁。

参数

to (str 或 WorkerInfo 或 int) – 目标工作节点的名称/rank/WorkerInfo。
func (Callable) – 一个可调用函数，例如 Python 可调用对象、内置运算符（例如 add()）和注解的 TorchScript 函数。
args (tuple) – 调用 func 的参数元组。
kwargs (dict) – 调用 func 的关键字参数字典。
timeout (float, optional) – 此远程调用的超时时间（秒）。如果在该超时时间内 to 节点上成功处理此 RRef 的创建，则下次尝试使用 RRef（例如 to_here()）时，将引发指示此失败的超时错误。值为 0 表示无限超时，即永远不会引发超时错误。如果未提供，将使用初始化期间或使用 _set_rpc_timeout 设置的默认值。

返回

用户 RRef 实例到结果值。使用阻塞 API torch.distributed.rpc.RRef.to_here() 在本地检索结果值。

警告

remote API 在发送张量参数到线路上之前不会复制它们的存储，这可能由不同的线程完成，具体取决于 RPC 后端类型。调用方应确保在返回的 RRef 被所有者确认之前，这些张量的内容保持不变，这可以使用 torch.distributed.rpc.RRef.confirmed_by_owner() API 进行检查。

警告

诸如 remote API 的超时之类的错误是尽力而为处理的。这意味着当 remote 启动的远程调用失败（例如，超时错误）时，我们会采取尽力而为的方法来处理错误。这意味着错误是异步处理和设置在结果 RRef 上的。如果 RRef 在此处理之前尚未被应用程序使用（例如 to_here 或 fork 调用），那么未来使用 RRef 将适当地引发错误。但是，用户应用程序可能会在错误处理之前使用 RRef。在这种情况下，错误可能不会引发，因为它们尚未得到处理。

示例

Make sure that ``MASTER_ADDR`` and ``MASTER_PORT`` are set properly
on both workers. Refer to :meth:`~torch.distributed.init_process_group`
API for more details. For example,

export MASTER_ADDR=localhost
export MASTER_PORT=5678

Then run the following code in two different processes:

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> rref1 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
>>> rref2 = rpc.remote("worker1", torch.add, args=(torch.ones(2), 1))
>>> x = rref1.to_here() + rref2.to_here()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

Below is an example of running a TorchScript function using RPC.

>>> # On both workers:
>>> @torch.jit.script
>>> def my_script_add(tensor: torch.Tensor, scalar: int):
>>>    return torch.add(tensor, scalar)

>>> # On worker 0:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> rref = rpc.remote("worker1", my_script_add, args=(torch.ones(2), 3))
>>> rref.to_here()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

torch.distributed.rpc.get_worker_info(worker_name=None)[source]#

获取给定工作节点名称的 WorkerInfo。使用此 WorkerInfo 以避免每次调用时传递昂贵的字符串。

参数: worker_name (str) – 工作节点的字符串名称。如果为 None，则返回当前工作节点的 ID。（默认 None）
返回: 给定 worker_name 的 WorkerInfo 实例，或者如果 worker_name 为 None，则为当前工作节点的 WorkerInfo。

torch.distributed.rpc.shutdown(graceful=True, timeout=0)[source]#

执行 RPC 代理的关闭，然后销毁 RPC 代理。这会阻止本地代理接受未完成的请求，并通过终止所有 RPC 线程来关闭 RPC 框架。如果 graceful=True，此方法将阻塞，直到所有本地和远程 RPC 进程都到达此方法并等待所有未完成的工作完成。否则，如果 graceful=False，则为本地关闭，并且不会等待其他 RPC 进程到达此方法。

警告

对于 Future 对象，由 rpc_async() 返回，不应在 shutdown() 之后调用 future.wait()。

参数: graceful (bool) – 是否执行优雅关闭。如果为 True，将 1) 等待直到没有待处理的 UserRRefs 系统消息并删除它们；2) 阻塞直到所有本地和远程 RPC 进程都到达此方法并等待所有未完成的工作完成。

示例：

请确保在两个工作节点上都正确设置了 MASTER_ADDR 和 MASTER_PORT。有关更多详细信息，请参阅 init_process_group() API。例如，

export MASTER_ADDR=localhost export MASTER_PORT=5678

然后，在两个不同的进程中运行以下代码

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> # do some work
>>> result = rpc.rpc_sync("worker1", torch.add, args=(torch.ones(1), 1))
>>> # ready to shutdown
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch.distributed.rpc as rpc
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> # wait for worker 0 to finish work, and then shutdown.
>>> rpc.shutdown()

class torch.distributed.rpc.WorkerInfo#

封装系统中工作节点信息的结构。包含工作节点的名称和 ID。此类不应直接实例化，而是可以通过 get_worker_info() 获取实例，并将结果传递给 rpc_sync()、rpc_async()、remote() 等函数，以避免每次调用时复制字符串。

property id#: 用于标识工作节点的全局唯一 ID。

property name#: 工作节点的名称。

RPC 包还提供了装饰器，允许应用程序指定给定函数在被调用方应如何处理。

torch.distributed.rpc.functions.async_execution(fn)[source]#

一个函数装饰器，表示函数的返回值保证是 Future 对象，并且该函数可以在 RPC 被调用方上异步运行。更具体地说，被调用方提取被包装函数返回的 Future，并将后续处理步骤安装为该 Future 的回调。安装的回调将在 Future 完成后从中读取值，并将该值作为 RPC 响应发送。这意味着返回的 Future 仅存在于被调用方，并且永不通过 RPC 发送。当被包装函数（fn）的执行需要暂停和恢复时，此装饰器非常有用，例如，当它包含 rpc_async() 或等待其他信号时。

注意

要启用异步执行，应用程序必须将此装饰器返回的函数对象传递给 RPC API。如果 RPC 检测到此装饰器安装的属性，它将知道该函数返回 Future 对象，并将相应地进行处理。但是，这并不意味着在定义函数时此装饰器必须是最外层的。例如，当与 @staticmethod 或 @classmethod 结合使用时，@rpc.functions.async_execution 需要成为内部装饰器，以便目标函数被识别为静态或类函数。此目标函数仍然可以异步执行，因为在访问时，静态或类方法会保留 @rpc.functions.async_execution 安装的属性。

示例：

返回的 Future 对象可以来自 rpc_async()、then() 或 Future 构造函数。下面的示例展示了直接使用 then() 返回的 Future。

>>> from torch.distributed import rpc
>>>
>>> # omitting setup and shutdown RPC
>>>
>>> # On all workers
>>> @rpc.functions.async_execution
>>> def async_add_chained(to, x, y, z):
>>>     # This function runs on "worker1" and returns immediately when
>>>     # the callback is installed through the `then(cb)` API. In the
>>>     # mean time, the `rpc_async` to "worker2" can run concurrently.
>>>     # When the return value of that `rpc_async` arrives at
>>>     # "worker1", "worker1" will run the lambda function accordingly
>>>     # and set the value for the previously returned `Future`, which
>>>     # will then trigger RPC to send the result back to "worker0".
>>>     return rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>         lambda fut: fut.wait() + z
>>>     )
>>>
>>> # On worker0
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     async_add_chained,
>>>     args=("worker2", torch.ones(2), 1, 1)
>>> )
>>> print(ret)  # prints tensor([3., 3.])

与 TorchScript 装饰器结合使用时，此装饰器必须是最外层的。

>>> from torch import Tensor
>>> from torch.futures import Future
>>> from torch.distributed import rpc
>>>
>>> # omitting setup and shutdown RPC
>>>
>>> # On all workers
>>> @torch.jit.script
>>> def script_add(x: Tensor, y: Tensor) -> Tensor:
>>>     return x + y
>>>
>>> @rpc.functions.async_execution
>>> @torch.jit.script
>>> def async_add(to: str, x: Tensor, y: Tensor) -> Future[Tensor]:
>>>     return rpc.rpc_async(to, script_add, (x, y))
>>>
>>> # On worker0
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     async_add,
>>>     args=("worker2", torch.ones(2), 1)
>>> )
>>> print(ret)  # prints tensor([2., 2.])

与静态方法或类方法结合使用时，此装饰器必须是内部装饰器。

>>> from torch.distributed import rpc
>>>
>>> # omitting setup and shutdown RPC
>>>
>>> # On all workers
>>> class AsyncExecutionClass:
>>>
>>>     @staticmethod
>>>     @rpc.functions.async_execution
>>>     def static_async_add(to, x, y, z):
>>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>             lambda fut: fut.wait() + z
>>>         )
>>>
>>>     @classmethod
>>>     @rpc.functions.async_execution
>>>     def class_async_add(cls, to, x, y, z):
>>>         ret_fut = torch.futures.Future()
>>>         rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>             lambda fut: ret_fut.set_result(fut.wait() + z)
>>>         )
>>>         return ret_fut
>>>
>>>     @rpc.functions.async_execution
>>>     def bound_async_add(self, to, x, y, z):
>>>         return rpc.rpc_async(to, torch.add, args=(x, y)).then(
>>>             lambda fut: fut.wait() + z
>>>         )
>>>
>>> # On worker0
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     AsyncExecutionClass.static_async_add,
>>>     args=("worker2", torch.ones(2), 1, 2)
>>> )
>>> print(ret)  # prints tensor([4., 4.])
>>>
>>> ret = rpc.rpc_sync(
>>>     "worker1",
>>>     AsyncExecutionClass.class_async_add,
>>>     args=("worker2", torch.ones(2), 1, 2)
>>> )
>>> print(ret)  # prints tensor([4., 4.])

此装饰器也适用于 RRef 助手，即 . torch.distributed.rpc.RRef.rpc_sync()、torch.distributed.rpc.RRef.rpc_async() 和 torch.distributed.rpc.RRef.remote()。

>>> from torch.distributed import rpc
>>>
>>> # reuse the AsyncExecutionClass class above
>>> rref = rpc.remote("worker1", AsyncExecutionClass)
>>> ret = rref.rpc_sync().static_async_add("worker2", torch.ones(2), 1, 2)
>>> print(ret)  # prints tensor([4., 4.])
>>>
>>> rref = rpc.remote("worker1", AsyncExecutionClass)
>>> ret = rref.rpc_async().static_async_add("worker2", torch.ones(2), 1, 2).wait()
>>> print(ret)  # prints tensor([4., 4.])
>>>
>>> rref = rpc.remote("worker1", AsyncExecutionClass)
>>> ret = rref.remote().static_async_add("worker2", torch.ones(2), 1, 2).to_here()
>>> print(ret)  # prints tensor([4., 4.])

后端#

RPC 模块可以利用不同的后端在节点之间执行通信。可以在 init_rpc() 函数中指定要使用的后端，方法是传递 BackendType 枚举的特定值。无论使用哪个后端，其余的 RPC API 都不会改变。每个后端还定义了自己的 RpcBackendOptions 类子类，其实例也可以传递给 init_rpc() 以配置后端行为。

class torch.distributed.rpc.BackendType(value)#

可用后端枚举类。

PyTorch 提供了内置的 BackendType.TENSORPIPE 后端。可以使用 register_backend() 函数注册其他后端。

class torch.distributed.rpc.RpcBackendOptions#

封装传递给 RPC 后端的选项的抽象结构。此类的实例可以传递给 init_rpc()，以便使用特定配置初始化 RPC，例如 RPC 超时和要使用的 init_method。

property init_method#: 用于初始化进程组的 URL。默认为 env://

property rpc_timeout#: 一个浮点数，表示所有 RPC 的超时时间。如果 RPC 在此时间范围内未完成，它将以指示超时的异常完成。

TensorPipe 后端#

TensorPipe 代理是默认的，它利用 TensorPipe 库，该库提供了本地点对点通信原语，专门适用于机器学习，从根本上解决了 Gloo 的一些局限性。与 Gloo 相比，它的优点是异步的，允许大量传输同时发生，每个传输以自己的速度进行，而不会互相阻塞。它仅在需要时按需在节点对之间打开管道，并且当一个节点失败时，只有其相关的管道会被关闭，而所有其他管道将正常工作。此外，它能够支持多种不同的传输（当然是 TCP，但也有共享内存、NVLink、InfiniBand 等），并可以自动检测它们的可用性并协商为每个管道使用的最佳传输。

TensorPipe 后端带有基于 TCP 的传输，就像 Gloo 一样。它还能够自动分块和多路复用大张量到多个套接字和线程，以实现非常高的带宽。代理将能够自行选择最佳传输，无需任何干预。

示例

import os
from torch.distributed import rpc
os.environ['MASTER_ADDR'] = 'localhost'
os.environ['MASTER_PORT'] = '29500'

rpc.init_rpc(
    "worker1",
    rank=0,
    world_size=2,
    rpc_backend_options=rpc.TensorPipeRpcBackendOptions(
        num_worker_threads=8,
        rpc_timeout=20 # 20 second timeout
    )
)

# omitting init_rpc invocation on worker2

class torch.distributed.rpc.TensorPipeRpcBackendOptions(*, num_worker_threads=16, rpc_timeout=60.0, init_method='env://', device_maps=None, devices=None, _transports=None, _channels=None)[source]#

TensorPipeAgent 的后端选项，继承自 RpcBackendOptions。

参数

num_worker_threads (int, optional) – TensorPipeAgent 用于执行请求的线程池中的线程数（默认值：16）。
rpc_timeout (float, optional) – RPC 请求的默认超时时间（秒）（默认值：60 秒）。如果在该时间内 RPC 未完成，将引发相应的异常。调用者可以在 rpc_sync() 和 rpc_async() 中为单个 RPC 覆盖此超时时间（如有必要）。
init_method (str, optional) – 用于初始化用于协调的分布式存储的 URL。它可以接受 init_process_group() 的相同参数的任何值（默认值：env://）。
device_maps (Dict[str, Dict], optional) – 从此工作节点到被调用方的设备放置映射。键是被调用方的工作节点名称，值是字典（Dict of int、str 或 torch.device），将此工作节点的设备映射到被调用方工作节点的设备。（默认值：None）
devices (List[int, str, 或 torch.device], optional) – RPC 代理使用的所有本地 CUDA 设备。默认情况下，它将初始化为来自其自身 device_maps 的所有本地设备以及来自其对等方 device_maps 的相应设备。在处理 CUDA RPC 请求时，代理将正确同步此 List 中所有设备的 CUDA 流。

property device_maps#: 设备映射位置。

property devices#: 本地代理使用的所有设备。

property init_method#: 用于初始化进程组的 URL。默认为 env://

property num_worker_threads#: TensorPipeAgent 用于执行请求的线程池中的线程数。

property rpc_timeout#: 一个浮点数，表示所有 RPC 的超时时间。如果 RPC 在此时间范围内未完成，它将以指示超时的异常完成。

set_device_map(to, device_map)[source]#

设置每个 RPC 调用方和被调用方对之间的设备映射。此函数可以多次调用以增量添加设备放置配置。

参数

to (str) – 被调用方名称。
device_map (Dict of int, str, or torch.device) – 从此工作节点到被调用方的设备放置映射。此映射必须是可逆的。

示例

>>> # both workers
>>> def add(x, y):
>>>     print(x)  # tensor([1., 1.], device='cuda:1')
>>>     return x + y, (x + y).to(2)
>>>
>>> # on worker 0
>>> options = TensorPipeRpcBackendOptions(
>>>     num_worker_threads=8,
>>>     device_maps={"worker1": {0: 1}}
>>> # maps worker0's cuda:0 to worker1's cuda:1
>>> )
>>> options.set_device_map("worker1", {1: 2})
>>> # maps worker0's cuda:1 to worker1's cuda:2
>>>
>>> rpc.init_rpc(
>>>     "worker0",
>>>     rank=0,
>>>     world_size=2,
>>>     backend=rpc.BackendType.TENSORPIPE,
>>>     rpc_backend_options=options
>>> )
>>>
>>> x = torch.ones(2)
>>> rets = rpc.rpc_sync("worker1", add, args=(x.to(0), 1))
>>> # The first argument will be moved to cuda:1 on worker1. When
>>> # sending the return value back, it will follow the invert of
>>> # the device map, and hence will be moved back to cuda:0 and
>>> # cuda:1 on worker0
>>> print(rets[0])  # tensor([2., 2.], device='cuda:0')
>>> print(rets[1])  # tensor([2., 2.], device='cuda:1')

set_devices(devices)[source]#

设置 TensorPipe RPC 代理使用的本地设备。在处理 CUDA RPC 请求时，TensorPipe RPC 代理将正确同步此 List 中所有设备的 CUDA 流。

参数: devices (List of int, str, or torch.device) – TensorPipe RPC 代理使用的本地设备。

注意

RPC 框架不会自动重试任何 rpc_sync()、rpc_async() 和 remote() 调用。原因是 RPC 框架无法确定操作是否是幂等的，以及是否可以安全地重试。因此，应用程序有责任处理故障并根据需要重试。RPC 通信基于 TCP，因此故障可能由网络故障或间歇性网络连接问题引起。在这种情况下，应用程序需要以合理的退避机制进行适当重试，以确保网络不会因激进的重试而过载。

RRef#

警告

使用 CUDA 张量时，目前不支持 RRefs。

RRef（远程引用）是对远程工作节点上某个类型 T（例如 Tensor）的值的引用。此句柄会使所有者处的被引用远程值保持活动状态，但并不意味着该值将来会被传输到本地工作节点。RRefs 可用于多机训练，通过持有存在于其他工作节点上的 nn.Modules 的引用，并在训练期间调用适当的函数来检索或修改其参数。有关更多详细信息，请参阅远程引用协议。

class torch.distributed.rpc.PyRRef(RRef)#

封装远程工作节点上某个类型的值的引用的类。此句柄将使被引用的远程值在工作节点上保持活动状态。UserRRef 将在以下情况下被删除：1) 应用程序代码和本地 RRef 上下文中不再有对其的引用；或 2) 应用程序已调用了优雅关机。调用已删除 RRef 上的方法会导致未定义行为。RRef 实现仅提供尽力错误检测，并且应用程序不应在 rpc.shutdown() 之后使用 UserRRefs。

警告

RRefs 只能由 RPC 模块序列化和反序列化。在没有 RPC 的情况下序列化和反序列化 RRefs（例如，Python pickle、torch save() / load()、JIT save() / load() 等）将导致错误。

参数

value (object) – 要由此 RRef 包装的值。
type_hint (Type, optional) – 应传递给 TorchScript 编译器作为 value 的类型提示的 Python 类型。

示例：

以下示例为简洁起见省略了 RPC 初始化和关闭代码。有关这些详细信息，请参阅 RPC 文档。

使用 rpc.remote 创建 RRef

>>> import torch
>>> import torch.distributed.rpc as rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.ones(2), 3))
>>> # get a copy of value from the RRef
>>> x = rref.to_here()

从本地对象创建 RRef

>>> import torch
>>> from torch.distributed.rpc import RRef
>>> x = torch.zeros(2, 2)
>>> rref = RRef(x)

与其它工作节点共享 RRef

>>> # On both worker0 and worker1:
>>> def f(rref):
>>>   return rref.to_here() + 1

>>> # On worker0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> from torch.distributed.rpc import RRef
>>> rref = RRef(torch.zeros(2, 2))
>>> # the following RPC shares the rref with worker1, reference
>>> # count is automatically updated.
>>> rpc.rpc_sync("worker1", f, args=(rref,))

backward(self: torch._C._distributed_rpc.PyRRef, dist_autograd_ctx_id: SupportsInt = -1, retain_graph: bool = False) → None#

使用 RRef 作为反向传播的根节点运行反向传播。如果提供了 dist_autograd_ctx_id，我们将使用提供的 ctx_id 从 RRef 的所有者开始执行分布式反向传播。在这种情况下，应使用 get_gradients() 来检索梯度。如果 dist_autograd_ctx_id 为 None，则假定这是一个本地自动微分图，我们只执行本地反向传播。在本地情况下，调用此 API 的节点必须是 RRef 的所有者。RRef 的值预计是标量 Tensor。

参数

dist_autograd_ctx_id (int, optional) – 我们应从中检索梯度的分布式自动微分上下文 ID（默认值：-1）。
retain_graph (bool, optional) – 如果为 False，则用于计算梯度的图将被释放。请注意，在几乎所有情况下，将此选项设置为 True 都是不必要的，并且通常可以通过更有效的方式来解决。通常，您需要将其设置为 True 以多次运行反向传播（默认值：False）。

示例：

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     rref.backward(context_id)

confirmed_by_owner(self: torch._C._distributed_rpc.PyRRef) → bool#: 返回此 RRef 是否已由所有者确认。OwnerRRef 始终返回 true，而 UserRRef 仅在所有者了解此 UserRRef 时才返回 true。

is_owner(self: torch._C._distributed_rpc.PyRRef) → bool#: 返回当前节点是否是此 RRef 的所有者。

local_value(self: torch._C._distributed_rpc.PyRRef) → object#: 如果当前节点是所有者，则返回本地值的引用。否则，抛出异常。

owner(self: torch._C._distributed_rpc.PyRRef) → torch._C._distributed_rpc.WorkerInfo#: 返回拥有此 RRef 的节点的工作节点信息。

owner_name(self: torch._C._distributed_rpc.PyRRef) → str#: 返回拥有此 RRef 的节点的工作节点名称。

remote(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat = -1.0) → object#

创建一个辅助代理，以便轻松启动一个 remote，并使用 RRef 的所有者作为目标，在该对象上运行函数。更具体地说，rref.remote().func_name(*args, **kwargs) 与以下内容相同

>>> def run(rref, func_name, args, kwargs):
>>>   return getattr(rref.local_value(), func_name)(*args, **kwargs)
>>>
>>> rpc.remote(rref.owner(), run, args=(rref, func_name, args, kwargs))

参数: timeout (float, optional) – rref.remote() 的超时时间。如果在超时时间内未能成功创建此 RRef，则下次尝试使用 RRef（例如 to_here）时将引发超时。如果未提供，将使用默认的 RPC 超时。有关 RRef 的具体超时语义，请参阅 rpc.remote()。

示例：

>>> from torch.distributed import rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.zeros(2, 2), 1))
>>> rref.remote().size().to_here()  # returns torch.Size([2, 2])
>>> rref.remote().view(1, 4).to_here()  # returns tensor([[1., 1., 1., 1.]])

rpc_async(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat = -1.0) → object#

创建一个辅助代理，以便轻松启动一个 rpc_async，并使用 RRef 的所有者作为目标，在该对象上运行函数。更具体地说，rref.rpc_async().func_name(*args, **kwargs) 与以下内容相同

>>> def run(rref, func_name, args, kwargs):
>>>   return getattr(rref.local_value(), func_name)(*args, **kwargs)
>>>
>>> rpc.rpc_async(rref.owner(), run, args=(rref, func_name, args, kwargs))

参数: timeout (float, optional) – rref.rpc_async() 的超时时间。如果在该时间范围内调用未完成，将引发指示此情况的异常。如果未提供此参数，将使用默认的 RPC 超时。

示例：

>>> from torch.distributed import rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.zeros(2, 2), 1))
>>> rref.rpc_async().size().wait()  # returns torch.Size([2, 2])
>>> rref.rpc_async().view(1, 4).wait()  # returns tensor([[1., 1., 1., 1.]])

rpc_sync(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat = -1.0) → object#

创建一个辅助代理，以便轻松启动一个 rpc_sync，并使用 RRef 的所有者作为目标，在该对象上运行函数。更具体地说，rref.rpc_sync().func_name(*args, **kwargs) 与以下内容相同

>>> def run(rref, func_name, args, kwargs):
>>>   return getattr(rref.local_value(), func_name)(*args, **kwargs)
>>>
>>> rpc.rpc_sync(rref.owner(), run, args=(rref, func_name, args, kwargs))

参数: timeout (float, optional) – rref.rpc_sync() 的超时时间。如果在该时间范围内调用未完成，将引发指示此情况的异常。如果未提供此参数，将使用默认的 RPC 超时。

示例：

>>> from torch.distributed import rpc
>>> rref = rpc.remote("worker1", torch.add, args=(torch.zeros(2, 2), 1))
>>> rref.rpc_sync().size()  # returns torch.Size([2, 2])
>>> rref.rpc_sync().view(1, 4)  # returns tensor([[1., 1., 1., 1.]])

to_here(self: torch._C._distributed_rpc.PyRRef, timeout: SupportsFloat = -1.0) → object#

阻塞调用，将 RRef 的值从所有者复制到本地节点并返回。如果当前节点是所有者，则返回本地值的引用。

参数: timeout (float, optional) – to_here 的超时时间。如果在该时间范围内调用未完成，将引发指示此情况的异常。如果未提供此参数，将使用默认的 RPC 超时（60 秒）。

有关 RRef 的更多信息

远程引用协议

RemoteModule#

警告

目前在使用 CUDA 张量时不支持 RemoteModule

RemoteModule 是在不同进程中远程创建 nn.Module 的简便方法。实际模块驻留在远程主机上，但本地主机拥有该模块的句柄，并可以像普通 nn.Module 一样调用该模块。调用会产生到远程端的 RPC 调用，并且可以通过 RemoteModule 支持的附加 API 异步执行。

class torch.distributed.nn.api.remote_module.RemoteModule(*args, **kwargs)[source]#

只有在 RPC 初始化后才能创建 RemoteModule 实例。

它会在指定的远程节点上创建一个用户指定的模块。它的行为类似于普通的 nn.Module，不同之处在于 forward 方法在远程节点上执行。它负责 autograd 记录，以确保反向传播将梯度传播回相应的远程模块。

它根据 module_cls 的 forward 方法的签名，生成两个方法：forward_async 和 forward。forward_async 异步运行并返回一个 Future。 forward_async 和 forward 的参数与 module_cls 返回的模块的 forward 方法相同。

例如，如果 module_cls 返回一个 nn.Linear 的实例，其 forward 方法签名如下：def forward(input: Tensor) -> Tensor:，则生成的 RemoteModule 将具有 2 个方法，签名如下：

def forward(input: Tensor) -> Tensor:

def forward_async(input: Tensor) -> Future[Tensor]:

参数

remote_device (str) – 我们希望将此模块放置在目标工作进程上的设备。格式应为“<workername>/<device>”，其中 device 字段可以解析为 torch.device 类型。例如，“trainer0/cpu”、“trainer0”、“ps0/cuda:0”。此外，device 字段是可选的，默认值为“cpu”。

module_cls (nn.Module) –

要远程创建的模块的类。例如，

>>> class MyModule(nn.Module):
>>>     def forward(input):
>>>         return input + 1
>>>
>>> module_cls = MyModule

args (Sequence, optional) – 要传递给 module_cls 的参数。
kwargs (Dict, optional) – 要传递给 module_cls 的关键字参数。

返回

一个远程模块实例，它包装了由用户提供的 module_cls 创建的 Module，它具有一个阻塞的 forward 方法和一个异步的 forward_async 方法，该方法返回远程端用户提供的模块上 forward 调用的 Future。

示例：

在两个不同的进程中运行以下代码

>>> # On worker 0:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>> from torch import nn, Tensor
>>> from torch.distributed.nn.api.remote_module import RemoteModule
>>>
>>> rpc.init_rpc("worker0", rank=0, world_size=2)
>>> remote_linear_module = RemoteModule(
>>>     "worker1/cpu", nn.Linear, args=(20, 30),
>>> )
>>> input = torch.randn(128, 20)
>>> ret_fut = remote_linear_module.forward_async(input)
>>> ret = ret_fut.wait()
>>> rpc.shutdown()

>>> # On worker 1:
>>> import torch
>>> import torch.distributed.rpc as rpc
>>>
>>> rpc.init_rpc("worker1", rank=1, world_size=2)
>>> rpc.shutdown()

此外，可以在此教程中找到与 DistributedDataParallel (DDP) 结合使用的更实用的示例。

get_module_rref()[source]#

返回一个指向远程模块的 RRef（RRef[nn.Module]）。

返回类型: RRef[Module]

remote_parameters(recurse=True)[source]#

返回指向远程模块参数的 RRef 列表。

这通常可以与 DistributedOptimizer 结合使用。

参数: recurse (bool) – 如果为 True，则返回远程模块及其所有子模块的参数。否则，仅返回直接属于远程模块的参数。
返回: 指向远程模块参数的 RRef 列表（List[RRef[nn.Parameter]]）。
返回类型: list[torch.distributed.rpc.api.RRef[torch.nn.parameter.Parameter]]

分布式自动微分框架#

警告

目前在使用 CUDA 张量时不支持分布式自动微分

此模块提供了一个基于 RPC 的分布式自动微分框架，可用于模型并行训练等应用。简而言之，应用程序可以通过 RPC 发送和接收梯度记录张量。在正向传播中，我们记录梯度记录张量何时通过 RPC 发送，在反向传播过程中，我们利用此信息通过 RPC 执行分布式反向传播。有关更多详细信息，请参阅分布式自动微分设计。

torch.distributed.autograd.backward(context_id: int, roots: List[Tensor], retain_graph=False) → None#

使用提供的根节点启动分布式反向传播。目前，它实现了 FAST 模式算法，该算法假定在不同工作节点之间同一分布式自动微分上下文中的所有 RPC 消息都将作为自动微分图的一部分在反向传播过程中进行处理。

我们使用提供的根节点来发现自动微分图并计算适当的依赖项。此方法将阻塞直到整个自动微分计算完成。

我们在每个节点上将梯度累积到适当的 torch.distributed.autograd.context 中。要使用的自动微分上下文是根据调用 torch.distributed.autograd.backward() 时传入的 context_id 查找的。如果不存在与给定 ID 对应的有效自动微分上下文，则会引发错误。您可以使用 get_gradients() API 来检索累积的梯度。

参数

context_id (int) – 应从中检索梯度的自动微分上下文 ID。
roots (list) – 代表自动微分计算根的张量。所有张量都应为标量。
retain_graph (bool, optional) – 如果为 False，则用于计算梯度的图将被释放。请注意，在几乎所有情况下，将此选项设置为 True 都是不必要的，并且通常可以通过更有效的方式来解决。通常，您需要将其设置为 True 才能多次运行 backward。

示例：

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     pred = model.forward()
>>>     loss = loss_func(pred, loss)
>>>     dist_autograd.backward(context_id, loss)

class torch.distributed.autograd.context[source]#

用于包装正向和反向传播的上下文对象，当使用分布式自动微分时。with 语句中生成的 context_id 对于在所有工作节点上唯一标识分布式反向传播至关重要。每个工作节点都存储与此 context_id 关联的元数据，这对于正确执行分布式自动微分传播至关重要。

示例：

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     t1 = torch.rand((3, 3), requires_grad=True)
>>>     t2 = torch.rand((3, 3), requires_grad=True)
>>>     loss = rpc.rpc_sync("worker1", torch.add, args=(t1, t2)).sum()
>>>     dist_autograd.backward(context_id, [loss])

torch.distributed.autograd.get_gradients(context_id: int) → Dict[Tensor, Tensor]#

检索一个映射，该映射将张量与其在提供的上下文（对应于给定 context_id）中作为分布式自动微分反向传播一部分累积的相应梯度关联起来。

参数: context_id (int) – 应从中检索梯度的自动微分上下文 ID。
返回: 一个映射，其中键是张量，值是该张量相关的梯度。

示例：

>>> import torch.distributed.autograd as dist_autograd
>>> with dist_autograd.context() as context_id:
>>>     t1 = torch.rand((3, 3), requires_grad=True)
>>>     t2 = torch.rand((3, 3), requires_grad=True)
>>>     loss = t1 + t2
>>>     dist_autograd.backward(context_id, [loss.sum()])
>>>     grads = dist_autograd.get_gradients(context_id)
>>>     print(grads[t1])
>>>     print(grads[t2])

有关 RPC 自动微分的更多信息

分布式 Autograd 设计

分布式优化器#

有关分布式优化器的文档，请参阅 torch.distributed.optim 页面。

设计说明#

分布式自动微分设计说明涵盖了基于 RPC 的分布式自动微分框架的设计，该框架对于模型并行训练等应用非常有用。

分布式 Autograd 设计

RRef 设计说明涵盖了框架用于引用远程工作节点上的值的 RRef（远程引用）协议的设计。

远程引用协议

教程#

RPC 教程向用户介绍了 RPC 框架，提供了使用 torch.distributed.rpc API 的几个示例应用程序，并演示了如何使用 profiler 来分析基于 RPC 的工作负载。

分布式 RPC 框架入门
使用分布式 RPC 框架实现参数服务器
将 Distributed DataParallel 与分布式 RPC 框架结合使用（也涵盖 **RemoteModule**）
实现批处理 RPC

分布式 RPC 框架#

基础#

RPC#

后端#

TensorPipe 后端#

RRef#

RemoteModule#

分布式自动微分框架#

分布式优化器#

设计说明#

教程#

文档

教程

资源