pytorch-fsdp

Pytorch-Fsdp 技能

为 pytorch-fsdp 开发提供全面协助，内容依据官方文档生成。

何时使用此技能

此技能应在以下情况下触发：

• 使用 pytorch-fsdp 时
• 询问 pytorch-fsdp 的功能或 API 时
• 实现 pytorch-fsdp 解决方案时
• 调试 pytorch-fsdp 代码时
• 学习 pytorch-fsdp 最佳实践时

快速参考

常见模式

模式 1： 通用加入上下文管理器# 创建日期：2025 年 6 月 6 日 | 最后更新日期：2025 年 6 月 6 日 通用加入上下文管理器用于在不均匀输入下进行分布式训练。本页概述了相关类 Join、Joinable 和 JoinHook 的 API。有关教程，请参阅使用 Join 上下文管理器进行不均匀输入的分布式训练。class torch.distributed.algorithms.Join(joinables, enable=True, throw_on_early_termination=False, **kwargs)[source]# 该类定义了通用加入上下文管理器，它允许在进程加入后调用自定义钩子。这些钩子应遮蔽未加入进程的集合通信，以防止挂起和报错，并确保算法的正确性。有关钩子定义的详细信息，请参阅 JoinHook。警告 上下文管理器要求每个参与的 Joinable 在其各自的每轮迭代集合通信之前调用 notify_join_context() 方法，以确保正确性。警告 上下文管理器要求 JoinHook 对象中的所有 process_group 属性相同。如果有多个 JoinHook 对象，则使用第一个对象的设备。进程组和设备信息用于检查未加入的进程，以及在启用 throw_on_early_termination 时通知进程抛出异常，这两项操作都使用 all-reduce。参数 joinables (List[Joinable]) – 参与的 Joinable 列表；它们的钩子会按给定顺序进行迭代。enable (bool) – 启用不均匀输入检测的标志；设置为 False 会禁用上下文管理器的功能，仅当用户知道输入不会不均匀时才应设置（默认值：True）。throw_on_early_termination (bool) – 控制是否在检测到不均匀输入时抛出异常的标志（默认值：False）。示例：>>> import os >>> import torch >>> import torch.distributed as dist >>> import torch.multiprocessing as mp >>> import torch.nn.parallel.DistributedDataParallel as DDP >>> import torch.distributed.optim.ZeroRedundancyOptimizer as ZeRO >>> from torch.distributed.algorithms.join import Join >>> >>> # 在每个生成的 worker 上 >>> def worker(rank): >>> dist.init_process_group("nccl", rank=rank, world_size=2) >>> model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank]) >>> optim = ZeRO(model.parameters(), torch.optim.Adam, lr=0.01) >>> # Rank 1 比 rank 0 多获得一个输入 >>> inputs = [torch.tensor([1.]).to(rank) for _ in range(10 + rank)] >>> with Join([model, optim]): >>> for input in inputs: >>> loss = model(input).sum() >>> loss.backward() >>> optim.step() >>> # 所有 rank 都能到达此处，而不会挂起或报错 static notify_join_context(joinable)[source]# 通知加入上下文管理器调用进程尚未加入。然后，如果 throw_on_early_termination=True，则检查是否检测到不均匀输入（即是否已有进程加入），如果检测到则抛出异常。此方法应由 Joinable 对象在其每轮迭代集合通信之前调用。例如，应在 DistributedDataParallel 中前向传播的开头调用它。只有传入上下文管理器的第一个 Joinable 对象会在该方法中执行集合通信，而对于其他对象，该方法为空操作。参数 joinable (Joinable) – 调用此方法的 Joinable 对象。返回 如果 joinable 是传入上下文管理器的第一个对象，则返回用于通知上下文管理器该进程尚未加入的 all-reduce 的异步工作句柄；否则返回 None。class torch.distributed.algorithms.Joinable[source]# 该类为可加入类定义了抽象基类。可加入类（继承自 Joinable）除了实现返回设备和进程组信息的 join_device() 和 join_process_group() 之外，还应实现返回 JoinHook 实例的 join_hook()。abstract property join_device: device# 返回用于执行加入上下文管理器所需集合通信的设备。abstract join_hook(**kwargs)[source]# 为给定的 Joinable 返回一个 JoinHook 实例。参数 kwargs (dict) – 一个字典，包含用于在运行时修改加入钩子行为的任意关键字参数；所有共享同一个加入上下文管理器的 Joinable 实例都会转发相同的 kwargs 值。返回类型 JoinHook abstract property join_process_group: Any# 返回加入上下文管理器自身所需的集合通信的进程组。class torch.distributed.algorithms.JoinHook[source]# 该类定义了加入钩子，它在加入上下文管理器中提供两个切入点。切入点：一个主钩子，在存在未加入进程时反复调用；以及一个后置钩子，在所有进程都已加入后调用一次。要为通用加入上下文管理器实现加入钩子，请定义一个继承自 JoinHook 的类，并按需重写 main_hook() 和 post_hook()。main_hook()[source]# 在存在未加入进程时调用此钩子，以在训练迭代中遮蔽集合通信。训练迭代即一次前向传播、反向传播和优化器步骤。post_hook(is_last_joiner)[source]# 在所有进程都已加入后调用该钩子。它会被传入一个额外的布尔参数 is_last_joiner，用于指示该 rank 是否为最后加入的 rank 之一。参数 is_last_joiner (bool) – 如果该 rank 是最后加入的 rank 之一，则为 True；否则为 False。

Join

**模式 2：**分布式通信包 - torch.distributed# 创建时间：2017 年 7 月 12 日 | 最后更新时间：2025 年 9 月 4 日

Note 请参阅 PyTorch Distributed Overview，简要了解分布式训练相关的所有功能。

Backends# torch.distributed 支持四种内置后端，各自具有不同的功能。下表列出了每种后端在 CPU 或 GPU 上可用的函数。对于 NCCL，GPU 指 CUDA GPU；对于 XCCL，GPU 指 XPU GPU。只有在用于构建 PyTorch 的 MPI 实现支持 CUDA 时，MPI 才支持 CUDA。

后端	gloo	mpi	nccl	xccl
Device	CPU GPU	CPU GPU	CPU GPU	CPU GPU
send	✓ ✘	✓ ?	✘ ✓	✘ ✓
recv	✓ ✘	✓ ?	✘ ✓	✘ ✓
broadcast	✓ ✓	✓ ?	✘ ✓	✘ ✓
all_reduce	✓ ✓	✓ ?	✘ ✓	✘ ✓
reduce	✓ ✓	✓ ?	✘ ✓	✘ ✓
all_gather	✓ ✓	✓ ?	✘ ✓	✘ ✓
gather	✓ ✓	✓ ?	✘ ✓	✘ ✓
scatter	✓ ✓	✓ ?	✘ ✓	✘ ✓
reduce_scatter	✓ ✓	✘ ✘	✘ ✓	✘ ✓
all_to_all	✓ ✓	✓ ?	✘ ✓	✘ ✓
barrier	✓ ✘	✓ ?	✘ ✓	✘ ✓

PyTorch 内置的后端# PyTorch 分布式包支持 Linux（稳定版）、MacOS（稳定版）和 Windows（原型版）。默认情况下，Linux 版本会构建并包含 Gloo 和 NCCL 后端（只有在使用 CUDA 构建时才包含 NCCL）。MPI 是可选后端，只有从源代码构建 PyTorch 时才能包含。（例如，在安装了 MPI 的主机上构建 PyTorch。）

Note 从 PyTorch v1.8 起，Windows 支持除 NCCL 之外的所有集合通信后端。如果 init_process_group() 的 init_method 参数指向一个文件，则必须符合以下格式：

本地文件系统，init_method="file:///d:/tmp/some_file"

共享文件系统，init_method="file://////{machine_name}/{share_folder_name}/some_file"

与 Linux 平台相同，可通过设置环境变量 MASTER_ADDR 和 MASTER_PORT 来启用 TcpStore。

应该使用哪个后端？# 过去，我们经常被问到：“我应该使用哪个后端？”

经验法则

使用 NCCL 后端进行基于 CUDA GPU 的分布式训练。

使用 XCCL 后端进行基于 XPU GPU 的分布式训练。

使用 Gloo 后端进行基于 CPU 的分布式训练。

使用 InfiniBand 互连的 GPU 主机

使用 NCCL，因为它是目前唯一支持 InfiniBand 和 GPUDirect 的后端。

使用 Ethernet 互连的 GPU 主机

使用 NCCL，因为它目前能提供最佳的分布式 GPU 训练性能，尤其是多进程单节点或多节点分布式训练。如果遇到 NCCL 问题，请使用 Gloo 作为备用选项。（请注意，Gloo 在 GPU 上的运行速度目前比 NCCL 慢。）

使用 InfiniBand 互连的 CPU 主机

如果 InfiniBand 已启用 IP over IB，请使用 Gloo；否则，请改用 MPI。我们计划在即将发布的版本中为 Gloo 添加 InfiniBand 支持。

使用 Ethernet 互连的 CPU 主机

除非有特定理由使用 MPI，否则请使用 Gloo。

常见环境变量# 选择要使用的网络接口# 默认情况下，NCCL 和 Gloo 后端都会尝试找到合适的网络接口。如果自动检测到的接口不正确，可以使用以下环境变量覆盖设置（适用于相应的后端）：

NCCL_SOCKET_IFNAME，例如 export NCCL_SOCKET_IFNAME=eth0

GLOO_SOCKET_IFNAME，例如 export GLOO_SOCKET_IFNAME=eth0

如果使用 Gloo 后端，可以用逗号分隔多个接口，如下所示：export GLOO_SOCKET_IFNAME=eth0,eth1,eth2,eth3。后端会以轮询方式在这些接口上分发操作。所有进程都必须在此变量中指定相同数量的接口。

其他 NCCL 环境变量# 调试——如果 NCCL 失败，可以设置 NCCL_DEBUG=INFO 来打印明确的警告消息以及基本的 NCCL 初始化信息。还可以使用 NCCL_DEBUG_SUBSYS 获取有关 NCCL 特定方面的更多详细信息。例如，NCCL_DEBUG_SUBSYS=COLL 会打印集合调用的日志，这有助于调试卡住问题，尤其是由集合类型或消息大小不匹配导致的问题。如果拓扑检测失败，可以设置 NCCL_DEBUG_SUBSYS=GRAPH 来检查详细的检测结果，并保存为参考资料，以便在需要 NCCL 团队进一步协助时使用。

性能调优——NCCL 会根据拓扑检测结果自动调优，减少用户的调优工作。在某些基于 socket 的系统上，用户仍可尝试调优 NCCL_SOCKET_NTHREADS 和 NCCL_NSOCKS_PERTHREAD，以提高 socket 网络带宽。NCCL 已针对 AWS 或 GCP 等部分云服务提供商预先调优了这两个环境变量。完整的 NCCL 环境变量列表请参阅 NVIDIA NCCL 官方文档。

还可以使用 torch.distributed.ProcessGroupNCCL.NCCLConfig 和 torch.distributed.ProcessGroupNCCL.Options 进一步调优 NCCL 通信器。可在解释器中使用 help（例如 help(torch.distributed.ProcessGroupNCCL.NCCLConfig)）了解更多信息。

基础知识# torch.distributed 包为在一台或多台机器上运行的多个计算节点之间进行多进程并行提供 PyTorch 支持和通信原语。torch.nn.parallel.DistributedDataParallel() 类基于此功能，通过包装任意 PyTorch 模型来提供同步分布式训练。与 Multiprocessing 包 torch.multiprocessing 和 torch.nn.DataParallel() 提供的并行方式不同，它支持多台通过网络连接的机器，并且用户必须为每个进程显式启动一份主训练脚本。

在单机同步场景中，torch.distributed 或 torch.nn.parallel.DistributedDataParallel() 包装器仍可能优于其他数据并行方法，包括 torch.nn.DataParallel()：

每个进程都会维护自己的优化器，并在每次迭代中执行完整的优化步骤。虽然这看起来是多余的，因为梯度已在各进程间收集并求平均，因此每个进程上的梯度都相同，但这样就不需要广播参数，从而减少了节点间传输张量的时间。

每个进程都包含独立的 Python 解释器，消除了由单个 Python 进程驱动多个执行线程、模型副本或 GPU 所产生的额外解释器开销和“GIL 争用”。对于大量使用 Python 运行时的模型，这一点尤为重要，包括含有循环层或许多小组件的模型。

初始化# 在调用其他任何方法之前，必须使用 torch.distributed.init_process_group() 或 torch.distributed.device_mesh.init_device_mesh() 函数初始化此包。两者都会阻塞，直到所有进程都加入。

Warning 初始化不是线程安全的。进程组创建应由单个线程执行，以避免不同 rank 之间的“UUID”分配不一致，并防止初始化期间发生可能导致卡住的竞争条件。

torch.distributed.is_available()[source]# 如果分布式包可用，则返回 True。否则，torch.distributed 不会公开任何其他 API。目前，torch.distributed 在 Linux、MacOS 和 Windows 上可用。从源代码构建 PyTorch 时，设置 USE_DISTRIBUTED=1 可启用此功能。目前，Linux 和 Windows 上 USE_DISTRIBUTED=1 为默认值，MacOS 上 USE_DISTRIBUTED=0 为默认值。

返回类型 bool

torch.distributed.init_process_group(backend=None, init_method=None, timeout=None, world_size=-1, rank=-1, store=None, group_name='', pg_options=None, device_id=None)[source]# 初始化默认分布式进程组。这也会初始化分布式包。

初始化进程组主要有两种方法：

显式指定 store、rank 和 world_size。

指定 init_method（URL 字符串），指示如何以及在哪里发现对等进程。可以选择指定 rank 和 world_size，或在 URL 中编码所有必需参数并省略它们。如果两者均未指定，则默认将 init_method 设为“env://”。

参数

backend（str 或 Backend，可选）——要使用的后端。根据构建时配置，有效值包括 mpi、gloo、nccl、ucc、xccl，或第三方插件注册的后端。自 2.6 起，如果未提供 backend，c10d 将使用为 device_id 关键字参数所指示的设备类型注册的后端（如果提供了 device_id）。目前已知的默认注册项为：cuda 对应 nccl，cpu 对应 gloo，xpu 对应 xccl。如果既未提供 backend，也未提供 device_id，c10d 会检测运行时机器上的加速器，并使用为检测到的加速器（或 cpu）注册的后端。此字段可使用小写字符串（例如 "gloo"），也可通过 Backend 属性访问（例如 Backend.GLOO）。如果使用 nccl 后端在每台机器上运行多个进程，每个进程都必须独占访问其使用的每个 GPU，因为多个进程共享 GPU 可能导致死锁或 NCCL 使用无效。ucc 后端为实验性功能。可使用 get_default_backend_for_device() 查询设备的默认后端。

init_method（str，可选）——指定如何初始化进程组的 URL。如果未指定 init_method 或 store，则默认值为“env://”。不能与 store 同时使用。

world_size（int，可选）——参与作业的进程数。指定 store 时为必填项。

rank（int，可选）——当前进程的 rank（应为 0 到 world_size-1 之间的数字）。指定 store 时为必填项。

store（Store，可选）——所有工作进程都可访问的键值存储，用于交换连接/地址信息。不能与 init_method 同时使用。

timeout（timedelta，可选）——进程组操作的超时时间。NCCL 的默认值为 10 分钟，其他后端为 30 分钟。超过此时间后，集合操作将被异步中止，进程也会崩溃。这是因为 CUDA 执行是异步的，一旦异步 NCCL 操作失败，后续 CUDA 操作可能会在损坏的数据上运行，因此继续执行用户代码已不再安全。设置 TORCH_NCCL_BLOCKING_WAIT 时，进程会阻塞并等待此超时时间。

group_name（str，可选，已弃用）——组名。此参数会被忽略。

pg_options（ProcessGroupOptions，可选）——进程组选项，用于指定构造特定进程组时要传入的附加选项。目前，我们支持的唯一选项是 nccl 后端的 ProcessGroupNCCL.Options；可设置 is_high_priority_stream，使 nccl 后端在有计算内核等待时选择高优先级 cuda 流。有关 nccl 的其他可用配置选项，请参阅 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t

device_id（torch.device | int，可选）——此进程将使用的单个特定设备，可启用特定于后端的优化。目前此参数仅对 NCCL 有两种效果：立即创建通信器（立即调用 ncclCommInit*，而不是通常的延迟调用）；在可能的情况下，子组会使用 ncclCommSplit，避免不必要的组创建开销。若要尽早了解 NCCL 初始化错误，也可使用此字段。如果提供的是 int，API 会假定使用编译时的加速器类型。

