分布式数据并行¶

警告

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 还在构造期间注册自动梯度钩子，每个参数一个钩子。当梯度准备好时，这些钩子将在反向传播期间触发。
前向传播：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: 帮助进程组实例的约会服务相互查找。

DistributedDataParallel¶

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

TorchDynamo DDPOptimizer¶

DDP 的性能优势来自在反向传播期间将所有减少集合与计算重叠。当与 TorchDynamo 一起使用以编译整个正向和整个反向图时，AotAutograd 会阻止这种重叠，因为所有减少操作是在整个优化的反向计算完成后由自动微分钩子启动的。

TorchDynamo 的 DDPOptimizer 通过在反向传播期间将正向图分解为 DDP 的所有减少桶的逻辑边界来提供帮助。注意：目标是在反向传播期间分解图，最简单的实现是分解正向图，然后对每个部分调用 AotAutograd 和编译。这允许 DDP 的所有减少钩子在反向传播部分之间触发，并安排通信以与计算重叠。

有关更深入的解释和实验结果，请参阅这篇博文，或阅读 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 的情况下正常工作，但性能会下降。