PyTorch 分布式概述

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

作者: Will Constable

注意

在github上查看和编辑本教程。

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

介绍

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

并行API

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

• 分布式数据并行 (DDP)

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

• Tensor Parallel (TP)

• Pipeline Parallel (PP)

分片原语

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

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

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

通信API

The PyTorch distributed communication layer (C10D) offers both collective communication APIs (e.g., all_reduce

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

启动器

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

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

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

当模型无法适应GPU时，需要模型并行技术（或分片数据并行技术），并且可以将它们组合在一起形成多维（N-D）并行技术。

在为您的模型选择并行技术时，请使用以下常见指南：

1. 如果您的模型适合单个GPU，但您希望使用多个GPU轻松扩展训练，请使用DistributedDataParallel (DDP)。

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

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

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

   • 另请参阅：FSDP入门

3. 如果您在使用FSDP时遇到扩展限制，请使用Tensor Parallel (TP)和/或Pipeline Parallel (PP)。

   • 试试我们的Tensor Parallelism Tutorial

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

注意

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

PyTorch 分布式开发者

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