Note 要启用 backend == Backend.MPI，必须在支持 MPI 的系统上从源代码构建 PyTorch。

Note 多后端支持目前为实验性功能。当前未指定后端时，会同时创建 gloo 和 nccl 后端。CPU 张量的集合操作将使用 gloo 后端，CUDA 张量的集合操作将使用 nccl 后端。可传入格式为“<device_type>:<backend_name>,<device_type>:<backend_name>”的字符串来指定自定义后端，例如“cpu:gloo,cuda:custom_backend”。

torch.distributed.device_mesh.init_device_mesh(device_type, mesh_shape, *, mesh_dim_names=None, backend_override=None)[source]# 根据 device_type、mesh_shape 和 mesh_dim_names 参数初始化 DeviceMesh。此函数会创建一个 n 维数组布局的 DeviceMesh，其中 n 为 mesh_shape 的长度。如果提供了 mesh_dim_names，则每个维度都会标记为 mesh_dim_names[i]。

Note init_device_mesh 遵循 SPMD 编程模型，这意味着集群中的所有进程/rank 都会运行同一个 PyTorch Python 程序。请确保所有 rank 上的 mesh_shape（描述设备布局的 nD 数组维度）都相同。mesh_shape 不一致可能导致卡住。

Note 如果找不到进程组，init_device_mesh 会在后台初始化分布式通信所需的进程组。

参数 device_type（str）——网格的设备类型。目前支持：“cpu”、“cuda/cuda-like”、“xpu”。不允许传入带 GPU 索引的设备类型，例如“cuda:0”。

mesh_shape（Tuple[int]）——定义多维数组维度的元组，用于描述设备布局。

mesh_dim_names（Tuple[str]，可选）——网格维度名称元组，用于为描述设备布局的多维数组的每个维度命名。其长度必须与 mesh_shape 的长度一致。mesh_dim_names 中的每个字符串都必须唯一。

backend_override（Dict[int | str, tuple[str, Options] | str | Options]，可选）——覆盖为每个网格维度创建的部分或全部 ProcessGroup 的设置。每个键可以是维度索引，也可以是维度名称（如果提供了 mesh_dim_names）。每个值可以是包含后端名称及其选项的元组，也可以只提供这两个部分之一（此时另一部分将使用默认值）。

返回值 一个表示设备布局的 DeviceMesh 对象。

返回类型 DeviceMesh

示例：>>> from torch.distributed.device_mesh import init_device_mesh >>> >>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,)) >>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp"))

torch.distributed.is_initialized()[source]# 检查默认进程组是否已初始化。

返回类型 bool

torch.distributed.is_mpi_available()[source]# 检查 MPI 后端是否可用。

返回类型 bool

torch.distributed.is_nccl_available()[source]# 检查 NCCL 后端是否可用。

返回类型 bool

torch.distributed.is_gloo_available()[source]# 检查 Gloo 后端是否可用。

返回类型 bool

torch.distributed.distributed_c10d.is_xccl_available()[source]# 检查 XCCL 后端是否可用。

返回类型 bool

torch.distributed.is_torchelastic_launched()[source]# 检查此进程是否通过 torch.distributed.elastic（又称 torchelastic）启动。此函数使用 TORCHELASTIC_RUN_ID 环境变量是否存在，作为判断当前进程是否由 torchelastic 启动的依据。这是合理的近似判断，因为 TORCHELASTIC_RUN_ID 对应 rendezvous ID，该值始终非空，表示用于发现对等进程的作业 ID。

返回类型 bool

torch.distributed.get_default_backend_for_device(device)[source]# 返回给定设备的默认后端。

参数 device（Union[str, torch.device]）——要查询默认后端的设备。

返回值 给定设备对应的默认后端，以小写字符串形式返回。

返回类型 str

目前支持三种初始化方法：

TCP 初始化# 使用 TCP 初始化有两种方法，两者都需要一个所有进程都可访问的网络地址，以及指定的 world_size。第一种方法要求指定属于 rank 0 进程的地址。此初始化方法要求所有进程手动指定 rank。请注意，最新的分布式包已不再支持多播地址。group_name 也已弃用。

import torch.distributed as dist # 使用其中一台机器的地址 dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4)

共享文件系统初始化# 另一种初始化方法是使用由组内所有机器共享并可见的文件系统，并指定 world_size。URL 应以 file:// 开头，并包含共享文件系统上一个不存在的文件（位于现有目录中）的路径。文件系统初始化会在文件不存在时自动创建该文件，但不会删除它。因此，在下一次对相同文件路径/名称调用 init_process_group() 之前，你有责任确保清理该文件。请注意，最新的分布式包已不再支持自动分配 rank，group_name 也已弃用。

Warning 此方法假定文件系统支持使用 fcntl 锁定——大多数本地系统和 NFS 都支持。

Warning 此方法始终会创建文件，并尽力在程序结束时清理和删除文件。换句话说，使用文件 init 方法的每次初始化都需要一个全新的空文件才能成功。如果再次使用上一次初始化所用的同一个文件（该文件可能未能清理），则会出现意外行为，并且常常导致死锁和失败。因此，即使此方法会尽力清理文件，如果自动删除失败，你仍有责任确保在训练结束时删除文件，以防下次再次使用该文件。如果计划使用相同文件名多次调用 init_process_group()，这一点尤其重要。也就是说，如果文件未删除/清理，而你再次使用该文件调用 init_process_group()，则预期会发生故障。经验法则是：每次调用 init_process_group() 时，都要确保文件不存在或为空。

import torch.distributed as dist # 始终应指定 rank dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile', world_size=4, rank=args.rank)

环境变量初始化# 此方法会从环境变量读取配置，从而允许用户完全自定义信息的获取方式。需要设置的变量如下：

MASTER_PORT - 必需；必须是 rank 0 所在机器上的空闲端口

MASTER_ADDR - 必需（rank 0 除外）；rank 0 节点的地址

WORLD_SIZE - 必需；可在此处设置，也可在调用初始化函数时设置

RANK - 必需；可在此处设置，也可在调用初始化函数时设置

rank 0 所在机器将用于建立所有连接。这是默认方法，因此无需指定 init_method（也可以指定为 env://）。

缩短初始化时间# TORCH_GLOO_LAZY_INIT——按需建立连接，而不使用全连接网格；这可以大幅缩短非 all2all 操作的初始化时间。

初始化后# 执行 torch.distributed.init_process_group() 后，可以使用以下函数。要检查进程组是否已初始化，请使用 torch.distributed.is_initialized()。

class torch.distributed.Backend(name)[source]# 类似枚举的后端类。可用后端包括：GLOO、NCCL、UCC、MPI、XCCL，以及其他已注册的后端。此类的值为小写字符串，例如 "gloo"。可通过属性访问，例如 Backend.NCCL。也可以直接调用此类来解析字符串，例如 Backend(backend_str) 会检查 backend_str 是否有效；如果有效，则返回解析后的小写字符串。也接受大写字符串，例如 Backend("GLOO") 会返回 "gloo"。

Note Backend.UNDEFINED 条目确实存在，但仅用作某些字段的初始值。用户不应直接使用它，也不应假定它必然存在。

classmethod register_backend(name, func, extended_api=False, devices=None)[source]# 使用给定名称和实例化函数注册新后端。第三方 ProcessGroup 扩展使用此类方法注册新后端。

参数 name（str）——ProcessGroup 扩展的后端名称。它应与 init_process_group() 中的名称一致。

func（function）——用于实例化后端的函数处理程序。该函数应在后端扩展中实现，并接受四个参数：store、rank、world_size 和 timeout。

extended_api（bool，可选）——后端是否支持扩展参数结构。默认值：False。如果设为 True，后端将获得 c10d::DistributedBackendOptions 实例，以及后端实现所定义的进程组选项对象。

device（str 或 str 列表，可选）——此后端支持的设备类型，例如“cpu”、“cuda”等。如果为 None，则假设同时支持“cpu”和“cuda”。

Note 对第三方后端的支持目前为实验性功能，可能会发生变化。

torch.distributed.get_backend(group=None)[source]# 返回给定进程组的后端。

参数 group（ProcessGroup，可选）——要使用的进程组。默认值为通用主进程组。如果指定了其他特定组，调用进程必须是该组的成员。

返回值 给定进程组的后端，以小写字符串形式返回。

返回类型 Backend

torch.distributed.get_rank(group=None)[source]# 返回当前进程在所提供组中的 rank，否则返回默认组中的 rank。Rank 是分配给分布式进程组中每个进程的唯一标识符。其值始终为从 0 到 world_size 的连续整数。

参数 group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

返回值 进程组中的 rank；如果不是组成员，则为 -1。

返回类型 int

torch.distributed.get_world_size(group=None)[source]# 返回当前进程组中的进程数。

参数 group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

返回值 进程组的 world size；如果不是组成员，则为 -1。

返回类型 int

关闭# 退出时调用 destroy_process_group() 清理资源很重要。最简单的做法是在训练脚本不再需要通信时（通常接近 main() 末尾），使用 group 参数的默认值 None 调用 destroy_process_group()，销毁每个进程组和后端。每个训练器进程都应调用一次，而不是在外层进程启动器级别调用。如果在超时时间内，某个 pg 中的所有 rank 未调用 destroy_process_group()，尤其是在应用中有多个进程组（例如用于 N 维并行）时，退出时可能会卡住。这是因为 ProcessGroupNCCL 的析构函数会调用 ncclCommAbort，而该调用必须由所有进程共同执行；但如果由 Python 的 GC 调用，ProcessGroupNCCL 析构函数的调用顺序并不确定。调用 destroy_process_group() 有助于确保所有 rank 以一致顺序调用 ncclCommAbort，并避免在 ProcessGroupNCCL 的析构函数中调用 ncclCommAbort。

重新初始化# destroy_process_group 也可用于销毁单个进程组。一个用例是容错训练：运行期间可能销毁一个进程组，然后初始化一个新的进程组。在这种情况下，销毁之后、随后重新初始化之前，必须使用 torch.distributed 原语之外的某种方式同步训练器进程。由于难以实现这种同步，目前不支持/未测试此行为，并将其视为已知问题。如果这是阻碍你使用的场景，请提交 github issue 或 RFC。

组# 默认情况下，集合操作会在默认组（也称 world）上运行，并要求所有进程都调用分布式函数。不过，某些工作负载可以从更细粒度的通信中受益。这时就需要使用分布式组。new_group() 函数可用于创建新组，并从所有进程中任意选择子集。它会返回一个不透明的组句柄，可作为 group 参数传递给所有集合操作（集合操作是用于按某些常见编程模式交换信息的分布式函数）。

torch.distributed.new_group(ranks=None, timeout=None, backend=None, pg_options=None, use_local_synchronization=False, group_desc=None, device_id=None)[source]# 创建新的分布式组。主组中的所有进程（即属于分布式作业的所有进程）都必须调用此函数，即使它们不会成为该组的成员。此外，所有进程都应按相同顺序创建组。

Warning 并发使用安全性：使用 NCCL 后端的多个进程组时，用户必须确保所有 rank 上的集合操作执行顺序全局一致。如果进程中的多个线程发出集合操作，则需要显式同步，以确保顺序一致。使用 torch.distributed 通信 API 的异步变体时，会返回一个 work 对象，并将通信内核排入单独的 CUDA 流，以便通信与计算重叠。在某个进程组上发出一个或多个异步操作后，必须调用 work.wait()，使它们与其他 cuda 流同步，然后才能使用另一个进程组。详情请参阅 Using multiple NCCL communicators concurrently https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently。

参数 ranks（list[int]）——组成员的 rank 列表。如果为 None，则设置为所有 rank。默认值为 None。

timeout（timedelta，可选）——详情和默认值请参阅 init_process_group。

backend（str 或 Backend，可选）——要使用的后端。根据构建时配置，有效值为 gloo 和 nccl。默认使用与全局组相同的后端。此字段应使用小写字符串（例如 "gloo"），也可通过 Backend 属性访问（例如 Backend.GLOO）。如果传入 None，则使用默认进程组对应的后端。默认值为 None。

pg_options（ProcessGroupOptions，可选）——进程组选项，用于指定构造特定进程组时要传入的附加选项。例如，对于 nccl 后端，可设置 is_high_priority_stream，使进程组能够使用高优先级 cuda 流。有关 nccl 的其他可用配置选项，请参阅 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t

use_local_synchronization（bool，可选）：在创建进程组结束时执行组内屏障。不同之处在于，非成员 rank 不需要调用此 API，也不会加入屏障。

group_desc（str，可选）——用于描述进程组的字符串。

device_id（torch.device，可选）——将此进程“绑定”到的单个特定设备。如果提供了此字段，new_group 调用会尝试立即为该设备初始化通信后端。

返回值 可传递给集合调用的分布式组句柄；如果该 rank 不属于 ranks，则返回 GroupMember.NON_GROUP_MEMBER。

N.B. use_local_synchronization 不适用于 MPI。

N.B. 虽然 use_local_synchronization=True 在较大的集群和较小的进程组中可能显著更快，但必须谨慎使用，因为它会改变集群行为：非成员 rank 不会加入组屏障。

N.B. 当每个 rank 创建多个重叠进程组时，use_local_synchronization=True 可能导致死锁。为避免这种情况，请确保所有 rank 遵循相同的全局创建顺序。

torch.distributed.get_group_rank(group, global_rank)[source]# 将全局 rank 转换为组内 rank。global_rank 必须属于 group，否则会引发 RuntimeError。

参数 group（ProcessGroup）——要查询相对 rank 的 ProcessGroup。

global_rank（int）——要查询的全局 rank。

返回值 global_rank 相对于 group 的组内 rank。

返回类型 int

N.B. 在默认进程组上调用此函数会返回原值。

torch.distributed.get_global_rank(group, group_rank)[source]# 将组内 rank 转换为全局 rank。group_rank 必须属于 group，否则会引发 RuntimeError。

参数 group（ProcessGroup）——用于获取全局 rank 的 ProcessGroup。

group_rank（int）——要查询的组内 rank。

返回值 group_rank 相对于 group 的全局 rank。

返回类型 int

N.B. 在默认进程组上调用此函数会返回原值。

torch.distributed.get_process_group_ranks(group)[source]# 获取与 group 关联的所有 rank。

参数 group（Optional[ProcessGroup]）——用于获取所有 rank 的 ProcessGroup。如果为 None，则使用默认进程组。

返回值 按组内 rank 排序的全局 rank 列表。

返回类型 list[int]

DeviceMesh# DeviceMesh 是一种更高层的抽象，用于管理进程组（或 NCCL 通信器）。它可以让用户轻松创建节点间和节点内的进程组，而无需操心如何为不同子进程组正确设置 rank，并有助于轻松管理这些分布式进程组。init_device_mesh() 函数可用于创建新的 DeviceMesh，并通过网格形状描述设备拓扑。

class torch.distributed.device_mesh.DeviceMesh(device_type, mesh, *, mesh_dim_names=None, backend_override=None, _init_backend=True)[source]# DeviceMesh 表示设备网格。设备布局可以表示为 n 维数组，n 维数组中的每个值都是默认进程组的全局 rank ID。DeviceMesh 可用于设置集群中的 N 维设备连接，并管理 N 维并行所需的 ProcessGroup。通信可以分别在 DeviceMesh 的每个维度上进行。

DeviceMesh 会遵循用户已选择的设备（即，如果用户在初始化 DeviceMesh 之前调用 torch.cuda.set_device），如果用户事先未设置设备，则会为当前进程选择/设置设备。请注意，必须在初始化 DeviceMesh 之前手动选择设备。与 DTensor API 配合使用时，DeviceMesh 也可用作上下文管理器。

Note DeviceMesh 遵循 SPMD 编程模型，这意味着集群中的所有进程/rank 都运行同一个 PyTorch Python 程序。因此，用户必须确保所有 rank 上的网格数组（用于描述设备布局）完全相同。网格不一致会导致无提示卡住。

参数 device_type（str）——网格的设备类型。目前支持：“cpu”、“cuda/cuda-like”。

mesh（ndarray）——描述设备布局的多维数组或整数张量，其中 ID 是默认进程组的全局 ID。

返回值 表示设备布局的 DeviceMesh 对象。

返回类型 DeviceMesh

以下程序以 SPMD 方式在每个进程/rank 上运行。此示例中有 2 台主机，每台主机有 4 个 GPU。对 mesh 的第一个维度执行归约，会跨列归约 (0, 4)、……和 (3, 7)；对 mesh 的第二个维度执行归约，会跨行归约 (0, 1, 2, 3) 和 (4, 5, 6, 7)。

