⚡

PyTorchFSDP分布式训练技能Skill pytorch-fsdp

这个技能专为PyTorch FSDP（全分片数据并行训练）设计，提供参数分片、混合精度、CPU卸载和FSDP2等功能的详细指导，适用于大规模深度学习模型的分布式训练。关键词：PyTorch, FSDP, 分布式训练, 深度学习, 参数分片, 混合精度, CPU卸载, 人工智能, 机器学习。

深度学习 0 次安装 0 次浏览 更新于 3/21/2026

名称: pytorch-fsdp 描述: PyTorch FSDP全分片数据并行训练的专家指导 - 参数分片、混合精度、CPU卸载、FSDP2 版本: 1.0.0 作者: Orchestra Research 许可证: MIT 标签: [分布式训练, PyTorch, FSDP, 数据并行, 分片, 混合精度, CPU卸载, FSDP2, 大规模训练] 依赖: [torch>=2.0, transformers]

PyTorch FSDP 技能

提供PyTorch FSDP开发的全面帮助，基于官方文档生成。

何时使用此技能

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

  • 使用PyTorch FSDP时
  • 询问PyTorch FSDP特性或API时
  • 实现PyTorch FSDP解决方案时
  • 调试PyTorch FSDP代码时
  • 学习PyTorch FSDP最佳实践时

快速参考

常见模式

模式1： 通用连接上下文管理器 创建于：2025年6月6日 | 最后更新于：2025年6月6日 通用连接上下文管理器便于在不均匀输入上进行分布式训练。此页面概述了相关类的API：Join、Joinable和JoinHook。有关教程，请参阅使用连接上下文管理器进行不均匀输入的分布式训练。

类 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

在每个衍生的工作进程上

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) # 排名1比排名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() # 所有排名到达此处而不挂起/错误 static notify_join_context(joinable)[source] 通知连接上下文管理器调用进程尚未连接。然后，如果 throw_on_early_termination=True，检查是否检测到不均匀输入（即，如果一个进程已连接）并在这种情况下抛出异常。此方法应在Joinable对象每次迭代集体通信之前调用。例如，这应在DistributedDataParallel的前向传递开始时调用。仅传递给上下文管理器的第一个Joinable对象在此方法中执行集体通信，对于其他对象，此方法是空操作。 参数 joinable (Joinable) – 调用此方法的Joinable对象。 返回 一个异步工作句柄，用于通知上下文管理器进程尚未连接的all-reduce，如果 joinable 是传递给上下文管理器的第一个对象；否则返回None。 类 torch.distributed.algorithms.Joinable[source] 此类定义了可连接类的抽象基类。可连接类（继承自Joinable）应实现 join_hook()，返回一个JoinHook实例，以及 join_device() 和 join_process_group()，分别返回设备信息和进程组信息。 abstract property join_device: device 返回用于连接上下文管理器所需集体通信的设备。 abstract join_hook(**kwargs)[source] 返回给定Joinable的JoinHook实例。 参数 kwargs (dict) – 包含任何关键字参数的字典，用于在运行时修改连接钩子的行为；共享相同连接上下文管理器的所有Joinable实例都转发相同的kwargs值。 返回类型 JoinHook abstract property join_process_group: Any 返回连接上下文管理器本身所需集体通信的进程组。 类 torch.distributed.algorithms.JoinHook[source] 此类定义了连接钩子，在连接上下文管理器中提供两个入口点。 入口点：一个主钩子，在存在非连接进程时重复调用，和一个后钩子，在所有进程连接后调用一次。 要为通用连接上下文管理器实现连接钩子，定义一个继承自JoinHook的类，并适当重写 main_hook() 和 post_hook()。 main_hook()[source] 在存在非连接进程时调用此钩子以在训练迭代中屏蔽集体通信。训练迭代即，一次前向传递、后向传递和优化器步骤。 post_hook(is_last_joiner)[source] 在所有进程连接后调用钩子。它传递一个额外的布尔参数 is_last_joiner，指示该排名是否是最后连接的之一。 参数 is_last_joiner (bool) – 如果该排名是最后连接的之一则为True；否则为False。

Join

