分布式数据并行¶

警告

的实现会随着时间的推移而发展。本设计说明是根据 v1.4 的状态编写的。

（DDP）以透明方式执行分布式数据并行训练。本页介绍了它的工作原理并揭示了实现细节。

例¶

让我们从一个简单的例子开始。此示例使用 a 作为本地模型 wraps 它，然后运行一个前向传递、一个向后传递和一个优化器 Step on DDP 模型。之后，本地模型上的参数将为 updated，并且不同进程上的所有模型都应该完全相同。

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 存储桶大小应用（图中断优化）。（有关更多信息，请参阅 TorchDynamo DDPOptimizer。DDPOptimizer

ddp_model = DDP(model, device_ids=[rank])
ddp_model = torch.compile(ddp_model)

内部设计¶

本节通过深入研究一次迭代中的每一步。

前提条件：DDP 依赖 c10d 进行通信。因此，应用程序必须在构造之前创建实例 DDP 的。ProcessGroupProcessGroup
构造：DDP 构造函数引用本地模块并从排名为 0 的进程广播到所有其他进程进程，以确保所有模型副本都从完全相同的状态。然后，每个 DDP 进程都会创建一个本地，该 later 将在 backward 期间处理梯度同步通过。为了提高通信效率，organizes 参数 gradients 转换为 bucket，并一次减少一个 bucket。存储桶大小可以是通过在 DDP 构造函数中设置 bucket_cap_mb 参数进行配置。这从参数梯度到存储桶的映射是在构造时确定的时间，具体取决于存储桶大小限制和参数大小。模型参数为以（大致）与给定模型的相反顺序分配到存储桶中。使用反向的原因 order 是因为 DDP 期望梯度在向后按大致该顺序传递。下图显示了一个示例。注意 that、和都在中，而其他两个渐变位于中。当然，这种假设可能并不总是是真的，当这种情况发生时，它可能会损害 DDP 的倒退速度，因为它无法尽早启动通信。除了分桶之外，它还会在构造，每个参数一个钩子。这些钩子将在渐变准备就绪时的 backward pass。state_dict()ReducerReducerModel.parameters()grad0grad1bucket1bucket0ReducerReducer
Forward Pass：DDP 获取输入并将其传递给本地模型，然后分析本地模型的输出（如果设置为）。此模式允许运行在模型的子图上向后，DDP 找出哪些参数是通过从模型中遍历 autograd 图参与向后传递 output 并将所有未使用的参数标记为 ready for reduction。在 backward pass，则只会等待 unready 参数，但它仍会减少所有存储桶。将参数梯度标记为就绪不会目前帮助 DDP 跳过存储桶，但它会阻止 DDP 等待在向后传递期间永远没有梯度。请注意，遍历 Autograd Graph 会带来额外的开销，因此应用程序应仅在必要时设置为。find_unused_parametersTrueReducerfind_unused_parametersTrue
Backward Pass：函数在 loss 时直接调用，这超出了 DDP 的控制范围，DDP 使用了 autograd hook 在构建时注册以触发渐变同步。什么时候一个梯度变为 ready，其对应的 DDP 钩子位于该 grad 上 accumulator 将触发，然后 DDP 会将该参数梯度标记为准备减少。当一个桶中的梯度全部准备就绪时，该桶将启动异步计算所有过程的梯度平均值。当所有存储桶都准备就绪时，将阻止等待所有操作完成。完成此操作后，平均梯度将写入字段的所有参数。所以在向后传递之后，相同的 grad 字段不同 DDP 过程中的相应参数应相同。backward()TensorReducerallreduceReducerallreduceparam.grad
Optimizer Step：从优化器的角度来看，它正在优化本地型。所有 DDP 进程上的模型副本都可以保持同步，因为它们都从相同的状态开始，它们具有相同的平均梯度每次迭代。

注意

DDP 要求所有进程上的实例以完全相同的顺序调用，这是通过始终以 bucket 索引顺序而不是实际的 bucket ready 顺序运行来完成的。跨流程的顺序不匹配可能会导致错误的结果或 DDP 倒退挂。Reducerallreduceallreduceallreduce

实现¶

以下是指向 DDP 实施组件的指针。堆叠图显示代码的结构。

进程组¶

ProcessGroup.hpp 中：包含所有进程组实现的抽象 API。该库提供了 3 种开箱即用的实现，即 ProcessGroupGloo、ProcessGroupNCCL和 ProcessGroupMPI。用于发送在初始化期间，从秩为 0 的进程到其他进程的模型状态并对梯度求和。c10dDistributedDataParallelProcessGroup::broadcast()ProcessGroup::allreduce()
Store.hpp 中：帮助进程组实例的 Rendezvous 服务找到彼此。