示例：>>> from torch.distributed.device_mesh import DeviceMesh >>> >>> # 将设备网格初始化为 (2, 4)，表示拓扑结构 >>> # 主机间（dim 0）和主机内（dim 1）。 >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])

static from_group(group, device_type, mesh=None, *, mesh_dim_names=None)[source]# 使用现有 ProcessGroup 或现有 ProcessGroup 列表构造指定 device_type 的 DeviceMesh。构造的设备网格维度数等于传入的组数。例如，如果传入单个进程组，生成的 DeviceMesh 为 1D 网格；如果传入 2 个进程组的列表，生成的 DeviceMesh 为 2D 网格。如果传入多个组，则必须提供 mesh 和 mesh_dim_names 参数。传入进程组的顺序决定网格拓扑。例如，第一个进程组将成为 DeviceMesh 的第 0 个维度。传入的 mesh 张量必须与传入的进程组数量具有相同维度，且 mesh 张量中的维度顺序必须与进程组的顺序一致。

参数 group（ProcessGroup 或 list[ProcessGroup]）——现有 ProcessGroup 或现有 ProcessGroup 列表。

device_type（str）——网格的设备类型。目前支持：“cpu”、“cuda/cuda-like”。不允许传入带 GPU 索引的设备类型，例如“cuda:0”。

mesh（torch.Tensor 或 ArrayLike，可选）——描述设备布局的多维数组或整数张量，其中 ID 是默认进程组的全局 ID。默认值为 None。

mesh_dim_names（tuple[str]，可选）——网格维度名称元组，用于为描述设备布局的多维数组的每个维度命名。其长度必须与 mesh_shape 的长度一致。mesh_dim_names 中的每个字符串都必须唯一。默认值为 None。

返回值 表示设备布局的 DeviceMesh 对象。

返回类型 DeviceMesh

get_all_groups()[source]# 返回网格所有维度的 ProcessGroup 列表。

返回值 ProcessGroup 对象列表。

返回类型 list[torch.distributed.distributed_c10d.ProcessGroup]

get_coordinate()[source]# 返回此 rank 相对于网格所有维度的索引。如果此 rank 不属于网格，则返回 None。

返回类型 Optional[list[int]]

get_group(mesh_dim=None)[source]# 返回由 mesh_dim 指定的单个 ProcessGroup；如果未指定 mesh_dim 且 DeviceMesh 是一维，则返回网格中唯一的 ProcessGroup。

参数 mesh_dim（str/python:int，可选）——可以是网格维度的名称或索引。默认值为 None。

返回值 ProcessGroup 对象。

返回类型 ProcessGroup

get_local_rank(mesh_dim=None)[source]# 返回 DeviceMesh 给定 mesh_dim 的本地 rank。

参数 mesh_dim（str/python:int，可选）——可以是网格维度的名称或索引。默认值为 None。

返回值 表示本地 rank 的整数。

返回类型 int

以下程序以 SPMD 方式在每个进程/rank 上运行。此示例中有 2 台主机，每台主机有 4 个 GPU。在 rank 0、1、2、3 上调用 mesh_2d.get_local_rank(mesh_dim=0) 会返回 0。在 rank 4、5、6、7 上调用 mesh_2d.get_local_rank(mesh_dim=0) 会返回 1。在 rank 0、4 上调用 mesh_2d.get_local_rank(mesh_dim=1) 会返回 0。在 rank 1、5 上调用 mesh_2d.get_local_rank(mesh_dim=1) 会返回 1。在 rank 2、6 上调用 mesh_2d.get_local_rank(mesh_dim=1) 会返回 2。在 rank 3、7 上调用 mesh_2d.get_local_rank(mesh_dim=1) 会返回 3。

示例：>>> from torch.distributed.device_mesh import DeviceMesh >>> >>> # 将设备网格初始化为 (2, 4)，表示拓扑结构 >>> # 主机间（dim 0）和主机内（dim 1）。 >>> mesh = DeviceMesh(device_type="cuda", mesh=[[0, 1, 2, 3],[4, 5, 6, 7]])

get_rank()[source]# 返回当前全局 rank。

返回类型 int

点对点通信# torch.distributed.send(tensor, dst=None, group=None, tag=0, group_dst=None)[source]# 同步发送张量。

Warning NCCL 后端不支持 tag。

参数 tensor（Tensor）——要发送的张量。

dst（int）——全局进程组中的目标 rank（无论 group 参数为何）。目标 rank 不得与当前进程的 rank 相同。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

tag（int，可选）——用于将 send 与远程 recv 匹配的标记。

group_dst（int，可选）——组内目标 rank。不能同时指定 dst 和 group_dst。

torch.distributed.recv(tensor, src=None, group=None, tag=0, group_src=None)[source]# 同步接收张量。

Warning NCCL 后端不支持 tag。

参数 tensor（Tensor）——用于填入接收数据的张量。

src（int，可选）——全局进程组中的源 rank（无论 group 参数为何）。如果未指定，则从任意进程接收。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

tag（int，可选）——用于将 recv 与远程 send 匹配的标记。

group_src（int，可选）——组内目标 rank。不能同时指定 src 和 group_src。

返回值 发送方 rank；如果不是组成员，则为 -1。

返回类型 int

使用 isend() 和 irecv() 时会返回分布式请求对象。通常，此对象的类型未指定，因为不应手动创建它们，但保证支持以下两个方法：

is_completed()——如果操作已完成，则返回 True。

wait()——阻塞进程，直到操作完成。is_completed() 一旦返回，就保证返回 True。

torch.distributed.isend(tensor, dst=None, group=None, tag=0, group_dst=None)[source]# 异步发送张量。

Warning 请求完成前修改张量会导致未定义行为。

Warning NCCL 后端不支持 tag。

与会阻塞的 send 不同，isend 允许 src == dst rank，即发送给自身。

参数 tensor（Tensor）——要发送的张量。

dst（int）——全局进程组中的目标 rank（无论 group 参数为何）。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

tag（int，可选）——用于将 send 与远程 recv 匹配的标记。

group_dst（int，可选）——组内目标 rank。不能同时指定 dst 和 group_dst。

返回值 分布式请求对象；如果不是组成员，则为 None。

返回类型 Optional[Work]

torch.distributed.irecv(tensor, src=None, group=None, tag=0, group_src=None)[source]# 异步接收张量。

Warning NCCL 后端不支持 tag。

与会阻塞的 recv 不同，irecv 允许 src == dst rank，即从自身接收。

参数 tensor（Tensor）——用于填入接收数据的张量。

src（int，可选）——全局进程组中的源 rank（无论 group 参数为何）。如果未指定，则从任意进程接收。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

tag（int，可选）——用于将 recv 与远程 send 匹配的标记。

group_src（int，可选）——组内目标 rank。不能同时指定 src 和 group_src。

返回值 分布式请求对象；如果不是组成员，则为 None。

返回类型 Optional[Work]

torch.distributed.send_object_list(object_list, dst=None, group=None, device=None, group_dst=None, use_batch=False)[source]# 同步发送 object_list 中可 pickle 的对象。与 send() 类似，但可传递 Python 对象。请注意，object_list 中的所有对象都必须可 pickle 才能发送。

参数 object_list（List[Any]）——要发送的输入对象列表。每个对象都必须可 pickle。接收方必须提供大小相同的列表。

dst（int）——要向其发送 object_list 的目标 rank。目标 rank 基于全局进程组（无论 group 参数为何）。

group（Optional[ProcessGroup]）——（ProcessGroup，可选）：要使用的进程组。如果为 None，则使用默认进程组。默认值为 None。

device（torch.device，可选）——如果不为 None，则对象会被序列化并转换为张量，然后移动到该设备再发送。默认值为 None。

group_dst（int，可选）——组内目标 rank。必须指定 dst 和 group_dst 之一，不能同时指定两者。

use_batch（bool，可选）——如果为 True，则使用批量 p2p 操作，而不是常规发送操作。这样可以避免初始化 2-rank 通信器，并使用现有的整个组通信器。用法和前提条件请参阅 batch_isend_irecv。默认值为 False。

返回值 None。

Note 对于基于 NCCL 的进程组，通信开始前，必须将对象的内部张量表示移动到 GPU 设备。在这种情况下，所用设备由 torch.cuda.current_device() 指定。用户有责任确保通过 torch.cuda.set_device() 为每个 rank 设置独立 GPU。

Warning 对象集合操作存在若干严重的性能和可扩展性限制。详情请参阅 Object collectives。

Warning send_object_list() 会隐式使用 pickle 模块，该模块已知存在安全风险。可以构造恶意 pickle 数据，在反 pickle 时执行任意代码。此函数只能处理可信数据。

Warning 使用 GPU 张量调用 send_object_list() 支持不佳且效率低下，因为张量会被 pickle，从而产生 GPU -> CPU 传输。请考虑改用 send()。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> # 假设后端不是 NCCL >>> device = torch.device("cpu") >>> if dist.get_rank() == 0: >>> # 假设 world_size 为 2。 >>> objects = ["foo", 12, {1: 2}] # 任意可 pickle 对象 >>> dist.send_object_list(objects, dst=1, device=device) >>> else: >>> objects = [None, None, None] >>> dist.recv_object_list(objects, src=0, device=device) >>> objects ['foo', 12, {1: 2}]

torch.distributed.recv_object_list(object_list, src=None, group=None, device=None, group_src=None, use_batch=False)[source]# 同步接收 object_list 中可 pickle 的对象。与 recv() 类似，但可以接收 Python 对象。

参数 object_list（List[Any]）——用于接收对象的列表。必须提供与发送列表大小相同的列表。

src（int，可选）——接收 object_list 的源 rank。源 rank 基于全局进程组（无论 group 参数为何）。如果设为 None，则从任意 rank 接收。默认值为 None。

group（Optional[ProcessGroup]）——（ProcessGroup，可选）：要使用的进程组。如果为 None，则使用默认进程组。默认值为 None。

device（torch.device，可选）——如果不为 None，则在此设备上接收。默认值为 None。

group_src（int，可选）——组内目标 rank。不能同时指定 src 和 group_src。

use_batch（bool，可选）——如果为 True，则使用批量 p2p 操作，而不是常规发送操作。这样可以避免初始化 2-rank 通信器，并使用现有的整个组通信器。用法和前提条件请参阅 batch_isend_irecv。默认值为 False。

返回值 发送方 rank。如果不是组成员，则为 -1。如果该 rank 是组成员，则 object_list 将包含来自 src rank 的已发送对象。

Note 对于基于 NCCL 的进程组，通信开始前，必须将对象的内部张量表示移动到 GPU 设备。在这种情况下，所用设备由 torch.cuda.current_device() 指定。用户有责任确保通过 torch.cuda.set_device() 为每个 rank 设置独立 GPU。

Warning 对象集合操作存在若干严重的性能和可扩展性限制。详情请参阅 Object collectives。

Warning recv_object_list() 会隐式使用 pickle 模块，该模块已知存在安全风险。可以构造恶意 pickle 数据，在反 pickle 时执行任意代码。此函数只能处理可信数据。

Warning 使用 GPU 张量调用 recv_object_list() 支持不佳且效率低下，因为张量会被 pickle，从而产生 GPU -> CPU 传输。请考虑改用 recv()。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> # 假设后端不是 NCCL >>> device = torch.device("cpu") >>> if dist.get_rank() == 0: >>> # 假设 world_size 为 2。 >>> objects = ["foo", 12, {1: 2}] # 任意可 pickle 对象 >>> dist.send_object_list(objects, dst=1, device=device) >>> else: >>> objects = [None, None, None] >>> dist.recv_object_list(objects, src=0, device=device) >>> objects ['foo', 12, {1: 2}]

torch.distributed.batch_isend_irecv(p2p_op_list)[source]# 异步发送或接收一批张量，并返回请求列表。处理 p2p_op_list 中的每个操作并返回对应请求。目前支持 NCCL、Gloo 和 UCC 后端。

参数 p2p_op_list（list[torch.distributed.distributed_c10d.P2POp]）——点对点操作列表（每个操作的类型为 torch.distributed.P2POp）。列表中 isend/irecv 的顺序很重要，必须与远端对应的 isend/irecv 顺序一致。

返回值 调用 op_list 中相应操作所返回的分布式请求对象列表。

返回类型 list[torch.distributed.distributed_c10d.Work]

示例 >>> send_tensor = torch.arange(2, dtype=torch.float32) + 2 * rank >>> recv_tensor = torch.randn(2, dtype=torch.float32) >>> send_op = dist.P2POp(dist.isend, send_tensor, (rank + 1) % world_size) >>> recv_op = dist.P2POp( ... dist.irecv, recv_tensor, (rank - 1 + world_size) % world_size ... ) >>> reqs = batch_isend_irecv([send_op, recv_op]) >>> for req in reqs: >>> req.wait() >>> recv_tensor tensor([2, 3]) # Rank 0 tensor([0, 1]) # Rank 1

Note 请注意，当此 API 与 NCCL PG 后端配合使用时，用户必须使用 torch.cuda.set_device 设置当前 GPU 设备，否则可能导致意外卡住。此外，如果此 API 是传递给 dist.P2POp 的组中的首次集合调用，则组内所有 rank 都必须参与此 API 调用；否则，行为未定义。如果此 API 不是组中的首次集合调用，则允许只涉及组内部分 rank 的批量 P2P 操作。

class torch.distributed.P2POp(op, tensor, peer=None, group=None, tag=0, group_peer=None)[source]# 用于构造 batch_isend_irecv 点对点操作的类。此类会构造 P2P 操作类型、通信缓冲区、对等 rank、Process Group 和 tag。此类的实例将传递给 batch_isend_irecv 以进行点对点通信。

参数 op（Callable）——向对等进程发送数据或从其接收数据的函数。op 的类型为 torch.distributed.isend 或 torch.distributed.irecv。

tensor（Tensor）——要发送或接收的张量。

peer（int，可选）——目标或源 rank。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

tag（int，可选）——用于将 send 与 recv 匹配的标记。

group_peer（int，可选）——目标或源 rank。

同步和异步集合操作# 每个集合操作函数都支持以下两种操作，具体取决于传入集合操作的 async_op 标志设置：

同步操作——默认模式，即 async_op 设为 False。函数返回时，可以保证集合操作已执行。对于 CUDA 操作，不能保证 CUDA 操作已完成，因为 CUDA 操作是异步的。对于 CPU 集合操作，任何使用集合调用输出的后续函数调用都会按预期运行。对于 CUDA 集合操作，在同一 CUDA 流上使用输出的函数调用会按预期运行。如果在不同流上运行，用户必须自行处理同步。有关流同步等 CUDA 语义的详情，请参阅 CUDA Semantics。以下脚本展示了 CPU 和 CUDA 操作在这些语义上的差异示例。

异步操作——将 async_op 设为 True 时。集合操作函数会返回分布式请求对象。通常无需手动创建该对象，并且它保证支持以下方法：

is_completed()——对于 CPU 集合操作，如果操作已完成，则返回 True。对于 CUDA 操作，如果操作已成功排入 CUDA 流，并且无需进一步同步即可在默认流上使用输出，则返回 True。

wait()——对于 CPU 集合操作，会阻塞进程，直到操作完成。对于 CUDA 操作，会阻塞当前活动的 CUDA 流，直到操作完成（但不会阻塞 CPU）。

get_future()——返回 torch._C.Future 对象。NCCL 支持此方法；GLOO 和 MPI 的大多数操作也支持，但不包括点对点操作。

Note：随着我们继续采用 Futures 并合并 API，get_future() 调用可能会变得多余。

示例 以下代码可作为使用分布式集合操作执行 CUDA 操作时语义的参考。它展示了在不同 CUDA 流上使用集合操作输出时明确进行同步的必要性：

代码在每个 rank 上运行。 dist.init_process_group("nccl", rank=rank, world_size=2) output = torch.tensor([rank]).cuda(rank) s = torch.cuda.Stream() handle = dist.all_reduce(output, async_op=True) # wait 可确保操作已排入队列，但不一定已经完成。 handle.wait() # 在非默认流上使用结果。 with torch.cuda.stream(s): s.wait_stream(torch.cuda.default_stream()) output.add_(100) if rank == 0: # 如果省略显式调用 wait_stream，下面的输出将不确定地为 1 或 101， # 取决于 allreduce 是否在 add 完成后覆盖该值。 print(output)

