PyTorch 分布式概览

创建日期：2020 年 7 月 28 日 | 最后更新：2024 年 10 月 8 日 | 最后验证：2024 年 11 月 5 日

作者: Will Constable

注意

在 github 中查看和编辑此教程。

这是 torch.distributed 包的概览页面。本页的目标是将文档按不同主题分类并简要描述每个主题。如果你是首次使用 PyTorch 构建分布式训练应用程序，建议使用本文档来查找最适合你的用例的技术。

简介

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

并行 API

这些并行模块提供高级功能，并可与现有模型组合使用

• 分布式数据并行 (DDP)

• 完全分片数据并行训练 (FSDP)

• 张量并行 (TP)

• 流水线并行 (PP)

分片原语

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

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

• DeviceMesh 将加速器设备通信器抽象为多维数组，用于在多维并行中管理底层 ProcessGroup 实例进行集体通信。尝试我们的 Device Mesh Recipe 以了解更多信息。

通信 API

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

和 all_gather）以及 P2P 通信 API（例如 send 和 isend），这些 API 在所有并行实现中底层使用。使用 PyTorch 编写分布式应用程序 展示了使用 c10d 通信 API 的示例。

启动器

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

应用并行来扩展你的模型

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

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

在决定为你的模型选择哪种并行技术时，请遵循这些常用指南

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

   • 如果你使用多个节点，请使用 torchrun 来启动多个 PyTorch 进程。

   • 另请参阅：分布式数据并行入门

2. 当你的模型无法放入一个 GPU 时，使用 FullyShardedDataParallel (FSDP)。

   • 另请参阅：FSDP 入门

3. 如果使用 FSDP 达到扩展限制，则使用 张量并行 (TP) 和/或 流水线并行 (PP)。

   • 尝试我们的 张量并行教程

   • 另请参阅：TorchTitan 3D 并行端到端示例

注意

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

PyTorch 分布式开发者

如果你想为 PyTorch 分布式做贡献，请参考我们的 开发者指南。