beginner/dist_overview

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

PyTorch 分布式概述

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

作者: Will Constable, Wei Feng

注意

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

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

简介

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

并行化 API

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

• 分布式数据并行 (DDP)

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

• 张量并行 (TP)

• 流水线并行 (PP)

分片原语

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

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

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

通信 API

PyTorch 分布式通信层 (C10D) 提供了集合通信 API（例如 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 (FSDP2)。

   • 另请参阅：FSDP2 入门

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

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

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

注意

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

PyTorch 分布式开发者

如果您想为 PyTorch Distributed 做出贡献，请参阅我们的 开发者指南。