模式2： 分布式通信包 - torch.distributed 创建于：2017年7月12日 | 最后更新于：2025年9月4日 注意 请参阅PyTorch分布式概述以获取分布式训练相关所有特性的简要介绍。 后端 torch.distributed支持四个内置后端，每个后端具有不同的能力。下表显示了每个后端在CPU或GPU上可用的函数。对于NCCL，GPU指CUDA GPU，而对于XCCL指XPU GPU。MPI仅支持CUDA，如果用于构建PyTorch的实现支持它。 后端 gloo mpi nccl xccl 设备 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后端被构建并包含在PyTorch分布式中（NCCL仅在构建时带有CUDA时包含）。MPI是一个可选后端，仅当您从源代码构建PyTorch时才能包含。（例如，在安装了MPI的主机上构建PyTorch。） 注意 自PyTorch v1.8起，Windows支持所有集体通信后端，但NCCL除外，如果 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。 应使用哪个后端？ 过去，我们常被问：“应使用哪个后端？”。 经验法则：对于使用CUDA GPU的分布式训练，使用NCCL后端。对于使用XPU GPU的分布式训练，使用XCCL后端。对于使用CPU的分布式训练，使用Gloo后端。 具有InfiniBand互连的GPU主机：使用NCCL，因为它是当前唯一支持InfiniBand和GPUDirect的后端。 具有以太网互连的GPU主机：使用NCCL，因为它当前提供最佳的分布式GPU训练性能，特别是对于多进程单节点或多节点分布式训练。如果遇到NCCL的任何问题，使用Gloo作为回退选项。（注意，Gloo当前对于GPU运行比NCCL慢。） 具有InfiniBand互连的CPU主机：如果您的InfiniBand启用了IP over IB，使用Gloo，否则使用MPI。我们计划在即将发布的版本中添加Gloo的InfiniBand支持。 具有以太网互连的CPU主机：使用Gloo，除非您有特定原因使用MPI。 常见环境变量 选择要使用的网络接口 默认情况下，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基于其拓扑检测进行自动调优以节省用户的调优努力。在一些基于套接字的系统上，用户仍可能尝试调优 NCCL_SOCKET_NTHREADS 和 NCCL_NSOCKS_PERTHREAD 以增加套接字网络带宽。这两个环境变量已由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模型的包装器提供同步分布式训练。这与多处理包 - torch.multiprocessing 和 torch.nn.DataParallel() 提供的并行类型不同，因为它支持多台网络连接的机器，并且用户必须为每个进程显式启动主训练脚本的单独副本。 在单机同步情况下，torch.distributed 或 torch.nn.parallel.DistributedDataParallel() 包装器可能仍然比其他数据并行方法具有优势，包括 torch.nn.DataParallel()： 每个进程维护其自己的优化器，并在每次迭代中执行完整的优化步骤。虽然这看起来冗余，因为梯度已经跨进程收集和平均，因此对于每个进程都相同，但这意味着不需要参数广播步骤，减少了在节点之间传输张量的时间。 每个进程包含独立的Python解释器，消除了从单个Python进程驱动多个执行线程、模型副本或GPU所带来的额外解释器开销和“GIL-thrashing”。这对于大量使用Python运行时的模型尤其重要，包括具有循环层或许多小组件的模型。 初始化 包需要在调用任何其他方法之前使用 torch.distributed.init_process_group() 或 torch.distributed.device_mesh.init_device_mesh() 函数进行初始化。两者都阻塞直到所有进程加入。 警告 初始化不是线程安全的。进程组创建应从单个线程执行，以防止跨排名的不一致’UUID’分配，并防止初始化期间的竞争导致挂起。 torch.distributed.is_available()[source] 如果分布式包可用则返回True。否则，torch.distributed 不暴露任何其他API。当前，torch.distributed 在Linux、MacOS和Windows上可用。从源代码构建PyTorch时设置 USE_DISTRIBUTED=1 以启用它。当前，默认值为 USE_DISTRIBUTED=1 对于Linux和Windows，USE_DISTRIBUTED=0 对于MacOS。 返回类型 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 kwarg（如果提供）指示的设备类型注册的后端。当前已知的默认注册有：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。默认是 “env://” 如果未指定 init_method 或 store。与 store 互斥。 world_size (int, 可选) – 参与作业的进程数。如果指定 store 则需要。 rank (int, 可选) – 当前进程的排名（应为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, 可选) – 进程组选项，指定在特定进程组构建期间需要传递哪些额外选项。截至目前，我们支持的唯一选项是 ProcessGroupNCCL.Options 用于 nccl 后端，可以指定 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假设将使用编译时的加速器类型。 注意 要启用 backend == Backend.MPI，PyTorch需要从支持MPI的系统上的源代码构建。 注意 对多个后端的支持是实验性的。当前当未指定后端时，将创建 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]。 注意 init_device_mesh 遵循SPMD编程模型，意味着相同的PyTorch Python程序在集群中的所有进程/排名上运行。确保 mesh_shape（描述设备布局的nD数组维度）在所有排名上相同。不一致的 mesh_shape 可能导致挂起。 注意 如果未找到进程组，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], 可选) – 为将为每个网格维度创建的某些或所有 ProcessGroups 的覆盖。每个键可以是维度的索引或其名称（如果提供 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 映射到集合点ID，该ID始终是一个非空值，指示对等体发现的作业ID。 返回类型 bool torch.distributed.get_default_backend_for_device(device)[source] 返回给定设备的默认后端。 参数 device (Union[str, torch.device]) – 要获取默认后端的设备。 返回 给定设备的默认后端作为小写字符串。 返回类型 str 当前支持三种初始化方法： TCP初始化 有两种方式使用TCP初始化，都需要所有进程可访问的网络地址和所需的 world_size。 第一种方式要求指定属于排名0进程的地址。此初始化方法要求所有进程手动指定排名。注意，组播地址在最新分布式包中不再支持。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() 调用相同文件路径/名称之前清理该文件。注意，自动排名分配在最新分布式包中不再支持，且 group_name 也已弃用。 警告 此方法假设文件系统支持使用 fcntl 锁定 - 大多数本地系统和NFS支持它。 警告 此方法将始终创建该文件，并尽力在程序结束时清理和删除该文件。换句话说，每次使用文件初始化方法初始化都需要一个全新的空文件才能使初始化成功。如果再次使用先前初始化使用的相同文件（恰好未清理），这是意外行为，通常可能导致死锁和失败。因此，即使此方法会尽力清理文件，如果自动删除不成功，您有责任确保在训练结束时移除该文件，以防止下次再次重用同一文件。换句话说，如果文件未移除/清理，并且您再次调用 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 - 必需；必须是排名0机器上的空闲端口。 MASTER_ADDR - 必需（排名0除外）；排名0节点的地址。 WORLD_SIZE - 必需；可以在此设置，或在 init 函数调用中设置。 RANK - 必需；可以在此设置，或在 init 函数调用中设置。 排名0的机器将用于建立所有连接。这是默认方法，意味着 init_method 不必指定（或可以是 env://）。 改进初始化时间 TORCH_GLOO_LAZY_INIT - 按需建立连接而不是使用完全网格，这可以大大改进非 all2all 操作的初始化时间。

