PyTorch 分布式概述

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

简介

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

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

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

DTensor 表示一个被分片和/或复制的张量，并根据操作的需要自动通信以重新分片张量。
DeviceMesh 将加速器设备的通信器抽象为一个多维数组，用于管理在多维并行中用于集体通信的底层 ProcessGroup 实例。尝试我们的 Device Mesh 实用技巧以了解更多。

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

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

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

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

在为您的模型选择并行技术时，请遵循以下通用准则：

如果您的模型可以放入单个 GPU，但您希望轻松地使用多个 GPU 扩展训练，请使用 DistributedDataParallel (DDP)。
- 如果您使用多个节点，请使用 torchrun 来启动多个 pytorch 进程。
- 另请参阅：分布式数据并行入门指南
当您的模型无法适应单个 GPU 时，请使用 FullyShardedDataParallel (FSDP)。
- 另请参阅：FSDP 入门指南
如果使用 FSDP 遇到扩展限制，请使用 Tensor Parallel (TP) 和/或 Pipeline Parallel (PP)。
- 尝试我们的 Tensor 并行教程
- 另请参阅：TorchTitan 3D 并行端到端示例

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

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