Distributed Data Parallel 入门-未完纪

作者：Shen Li。编辑：Joe Zhu、Chirag Pandya。原页创建于 2019-04-23，更新于 2025-09-23，原页标注最近验证日期为 2024-11-05。

先决知识与工作原理

建议先阅读 PyTorch Distributed 概览、DistributedDataParallel API 与 DDP 设计说明。

DistributedDataParallel（DDP）是 PyTorch 中用于跨多台机器并行训练模型的模块。使用时需要启动多个进程，每个进程创建一个 DDP 实例。它利用 torch.distributed 的集合通信，在所有进程之间同步梯度与缓冲区。每个进程拥有自己的模型副本，共同参与训练。

DDP 为模型的每个参数注册 autograd hook。反向传播时，hook 被触发，启动进程之间的梯度同步。各进程使用同步后的梯度更新模型。更详细的设计与使用方式参见 DDP 设计说明。

建议为每个模型副本启动一个进程；一个副本可以跨多个设备。进程可以位于同一机器，也可以跨机器。不同 DDP 进程不能共享 GPU；单 GPU 副本情况下，一个进程独占一个 GPU。

本教程从基本用例开始，再介绍检查点和与模型并行结合的用例。原作者说明代码在 8 GPU 服务器上运行，可以推广到其他环境；本次未使用或验证该硬件环境。

DataParallel 与 DistributedDataParallel 的比较

DataParallel 是单进程、多线程实现，只支持单机。DDP 采用多进程，同时支持单机与多机。由于线程间 GIL 竞争、每次迭代复制模型，以及输入分发、输出汇集的额外开销，即便在单机上，DataParallel 通常也慢于 DDP。
如果模型太大，单 GPU 无法容纳，需要模型并行把模型拆到多个 GPU。原文说明 DDP 可以与模型并行组合，而 DataParallel 在该语境下不能。组合时，每个进程内部采用模型并行，所有进程之间再采用数据并行。

基础用例

创建 DDP 模块前，先正确设置进程组。更多细节见编写分布式应用。

import os
import sys
import tempfile
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
import torch.multiprocessing as mp

from torch.nn.parallel import DistributedDataParallel as DDP

# On Windows platform, the torch.distributed package only
# supports Gloo backend, FileStore and TcpStore.
# For FileStore, set init_method parameter in init_process_group
# to a local file. Example as follow:
# init_method="file:///f:/libtmp/some_file"
# dist.init_process_group(
#    "gloo",
#    rank=rank,
#    init_method=init_method,
#    world_size=world_size)
# For TcpStore, same way as on Linux.

def setup(rank, world_size):
    os.environ['MASTER_ADDR'] = 'localhost'
    os.environ['MASTER_PORT'] = '12355'

    # We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
    # such as CUDA, MPS, MTIA, or XPU.
    acc = torch.accelerator.current_accelerator()
    backend = torch.distributed.get_default_backend_for_device(acc)
    # initialize the process group
    dist.init_process_group(backend, rank=rank, world_size=world_size)

def cleanup():
    dist.destroy_process_group()

接下来创建简单模型，用 DDP 包装，并传入模拟数据。DDP 构造函数会把 rank 0 的模型状态广播给其他进程，因此无须担心各进程随机初始化的参数不同。

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic(rank, world_size):
    print(f"Running basic DDP example on rank {rank}.")
    setup(rank, world_size)

    # create model and move it to GPU with id rank
    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()
    print(f"Finished running basic DDP example on rank {rank}.")


def run_demo(demo_fn, world_size):
    mp.spawn(demo_fn,
             args=(world_size,),
             nprocs=world_size,
             join=True)

DDP 包装了底层分布式通信，提供类似本地模型的接口。梯度同步发生在反向传播期间，并与反向计算重叠。backward() 返回时，param.grad 已包含同步后的梯度张量。基础用例只需要增加少量进程组设置代码；更高级的用法需要注意下面的边界。

处理速度不均衡

DDP 构造函数、前向传播和反向传播都是分布式同步点。各进程预期启动同样次数的同步，按相同顺序到达同步点，并在大致相同的时间进入它们。否则较快进程可能提前到达，在等待较慢进程时超时。用户需要平衡不同进程的工作负载。

网络延迟、资源竞争或不可预测的工作量峰值有时会造成不可避免的速度差异。为减少这种情况下的超时，在调用 init_process_group 时设置足够大的 timeout。

保存与加载检查点

训练期间通常使用 torch.save 与 torch.load 保存、恢复模块，详见保存和加载模型。DDP 的一种优化是只在一个进程保存，再由所有进程加载，减少写入开销。这之所以可行，是因为各进程从相同参数开始，反向传播同步梯度，优化器应将参数更新为相同值。

采用这种方式时，必须保证保存完成后才开始加载。加载时还应提供正确的 map_location，避免进程进入其他进程的设备。没有 map_location 时，torch.load 会先把模块加载到 CPU，再把参数复制到原先保存的设备，可能导致同机进程都使用相同的一组设备。更高级的故障恢复和弹性支持参见 TorchElastic。

