评价此页

PyTorch 分布式概述#

创建日期:2020年7月28日 | 最后更新:2025年7月20日 | 最后验证:2024年11月5日

作者Will Constable, Wei Feng

注意

editgithub 上查看并编辑本教程。

这是 torch.distributed 包的概览页面。本页面的目标是将文档归类为不同的主题并对其进行简要说明。如果您是第一次使用 PyTorch 构建分布式训练应用程序,建议使用本文档导航至最适合您用例的技术。

简介#

PyTorch 分布式库包含一组并行模块、一个通信层,以及用于启动和调试大型训练任务的基础设施。

并行 API#

这些并行模块提供了高级功能,并可与现有模型进行组合

分片原语#

DTensorDeviceMesh 是用于在 N 维进程组上构建分片或复制张量并行性的原语。

  • DTensor 表示一个已分片和/或已复制的张量,它会自动进行通信,以便根据操作的需要重新分片张量。

  • DeviceMesh 将加速器设备通信器抽象为一个多维数组,该数组管理用于多维并行性中集合通信的底层 ProcessGroup 实例。尝试我们的 Device Mesh 使用指南以了解更多信息。

通信 API#

PyTorch 分布式通信层 (C10D) 提供了集合通信 API(例如 all_reduce

all_gather)以及点对点 (P2P) 通信 API(例如 sendisend),它们在所有并行实现底层被使用。使用 PyTorch 编写分布式应用程序展示了使用 c10d 通信 API 的示例。

启动器#

torchrun 是一个广泛使用的启动器脚本,它在本地和远程机器上生成进程,用于运行分布式 PyTorch 程序。

应用并行性来扩展您的模型#

数据并行是一种被广泛采用的单程序多数据训练范式,其中模型在每个进程上进行复制,每个模型副本计算不同输入数据样本的局部梯度,在每次优化器步骤之前,梯度会在数据并行通信组内进行平均。

当模型无法放入 GPU 时,需要使用模型并行技术(或分片数据并行),并可以将它们组合起来形成多维 (N-D) 并行技术。

在决定为您的模型选择何种并行技术时,请参考以下通用指南

  1. 如果您的模型可以放入单个 GPU,但您希望通过使用多个 GPU 轻松扩展训练,请使用 DistributedDataParallel (DDP)

  2. 当您的模型无法放入单个 GPU 时,请使用 FullyShardedDataParallel (FSDP2)

  3. 如果您在使用 FSDP2 时达到扩展限制,请使用 张量并行 (TP) 和/或 流水线并行 (PP)

注意

数据并行训练也适用于 自动混合精度 (AMP)

PyTorch 分布式开发者#

如果您想为 PyTorch 分布式做出贡献,请参阅我们的 开发者指南