集合函数# torch.distributed.broadcast(tensor, src=None, group=None, async_op=False, group_src=None)[source]# 向整个组广播张量。参与集合操作的所有进程中的 tensor 必须具有相同数量的元素。

参数 tensor（Tensor）——如果 src 是当前进程的 rank，则为要发送的数据；否则为用于保存接收数据的张量。

src（int）——全局进程组中的源 rank（无论 group 参数为何）。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

group_src（int）——组内源 rank。必须指定 group_src 和 src 之一，不能同时指定两者。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

torch.distributed.broadcast_object_list(object_list, src=None, group=None, device=None, group_src=None)[source]# 向整个组广播 object_list 中可 pickle 的对象。与 broadcast() 类似，但可以传递 Python 对象。请注意，object_list 中的所有对象都必须可 pickle 才能广播。

参数 object_list（List[Any]）——要广播的输入对象列表。每个对象都必须可 pickle。只有 src rank 上的对象会被广播，但每个 rank 都必须提供大小相同的列表。

src（int）——广播 object_list 的源 rank。源 rank 基于全局进程组（无论 group 参数为何）。

group（Optional[ProcessGroup]）——（ProcessGroup，可选）：要使用的进程组。如果为 None，则使用默认进程组。默认值为 None。

device（torch.device，可选）——如果不为 None，则对象会被序列化并转换为张量，然后移动到该设备再广播。默认值为 None。

group_src（int）——组内源 rank。不能同时指定 group_src 和 src。

返回值 None。如果 rank 是组成员，object_list 将包含来自 src rank 的广播对象。

Note 对于基于 NCCL 的进程组，通信开始前，必须将对象的内部张量表示移动到 GPU 设备。在这种情况下，所用设备由 torch.cuda.current_device() 指定。用户有责任确保通过 torch.cuda.set_device() 为每个 rank 设置独立 GPU。

Note 请注意，此 API 与 broadcast() 集合操作略有不同：它不提供 async_op 句柄，因此是阻塞调用。

Warning 对象集合操作存在若干严重的性能和可扩展性限制。详情请参阅 Object collectives。

Warning broadcast_object_list() 会隐式使用 pickle 模块，该模块已知存在安全风险。可以构造恶意 pickle 数据，在反 pickle 时执行任意代码。此函数只能处理可信数据。

Warning 使用 GPU 张量调用 broadcast_object_list() 支持不佳且效率低下，因为张量会被 pickle，从而产生 GPU -> CPU 传输。请考虑改用 broadcast()。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> if dist.get_rank() == 0: >>> # 假设 world_size 为 3。 >>> objects = ["foo", 12, {1: 2}] # 任意可 pickle 对象 >>> else: >>> objects = [None, None, None] >>> # 假设后端不是 NCCL >>> device = torch.device("cpu") >>> dist.broadcast_object_list(objects, src=0, device=device) >>> objects ['foo', 12, {1: 2}]

torch.distributed.all_reduce(tensor, op=<RedOpType.SUM: 0>, group=None, async_op=False)[source]# 在所有机器之间归约张量数据，使所有机器都获得最终结果。调用后，所有进程中的 tensor 都将逐位完全相同。支持复数张量。

参数 tensor（Tensor）——集合操作的输入和输出。函数会原地操作。

op（可选）——torch.distributed.ReduceOp 枚举中的一个值。指定逐元素归约操作。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

示例 >>> # 以下所有张量均为 torch.int64 类型。 >>> # 有 2 个进程组、2 个 rank。 >>> device = torch.device(f"cuda:{rank}") >>> tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank >>> tensor tensor([1, 2], device='cuda:0') # Rank 0 tensor([3, 4], device='cuda:1') # Rank 1 >>> dist.all_reduce(tensor, op=ReduceOp.SUM) >>> tensor tensor([4, 6], device='cuda:0') # Rank 0 tensor([4, 6], device='cuda:1') # Rank 1

torch.distributed.reduce(tensor, dst=None, op=<RedOpType.SUM: 0>, group=None, async_op=False, group_dst=None)[source]# 在所有机器之间归约张量数据。只有 rank 为 dst 的进程会收到最终结果。

参数 tensor（Tensor）——集合操作的输入和输出。函数会原地操作。

dst（int）——全局进程组中的目标 rank（无论 group 参数为何）。

op（可选）——torch.distributed.ReduceOp 枚举中的一个值。指定逐元素归约操作。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

group_dst（int）——组内目标 rank。必须指定 group_dst 和 dst 之一，不能同时指定两者。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

torch.distributed.all_gather(tensor_list, tensor, group=None, async_op=False)[source]# 将整个组的张量收集到一个列表中。支持复数张量和大小不等的张量。

参数 tensor_list（list[Tensor]）——输出列表。应包含大小正确的张量，用作集合操作的输出。支持大小不等的张量。

tensor（Tensor）——要从当前进程广播的张量。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

示例 >>> # 以下所有张量均为 torch.int64 dtype。 >>> # 有 2 个进程组、2 个 rank。 >>> device = torch.device(f"cuda:{rank}") >>> tensor_list = [ ... torch.zeros(2, dtype=torch.int64, device=device) for _ in range(2) ... ] >>> tensor_list [tensor([0, 0], device='cuda:0'), tensor([0, 0], device='cuda:0')] # Rank 0 [tensor([0, 0], device='cuda:1'), tensor([0, 0], device='cuda:1')] # Rank 1 >>> tensor = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank >>> tensor tensor([1, 2], device='cuda:0') # Rank 0 tensor([3, 4], device='cuda:1') # Rank 1 >>> dist.all_gather(tensor_list, tensor) >>> tensor_list [tensor([1, 2], device='cuda:0'), tensor([3, 4], device='cuda:0')] # Rank 0 [tensor([1, 2], device='cuda:1'), tensor([3, 4], device='cuda:1')] # Rank 1

torch.distributed.all_gather_into_tensor(output_tensor, input_tensor, group=None, async_op=False)[source]# 从所有 rank 收集张量，并将其放入单个输出张量。此函数要求每个进程上的所有张量大小相同。

参数 output_tensor（Tensor）——用于容纳所有 rank 张量元素的输出张量。其大小必须正确，并采用以下形式之一：(i) 沿主维度拼接所有输入张量；有关“拼接”的定义，请参阅 torch.cat()；(ii) 沿主维度堆叠所有输入张量；有关“堆叠”的定义，请参阅 torch.stack()。下面的示例会更清楚地说明支持的输出形式。

input_tensor（Tensor）——要从当前 rank 收集的张量。与 all_gather API 不同，此 API 要求所有 rank 上的输入张量大小相同。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

示例 >>> # 以下所有张量均为 torch.int64 dtype，且位于 CUDA 设备上。 >>> # 有两个 rank。 >>> device = torch.device(f"cuda:{rank}") >>> tensor_in = torch.arange(2, dtype=torch.int64, device=device) + 1 + 2 * rank >>> tensor_in tensor([1, 2], device='cuda:0') # Rank 0 tensor([3, 4], device='cuda:1') # Rank 1 >>> # 拼接形式的输出 >>> tensor_out = torch.zeros(world_size * 2, dtype=torch.int64, device=device) >>> dist.all_gather_into_tensor(tensor_out, tensor_in) >>> tensor_out tensor([1, 2, 3, 4], device='cuda:0') # Rank 0 tensor([1, 2, 3, 4], device='cuda:1') # Rank 1 >>> # 堆叠形式的输出 >>> tensor_out2 = torch.zeros(world_size, 2, dtype=torch.int64, device=device) >>> dist.all_gather_into_tensor(tensor_out2, tensor_in) >>> tensor_out2 tensor([[1, 2], [3, 4]], device='cuda:0') # Rank 0 tensor([[1, 2], [3, 4]], device='cuda:1') # Rank 1

torch.distributed.all_gather_object(object_list, obj, group=None)[source]# 将整个组中可 pickle 的对象收集到一个列表中。与 all_gather() 类似，但可以传递 Python 对象。请注意，对象必须可 pickle 才能收集。

参数 object_list（list[Any]）——输出列表。其大小应正确，等于此集合操作所在组的大小，并用于存放输出。

obj（Any）——要从当前进程广播的可 pickle Python 对象。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。默认值为 None。

返回值 None。如果调用 rank 属于此组，集合操作的输出将填入 object_list。如果调用 rank 不属于此组，则传入的 object_list 不会被修改。

Note 请注意，此 API 与 all_gather() 集合操作略有不同：它不提供 async_op 句柄，因此是阻塞调用。

Note 对于基于 NCCL 的进程组，通信开始前，必须将对象的内部张量表示移动到 GPU 设备。在这种情况下，所用设备由 torch.cuda.current_device() 指定。用户有责任确保通过 torch.cuda.set_device() 为每个 rank 设置独立 GPU。

Warning 对象集合操作存在若干严重的性能和可扩展性限制。详情请参阅 Object collectives。

Warning all_gather_object() 会隐式使用 pickle 模块，该模块已知存在安全风险。可以构造恶意 pickle 数据，在反 pickle 时执行任意代码。此函数只能处理可信数据。

Warning 使用 GPU 张量调用 all_gather_object() 支持不佳且效率低下，因为张量会被 pickle，从而产生 GPU -> CPU 传输。请考虑改用 all_gather()。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> # 假设 world_size 为 3。 >>> gather_objects = ["foo", 12, {1: 2}] # 任意可 pickle 对象 >>> output = [None for _ in gather_objects] >>> dist.all_gather_object(output, gather_objects[dist.get_rank()]) >>> output ['foo', 12, {1: 2}]

torch.distributed.gather(tensor, gather_list=None, dst=None, group=None, async_op=False, group_dst=None)[source]# 在单个进程中收集张量列表。此函数要求每个进程上的所有张量大小相同。

参数 tensor（Tensor）——输入张量。

gather_list（list[Tensor]，可选）——用于存放收集数据的、大小适当且一致的张量列表（默认值为 None；必须在目标 rank 上指定）。

dst（int，可选）——全局进程组中的目标 rank（无论 group 参数为何）。（如果 dst 和 group_dst 都为 None，则默认为全局 rank 0。）

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

group_dst（int，可选）——组内目标 rank。不能同时指定 dst 和 group_dst。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

Note 请注意，gather_list 中的所有 Tensor 必须大小相同。

示例::>>> # 有 2 个进程组、2 个 rank。 >>> tensor_size = 2 >>> device = torch.device(f'cuda:{rank}') >>> tensor = torch.ones(tensor_size, device=device) + rank >>> if dist.get_rank() == 0: >>> gather_list = [torch.zeros_like(tensor, device=device) for i in range(2)] >>> else: >>> gather_list = None >>> dist.gather(tensor, gather_list, dst=0) >>> # Rank 0 获取收集到的数据。 >>> gather_list [tensor([1., 1.], device='cuda:0'), tensor([2., 2.], device='cuda:0')] # Rank 0 None # Rank 1

torch.distributed.gather_object(obj, object_gather_list=None, dst=None, group=None, group_dst=None)[source]# 在单个进程中收集整个组中可 pickle 的对象。与 gather() 类似，但可以传递 Python 对象。请注意，对象必须可 pickle 才能收集。

参数 obj（Any）——输入对象。必须可 pickle。

object_gather_list（list[Any]）——输出列表。在 dst rank 上，其大小应正确，等于此集合操作所在组的大小，并用于存放输出。在非 dst rank 上必须为 None。（默认值为 None。）

dst（int，可选）——全局进程组中的目标 rank（无论 group 参数为何）。（如果 dst 和 group_dst 都为 None，则默认为全局 rank 0。）

group（Optional[ProcessGroup]）——（ProcessGroup，可选）：要使用的进程组。如果为 None，则使用默认进程组。默认值为 None。

group_dst（int，可选）——组内目标 rank。不能同时指定 dst 和 group_dst。

返回值 None。在 dst rank 上，object_gather_list 将包含集合操作的输出。

Note 请注意，此 API 与 gather 集合操作略有不同：它不提供 async_op 句柄，因此是阻塞调用。

Note 对于基于 NCCL 的进程组，通信开始前，必须将对象的内部张量表示移动到 GPU 设备。在这种情况下，所用设备由 torch.cuda.current_device() 指定。用户有责任确保通过 torch.cuda.set_device() 为每个 rank 设置独立 GPU。

Warning 对象集合操作存在若干严重的性能和可扩展性限制。详情请参阅 Object collectives。

Warning gather_object() 会隐式使用 pickle 模块，该模块已知存在安全风险。可以构造恶意 pickle 数据，在反 pickle 时执行任意代码。此函数只能处理可信数据。

Warning 使用 GPU 张量调用 gather_object() 支持不佳且效率低下，因为张量会被 pickle，从而产生 GPU -> CPU 传输。请考虑改用 gather()。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> # 假设 world_size 为 3。 >>> gather_objects = ["foo", 12, {1: 2}] # 任意可 pickle 对象 >>> output = [None for _ in gather_objects] >>> dist.gather_object( ... gather_objects[dist.get_rank()], ... output if dist.get_rank() == 0 else None, ... dst=0 ... ) >>> # 在 rank 0 上 >>> output ['foo', 12, {1: 2}]

torch.distributed.scatter(tensor, scatter_list=None, src=None, group=None, async_op=False, group_src=None)[source]# 将张量列表分发给组内所有进程。每个进程都会接收一个张量，并将其数据存储在 tensor 参数中。支持复数张量。

参数 tensor（Tensor）——输出张量。

scatter_list（list[Tensor]）——要分发的张量列表（默认值为 None；必须在源 rank 上指定）。

src（int）——全局进程组中的源 rank（无论 group 参数为何）。（如果 src 和 group_src 都为 None，则默认为全局 rank 0。）

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

group_src（int，可选）——组内源 rank。不能同时指定 src 和 group_src。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

Note 请注意，scatter_list 中的所有 Tensor 必须大小相同。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> tensor_size = 2 >>> device = torch.device(f'cuda:{rank}') >>> output_tensor = torch.zeros(tensor_size, device=device) >>> if dist.get_rank() == 0: >>> # 假设 world_size 为 2。 >>> # 只能传入张量，且所有张量大小必须相同。 >>> t_ones = torch.ones(tensor_size, device=device) >>> t_fives = torch.ones(tensor_size, device=device) * 5 >>> scatter_list = [t_ones, t_fives] >>> else: >>> scatter_list = None >>> dist.scatter(output_tensor, scatter_list, src=0) >>> # Rank i 获取 scatter_list[i]。 >>> output_tensor tensor([1., 1.], device='cuda:0') # Rank 0 tensor([5., 5.], device='cuda:1') # Rank 1

torch.distributed.scatter_object_list(scatter_object_output_list, scatter_object_input_list=None, src=None, group=None, group_src=None)[source]# 将 scatter_object_input_list 中可 pickle 的对象分发给整个组。与 scatter() 类似，但可以传递 Python 对象。在每个 rank 上，分发的对象都会存储在 scatter_object_output_list 的第一个元素中。请注意，scatter_object_input_list 中的所有对象都必须可 pickle 才能分发。

参数 scatter_object_output_list（List[Any]）——非空列表，其第一个元素将存储分发给此 rank 的对象。

scatter_object_input_list（List[Any]，可选）——要分发的输入对象列表。每个对象都必须可 pickle。只有 src rank 上的对象会被分发；对于非 src rank，此参数可以为 None。

src（int）——分发 scatter_object_input_list 的源 rank。源 rank 基于全局进程组（无论 group 参数为何）。（如果 src 和 group_src 都为 None，则默认为全局 rank 0。）

group（Optional[ProcessGroup]）——（ProcessGroup，可选）：要使用的进程组。如果为 None，则使用默认进程组。默认值为 None。

group_src（int，可选）——组内源 rank。不能同时指定 src 和 group_src。

返回值 None。如果 rank 是组成员，scatter_object_output_list 的第一个元素将被设置为分发给此 rank 的对象。

Note 请注意，此 API 与 scatter 集合操作略有不同：它不提供 async_op 句柄，因此是阻塞调用。

Warning 对象集合操作存在若干严重的性能和可扩展性限制。详情请参阅 Object collectives。

Warning scatter_object_list() 会隐式使用 pickle 模块，该模块已知存在安全风险。可以构造恶意 pickle 数据，在反 pickle 时执行任意代码。此函数只能处理可信数据。