torch.distributed.init_process_group()

模式3： 初始化 包需要在调用任何其他方法之前使用 torch.distributed.init_process_group() 或 torch.distributed.device_mesh.init_device_mesh() 函数进行初始化。两者都阻塞直到所有进程加入。 警告 初始化不是线程安全的。进程组创建应从单个线程执行，以防止跨排名的不一致’UUID’分配，并防止初始化期间的竞争导致挂起。 torch.distributed.is_available()[source] 如果分布式包可用则返回True。否则，torch.distributed 不暴露任何其他API。当前，torch.distributed 在Linux、MacOS和Windows上可用。从源代码构建PyTorch时设置 USE_DISTRIBUTED=1 以启用它。当前，默认值为 USE_DISTRIBUTED=1 对于Linux和Windows，USE_DISTRIBUTED=0 对于MacOS。 返回类型 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 kwarg（如果提供）指示的设备类型注册的后端。当前已知的默认注册有：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。默认是 “env://” 如果未指定 init_method 或 store。与 store 互斥。 world_size (int, 可选) – 参与作业的进程数。如果指定 store 则需要。 rank (int, 可选) – 当前进程的排名（应为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, 可选) – 进程组选项，指定在特定进程组构建期间需要传递哪些额外选项。截至目前，我们支持的唯一选项是 ProcessGroupNCCL.Options 用于 nccl 后端，可以指定 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假设将使用编译时的加速器类型。 注意 要启用 backend == Backend.MPI，PyTorch需要从支持MPI的系统上的源代码构建。 注意 对多个后端的支持是实验性的。当前当未指定后端时，将创建 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]。 注意 init_device_mesh 遵循SPMD编程模型，意味着相同的PyTorch Python程序在集群中的所有进程/排名上运行。确保 mesh_shape（描述设备布局的nD数组维度）在所有排名上相同。不一致的 mesh_shape 可能导致挂起。 注意 如果未找到进程组，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], 可选) – 为将为每个网格维度创建的某些或所有 ProcessGroups 的覆盖。每个键可以是维度的索引或其名称（如果提供 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.init_process_group()

