分布式数据并行#

创建于：2020年1月15日 | 最后更新于：2024年1月25日

警告

torch.nn.parallel.DistributedDataParallel 的实现会随着时间推移而不断演进。本文档基于 v1.4 版本时的状态编写。

torch.nn.parallel.DistributedDataParallel (DDP) 可透明地执行分布式数据并行训练。本页面将介绍其工作原理并揭示实现细节。

示例#

让我们从一个简单的 torch.nn.parallel.DistributedDataParallel 示例开始。该示例使用 torch.nn.Linear 作为本地模型，用 DDP 包装它，然后对 DDP 模型执行一次前向传播、一次后向传播和一次优化器步骤。之后，本地模型上的参数将得到更新，并且不同进程上的所有模型应完全相同。

import torch
import torch.distributed as dist
import torch.multiprocessing as mp
import torch.nn as nn
import torch.optim as optim
import os
from torch.nn.parallel import DistributedDataParallel as DDP


def example(rank, world_size):
    # create default process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)
    # create local model
    model = nn.Linear(10, 10).to(rank)
    # construct DDP model
    ddp_model = DDP(model, device_ids=[rank])
    # define loss function and optimizer
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    # forward pass
    outputs = ddp_model(torch.randn(20, 10).to(rank))
    labels = torch.randn(20, 10).to(rank)
    # backward pass
    loss_fn(outputs, labels).backward()
    # update parameters
    optimizer.step()

def main():
    world_size = 2
    mp.spawn(example,
        args=(world_size,),
        nprocs=world_size,
        join=True)

if __name__=="__main__":
    # Environment variables which need to be
    # set when using c10d's default "env"
    # initialization mode.
    os.environ["MASTER_ADDR"] = "localhost"
    os.environ["MASTER_PORT"] = "29500"
    main()

DDP 可与 TorchDynamo 配合使用。当与 TorchDynamo 一起使用时，请在编译模型之前应用 DDP 模型包装器，以便 torchdynamo 可以基于 DDP 桶大小应用 DDPOptimizer（图中断优化）。（更多信息请参阅 TorchDynamo DDPOptimizer。）

ddp_model = DDP(model, device_ids=[rank])
ddp_model = torch.compile(ddp_model)

内部设计#

本节将通过深入探讨单次迭代中每个步骤的细节，揭示 torch.nn.parallel.DistributedDataParallel 引擎盖下的工作原理。

前提条件：DDP 依赖 c10d ProcessGroup 进行通信。因此，应用程序必须在构建 DDP 之前创建 ProcessGroup 实例。
构造：DDP 构造函数接收本地模块的引用，并将 rank 为 0 的进程的 state_dict() 广播到组内的所有其他进程，以确保所有模型副本从完全相同的状态开始。然后，每个 DDP 进程创建一个本地 Reducer，该 Reducer 稍后将在后向传播期间负责梯度同步。为了提高通信效率，Reducer 将参数梯度组织成桶（bucket），并一次减少一个桶。桶大小可以通过在 DDP 构造函数中设置 bucket_cap_mb 参数来配置。参数梯度到桶的映射是在构造时确定的，基于桶大小限制和参数大小。模型参数根据给定模型 Model.parameters() 的（大致）反向顺序分配到桶中。使用反向顺序的原因是 DDP 期望梯度在后向传播期间以大约该顺序准备好。下图展示了一个示例。请注意，grad0 和 grad1 位于 bucket1 中，而另外两个梯度位于 bucket0 中。当然，这种假设不一定总是成立，当这种情况发生时，可能会影响 DDP 后向传播的速度，因为 Reducer 无法在最早可能的时间开始通信。除了分桶之外，Reducer 在构造时还会注册 autograd 钩子，每个参数一个钩子。这些钩子将在后向传播期间，当梯度准备好时被触发。
前向传播：DDP 接收输入并将其传递给本地模型，然后分析本地模型的输出，如果 find_unused_parameters 设置为 True。此模式允许在模型子图上运行后向传播，DDP 通过从模型输出遍历 autograd 图并标记所有未使用的参数为可归约状态来找出哪些参数参与后向传播。在后向传播期间，Reducer 只会等待未准备好的参数，但它仍然会归约所有桶。将参数梯度标记为已准备好并不能帮助 DDP 跳过桶（目前是这样），但它可以防止 DDP 在后向传播期间永远等待缺失的梯度。请注意，遍历 autograd 图会引入额外的开销，因此应用程序应仅在必要时将 find_unused_parameters 设置为 True。
后向传播：backward() 函数直接在损失 Tensor 上调用，DDP 无法控制这一点，DDP 使用在构造时注册的 autograd 钩子来触发梯度同步。当一个梯度准备好后，它对应的 DDP 钩子会触发，DDP 随后会将该参数梯度标记为可归约状态。当一个桶中的所有梯度都准备好后，Reducer 将为该桶启动一个异步 allreduce 操作，以计算所有进程上的梯度均值。当所有桶都准备好后，Reducer 将阻塞等待所有 allreduce 操作完成。完成后，平均梯度将被写入所有参数的 param.grad 字段。因此，在后向传播之后，不同 DDP 进程上相同相应参数的 grad 字段应该相同。
优化器步骤：从优化器的角度来看，它正在优化一个本地模型。由于所有 DDP 进程上的模型副本都从相同的状态开始，并且在每次迭代中具有相同的平均梯度，因此它们可以保持同步。