Warning 使用 GPU 张量调用 scatter_object_list() 支持不佳且效率低下，因为张量会被 pickle，从而产生 GPU -> CPU 传输。请考虑改用 scatter()。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> if dist.get_rank() == 0: >>> # 假设 world_size 为 3。 >>> objects = ["foo", 12, {1: 2}] # 任意可 pickle 对象 >>> else: >>> # 非 src rank 可以提供任意列表，其元素不会被使用。 >>> objects = [None, None, None] >>> output_list = [None] >>> dist.scatter_object_list(output_list, objects, src=0) >>> # Rank i 获取 objects[i]。例如，在 rank 2 上： >>> output_list [{1: 2}]

torch.distributed.reduce_scatter(output, input_list, op=<RedOpType.SUM: 0>, group=None, async_op=False)[source]# 先归约，再将张量列表分发给组内所有进程。

参数 output（Tensor）——输出张量。

input_list（list[Tensor]）——要归约并分发的张量列表。

op（可选）——torch.distributed.ReduceOp 枚举中的一个值。指定逐元素归约操作。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

torch.distributed.reduce_scatter_tensor(output, input, op=<RedOpType.SUM: 0>, group=None, async_op=False)[source]# 先归约，再将张量分发给组内所有 rank。

参数 output（Tensor）——输出张量。所有 rank 上的大小都应相同。

input（Tensor）——要归约并分发的输入张量。其大小应为输出张量大小乘以 world size。输入张量可以采用以下形状之一：(i) 沿主维度拼接输出张量；(ii) 沿主维度堆叠输出张量。有关“拼接”的定义，请参阅 torch.cat()；有关“堆叠”的定义，请参阅 torch.stack()。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

示例 >>> # 以下所有张量均为 torch.int64 dtype，且位于 CUDA 设备上。 >>> # 有两个 rank。 >>> device = torch.device(f"cuda:{rank}") >>> tensor_out = torch.zeros(2, dtype=torch.int64, device=device) >>> # 拼接形式的输入 >>> tensor_in = torch.arange(world_size * 2, dtype=torch.int64, device=device) >>> tensor_in tensor([0, 1, 2, 3], device='cuda:0') # Rank 0 tensor([0, 1, 2, 3], device='cuda:1') # Rank 1 >>> dist.reduce_scatter_tensor(tensor_out, tensor_in) >>> tensor_out tensor([0, 2], device='cuda:0') # Rank 0 tensor([4, 6], device='cuda:1') # Rank 1

torch.distributed.all_to_all_single(output, input, output_split_sizes=None, input_split_sizes=None, group=None, async_op=False)[source]# 将输入张量拆分，然后将拆分后的列表分发给组内所有进程。之后，接收的张量会从组内所有进程拼接为单个输出张量并返回。支持复数张量。

参数 output（Tensor）——拼接后的收集输出张量。

input（Tensor）——要分发的输入张量。

output_split_sizes——（list[Int]，可选）：dim 0 的输出拆分大小。如果为 None 或空列表，则输出张量的 dim 0 必须能被 world_size 整除。

input_split_sizes——（list[Int]，可选）：dim 0 的输入拆分大小。如果为 None 或空列表，则输入张量的 dim 0 必须能被 world_size 整除。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

Warning all_to_all_single 是实验性功能，可能会发生变化。

示例 >>> input = torch.arange(4) + rank * 4 >>> input tensor([0, 1, 2, 3]) # Rank 0 tensor([4, 5, 6, 7]) # Rank 1 tensor([8, 9, 10, 11]) # Rank 2 tensor([12, 13, 14, 15]) # Rank 3 >>> output = torch.empty([4], dtype=torch.int64) >>> dist.all_to_all_single(output, input) >>> output tensor([0, 4, 8, 12]) # Rank 0 tensor([1, 5, 9, 13]) # Rank 1 tensor([2, 6, 10, 14]) # Rank 2 tensor([3, 7, 11, 15]) # Rank 3 >>> # 本质上，它类似于以下操作： >>> scatter_list = list(input.chunk(world_size)) >>> gather_list = list(output.chunk(world_size)) >>> for i in range(world_size): >>> dist.scatter(gather_list[i], scatter_list if i == rank else [], src = i)

torch.distributed.all_to_all(output_tensor_list, input_tensor_list, group=None, async_op=False)[source]# 将输入张量列表分发给组内所有进程，并返回输出列表中收集到的张量列表。支持复数张量。

参数 output_tensor_list（list[Tensor]）——要收集的张量列表，每个 rank 对应一个张量。

input_tensor_list（list[Tensor]）——要分发的张量列表，每个 rank 对应一个张量。

group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

Warning all_to_all 是实验性功能，可能会发生变化。

torch.distributed.barrier(group=None, async_op=False, device_ids=None)[source]# 同步所有进程。如果 async_op 为 False，此集合操作会阻塞进程，直到整个组都进入此函数；如果为 True，则会阻塞，直到在异步 work 句柄上调用 wait()。

参数 group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

async_op（bool，可选）——此操作是否应为异步操作。

device_ids（[int]，可选）——设备/GPU ID 列表。只应提供一个 ID。

返回值 如果 async_op 设为 True，则返回异步 work 句柄；如果 async_op 为 False 或当前进程不是组成员，则返回 None。

Note ProcessGroupNCCL 现在会阻塞 CPU 线程，直到屏障集合操作完成。

Note ProcessGroupNCCL 将屏障实现为对 1 个元素张量执行 all_reduce。必须为该张量分配设备。设备选择顺序如下：(1) 如果 barrier 的 device_ids 参数不为 None，则使用其中第一个设备；(2) 如果 init_process_group 指定了设备，则使用该设备；(3) 如果已执行过使用张量输入的其他集合操作，则使用此进程组首次使用的设备；(4) 使用全局 rank 对本地设备数量取模所对应的设备索引。

torch.distributed.monitored_barrier(group=None, timeout=None, wait_all_ranks=False)[source]# 类似 torch.distributed.barrier，但可配置超时时间来同步进程。它能够报告未在给定超时时间内通过此屏障的 rank。具体来说，非零 rank 会阻塞，直到处理完来自 rank 0 的 send/recv。Rank 0 会阻塞，直到处理完来自其他 rank 的所有 send/recv，并报告未能及时响应的 rank。如果某个 rank 未到达 monitored_barrier（例如因卡住），所有其他 rank 在 monitored_barrier 中都会失败。此集合操作会阻塞组内所有进程/rank，直到整个组成功退出函数，因此适用于调试和同步。不过，它可能影响性能，因此仅应在调试或需要主机端完全同步的场景中使用。为了调试，可以在应用的集合调用之前插入此屏障，以检查是否有 rank 不同步。

Note 此集合操作仅支持 GLOO 后端。

参数 group（ProcessGroup，可选）——要使用的进程组。如果为 None，则使用默认进程组。

timeout（datetime.timedelta，可选）——monitored_barrier 的超时时间。如果为 None，则使用默认进程组的超时时间。

wait_all_ranks（bool，可选）——是否收集所有失败的 rank。默认值为 False，rank 0 上的 monitored_barrier 会在遇到第一个失败的 rank 时立即抛出异常，以便尽快失败。设为 wait_all_ranks=True 时，monitored_barrier 会收集所有失败的 rank，并抛出包含所有失败 rank 信息的错误。

返回值 None。

示例::>>> # 注意：每个 rank 上均省略了进程组初始化。 >>> import torch.distributed as dist >>> if dist.get_rank() != 1: >>> dist.monitored_barrier() # 引发异常，指出 >>> # rank 1 未调用 monitored_barrier。 >>> # wait_all_ranks=True 的示例 >>> if dist.get_rank() == 0: >>> dist.monitored_barrier(wait_all_ranks=True) # 引发异常 >>> # 指出 rank 1、2、……world_size - 1 未调用 >>> # monitored_barrier。

class torch.distributed.Work# Work 对象表示 PyTorch 分布式包中待完成异步操作的句柄。非阻塞集合操作（例如 dist.all_reduce(tensor, async_op=True)）会返回此对象。

block_current_stream(self: torch._C._distributed_c10d.Work) → None# 阻塞当前活动 GPU 流，直到操作完成。对于基于 GPU 的集合操作，这等同于 synchronize。对于使用 Gloo 等方式发起的 CPU 集合操作，此方法会阻塞 CUDA 流，直到操作完成。所有情况下都会立即返回。要检查操作是否成功，应异步检查 Work 对象的结果。

boxed(self: torch._C._distributed_c10d.Work) → object#

exception(self: torch._C._distributed_c10d.Work) → std::__exception_ptr::exception_ptr#

get_future(self: torch._C._distributed_c10d.Work) → torch.Future#

返回值 与 Work 完成相关联的 torch.futures.Future 对象。例如，可以通过 fut = process_group.allreduce(tensors).get_future() 获取 Future 对象。

示例::以下示例展示了一个简单的 allreduce DDP 通信钩子，该钩子使用 get_future API 获取与 allreduce 完成相关联的 Future。>>> def allreduce(process_group: dist.ProcessGroup, bucket: dist.GradBucket): -> torch.futures.Future >>> group_to_use = process_group if process_group is not None else torch.distributed.group.WORLD >>> tensor = bucket.buffer().div_(group_to_use.size()) >>> return torch.distributed.all_reduce(tensor, group=group_to_use, async_op=True).get_future() >>> ddp_model.register_comm_hook(state=None, hook=allreduce)

Warning get_future API 支持 NCCL，并部分支持 GLOO 和 MPI 后端（不支持 send/recv 等点对点操作），并会返回 torch.futures.Future。以上示例中，allreduce 操作会通过 NCCL 后端在 GPU 上执行；调用 fut.wait() 后，会在适当的 NCCL 流与 PyTorch 当前设备流同步后返回，以确保能够执行异步 CUDA 操作，而不必等待 GPU 上的整个操作完成。请注意，CUDAFuture 不支持 TORCH_NCCL_BLOCKING_WAIT 标志或 NCCL 的 barrier()。此外，如果通过 fut.then() 添加了回调函数，它会等待 WorkNCCL 的 NCCL 流与 ProcessGroupNCCL 的专用回调流同步，并在回调流上运行回调后以内联方式调用该回调。fut.then() 会返回另一个 CUDAFuture，其中包含回调的返回值和记录了回调流的 CUDAEvent。对于 CPU work，fut.done() 在 work 完成且 value() 张量就绪时返回 true。对于 GPU work，fut.done() 仅在操作已排入队列时返回 true。对于混合 CPU-GPU work（例如使用 GLOO 发送 GPU 张量），fut.done() 会在张量到达相应节点时返回 true，但不一定已在相应 GPU 上同步（与 GPU work 类似）。

get_future_result(self: torch._C._distributed_c10d.Work) → torch.Future#

返回值 类型为 int 的 torch.futures.Future 对象，其值对应 WorkResult 枚举类型。例如，可以通过 fut = process_group.allreduce(tensor).get_future_result() 获取 Future 对象。

示例::用户可以通过 fut.wait() 阻塞等待 work 完成，并通过 fut.value() 获取 WorkResult。用户还可以通过 fut.then(call_back_func) 注册回调函数，使其在 work 完成时调用，而不阻塞当前线程。

Warning get_future_result API 支持 NCCL。

is_completed(self: torch._C._distributed_c10d.Work) → bool#

is_success(self: torch._C._distributed_c10d.Work) → bool#

result(self: torch._C._distributed_c10d.Work) → list[torch.Tensor]#

source_rank(self: torch._C._distributed_c10d.Work) → int#

synchronize(self: torch._C._distributed_c10d.Work) → None#

static unbox(arg0: object) → torch._C._distributed_c10d.Work#

wait(self: torch._C._distributed_c10d.Work, timeout: datetime.timedelta = datetime.timedelta(0)) → bool#

返回 true/false。

示例:: try:work.wait(timeout) except:# 处理方式

Warning 正常情况下，用户无需设置超时时间。调用 wait() 与调用 synchronize() 相同：让当前流阻塞，直到 NCCL work 完成。但是，如果设置了超时时间，CPU 线程将阻塞，直到 NCCL work 完成或超时。如果超时，则会抛出异常。

class torch.distributed.ReduceOp# 可用归约操作的类枚举：SUM、PRODUCT、MIN、MAX、BAND、BOR、BXOR 和 PREMUL_SUM。使用 NCCL 后端时不可用 BAND、BOR 和 BXOR 归约。AVG 会先将值除以 world size，再跨 rank 求和。AVG 仅适用于 NCCL 后端，且只适用于 NCCL 2.10 或更高版本。PREMUL_SUM 会在归约前在本地将输入乘以给定标量。PREMUL_SUM 仅适用于 NCCL 后端，且只适用于 NCCL 2.11 或更高版本。用户应使用 torch.distributed._make_nccl_premul_sum。此外，复数张量不支持 MAX、MIN 和 PRODUCT。可通过属性访问此类的值，例如 ReduceOp.SUM。它们用于指定归约集合操作的策略，例如 reduce()。此类不支持 members 属性。

class torch.distributed.reduce_op# 已弃用的归约操作类枚举：SUM、PRODUCT、MIN 和 MAX。建议使用 ReduceOp。

分布式键值存储# 分布式包提供了分布式键值存储，可用于在组内进程之间共享信息，也可用于初始化 torch.distributed.init_process_group() 中的分布式包（显式创建存储，作为指定 init_method 的替代方式）。键值存储有 3 种选择：TCPStore、FileStore 和 HashStore。

class torch.distributed.Store# 所有存储实现的基类，例如 PyTorch 分布式提供的 3 种实现：TCPStore、FileStore 和 HashStore。

init(self: torch._C._distributed_c10d.Store) → None#

add(self: torch._C._distributed_c10d.Store, arg0: str, arg1: SupportsInt) → int# 对于给定键，首次调用 add 会在存储中创建一个与该键关联的计数器，并将其初始化为 amount。之后使用相同键调用 add 时，会按指定数量递增计数器。如果对已通过 set() 在存储中设置的键调用 add()，则会引发异常。

参数 key（str）——存储中要递增计数器的键。

amount（int）——计数器递增的数量。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用其他存储类型 >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.add("first_key", 1) >>> store.add("first_key", 6) >>> # 应返回 7 >>> store.get("first_key")

append(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str) → None# 根据提供的键和值，将键值对追加到存储中。如果存储中不存在该键，则会创建该键。

参数 key（str）——要追加到存储中的键。

value（str）——要添加到存储中的键对应的值。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.append("first_key", "po") >>> store.append("first_key", "tato") >>> # 应返回 "potato" >>> store.get("first_key")

check(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) → bool# 检查给定键列表是否都已在存储中设置值。正常情况下，此调用会立即返回，但在某些边缘死锁情况下仍可能出问题，例如在 TCPStore 销毁后调用 check。调用 check() 时会检查指定的键是否已存储。

参数 keys（list[str]）——要查询是否存储的键。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用其他存储类型 >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.add("first_key", 1) >>> # 应返回 7 >>> store.check(["first_key"])

clone(self: torch._C._distributed_c10d.Store) → torch._C._distributed_c10d.Store# 克隆存储并返回一个指向相同底层存储的新对象。返回的存储可以与原对象并发使用。此功能旨在通过为每个线程克隆一个存储，为多线程使用存储提供安全方式。

compare_set(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str, arg2: str) → bytes# 根据提供的键插入键值对，并在插入前比较 expected_value 与 desired_value。只有在存储中已存在该键的 expected_value，或 expected_value 为空字符串时，才会设置 desired_value。

参数 key（str）——要在存储中检查的键。

expected_value（str）——插入前要检查的键对应的值。

desired_value（str）——要添加到存储中的键对应的值。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("key", "first_value") >>> store.compare_set("key", "first_value", "second_value") >>> # 应返回 "second_value" >>> store.get("key")

delete_key(self: torch._C._distributed_c10d.Store, arg0: str) → bool# 从存储中删除与 key 关联的键值对。如果成功删除键，则返回 true；否则返回 false。

Warning delete_key API 仅由 TCPStore 和 HashStore 支持。在 FileStore 上使用此 API 会引发异常。

参数 key（str）——要从存储中删除的键。

返回值 如果 key 已删除，则返回 True；否则返回 False。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用 HashStore >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key") >>> # 应返回 true >>> store.delete_key("first_key") >>> # 应返回 false >>> store.delete_key("bad_key")

get(self: torch._C._distributed_c10d.Store, arg0: str) → bytes# 获取存储中给定键对应的值。如果存储中不存在该键，函数会等待到存储初始化时定义的超时时间，然后引发异常。

参数 key（str）——函数将返回此键对应的值。

返回值 如果 key 存在于存储中，则返回其对应的值。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key", "first_value") >>> # 应返回 "first_value" >>> store.get("first_key")

has_extended_api(self: torch._C._distributed_c10d.Store) → bool# 如果存储支持扩展操作，则返回 true。