模式4： 示例：

>>> 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： 组 默认情况下，集体操作在默认组（也称为世界）上运行，并要求所有进程进入分布式函数调用。然而，一些工作负载可以受益于更细粒度的通信。这就是分布式组的用武之地。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后端时，用户必须确保跨排名的集体执行顺序全局一致。如果进程内的多个线程发出集体，需要显式同步以确保一致顺序。当使用 torch.distributed 通信API的异步变体时，返回一个工作对象，通信内核在单独的CUDA流上排队，允许通信和计算重叠。一旦在一个进程组上发出一个或多个异步操作，必须通过调用 work.wait() 与其他cuda流同步，然后再使用另一个进程组。有关更多详细信息，请参阅使用多个NCCL通信器并发。 参数 ranks (list[int]) – 组成员的排名列表。如果为None，将设置为所有排名。默认是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, 可选): 在进程组创建结束时执行组本地屏障。这与非成员排名不需要调用API且不加入屏障不同。 group_desc (str, 可选) – 描述进程组的字符串。 device_id (torch.device, 可选) – 此进程将“绑定”到的单个特定设备，如果给定此字段，new_group 调用将尝试立即为设备初始化通信后端。 返回 可以给集体调用的分布式组句柄，或 GroupMember.NON_GROUP_MEMBER 如果排名不是 ranks 的一部分。 注意 use_local_synchronization 不适用于MPI。 注意 虽然 use_local_synchronization=True 对于更大集群和小进程组可以显著更快，但必须小心，因为它改变了集群行为，因为非成员排名不加入组屏障()。 注意 use_local_synchronization=True 当每个排名创建多个重叠进程组时可能导致死锁。为避免这种情况，确保所有排名遵循相同的全局创建顺序。 torch.distributed.get_group_rank(group, global_rank)[source] 将全局排名转换为组排名。global_rank 必须是 group 的一部分，否则引发 RuntimeError。 参数 group (ProcessGroup) – 要查找相对排名的 ProcessGroup。 global_rank (int) – 要查询的全局排名。 返回 global_rank 相对于 group 的组排名。 返回类型 int 注意 在默认进程组上调用此函数返回身份。 torch.distributed.get_global_rank(group, group_rank)[source] 将组排名转换为全局排名。group_rank 必须是 group 的一部分，否则引发 RuntimeError。 参数 group (ProcessGroup) – 要从中查找全局排名的 ProcessGroup。 group_rank (int) – 要查询的组排名。 返回 group_rank 相对于 group 的全局排名。 返回类型 int 注意 在默认进程组上调用此函数返回身份。 torch.distributed.get_process_group_ranks(group)[source] 获取与 group 关联的所有排名。 参数 group (Optional[ProcessGroup]) – 要从中获取所有排名的 ProcessGroup。如果为None，将使用默认进程组。 返回 按组排名排序的全局排名列表。 返回类型 list[int]

new_group()

模式6： 警告 安全并发使用：当使用多个进程组与NCCL后端时，用户必须确保跨排名的集体执行顺序全局一致。如果进程内的多个线程发出集体，需要显式同步以确保一致顺序。当使用 torch.distributed 通信API的异步变体时，返回一个工作对象，通信内核在单独的CUDA流上排队，允许通信和计算重叠。一旦在一个进程组上发出一个或多个异步操作，必须通过调用 work.wait() 与其他cuda流同步，然后再使用另一个进程组。有关更多详细信息，请参阅使用多个NCCL通信器并发。

NCCL

模式7： 注意 如果您在使用 DistributedDataParallel 结合分布式RPC框架，应始终使用 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)

设置优化器

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) 有模型参数在前向函数之外。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/ 中的全面文档：

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

使用此技能

对于初学者

从入门或教程参考文件开始，了解基础概念。

对于特定功能

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

对于代码示例

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

资源

references/

从官方来源提取的组织文档。这些文件包含：

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

scripts/

在此添加常见自动化任务的助手脚本。

assets/

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

注意事项

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

更新

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

  1. 使用相同配置重新运行爬虫
  2. 技能将使用最新信息重建