def demo_checkpoint(rank, world_size):
    print(f"Running DDP checkpoint example on rank {rank}.")
    setup(rank, world_size)

    model = ToyModel().to(rank)
    ddp_model = DDP(model, device_ids=[rank])


    CHECKPOINT_PATH = tempfile.gettempdir() + "/model.checkpoint"
    if rank == 0:
        # All processes should see same parameters as they all start from same
        # random parameters and gradients are synchronized in backward passes.
        # Therefore, saving it in one process is sufficient.
        torch.save(ddp_model.state_dict(), CHECKPOINT_PATH)

    # Use a barrier() to make sure that process 1 loads the model after process
    # 0 saves it.
    dist.barrier()
    # We want to be able to train our model on an `accelerator <https://pytorch.org/docs/stable/torch.html#accelerators>`__
    # such as CUDA, MPS, MTIA, or XPU.
    acc = torch.accelerator.current_accelerator()
    # configure map_location properly
    map_location = {f'{acc}:0': f'{acc}:{rank}'}
    ddp_model.load_state_dict(
        torch.load(CHECKPOINT_PATH, map_location=map_location, weights_only=True))

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(rank)

    loss_fn(outputs, labels).backward()
    optimizer.step()

    # Not necessary to use a dist.barrier() to guard the file deletion below
    # as the AllReduce ops in the backward pass of DDP already served as
    # a synchronization.

    if rank == 0:
        os.remove(CHECKPOINT_PATH)

    cleanup()
    print(f"Finished running DDP checkpoint example on rank {rank}.")

与模型并行结合

DDP 同样支持跨多个 GPU 的模型。训练大型模型与大量数据时，包装多 GPU 模型尤其有用。

class ToyMpModel(nn.Module):
    def __init__(self, dev0, dev1):
        super(ToyMpModel, self).__init__()
        self.dev0 = dev0
        self.dev1 = dev1
        self.net1 = torch.nn.Linear(10, 10).to(dev0)
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to(dev1)
    def forward(self, x):
        x = x.to(self.dev0)
        x = self.relu(self.net1(x))
        x = x.to(self.dev1)
        return self.net2(x)

把多 GPU 模型传给 DDP 时，不能设置 device_ids 和 output_device。输入与输出应由应用程序或模型 forward() 方法放到合适的设备。

def demo_model_parallel(rank, world_size):
    print(f"Running DDP with model parallel example on rank {rank}.")
    setup(rank, world_size)

    # setup mp_model and devices for this process
    dev0 = rank * 2
    dev1 = rank * 2 + 1
    mp_model = ToyMpModel(dev0, dev1)
    ddp_mp_model = DDP(mp_model)

    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_mp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    # outputs will be on dev1
    outputs = ddp_mp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(dev1)
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()
    print(f"Finished running DDP with model parallel example on rank {rank}.")


if __name__ == "__main__":
    n_gpus = torch.accelerator.device_count()
    assert n_gpus >= 2, f"Requires at least 2 GPUs to run, but got {n_gpus}"
    world_size = n_gpus
    run_demo(demo_basic, world_size)
    run_demo(demo_checkpoint, world_size)
    world_size = n_gpus//2
    run_demo(demo_model_parallel, world_size)

上述模型在单个进程中使用 dev0 和 dev1 两个设备；主入口先按全部设备数运行基础与检查点示例，再把进程数设为设备数的一半运行模型并行示例。

使用 torch.distributed.run / torchrun 初始化 DDP

PyTorch Elastic 可以简化 DDP 初始化。仍使用 ToyModel，创建名为 elastic_ddp.py 的文件：

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim

from torch.nn.parallel import DistributedDataParallel as DDP

class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = nn.Linear(10, 10)
        self.relu = nn.ReLU()
        self.net2 = nn.Linear(10, 5)

    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))


def demo_basic():
    torch.accelerator.set_device_index(int(os.environ["LOCAL_RANK"]))
    acc = torch.accelerator.current_accelerator()
    backend = torch.distributed.get_default_backend_for_device(acc)
    dist.init_process_group(backend)
    rank = dist.get_rank()
    print(f"Start running basic DDP example on rank {rank}.")
    # create model and move it to GPU with id rank
    device_id = rank % torch.accelerator.device_count()
    model = ToyModel().to(device_id)
    ddp_model = DDP(model, device_ids=[device_id])
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(ddp_model.parameters(), lr=0.001)

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_id)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    dist.destroy_process_group()
    print(f"Finished running basic DDP example on rank {rank}.")

if __name__ == "__main__":
    demo_basic()

在所有节点上运行下面的 torchrun 命令来启动该任务：

torchrun --nnodes=2 --nproc_per_node=8 --rdzv_id=100 --rdzv_backend=c10d --rdzv_endpoint=$MASTER_ADDR:29400 elastic_ddp.py

这个示例使用两台主机，每台启动 8 个进程，总共使用 16 个 GPU。所有节点上的 MASTER_ADDR 必须相同。torchrun 在它所在的节点启动 8 个进程，并让每个进程调用 elastic_ddp.py；用户仍需借助 SLURM 等集群管理工具，在两台主机上实际执行这个命令。

例如，在支持 SLURM 的集群中，可以编写脚本执行上面的命令，并按下面的方式设置 MASTER_ADDR：

export MASTER_ADDR=$(scontrol show hostname ${SLURM_NODELIST} | head -n 1)

然后用 SLURM 命令 srun --nodes=2 ./torchrun_script.sh 运行脚本。这只是示例，也可选择其他调度工具来启动 torchrun 作业。更多 Elastic run 信息见快速入门。

文章版权归作者所有，未经允许请勿转载。

THE END

技术译文