评价此页
★ ★ ★ ★ ★

分布式 RPC 框架#

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

分布式 RPC 框架提供了一套原语，通过远程通信机制支持多机模型训练，并提供了一套更高级别的 API，用于自动微分跨多台机器拆分的模型。

警告

RPC 包中的 API 目前处于稳定且维护模式。

警告

CUDA 支持是一项 beta 功能。RPC 包的并非所有功能都与 CUDA 支持兼容，因此不建议在 CUDA 环境下使用。这些不支持的功能包括：RRefs、JIT 兼容性、分布式自动微分（dist autograd）、分布式优化器（dist optimizer）以及性能分析（profiling）。

注意

请参考 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 获取 Future（未来对象），并在调用端需要返回值时等待该 Future。当需求是远程创建某物但不需要将其获取到调用端时，remote() API 非常有用。设想一个驱动进程正在设置参数服务器和训练器的情况。驱动程序可以在参数服务器上创建一个嵌入表，然后将该嵌入表的引用共享给训练器，但驱动程序本身在本地并不会使用该嵌入表。在这种情况下，rpc_sync() 和 rpc_async() 就不再适用，因为它们总是意味着返回值会立即或在未来返回给调用端。

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

  3. 分布式自动微分 将涉及前向传递的所有工作节点上的本地自动微分引擎拼接在一起，并在反向传递期间自动联系它们以计算梯度。如果前向传递需要在执行分布式模型并行训练、参数服务器训练等时跨越多个机器，这将非常有帮助。有了这个功能，用户代码不再需要担心如何跨 RPC 边界发送梯度，也不需要担心以何种顺序启动本地自动微分引擎（在前向传递中存在嵌套和相互依赖的 RPC 调用时，这可能会变得非常复杂）。

  4. 分布式优化器 的构造函数接受一个 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 or WorkerInfo or int) – 目标工作节点的名称/rank/WorkerInfo。

  • func (Callable) – 可调用函数，例如 Python 可调用对象、内置运算符（如 add()）和带注释的 TorchScript 函数。

  • args (tuple) – func 调用的参数元组。

  • kwargs (dict) – func 调用的关键字参数字典。

  • timeout (float, optional) – 此 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 or WorkerInfo or int) – 目标工作节点的名称/rank/WorkerInfo。

  • func (Callable) – 可调用函数，例如 Python 可调用对象、内置运算符（如 add()）和带注释的 TorchScript 函数。

  • args (tuple) – func 调用的参数元组。

  • kwargs (dict) – func 调用的关键字参数字典。

  • timeout (float, optional) – 此 RPC 的超时时间（秒）。如果 RPC 在此时间内未完成，将引发指示超时的异常。值为 0 表示无限超时，即永远不会引发超时错误。如果未提供，则使用初始化期间或通过 _set_rpc_timeout 设置的默认值。

返回:

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

警告

不支持将 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 or WorkerInfo or 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 不会复制参数张量的存储，直到通过网络发送它们。调用者应确保这些张量的内容在返回的 RRef 被所有者确认之前保持完整，这可以通过 torch.distributed.rpc.RRef.confirmed_by_owner() API 进行检查。

警告

remote API 的超时等错误以尽力而为（best-effort）的方式处理。这意味着当由 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 进程到达此方法。

警告

对于由 rpc_async() 返回的 Future 对象，在 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, 可选) – TensorPipeAgent 用于执行请求的线程池中的线程数（默认值：16）。

  • rpc_timeout (float, 可选) – RPC 请求的默认超时时间（以秒为单位，默认值：60 秒）。如果 RPC 未在此时间范围内完成，将抛出异常。如有必要，调用者可以在 rpc_sync() 和 rpc_async() 中为单个 RPC 覆盖此超时设置。

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

  • device_maps (Dict[str, Dict], 可选) – 从当前工作节点到被调用节点的设备放置映射。键是被调用工作节点的名称，值是将当前工作节点设备映射到被调用工作节点设备的字典（Dict，由 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, 由int, str, 或 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, 由int, str, 或 torch.device 组成) – TensorPipe RPC 代理使用的本地设备。

注意

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

RRef#

警告

使用 CUDA 张量时当前不支持 RRef

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

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

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

警告

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: 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 的值应为标量张量。

参数:

  • 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: SupportsFloat = -1.0) → object#

创建一个辅助代理，以便轻松地使用 RRef 的所有者作为目的地来启动 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, 可选) – rref.remote() 的超时时间。如果此 RRef 的创建未能在超时时间内成功完成，则下次尝试使用该 RRef（例如 to_here）时，将引发超时。如果未提供，将使用默认 RPC 超时。请参阅 rpc.remote() 获取 RRef 的具体超时语义。

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

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

创建一个辅助代理，以便通过使用 RRef 的所有者作为目的地，轻松启动 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]#

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) – 希望放置此模块的目标 worker 上的设备。格式应为 “<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[RRef[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 以运行多次反向传播。

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

在使用分布式自动微分时，用于封装正向和反向传递的上下文对象。with 语句中生成的 context_id 是在所有 worker 上唯一标识分布式反向传递所必需的。每个 worker 都会存储与此 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 自动微分的更多信息

分布式优化器#

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

设计说明#

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

RRef 设计说明涵盖了框架用于引用远程 worker 上值的 RRef (远程引用) 协议的设计。

教程#

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

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