分布式数据并行¶

警告

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 将参数梯度组织到桶中，并一次减少一个桶。可以通过在 DDP 构造函数中设置 bucket_cap_mb 参数来配置桶大小。参数梯度到桶的映射是在构造时确定的，基于桶大小限制和参数大小。模型参数被分配到桶中，顺序（大致）与给定模型的 Model.parameters() 相反。使用相反顺序的原因是 DDP 期望梯度在反向传递期间大约按照该顺序准备好。下图显示了一个示例。请注意，grad0 和 grad1 在 bucket1 中，而其他两个梯度在 bucket0 中。当然，这种假设可能并不总是正确的，当这种情况发生时，它可能会损害 DDP 反向速度，因为 Reducer 无法在最早的可能时间开始通信。除了分桶之外，Reducer 还在构造期间注册 autograd 钩子，每个参数一个钩子。当梯度准备好时，这些钩子将在反向传递期间被触发。
**前向传递**：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() 在初始化期间将模型状态从 rank 0 的进程发送到其他进程，并使用 ProcessGroup::allreduce() 对梯度求和。
Store.hpp：协助进程组实例的会合服务找到彼此。