分布式数据并行#

创建于：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 构造函数接受对本地模块的引用，并将 state_dict() 从等级为 0 的进程广播到组中的所有其他进程，以确保所有模型副本从完全相同的状态开始。然后，每个 DDP 进程创建一个本地 Reducer，它稍后将在反向传播期间负责梯度同步。为了提高通信效率，Reducer 将参数梯度组织成桶，并一次减少一个桶。桶大小可以通过在 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.hpp：包含所有进程组实现的抽象 API。c10d 库提供了 3 种开箱即用的实现，即 ProcessGroupGloo、ProcessGroupNCCL 和 ProcessGroupMPI。DistributedDataParallel 使用 ProcessGroup::broadcast() 在初始化期间将模型状态从等级为 0 的进程发送给其他进程，并使用 ProcessGroup::allreduce() 来求和梯度。
Store.hpp：协助会合服务，以便进程组实例可以相互找到。