分布式数据并行¶

警告

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 为了提高通信效率，将参数梯度组织成桶，并一次减少一个桶。桶大小可以通过在 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：包含所有进程组实现的抽象 API。c10d 库开箱即用地提供了 3 种实现，即 ProcessGroupGloo、ProcessGroupNCCL 和 ProcessGroupMPI。DistributedDataParallel 使用 ProcessGroup::broadcast() 在初始化期间将模型状态从 rank 0 进程发送到其他进程，并使用 ProcessGroup::allreduce() 来求和梯度。
Store.hpp：协助进程组实例的 rendezvous 服务相互查找。

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() 在 distributed.py 中的 DDP 前向传递结束时调用。当 find_unused_parameters 在 DDP 构造函数中设置为 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 的情况下仍应正常工作，但性能会下降。