multi_get(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) → list[bytes]# 获取 keys 中所有键的值。如果 keys 中任意键不存在于存储中，函数会等待至超时。

参数 keys（List[str]）——要从存储中获取的键。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key", "po") >>> store.set("second_key", "tato") >>> # 应返回 [b"po", b"tato"] >>> store.multi_get(["first_key", "second_key"])

multi_set(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str], arg1: collections.abc.Sequence[str]) → None# 根据提供的键和值列表，将键值对列表插入存储。

参数 keys（List[str]）——要插入的键。

values（List[str]）——要插入的值。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.multi_set(["first_key", "second_key"], ["po", "tato"]) >>> # 应返回 b"po" >>> store.get("first_key")

num_keys(self: torch._C._distributed_c10d.Store) → int# 返回存储中已设置的键数量。请注意，此数量通常比通过 set() 和 add() 添加的键数量多 1，因为其中有一个键用于协调使用此存储的所有工作进程。

Warning 与 TCPStore 配合使用时，num_keys 返回写入底层文件的键数量。如果存储被销毁后又使用同一文件创建另一个存储，则原有键会保留。

返回值 存储中的键数量。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用其他存储类型 >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key", "first_value") >>> # 应返回 2 >>> store.num_keys()

queue_len(self: torch._C._distributed_c10d.Store, arg0: str) → int# 返回指定队列的长度。如果队列不存在，则返回 0。更多信息请参阅 queue_push。

参数 key（str）——要获取长度的队列的键。

queue_pop(self: torch._C._distributed_c10d.Store, key: str, block: bool = True) → bytes# 从指定队列弹出一个值，或在队列为空时等待至超时。更多信息请参阅 queue_push。如果 block 为 False 且队列为空，则会引发 dist.QueueEmptyError。

参数 key（str）——要从中弹出元素的队列键。

block（bool）——是否阻塞并等待键，或立即返回。

queue_push(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str) → None# 向指定队列推送一个值。将同一键用于队列和 set/get 操作可能导致意外行为。队列支持 wait/check 操作。对队列执行 wait 时，只会唤醒一个等待中的工作进程，而不是全部。

参数 key（str）——要推送到的队列键。

value（str）——要推送到队列中的值。

set(self: torch._C._distributed_c10d.Store, arg0: str, arg1: str) → None# 根据提供的键和值，将键值对插入存储。如果该键已存在，则会用新提供的值覆盖旧值。

参数 key（str）——要添加到存储中的键。

value（str）——要添加到存储中的键对应的值。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set("first_key", "first_value") >>> # 应返回 "first_value" >>> store.get("first_key")

set_timeout(self: torch._C._distributed_c10d.Store, arg0: datetime.timedelta) → None# 设置存储的默认超时时间。此超时时间用于初始化以及 wait() 和 get()。

参数 timeout（timedelta）——要在存储中设置的超时时间。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用其他存储类型 >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> store.set_timeout(timedelta(seconds=10)) >>> # 10 秒后将引发异常 >>> store.wait(["bad_key"])

property timeout# 获取存储的超时时间。

wait(*args, **kwargs)# 重载函数。

wait(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str]) -> None

等待 keys 中的每个键被添加到存储中。如果在超时前（超时在存储初始化时设置）并非所有键都已设置，则 wait 会引发异常。

参数 keys（list）——要等待其在存储中被设置的键列表。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用其他存储类型 >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> # 30 秒后将引发异常 >>> store.wait(["bad_key"])

wait(self: torch._C._distributed_c10d.Store, arg0: collections.abc.Sequence[str], arg1: datetime.timedelta) -> None

等待 keys 中的每个键被添加到存储中。如果在提供的超时时间内未设置这些键，则会引发异常。

参数 keys（list）——要等待其在存储中被设置的键列表。

timeout（timedelta）——等待键被添加的时长，超时后将引发异常。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 以 TCPStore 为例，也可以使用其他存储类型 >>> store = dist.TCPStore("127.0.0.1", 0, 1, True, timedelta(seconds=30)) >>> # 10 秒后将引发异常 >>> store.wait(["bad_key"], timedelta(seconds=10))

class torch.distributed.TCPStore# 基于 TCP 的分布式键值存储实现。服务器存储负责保存数据，而客户端存储可通过 TCP 连接到服务器存储，并执行 set()（插入键值对）、get()（获取键值对）等操作。始终应初始化一个服务器存储，因为客户端存储会等待服务器建立连接。

参数 host_name（str）——服务器存储运行所在的主机名或 IP 地址。

port（int）——服务器存储监听传入请求的端口。

world_size（int，可选）——存储用户总数（客户端数量 + 服务器数量 1）。默认值为 None（None 表示存储用户数量不固定）。

is_master（bool，可选）——初始化服务器存储时为 True，初始化客户端存储时为 False。默认值为 False。

timeout（timedelta，可选）——存储初始化以及 get() 和 wait() 等方法使用的超时时间。默认值为 timedelta(seconds=300)。

wait_for_workers（bool，可选）——是否等待所有工作进程连接到服务器存储。仅当 world_size 为固定值时适用。默认值为 True。

multi_tenant（bool，可选）——如果为 True，则当前进程中所有 host/port 相同的 TCPStore 实例都会使用同一个底层 TCPServer。默认值为 False。

master_listen_fd（int，可选）——如果指定，底层 TCPServer 将使用此文件描述符进行监听；该文件描述符必须是已绑定到端口的 socket。要绑定临时端口，建议将 port 设为 0，然后读取 .port。默认值为 None（表示服务器会创建新 socket 并尝试将其绑定到 port）。

use_libuv（bool，可选）——如果为 True，则 TCPServer 后端使用 libuv。默认值为 True。

示例::>>> import torch.distributed as dist >>> from datetime import timedelta >>> # 在进程 1 上运行（服务器） >>> server_store = dist.TCPStore("127.0.0.1", 1234, 2, True, timedelta(seconds=30)) >>> # 在进程 2 上运行（客户端） >>> client_store = dist.TCPStore("127.0.0.1", 1234, 2, False) >>> # 初始化后，可从客户端或服务器调用任意存储方法 >>> server_store.set("first_key", "first_value") >>> client_store.get("first_key")

init(self: torch._C._distributed_c10d.TCPStore, host_name: str, port: SupportsInt, world_size: SupportsInt | None = None, is_master: bool = False, timeout: datetime.timedelta = datetime.timedelta(seconds=300), wait_for_workers: bool = True, multi_tenant: bool = False, master_listen_fd: SupportsInt | None = None, use_libuv: bool = True) → None# 创建新的 TCPStore。

property host# 获取存储监听请求的主机名。

property libuvBackend# 如果使用 libuv 后端，则返回 True。

property port# 获取存储监听请求的端口号。

class torch.distributed.HashStore# 基于底层 hashmap 的线程安全存储实现。此存储可在同一进程内使用（例如由其他线程使用），但不能跨进程使用。

示例::>>> import torch.distributed as dist >>> store = dist.HashStore() >>> # 其他线程也可以使用 store >>> # 初始化后可调用任意存储方法 >>> store.set("first_key", "first_value")

init(self: torch._C._distributed_c10d.HashStore) → None# 创建新的 HashStore。

class torch.distributed.FileStore# 使用文件存储底层键值对的存储实现。

参数 file_name（str）——用于存储键值对的文件路径。

world_size（int，可选）——使用此存储的进程总数。默认值为 -1（负值表示存储用户数量不固定）。

示例::>>> import torch.distributed as dist >>> store1 = dist.FileStore("/tmp/filestore", 2) >>> store2 = dist.FileStore("/tmp/filestore", 2) >>> # 初始化后，可从客户端或服务器调用任意存储方法 >>> store1.set("first_key", "first_value") >>> store2.get("first_key")

init(self: torch._C._distributed_c10d.FileStore, file_name: str, world_size: SupportsInt = -1) → None# 创建新的 FileStore。

property path# 获取 FileStore 用于存储键值对的文件路径。

class torch.distributed.PrefixStore# 对 3 种键值存储（TCPStore、FileStore 和 HashStore）中的任一种进行包装，并为插入存储的每个键添加前缀。

参数 prefix（str）——在每个键插入存储前添加的前缀字符串。

store（torch.distributed.store）——构成底层键值存储的存储对象。

init(self: torch._C._distributed_c10d.PrefixStore, prefix: str, store: torch._C._distributed_c10d.Store) → None# 创建新的 PrefixStore。

property underlying_store# 获取 PrefixStore 包装的底层存储对象。

分析集合通信性能# 请注意，可以使用 torch.profiler（推荐，仅在 1.8.1 及更高版本中提供）或 torch.autograd.profiler 对此处提及的集合通信和点对点通信 API 进行性能分析。所有开箱即用的后端（gloo、nccl、mpi）都受支持，集合通信的使用情况也会按预期呈现在性能分析输出/跟踪中。对代码进行性能分析的方式与普通 torch 运算符相同：

import torch import torch.distributed as dist with torch.profiler(): tensor = torch.randn(20, 10) dist.all_reduce(tensor)

完整了解性能分析器功能，请参阅性能分析器文档。

多 GPU 集合函数# Warning 多 GPU 函数（指每个 CPU 线程使用多个 GPU）已弃用。截至目前，PyTorch Distributed 首选的编程模型是每个线程一个设备，本文档中的 API 即为示例。如果你是后端开发者，并希望支持每个线程使用多个设备，请联系 PyTorch Distributed 的维护者。

对象集合操作# Warning 对象集合操作存在若干严重限制。请继续阅读，判断它们是否适合你的用例。

对象集合操作是一组类似集合操作的函数，可处理任意 Python 对象，只要这些对象能够 pickle。实现了多种集合模式（例如 broadcast、all_gather、……），但它们大致遵循以下模式：

将输入对象转换为 pickle（原始字节），然后放入字节张量中

将此字节张量的大小传达给对等进程（第一个集合操作）

分配大小合适的张量以执行实际集合操作

传递对象数据（第二个集合操作）

将原始数据转换回 Python 对象（反 pickle）

对象集合操作的性能或内存特性有时会令人意外，导致运行时间过长或内存溢出，因此应谨慎使用。以下是一些常见问题。

pickle/unpickle 时间不对称——pickle 对象可能很慢，具体取决于对象的数量、类型和大小。当集合操作采用扇入模式（例如 gather_object）时，接收 rank 必须反 pickle 的对象数量是发送 rank 进行 pickle 的数量的 N 倍，这可能导致其他 rank 在下一次集合操作中超时。

张量通信效率低——张量应通过常规集合 API 发送，而不是对象集合 API。虽然可以通过对象集合 API 发送 Tensor，但它们会被序列化和反序列化（对于非 CPU 张量，还包括 CPU 同步和设备到主机的复制）。除了调试或排查问题之外，几乎所有情况下，重构代码以使用非对象集合操作都值得付出额外工作。

意外的张量设备——如果仍要通过对象集合操作发送张量，还需考虑 cuda（以及可能的其他加速器）张量特有的另一个问题。如果 pickle 当前位于 cuda:3 上的张量，然后再反 pickle，则无论当前进程是什么，也无论该进程的“默认”CUDA 设备是什么，都会得到另一个位于 cuda:3 上的张量。使用常规张量集合 API 时，“输出张量”始终位于同一本地设备上，这通常符合预期。如果这是进程首次使用 GPU，反 pickle 张量会隐式激活 CUDA 上下文，从而可能浪费大量 GPU 内存。可通过在将张量作为对象集合操作的输入传递前，将其移动到 CPU 来避免此问题。

第三方后端# 除了内置的 GLOO/MPI/NCCL 后端，PyTorch 分布式还通过运行时注册机制支持第三方后端。有关如何通过 C++ 扩展开发第三方后端的参考资料，请参阅 Tutorials - Custom C++ and CUDA Extensions 和 test/cpp_extensions/cpp_c10d_extension.cpp。第三方后端的能力由其自身实现决定。新后端从 c10d::ProcessGroup 派生，并在导入时通过 torch.distributed.Backend.register_backend() 注册后端名称及实例化接口。手动导入此后端，并使用相应的后端名称调用 torch.distributed.init_process_group() 后，torch.distributed 包便会使用新后端。

Warning 对第三方后端的支持目前为实验性功能，可能会发生变化。

启动工具# torch.distributed 包还在 torch.distributed.launch 中提供启动工具。此辅助工具可用于为分布式训练启动每个节点上的多个进程。

Module torch.distributed.launch. torch.distributed.launch 是一个模块，可在每个训练节点上启动多个分布式训练进程。

Warning 此模块将被弃用，改用 torchrun。

此工具可用于单节点分布式训练，在每个节点上启动一个或多个进程。此工具既可用于 CPU 训练，也可用于 GPU 训练。如果用于 GPU 训练，每个分布式进程将使用单个 GPU。这可以显著提升单节点训练性能。此工具也可用于多节点分布式训练，通过在每个节点上启动多个进程来显著提升多节点分布式训练性能。对于具有多个支持直接 GPU 访问的 InfiniBand 接口的系统，尤其有益，因为所有接口都可以用于汇总通信带宽。

无论是单节点还是多节点分布式训练，此工具都会启动指定数量的每节点进程（--nproc-per-node）。如果用于 GPU 训练，此数量必须小于或等于当前系统中的 GPU 数量（nproc_per_node），并且每个进程将使用一个 GPU，范围从 GPU 0 到 GPU (nproc_per_node - 1)。

如何使用此模块：

单节点多进程分布式训练

python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 和训练脚本的所有其他参数)

多节点多进程分布式训练：（例如两个节点）

节点 1：（IP：192.168.1.1，且有一个空闲端口：1234）

python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE --nnodes=2 --node-rank=0 --master-addr="192.168.1.1" --master-port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 和训练脚本的所有其他参数)

节点 2：

python -m torch.distributed.launch --nproc-per-node=NUM_GPUS_YOU_HAVE --nnodes=2 --node-rank=1 --master-addr="192.168.1.1" --master-port=1234 YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 和训练脚本的所有其他参数)

查看此模块提供的可选参数：

python -m torch.distributed.launch --help

重要提示：

1. 对于 GPU 训练，此工具和多进程分布式训练（单节点或多节点）目前只有使用 NCCL 分布式后端才能达到最佳性能。因此，建议 GPU 训练使用 NCCL 后端。

2. 在训练程序中，必须解析命令行参数：--local-rank=LOCAL_PROCESS_RANK，此参数由此模块提供。如果训练程序使用 GPU，应确保代码仅在 LOCAL_PROCESS_RANK 对应的 GPU 设备上运行。可通过以下方式实现：

解析 local_rank 参数

import argparse >>> parser = argparse.ArgumentParser() >>> parser.add_argument("--local-rank", "--local_rank", type=int) >>> args = parser.parse_args()

使用 local rank 设置设备，可采用以下任一方式

torch.cuda.set_device(args.local_rank) # 在代码运行前设置

或

with torch.cuda.device(args.local_rank): >>> # 要运行的代码 >>> ...

版本 2.0.0 中的更改：启动器会将 --local-rank= 参数传递给脚本。从 PyTorch 2.0.0 起，优先使用带连字符的 --local-rank，而不是此前使用的带下划线的 --local_rank。为了向后兼容，用户可能需要在参数解析代码中同时处理这两种情况。这意味着应在参数解析器中同时包含 "--local-rank" 和 "--local_rank"。如果只提供 "--local_rank"，启动器会引发错误：“error: unrecognized arguments: –local-rank=”。对于仅支持 PyTorch 2.0.0 及更高版本的训练代码，只需包含 "--local-rank" 即可。

3. 在训练程序中，应在开始时调用以下函数以启动分布式后端。强烈建议使用 init_method=env://。其他初始化方法（例如 tcp://）可能可用，但 env:// 是此模块正式支持的方法。

torch.distributed.init_process_group(backend='YOUR BACKEND', >>> init_method='env://')

4. 在训练程序中，可以使用常规分布式函数，也可以使用 torch.nn.parallel.DistributedDataParallel() 模块。如果训练程序使用 GPU 进行训练，并且希望使用 torch.nn.parallel.DistributedDataParallel() 模块，可按以下方式配置。

model = torch.nn.parallel.DistributedDataParallel(model, >>> device_ids=[args.local_rank], >>> output_device=args.local_rank)

请确保将 device_ids 参数设置为代码将使用的唯一 GPU 设备 ID。通常这是进程的本地 rank。换言之，要使用此工具，device_ids 必须设为 [args.local_rank]，output_device 必须设为 args.local_rank。

