评价此页
★ ★ ★ ★ ★

分布式 RPC 框架#

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

分布式 RPC 框架通过一组原语提供多机模型训练机制，以实现远程通信，以及一个更高级别的 API 来自动区分跨多台机器拆分模型。

警告

RPC 包中的 API 是稳定的。目前正在进行多项工作以改进性能和错误处理，这些将在未来版本中发布。

警告

PyTorch 1.9 中引入了 CUDA 支持，并且仍处于 beta 功能阶段。RPC 包并非所有功能都与 CUDA 支持兼容，因此不建议使用它们。这些不支持的功能包括：RRefs、JIT 兼容性、分布式自动求导和分布式优化器以及性能分析。这些缺点将在未来版本中解决。

注意

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

基础知识#

分布式 RPC 框架使得远程运行函数、无需复制实际数据即可引用远程对象以及提供自动求导和优化器 API 以透明地跨 RPC 边界运行反向传播和更新参数变得容易。这些功能可以分为四组 API。

  1. 远程过程调用 (RPC) 支持在指定的目标工作器上使用给定参数运行函数，并获取返回值或创建对返回值的引用。有三个主要的 RPC API：rpc_sync()（同步）、rpc_async()（异步）和 remote()（异步并返回对远程返回值的引用）。如果用户代码在没有返回值的情况下无法继续，请使用同步 API。否则，使用异步 API 获取未来对象，并在调用者需要返回值时等待该未来对象。remote() API 在需要远程创建某些内容但从不需要将其提取到调用者时很有用。想象一下驱动程序进程正在设置参数服务器和训练器的情况。驱动程序可以在参数服务器上创建嵌入表，然后与训练器共享对嵌入表的引用，但它本身永远不会在本地使用嵌入表。在这种情况下，rpc_sync() 和 rpc_async() 不再适用，因为它们总是意味着返回值将立即或将来返回给调用者。

  2. 远程引用 (RRef) 作为本地或远程对象的分布式共享指针。它可以与其他工作器共享，并且引用计数将透明处理。每个 RRef 只有一个所有者，并且对象仅存在于该所有者上。持有 RRef 的非所有者工作器可以通过显式请求从所有者处获取对象的副本。当工作器需要访问某些数据对象，但它本身既不是创建者（remote() 的调用者）也不是对象的所有者时，这很有用。分布式优化器，正如我们将在下面讨论的那样，就是这种用例的一个例子。

  3. 分布式自动求导 将正向传播中涉及的所有工作器上的本地自动求导引擎连接起来，并在反向传播期间自动联系它们以计算梯度。这在正向传播需要跨多台机器进行时特别有用，例如，分布式模型并行训练、参数服务器训练等。有了此功能，用户代码不再需要担心如何跨 RPC 边界发送梯度以及以何种顺序启动本地自动求导引擎，这在正向传播中存在嵌套和相互依赖的 RPC 调用时会变得非常复杂。

  4. 分布式优化器 的构造函数接受一个 Optimizer()（例如，SGD()、Adagrad() 等）和参数 RRefs 列表，在每个不同的 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, 可选) – RPC 后端实现类型。支持的值为 BackendType.TENSORPIPE（默认）。有关更多信息，请参阅 后端。

  • rank (int) – 此节点的全局唯一 ID/排名。

  • world_size (int) – 组中工作器的数量。

  • rpc_backend_options (RpcBackendOptions, 可选) – 传递给 RpcAgent 构造函数的选项。它必须是 RpcBackendOptions 的代理特定子类，并包含代理特定的初始化配置。默认情况下，对于所有代理，它将默认超时设置为 60 秒，并使用 init_method = "env://" 初始化的底层进程组执行会合，这意味着需要正确设置环境变量 MASTER_ADDR 和 MASTER_PORT。有关更多信息以及可用的选项，请参阅 后端。

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

警告

RPC 中的 TorchScript 支持是一个原型功能，可能会更改。自 v1.5.0 起，torch.distributed.rpc 支持将 TorchScript 函数作为 RPC 目标函数调用，这将有助于提高被调用者端的并行性，因为执行 TorchScript 函数不需要 GIL。

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, 可选) – 此 RPC 的超时时间（秒）。如果 RPC 在此时间内未完成，将引发指示超时的异常。值为 0 表示无限超时，即永远不会引发超时错误。如果未提供，则使用初始化期间或使用 _set_rpc_timeout 设置的默认值。

返回

返回使用 args 和 kwargs 运行 func 的结果。

示例：

确保在两个工作器上都正确设置了 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, 可选) – 此 RPC 的超时时间（秒）。如果 RPC 在此时间内未完成，将引发指示超时的异常。值为 0 表示无限超时，即永远不会引发超时错误。如果未提供，则使用初始化期间或使用 _set_rpc_timeout 设置的默认值。

返回

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

警告