注意

DDP 要求所有进程上的 Reducer 实例以完全相同的顺序调用 allreduce，这是通过始终按桶索引顺序而不是实际的桶就绪顺序运行 allreduce 来实现的。进程间 allreduce 顺序不匹配可能导致错误结果或 DDP 后向传播挂起。

实现#

以下是指向 DDP 实现组件的链接。堆叠图显示了代码的结构。

ProcessGroup#

ProcessGroup.hpp：包含所有进程组实现。 c10d 库开箱即用提供了 3 种实现，即 ProcessGroupGloo、ProcessGroupNCCL 和 ProcessGroupMPI。 DistributedDataParallel 在初始化期间使用 ProcessGroup::broadcast() 将模型状态从 rank 为 0 的进程发送到其他进程，并在后向传播期间使用 ProcessGroup::allreduce() 对梯度进行求和。
Store.hpp：协助进程组实例相互查找的集合服务。

DistributedDataParallel#

distributed.py：DDP 的 Python 入口点。它实现了 nn.parallel.DistributedDataParallel 模块的初始化步骤和 forward 函数，这些函数调用 C++ 库。其 _sync_param 函数在 DDP 进程处理多个设备时执行进程内参数同步，并且还将模型缓冲区从 rank 为 0 的进程广播到所有其他进程。进程间参数同步发生在 Reducer.cpp 中。
comm.h：实现了合并广播辅助函数，该函数在初始化期间用于广播模型状态，并在前向传播之前同步模型缓冲区。
reducer.h：提供了后向传播中梯度同步的核心实现。它有三个入口点函数：
- Reducer：构造函数在 distributed.py 中调用，该函数将 Reducer::autograd_hook() 注册到梯度累加器。
- autograd_hook() 函数将在梯度准备就绪时由 autograd 引擎调用。
- prepare_for_backward() 在 DDP 前向传播结束时在 distributed.py 中调用。当 DDP 构造函数中设置 find_unused_parameters 为 True 时，它会遍历 autograd 图以查找未使用的参数。

TorchDynamo DDPOptimizer#

DDP 的性能优势来自于在后向传播期间将 allreduce 集合操作与计算重叠。当与 TorchDynamo 一起用于编译整个前向和后向图时，AotAutograd 会阻止这种重叠，因为 allreduce 操作是在整个优化后的后向计算完成后，由 autograd 钩子 _之后_ 启动的。

TorchDynamo 的 DDPOptimizer 通过在后向传播期间的 DDP allreduce 桶的逻辑边界处中断前向图来提供帮助。注意：目标是在后向传播期间中断图，最简单的实现是中断前向图，然后在每个部分上调用 AotAutograd 和编译。这使得 DDP 的 allreduce 钩子可以在后向传播的各个部分之间触发，并调度通信与计算重叠。

有关更深入的解释和实验结果，请参阅这篇博文，或在 torch/_dynamo/optimizations/distributed.py 阅读文档和代码。

要调试 DDPOptimizer，请设置 TORCH_LOGS=’ddp_graphs’ 以获取完整的图转储。要获取不带图的日志，请将 ‘dynamo’、‘distributed’ 或 ‘dist_ddp’ 添加到 TORCH_LOGS（用于关于桶边界的基本信息）。要禁用 DDPOptimizer，请设置 torch._dynamo.config.optimize_ddp=False。DDP 和 TorchDynamo 在没有 DDPOptimizer 的情况下仍能正常工作，但性能会下降。

分布式数据并行#

示例#

内部设计#

实现#

ProcessGroup#

DistributedDataParallel#

TorchDynamo DDPOptimizer#

文档

教程

资源