5. 也可以通过环境变量 LOCAL_RANK 将 local_rank 传递给子进程。使用 --use-env=True 启动脚本时会启用此行为。必须调整上述子进程示例，将 args.local_rank 替换为 os.environ['LOCAL_RANK']；指定此标志后，启动器不会传递 --local-rank。

Warning local_rank 不是全局唯一值：它只在每台机器上的进程之间唯一。因此，不要用它来决定是否应执行某些操作，例如写入网络文件系统。请参阅 pytorch/pytorch#12042，了解未正确处理时可能出现问题的示例。

Spawn 工具# Multiprocessing 包 torch.multiprocessing 也在 torch.multiprocessing.spawn() 中提供 spawn 函数。此辅助函数可用于启动多个进程。它通过传入要运行的函数来启动 N 个进程并运行该函数。此函数也可用于多进程分布式训练。有关用法参考，请参阅 PyTorch example - ImageNet implementation。请注意，此函数要求 Python 3.4 或更高版本。

调试 torch.distributed 应用# 分布式应用的调试可能很困难，因为卡住、崩溃或不同 rank 之间行为不一致的问题往往难以理解。torch.distributed 提供了一套工具，帮助用户自行调试训练应用：

Python 断点# 在分布式环境中使用 Python 调试器非常方便，但由于它开箱即用时无法正常工作，许多人根本不会使用它。PyTorch 提供了 pdb 的自定义包装器，使这一过程更加简单。torch.distributed.breakpoint 可以轻松完成此操作。它在内部通过两种方式自定义 pdb 的断点行为，其他方面与普通 pdb 相同。

仅在一个 rank 上附加调试器（由用户指定）。

使用 torch.distributed.barrier() 确保所有其他 rank 停止；调试 rank 发出 continue 后，屏障才会释放。

重定向子进程的 stdin，使其连接到终端。

要使用此功能，只需在所有 rank 上调用 torch.distributed.breakpoint(rank)，并确保每种情况下的 rank 值相同。

监控屏障# 从 v1.10 起，torch.distributed.monitored_barrier() 可作为 torch.distributed.barrier() 的替代方法。当有 rank 可能出错时（即并非所有 rank 都在给定超时时间内调用 torch.distributed.monitored_barrier()），它会失败并提供有关故障 rank 的有用信息。torch.distributed.monitored_barrier() 使用进程间的 send/recv 通信原语在主机端实现屏障，类似于确认机制，使 rank 0 能够报告哪些 rank 未及时确认屏障。例如，考虑以下函数：rank 1 未调用 torch.distributed.monitored_barrier()（在实际场景中，这可能是由于应用程序错误或之前的集合操作卡住所致）：

import os from datetime import timedelta import torch import torch.distributed as dist import torch.multiprocessing as mp def worker(rank): dist.init_process_group("nccl", rank=rank, world_size=2) # monitored barrier 需要 gloo 进程组执行主机端同步。 group_gloo = dist.new_group(backend="gloo") if rank not in [1]: dist.monitored_barrier(group=group_gloo, timeout=timedelta(seconds=2)) if name == "main": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29501" mp.spawn(worker, nprocs=2, args=())

rank 0 上会产生以下错误消息，帮助用户确定可能出错的 rank 并进一步调查：

RuntimeError: Rank 1 failed to pass monitoredBarrier in 2000 ms Original exception: [gloo/transport/tcp/pair.cc:598] Connection closed by peer [2401:db00:eef0:1100:3560:0:1c05:25d]:8594

TORCH_DISTRIBUTED_DEBUG# 配合 TORCH_CPP_LOG_LEVEL=INFO 使用时，可通过环境变量 TORCH_DISTRIBUTED_DEBUG 触发额外的实用日志和集合同步检查，以确保所有 rank 正确同步。根据所需调试级别，TORCH_DISTRIBUTED_DEBUG 可设为 OFF（默认值）、INFO 或 DETAIL。请注意，最详细的选项 DETAIL 可能会影响应用性能，因此只应在调试问题时使用。将 TORCH_DISTRIBUTED_DEBUG=INFO 设为 INFO 时，使用 torch.nn.parallel.DistributedDataParallel() 训练的模型初始化时会输出额外调试日志；将其设为 DETAIL 时，还会在指定的若干次迭代期间记录运行时性能统计信息。这些运行时统计信息包括前向时间、反向时间、梯度通信时间等。

例如，给定以下应用：

import os import torch import torch.distributed as dist import torch.multiprocessing as mp class TwoLinLayerNet(torch.nn.Module): def init(self): super().init() self.a = torch.nn.Linear(10, 10, bias=False) self.b = torch.nn.Linear(10, 1, bias=False) def forward(self, x): a = self.a(x) b = self.b(x) return (a, b) def worker(rank): dist.init_process_group("nccl", rank=rank, world_size=2) torch.cuda.set_device(rank) print("init model") model = TwoLinLayerNet().cuda() print("init ddp") ddp_model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[rank]) inp = torch.randn(10, 10).cuda() print("train") for _ in range(20): output = ddp_model(inp) loss = output[0] + output[1] loss.sum().backward() if name == "main": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29501" os.environ["TORCH_CPP_LOG_LEVEL"]="INFO" os.environ[ "TORCH_DISTRIBUTED_DEBUG" ] = "DETAIL" # 设为 DETAIL 可启用运行时日志。 mp.spawn(worker, nprocs=2, args=())

初始化时会输出以下日志：

I0607 16:10:35.739390 515217 logger.cpp:173] [Rank 0]: DDP Initialized with: broadcast_buffers: 1 bucket_cap_bytes: 26214400 find_unused_parameters: 0 gradient_as_bucket_view: 0 is_multi_device_module: 0 iteration: 0 num_parameter_tensors: 2 output_device: 0 rank: 0 total_parameter_size_bytes: 440 world_size: 2 backend_name: nccl bucket_sizes: 440 cuda_visible_devices: N/A device_ids: 0 dtypes: float master_addr: localhost master_port: 29501 module_name: TwoLinLayerNet nccl_async_error_handling: N/A nccl_blocking_wait: N/A nccl_debug: WARN nccl_ib_timeout: N/A nccl_nthreads: N/A nccl_socket_ifname: N/A torch_distributed_debug: INFO

运行期间（设置 TORCH_DISTRIBUTED_DEBUG=DETAIL 时）会输出以下日志：

I0607 16:18:58.085681 544067 logger.cpp:344] [Rank 1 / 2] Training TwoLinLayerNet unused_parameter_size=0 Avg forward compute time: 40838608 Avg backward compute time: 5983335 Avg backward comm. time: 4326421 Avg backward comm/comp overlap time: 4207652 I0607 16:18:58.085693 544066 logger.cpp:344] [Rank 0 / 2] Training TwoLinLayerNet unused_parameter_size=0 Avg forward compute time: 42850427 Avg backward compute time: 3885553 Avg backward comm. time: 2357981 Avg backward comm/comp overlap time: 2234674

此外，TORCH_DISTRIBUTED_DEBUG=INFO 还能增强 torch.nn.parallel.DistributedDataParallel() 中因模型存在未使用参数导致的崩溃日志。目前，如果前向传播中可能存在未使用的参数，则初始化 torch.nn.parallel.DistributedDataParallel() 时必须传入 find_unused_parameters=True；自 v1.10 起，所有模型输出都必须参与损失计算，因为 torch.nn.parallel.DistributedDataParallel() 不支持反向传播中存在未使用参数。这些限制对大型模型尤其棘手，因此发生错误崩溃时，torch.nn.parallel.DistributedDataParallel() 会记录所有未使用参数的完整限定名称。例如，在上述应用中，如果将 loss 改为 loss = output[1]，则 TwoLinLayerNet.a 在反向传播中不会收到梯度，因此会导致 DDP 失败。发生崩溃时，用户会收到未使用参数的信息；对于大型模型，手动查找这些参数可能很困难：

RuntimeError: Expected to have finished reduction in the prior iteration before starting a new one. This error indicates that your module has parameters that were not used in producing loss. You can enable unused parameter detection by passing the keyword argument find_unused_parameters=True to torch.nn.parallel.DistributedDataParallel, and by making sure all forward function outputs participate in calculating loss. If you already have done the above, then the distributed data parallel module wasn't able to locate the output tensors in the return value of your module's forward function. Please include the loss function and the structure of the return va lue of forward of your module when reporting this issue (e.g. list, dict, iterable). Parameters which did not receive grad for rank 0: a.weight Parameter indices which did not receive grad for rank 0: 0

设置 TORCH_DISTRIBUTED_DEBUG=DETAIL 会对用户直接或间接发起的每次集合调用（例如 DDP allreduce）触发额外的一致性和同步检查。这通过创建一个包装所有进程组的包装进程组实现，这些进程组由 torch.distributed.init_process_group() 和 torch.distributed.new_group() API 返回。因此，这些 API 会返回一个包装进程组，其用法与常规进程组完全相同，但会在向底层进程组分发集合操作前执行一致性检查。目前，这些检查包括 torch.distributed.monitored_barrier()，用于确保所有 rank 完成待处理的集合调用，并报告卡住的 rank。接下来会检查集合操作的一致性，确保所有集合函数匹配，并使用一致的张量形状调用。如果不一致，应用崩溃时会包含详细错误报告，而不是卡住或显示无信息的错误消息。

例如，考虑以下函数，其向 torch.distributed.all_reduce() 传入的输入形状不匹配：

import torch import torch.distributed as dist import torch.multiprocessing as mp def worker(rank): dist.init_process_group("nccl", rank=rank, world_size=2) torch.cuda.set_device(rank) tensor = torch.randn(10 if rank == 0 else 20).cuda() dist.all_reduce(tensor) torch.cuda.synchronize(device=rank) if name == "main": os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = "29501" os.environ["TORCH_CPP_LOG_LEVEL"]="INFO" os.environ["TORCH_DISTRIBUTED_DEBUG"] = "DETAIL" mp.spawn(worker, nprocs=2, args=())

使用 NCCL 后端时，此类应用很可能会卡住，在复杂场景中很难找出根本原因。如果用户启用 TORCH_DISTRIBUTED_DEBUG=DETAIL 并重新运行应用，以下错误消息会揭示根本原因：

work = default_pg.allreduce([tensor], opts) RuntimeError: Error when verifying shape tensors for collective ALLREDUCE on rank 0. This likely indicates that input shapes into the collective are mismatched across ranks. Got shapes: 10 [ torch.LongTensor{1} ]

Note 要在运行时精细控制调试级别，也可使用函数 torch.distributed.set_debug_level()、torch.distributed.set_debug_level_from_env() 和 torch.distributed.get_debug_level()。此外，可以结合 TORCH_SHOW_CPP_STACKTRACES=1 使用 TORCH_DISTRIBUTED_DEBUG=DETAIL，以便在检测到集合操作不同步时记录完整调用栈。这些集合不同步检查适用于所有使用 c10d 集合调用的应用，这些集合调用由通过 torch.distributed.init_process_group() 和 torch.distributed.new_group() API 创建的进程组提供支持。

日志记录# 除了通过 torch.distributed.monitored_barrier() 和 TORCH_DISTRIBUTED_DEBUG 显式提供调试支持外，torch.distributed 的底层 C++ 库还会输出不同级别的日志消息。这些消息有助于了解分布式训练作业的执行状态，以及排查网络连接失败等问题。下表显示了如何通过组合 TORCH_CPP_LOG_LEVEL 和 TORCH_DISTRIBUTED_DEBUG 环境变量来调整日志级别。

TORCH_CPP_LOG_LEVEL	TORCH_DISTRIBUTED_DEBUG	有效日志级别
ERROR	ignored	Error
WARNING	ignored	Warning
INFO	ignored	Info
INFO	INFO	Debug
INFO	DETAIL	Trace（又称 All）

分布式组件会抛出从 RuntimeError 派生的自定义 Exception 类型：

torch.distributed.DistError：所有分布式异常的基类。

torch.distributed.DistBackendError：发生特定于后端的错误时抛出此异常。例如，使用 NCCL 后端时，用户尝试使用 NCCL 库不可用的 GPU。

torch.distributed.DistNetworkError：网络库遇到错误时抛出此异常（例如：对等方重置连接）。

torch.distributed.DistStoreError：Store 遇到错误时抛出此异常（例如：TCPStore 超时）。

class torch.distributed.DistError# 分布式库发生错误时引发的异常。

class torch.distributed.DistBackendError# 分布式后端发生错误时引发的异常。

class torch.distributed.DistNetworkError# 分布式网络发生错误时引发的异常。

class torch.distributed.DistStoreError# 分布式存储发生错误时引发的异常。

如果运行的是单节点训练，可以方便地对脚本设置交互式断点。我们提供了一种便捷方式，可在单个 rank 上设置断点：

torch.distributed.breakpoint(rank=0, skip=0, timeout_s=3600)[source]# 设置断点，但仅在单个 rank 上生效。所有其他 rank 都会等待你完成断点调试后再继续。

参数 rank（int）——在哪个 rank 上设置断点。默认值：0。

skip（int）——跳过对此断点的前几次调用。默认值：0。

torch.distributed

