分布式数据并行¶

警告

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 构造函数获取对本地模块的引用，并将来自秩为 0 的进程的 state_dict() 广播到组中的所有其他进程，以确保所有模型副本都从完全相同的状态开始。然后，每个 DDP 进程创建一个本地 Reducer，该 Reducer 稍后将在反向传递期间负责梯度同步。为了提高通信效率，Reducer 将参数梯度组织成桶，并一次减少一个桶。可以通过在 DDP 构造函数中设置 bucket_cap_mb 参数来配置桶的大小。参数梯度到桶的映射是在构造时确定的，基于桶大小限制和参数大小。模型参数按给定模型的 Model.parameters() 的（大致）反序分配到桶中。使用反序的原因是 DDP 预计梯度在反向传递期间将大致按该顺序准备好。下图显示了一个示例。请注意，grad0 和 grad1 位于 bucket1 中，其他两个梯度位于 bucket0 中。当然，此假设可能并不总是正确的，当这种情况发生时，它可能会损害 DDP 反向速度，因为 Reducer 无法在最早可能的时间启动通信。除了桶化之外，Reducer 还在构造期间注册自动梯度挂钩，每个参数一个挂钩。当梯度准备好时，这些挂钩将在反向传递期间触发。
**前向传递**：DDP 获取输入并将其传递给本地模型，然后如果 find_unused_parameters 设置为 True，则分析本地模型的输出。此模式允许在模型的子图上运行反向传播，DDP 通过从模型输出遍历自动梯度图并将所有未使用的参数标记为准备好进行归约来找出哪些参数参与了反向传播。在反向传递期间，Reducer 仅等待未准备好的参数，但它仍会减少所有桶。目前，将参数梯度标记为已准备好无助于 DDP 跳过桶，但它将防止 DDP 在反向传递期间永远等待不存在的梯度。请注意，遍历自动梯度图会引入额外的开销，因此应用程序应仅在必要时将 find_unused_parameters 设置为 True。
反向传播：backward() 函数直接作用于损失 Tensor，这超出了 DDP 的控制范围。DDP 使用在构造时注册的自动微分钩子来触发梯度同步。当一个梯度准备好时，其在该梯度累加器上的相应 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：协助进程组实例的会合服务找到彼此。