不支持使用 GPU 张量作为 func 的参数或返回值，因为我们不支持通过网络发送 GPU 张量。在使用 GPU 张量作为 func 的参数或返回值之前，需要显式地将其复制到 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, 可选) – 此远程调用的超时时间（秒）。如果在此超时时间内，此 RRef 在工作器 to 上的创建未在此工作器上成功处理，则下次尝试使用 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 后端已在 PyTorch v1.6 中引入，并正在积极开发中。目前，它只支持 CPU 张量，即将支持 GPU。它带有一个基于 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, 可选) – TensorPipeAgent 用于执行请求的线程池中的线程数（默认值：16）。

  • rpc_timeout (float, 可选) – RPC 请求的默认超时时间（秒）（默认值：60 秒）。如果 RPC 未在此时间段内完成，将引发指示此情况的异常。调用者可以在 rpc_sync() 和 rpc_async() 中根据需要覆盖单个 RPC 的此超时。

  • init_method (str, 可选) – 初始化用于协调的分布式存储的 URL。它接受与 init_process_group() 相同参数的任何值（默认值：env://）。

  • device_maps (Dict[str, Dict], 可选) – 从此工作器到被调用方的设备放置映射。键是被调用方工作器名称，值是字典（Dict of int、str 或 torch.device），它将此工作器的设备映射到被调用方工作器的设备。（默认值：None）

  • devices (List[int, str, 或 torch.device], 可选) – 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 张量时目前不支持 RRef

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

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

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

警告

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

参数

  • value (object) – 将由此 RRef 包装的值。

  • type_hint (Type, 可选) – 应作为 value 的类型提示传递给 TorchScript 编译器的 Python 类型。

示例：

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

  1. 使用 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()

  1. 从本地对象创建 RRef

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

  1. 与其他工作器共享 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: int = -1, retain_graph: bool = False) → None#

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

参数

  • dist_autograd_ctx_id (int, 可选) – 我们应为其检索梯度的分布式自动求导上下文 ID（默认值：-1）。

  • retain_graph (bool, 可选) – 如果为 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: float = -1.0) → object#

创建一个辅助代理，以便轻松启动一个 remote，使用 RRef 的所有者作为目的地，对该 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, 可选) – 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: float = -1.0) → object#

创建一个辅助代理，以便轻松启动一个 rpc_async，使用 RRef 的所有者作为目的地，对该 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, 可选) – 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: float = -1.0) → object#

创建一个辅助代理，以便轻松启动一个 rpc_sync，使用 RRef 的所有者作为目的地，对该 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, 可选) – 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: float = -1.0) → object#

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

参数

timeout (float, 可选) – 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]#

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

它在指定的远程节点上创建一个用户指定的模块。它的行为类似于常规的 nn.Module，只是 forward 方法在远程节点上执行。它负责自动求导记录，以确保反向传播将梯度传播回相应的远程模块。

它根据 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) – 我们希望放置此模块的目标工作器上的设备。格式应为“<工作器名称>/<设备>”，其中设备字段可解析为 torch.device 类型。例如，“trainer0/cpu”、“trainer0”、“ps0/cuda:0”。此外，设备字段可以可选，默认值为“cpu”。

  • module_cls (nn.Module) –

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

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

  • args (Sequence, 可选) – 传递给 module_cls 的参数。

  • kwargs (Dict, 可选) – 传递给 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#

使用提供的根启动分布式反向传播。这目前实现了 快速模式算法，该算法假定在工作器之间相同分布式自动求导上下文中发送的所有 RPC 消息都将作为反向传播期间自动求导图的一部分。

我们使用提供的根来发现自动求导图并计算适当的依赖项。此方法会阻塞，直到整个自动求导计算完成。

我们将梯度累积到每个节点上相应 torch.distributed.autograd.context 中。要使用的自动求导上下文将根据调用 torch.distributed.autograd.backward() 时传入的 context_id 进行查找。如果不存在与给定 ID 对应的有效自动求导上下文，我们将抛出错误。您可以使用 get_gradients() API 检索累积的梯度。

参数

  • context_id (int) – 我们应为其检索梯度的自动求导上下文 ID。

  • roots (list) – 表示自动求导计算根的张量。所有张量都应为标量。

  • retain_graph (bool, 可选) – 如果为 False，则用于计算梯度的图将被释放。请注意，在几乎所有情况下，将此选项设置为 True 都不需要，并且通常可以通过更高效的方式解决。通常，您需要将此设置为 True 才能多次运行反向传播。

示例：
>>> 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]#

检索从 Tensor 到在给定 context_id 对应的上下文中作为分布式自动求导反向传播的一部分累积的该 Tensor 的相应梯度的映射。

参数

context_id (int) – 我们应为其检索梯度的自动求导上下文 ID。

返回

一个映射，其中键是 Tensor，值是该 Tensor 的关联梯度。

示例：
>>> 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 自动求导的更多信息

分布式优化器#

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

设计说明#

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

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

教程#

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

在此页面上
在 GitHub 上编辑
显示源代码
PyTorch 库