模式 3： 初始化# 在调用任何其他方法之前，需要使用 torch.distributed.init_process_group() 或 torch.distributed.device_mesh.init_device_mesh() 函数初始化该包。两者都会阻塞，直到所有进程都已加入。Warning 初始化不是线程安全的。进程组的创建应在单个线程中执行，以防止各 rank 之间出现不一致的 'UUID' 分配，并防止初始化期间的竞态条件导致挂起。 torch.distributed.is_available()[source]# 如果分布式包可用则返回 True。否则，torch.distributed 不会暴露任何其他 API。目前，torch.distributed 在 Linux、MacOS 和 Windows 上可用。从源码构建 PyTorch 时，设置 USE_DISTRIBUTED=1 以启用它。目前，Linux 和 Windows 的默认值为 USE_DISTRIBUTED=1，MacOS 的默认值为 USE_DISTRIBUTED=0。 Return type bool torch.distributed.init_process_group(backend=None, init_method=None, timeout=None, world_size=-1, rank=-1, store=None, group_name='', pg_options=None, device_id=None)[source]# 初始化默认的分布式进程组。这也会初始化分布式包。 有 2 种主要方式初始化进程组： 显式指定 store、rank 和 world_size。 指定 init_method（一个 URL 字符串），它指示在何处/如何发现对等节点。可以选择指定 rank 和 world_size，或者将所有必需参数编码在 URL 中并省略它们。 如果两者都未指定，则 init_method 假定为 “env://”。 Parameters backend (str or Backend, optional) – 要使用的后端。取决于构建时配置，有效值包括 mpi、gloo、nccl、ucc、xccl 或由第三方插件注册的后端。自 2.6 起，如果未提供 backend，c10d 将使用为 device_id 关键字参数（如果提供）所指示的设备类型注册的后端。目前已知的默认注册为：cuda 使用 nccl，cpu 使用 gloo，xpu 使用 xccl。如果既未提供 backend 也未提供 device_id，c10d 将检测运行时机器上的加速器，并使用为该检测到的加速器（或 cpu）注册的后端。此字段可以以小写字符串形式给出（例如 "gloo"），也可以通过 Backend 属性访问（例如 Backend.GLOO）。如果使用 nccl 后端在每台机器上运行多个进程，则每个进程必须对其使用的每个 GPU 具有独占访问权限，因为进程之间共享 GPU 可能导致死锁或 NCCL 无效用法。ucc 后端是实验性的。设备的默认后端可以通过 get_default_backend_for_device() 查询。 init_method (str, optional) – 指定如何初始化进程组的 URL。如果未指定 init_method 或 store，则默认为 “env://”。与 store 互斥。 world_size (int, optional) – 参与该作业的进程数。如果指定了 store，则为必需。 rank (int, optional) – 当前进程的 rank（应为 0 到 world_size-1 之间的数字）。如果指定了 store，则为必需。 store (Store, optional) – 所有 worker 都可访问的键/值存储，用于交换连接/地址信息。与 init_method 互斥。 timeout (timedelta, optional) – 针对进程组执行的操作的超时时间。NCCL 的默认值为 10 分钟，其他后端为 30 分钟。这是集合操作将被异步中止且进程将崩溃的持续时间。这样做是因为 CUDA 执行是异步的，并且继续执行用户代码不再安全，因为失败的异步 NCCL 操作可能导致后续 CUDA 操作在损坏的数据上运行。当设置 TORCH_NCCL_BLOCKING_WAIT 时，进程将阻塞并等待此超时。 group_name (str, optional, deprecated) – 组名称。此参数被忽略 pg_options (ProcessGroupOptions, optional) – 进程组选项，指定在构造特定进程组时需要传入哪些额外选项。截至目前，我们唯一支持的选项是用于 nccl 后端的 ProcessGroupNCCL.Options，可以指定 is_high_priority_stream，以便 nccl 后端在有计算内核等待时能够拾取高优先级 cuda 流。有关配置 nccl 的其他可用选项，请参阅 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t device_id (torch.device | int, optional) – 此进程将在其上工作的单个特定设备，允许进行后端特定的优化。目前这有两个效果，仅在 NCCL 下：通信器会立即形成（立即调用 ncclCommInit* 而不是通常的惰性调用），子组在可能的情况下将使用 ncclCommSplit 以避免不必要的组创建开销。如果你想尽早了解 NCCL 初始化错误，也可以使用此字段。如果提供的是 int，则 API 假定将使用编译时的加速器类型。 Note 要启用 backend == Backend.MPI，PyTorch 需要在支持 MPI 的系统上从源码构建。 Note 对多个后端的支持是实验性的。目前当未指定后端时，将同时创建 gloo 和 nccl 后端。gloo 后端将用于 CPU 张量的集合操作，nccl 后端将用于 CUDA 张量的集合操作。可以通过传入格式为 “<device_type>:<backend_name>,<device_type>:<backend_name>” 的字符串来指定自定义后端，例如 “cpu:gloo,cuda:custom_backend”。 torch.distributed.device_mesh.init_device_mesh(device_type, mesh_shape, *, mesh_dim_names=None, backend_override=None)[source]# 根据 device_type、mesh_shape 和 mesh_dim_names 参数初始化 DeviceMesh。 这会创建一个具有 n 维数组布局的 DeviceMesh，其中 n 是 mesh_shape 的长度。如果提供了 mesh_dim_names，则每个维度被标记为 mesh_dim_names[i]。 Note init_device_mesh 遵循 SPMD 编程模型，这意味着相同的 PyTorch Python 程序在集群中的所有进程/rank 上运行。确保 mesh_shape（描述设备布局的 nD 数组的维度）在所有 rank 上完全相同。不一致的 mesh_shape 可能导致挂起。 Note 如果未找到进程组，init_device_mesh 将在后台初始化分布式通信所需的分布式进程组。 Parameters device_type (str) – mesh 的设备类型。目前支持：“cpu”、“cuda/cuda-like”、“xpu”。不允许传入带有 GPU 索引的设备类型，例如 “cuda:0”。 mesh_shape (Tuple[int]) – 一个元组，定义描述设备布局的多维数组的维度。 mesh_dim_names (Tuple[str], optional) – 一个 mesh 维度名称元组，用于分配给描述设备布局的多维数组的每个维度。其长度必须与 mesh_shape 的长度匹配。mesh_dim_names 中的每个字符串必须唯一。 backend_override (Dict[int | str, tuple[str, Options] | str | Options], optional) – 对将为每个 mesh 维度创建的部分或全部 ProcessGroup 的覆盖。每个键可以是某个维度的索引或其名称（如果提供了 mesh_dim_names）。每个值可以是一个包含后端名称及其选项的元组，或仅包含这两部分之一（在这种情况下，另一部分将被设置为其默认值）。 Returns 一个表示设备布局的 DeviceMesh 对象。 Return type DeviceMesh Example: >>> from torch.distributed.device_mesh import init_device_mesh >>> >>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,)) >>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp")) torch.distributed.is_initialized()[source]# 检查默认进程组是否已初始化。 Return type bool torch.distributed.is_mpi_available()[source]# 检查 MPI 后端是否可用。 Return type bool torch.distributed.is_nccl_available()[source]# 检查 NCCL 后端是否可用。 Return type bool torch.distributed.is_gloo_available()[source]# 检查 Gloo 后端是否可用。 Return type bool torch.distributed.distributed_c10d.is_xccl_available()[source]# 检查 XCCL 后端是否可用。 Return type bool torch.distributed.is_torchelastic_launched()[source]# 检查此进程是否使用 torch.distributed.elastic（又名 torchelastic）启动。TORCHELASTIC_RUN_ID 环境变量的存在被用作确定当前进程是否使用 torchelastic 启动的代理。这是一个合理的代理，因为 TORCHELASTIC_RUN_ID 映射到 rendezvous id，后者始终是一个非空值，指示用于对等发现目的的作业 id。 Return type bool torch.distributed.get_default_backend_for_device(device)[source]# 返回给定设备的默认后端。 Parameters device (Union[str, torch.device]) – 要获取其默认后端的设备。 Returns 给定设备的默认后端，以小写字符串形式表示。 Return type str 目前支持三种初始化方法： TCP 初始化# 有两种使用 TCP 初始化的方式，两者都需要一个所有进程都可访问的网络地址以及所需的 world_size。第一种方式需要指定一个属于 rank 0 进程的地址。此初始化方法要求所有进程都手动指定 rank。请注意，最新的分布式包已不再支持组播地址。group_name 也已被弃用。 import torch.distributed as dist # Use address of one of the machines dist.init_process_group(backend, init_method='tcp://10.1.1.20:23456', rank=args.rank, world_size=4) 共享文件系统初始化# 另一种初始化方法利用组中所有机器都可共享并可见的文件系统，以及所需的 world_size。URL 应以 file:// 开头，并包含共享文件系统上一个不存在的文件（位于已存在的目录中）的路径。文件系统初始化会在该文件不存在时自动创建它，但不会删除该文件。因此，你有责任确保在下一次对同一文件路径/名称调用 init_process_group() 之前清理该文件。请注意，最新的分布式包已不再支持自动 rank 分配，group_name 也已被弃用。 Warning 此方法假定文件系统支持使用 fcntl 进行锁定——大多数本地系统和 NFS 都支持。 Warning 此方法将始终创建该文件，并尽最大努力在程序结束时清理和删除该文件。换句话说，使用文件初始化方法的每次初始化都需要一个全新的空文件，初始化才能成功。如果使用了前一次初始化所用过的同一文件（碰巧未被清理），这是意外的行为，通常会导致死锁和失败。因此，即使此方法会尽最大努力清理该文件，如果自动删除碰巧失败，你有责任确保在训练结束时删除该文件，以防止同一文件在下次被再次使用。如果你计划对同一文件名多次调用 init_process_group()，这一点尤为重要。换句话说，如果该文件未被删除/清理，而你又对该文件再次调用 init_process_group()，预计会失败。这里的经验法则是：确保每次调用 init_process_group() 时该文件不存在或为空。 import torch.distributed as dist # rank should always be specified dist.init_process_group(backend, init_method='file:///mnt/nfs/sharedfile', world_size=4, rank=args.rank) 环境变量初始化# 此方法将从环境变量中读取配置，允许完全自定义信息的获取方式。需要设置的变量有： MASTER_PORT - 必需；必须是 rank 0 机器上的空闲端口 MASTER_ADDR - 必需（rank 0 除外）；rank 0 节点的地址 WORLD_SIZE - 必需；可以在此处设置，也可以在调用 init 函数时设置 RANK - 必需；可以在此处设置，也可以在调用 init 函数时设置 rank 0 的机器将用于建立所有连接。 这是默认方法，意味着不必指定 init_method（或者可以是 env://）。 缩短初始化时间# TORCH_GLOO_LAZY_INIT - 按需建立连接，而不是使用全网格，这可以大大缩短非 all2all 操作的初始化时间。

torch.distributed.init_process_group()

Pattern 4: Example:

>>> from torch.distributed.device_mesh import init_device_mesh
>>>
>>> mesh_1d = init_device_mesh("cuda", mesh_shape=(8,))
>>> mesh_2d = init_device_mesh("cuda", mesh_shape=(2, 8), mesh_dim_names=("dp", "tp"))

模式 5： 组# 默认情况下，集合通信作用于默认组（也称为 world），并要求所有进程进入该分布式函数调用。然而，某些工作负载可以从更细粒度的通信中获益。这正是分布式组发挥作用的地方。可以使用 new_group() 函数创建新的组，其成员可以是所有进程中的任意子集。它返回一个不透明的组句柄，可作为 group 参数传递给所有集合通信（集合通信是用于以某些众所周知的编程模式交换信息的分布式函数）。torch.distributed.new_group(ranks=None, timeout=None, backend=None, pg_options=None, use_local_synchronization=False, group_desc=None, device_id=None)[source]# 创建一个新的分布式组。此函数要求主组中的所有进程（即属于该分布式作业的所有进程）都进入此函数，即使它们不会成为该组的成员。此外，所有进程中应以相同的顺序创建组。警告 安全的并发使用：当使用 NCCL 后端和多个进程组时，用户必须确保各 rank 之间集合通信的执行顺序全局一致。如果某个进程内的多个线程发起集合通信，则必须进行显式同步以确保顺序一致。当使用 torch.distributed 通信 API 的异步变体时，会返回一个 work 对象，通信内核会被排入单独的 CUDA 流，从而允许通信与计算重叠。一旦在某个进程组上发起了一个或多个异步操作，在改用另一个进程组之前，必须通过调用 work.wait() 与其他 cuda 流同步。更多详情请参见 Using multiple NCCL communicators concurrently https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently。参数 ranks (list[int]) – 组成员 rank 的列表。如果为 None，将设置为所有 rank。默认值为 None。timeout (timedelta, 可选) – 详情和默认值参见 init_process_group。backend (str 或 Backend, 可选) – 要使用的后端。取决于构建时配置，有效值为 gloo 和 nccl。默认使用与全局组相同的后端。此字段应以小写字符串形式给出（例如 "gloo"），也可以通过 Backend 属性访问（例如 Backend.GLOO）。如果传入 None，将使用与默认进程组对应的后端。默认值为 None。pg_options (ProcessGroupOptions, 可选) – 进程组选项，用于指定在构造特定进程组时所需传入的附加选项。例如对于 nccl 后端，可以指定 is_high_priority_stream，以便进程组可以获取高优先级 cuda 流。有关配置 nccl 的其他可用选项，请参见 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/types.html#ncclconfig-t use_local_synchronization (bool, 可选)：在进程组创建结束时执行组内屏障。其不同之处在于，非成员 rank 无需调用该 API，也不加入该屏障。group_desc (str, 可选) – 描述该进程组的字符串。device_id (torch.device, 可选) – 一个单一且特定的设备，用于将此进程“绑定”到该设备。如果给出了此字段，new_group 调用将尝试立即为该设备初始化通信后端。返回 一个分布式组的句柄，可传递给集合通信调用；如果该 rank 不属于 ranks，则返回 GroupMember.NON_GROUP_MEMBER。注意 use_local_synchronization 不适用于 MPI。注意 虽然 use_local_synchronization=True 在较大的集群和较小的进程组中可能显著更快，但必须谨慎，因为它会改变集群行为，因为非成员 rank 不加入组 barrier()。注意 当每个 rank 创建多个重叠的进程组时，use_local_synchronization=True 可能导致死锁。为避免这种情况，请确保所有 rank 遵循相同的全局创建顺序。torch.distributed.get_group_rank(group, global_rank)[source]# 将全局 rank 转换为组 rank。global_rank 必须是 group 的一部分，否则会引发 RuntimeError。参数 group (ProcessGroup) – 用于查找相对 rank 的 ProcessGroup。global_rank (int) – 要查询的全局 rank。返回 相对于 group 的 global_rank 的组 rank。返回类型 int 注意 在默认进程组上调用此函数会返回恒等值。torch.distributed.get_global_rank(group, group_rank)[source]# 将组 rank 转换为全局 rank。group_rank 必须是 group 的一部分，否则会引发 RuntimeError。参数 group (ProcessGroup) – 用于查找全局 rank 的 ProcessGroup。group_rank (int) – 要查询的组 rank。返回 相对于 group 的 group_rank 的全局 rank。返回类型 int 注意 在默认进程组上调用此函数会返回恒等值。torch.distributed.get_process_group_ranks(group)[source]# 获取与 group 关联的所有 rank。参数 group (Optional[ProcessGroup]) – 用于获取所有 rank 的 ProcessGroup。如果为 None，将使用默认进程组。返回 按组 rank 排序的全局 rank 列表。返回类型 list[int]

new_group()

模式 6： 警告 安全的并发使用：当使用 NCCL 后端与多个进程组时，用户必须确保跨 rank 的集合操作具有全局一致的执行顺序。如果进程内的多个线程发起集合操作，则需要显式同步以确保顺序一致。使用 torch.distributed 通信 API 的异步变体时，会返回一个 work 对象，并且通信内核会排队到单独的 CUDA 流上，从而允许通信与计算重叠。一旦在一个进程组上发起了一个或多个异步操作，在调用 work.wait() 与其他 CUDA 流进行同步之前，不得使用另一个进程组。有关更多详细信息，请参阅并发使用多个 NCCL 通信器 https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/communicators.html#using-multiple-nccl-communicators-concurrently。

NCCL

模式 7： 注意 如果你将 DistributedDataParallel 与 Distributed RPC Framework 一起使用，则应始终使用 torch.distributed.autograd.backward() 来计算梯度，并使用 torch.distributed.optim.DistributedOptimizer 来优化参数。示例：>>> import torch.distributed.autograd as dist_autograd >>> from torch.nn.parallel import DistributedDataParallel as DDP >>> import torch >>> from torch import optim >>> from torch.distributed.optim import DistributedOptimizer >>> import torch.distributed.rpc as rpc >>> from torch.distributed.rpc import RRef >>> >>> t1 = torch.rand((3, 3), requires_grad=True) >>> t2 = torch.rand((3, 3), requires_grad=True) >>> rref = rpc.remote("worker1", torch.add, args=(t1, t2)) >>> ddp_model = DDP(my_model) >>> >>> # Setup optimizer >>> optimizer_params = [rref] >>> for param in ddp_model.parameters(): >>> optimizer_params.append(RRef(param)) >>> >>> dist_optim = DistributedOptimizer( >>> optim.SGD, >>> optimizer_params, >>> lr=0.05, >>> ) >>> >>> with dist_autograd.context() as context_id: >>> pred = ddp_model(rref.to_here()) >>> loss = loss_func(pred, target) >>> dist_autograd.backward(context_id, [loss]) >>> dist_optim.step(context_id)

torch.distributed.autograd.backward()

模式 8： static_graph (bool) – 当设置为 True 时，DDP 知道被训练的图是静态的。静态图意味着 1）在整个训练循环期间，已使用和未使用的参数集合不会改变；在这种情况下，用户是否设置 find_unused_parameters = True 并不重要。2）在整个训练循环期间，图的训练方式不会改变（意味着不存在取决于迭代次数的控制流）。当 static_graph 设置为 True 时，DDP 将支持过去无法支持的情况：1）可重入的反向传播。2）多次激活检查点。3）当模型有未使用参数时的激活检查点。4）存在位于 forward 函数之外的模型参数。5）当存在未使用参数时可能提升性能，因为当 static_graph 设置为 True 时，DDP 不会在每次迭代中搜索图来检测未使用参数。要检查是否可以将 static_graph 设置为 True，一种方法是在你之前的模型训练结束时检查 ddp 日志数据，如果 ddp_logging_data.get("can_set_static_graph") == True，那么通常你也可以设置 static_graph = True。示例::>>> model_DDP = torch.nn.parallel.DistributedDataParallel(model) >>> # 训练循环 >>> ... >>> ddp_logging_data = model_DDP._get_ddp_logging_data() >>> static_graph = ddp_logging_data.get("can_set_static_graph")

True

参考文件

本技能在 references/ 中包含全面的文档：

• other.md - 其他文档

在需要详细信息时，使用 view 读取特定的参考文件。

使用本技能

面向初学者

从 getting_started 或 tutorials 参考文件开始，了解基础概念。

面向特定功能

使用适当的类别参考文件（api、guides 等）获取详细信息。

代码示例

上方的快速参考部分包含从官方文档中提取的常见模式。

资源

references/

从官方来源提取的整理后的文档。这些文件包含：

• 详细说明
• 带语言注释的代码示例
• 指向原始文档的链接
• 便于快速导航的目录

scripts/

在此处添加用于常见自动化任务的辅助脚本。

assets/

在此处添加模板、样板或示例项目。

说明

• 此技能是从官方文档自动生成的
• 参考文件保留了源文档的结构和示例
• 代码示例包含语言检测，以获得更好的语法高亮
• 快速参考模式是从文档中的常见用法示例中提取的

更新

要使用更新后的文档刷新此技能：

1. 使用相同配置重新运行抓取器
2. 此技能将使用最新信息重新构建