pytorch分布式数据并行DistributedDataParallel（DDP）

DistributedDataParallel（DDP）在module级别实现数据并行性。它使用torch.distributed包communication collectives来同步梯度，参数和缓冲区。并行性在单个进程内部和跨进程均有用。在一个进程中，DDP将input module 复制到device_ids指定的设备，相应地按batch维度分别扔进模型，并将输出收集到output_device，这与DataParallel相似。Across processes, DDP inserts necessary parameter synchronizations in forward passes and gradient synchronizations in backward passes. It is up to users to map processes to available resources, as long as processes do not share GPU devices.

推荐（通常是最快的方法）为每个module 副本创建一个进程，即在一个进程中不进行任何module复制。

DataParallel 和 DistributedDataParallel 区别

在深入探讨之前，让我们澄清一下，尽管增加了复杂性，但为什么要考虑在DataParallel上使用DistributedDataParallel：

1. 如果模型太大而无法容纳在单个GPU上，则必须使用 model parallel 将其拆分到多个GPU中。 DistributedDataParallel与模型并行工作； DataParallel目前不提供。
2. DataParallel是单进程，多线程，并且只能在单台计算机上运行，​​而DistributedDataParallel是多进程，并且可以在单机和分布式训练中使用。因此，即使在单机训练中，您的数据足够小以适合单机，DistributedDataParallel仍要比DataParallel更快。 DistributedDataParallel还可以预先复制模型，而不是在每次迭代时复制模型，并且可以避免PIL全局解释器锁定。
3. 如果数据和模型同时很大而无法用一个GPU训练，则可以将model parallel（与DistributedDataParallel结合使用。在这种情况下，每个DistributedDataParallel进程都可以model parallel，并且所有进程共同用数据并行

基本使用

import os
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


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

    # initialize the process group
    dist.init_process_group("gloo", rank=rank, world_size=world_size)

    # Explicitly setting seed to make sure that models created in two processes
    # start from same random weights and biases.
    torch.manual_seed(42)


def cleanup():
    dist.destroy_process_group()

现在，创建一个toy model，将其与DDP封装在一起，并提供一些虚拟输入数据。请注意，如果训练从随机参数开始，则可能要确保所有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(rank, world_size):
    setup(rank, world_size)

    # setup devices for this process, rank 1 uses GPUs [0, 1, 2, 3] and
    # rank 2 uses GPUs [4, 5, 6, 7].
    n = torch.cuda.device_count() // world_size
    device_ids = list(range(rank * n, (rank + 1) * n))

    # create model and move it to device_ids[0]
    model = ToyModel().to(device_ids[0])
    # output_device defaults to device_ids[0]
    ddp_model = DDP(model, device_ids=device_ids)

    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_ids[0])
    loss_fn(outputs, labels).backward()
    optimizer.step()

    cleanup()


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

DDP封装了 lower level distributed communication details，并提供了干净的API，就好像它是本地模型一样。对于基本用例，DDP仅需要几个LoCs来设置 process group。在将DDP应用到更高级的用例时，需要注意一些警告。

处理速度不同步时

在DDP中，Model, forward method 和 differentiation of the outputs是分布式的同步点。期望不同的过程以相同的顺序到达同步点，并在大致相同的时间进入每个同步点。否则，快速流程可能会提早到达，并在等待时超时。因此，用户负责进程之间的工作负载分配。有时，由于例如网络延迟，资源争用，不可预测的工作量峰值，不可避免地会出现不同步的处理速度。为了避免在这些情况下超时，请确保在调用init_process_group时传递足够大timeoutvalue

保存和加载 Checkpoints

在训练期间，通常使用torch.save 和torch.load 来保存和加载 Checkpoints，有关更多详细信息，请参见 SAVING AND LOADING MODELS，使用DDP时，一种优化方法是仅在一个进程中保存模型，然后将其加载到所有进程中，从而减少写开销，这是正确的，因为所有过程都从相同的参数开始，并且梯度在反向传播中同步，因此优化程序应将参数设置为相同的值。如果使用此优化，请确保在保存完成之前不要启动所有进程。此外，在加载模块时，您需要提供适当的 map_location 参数，以防止进程进入其他人的设备。如果缺少map_location，torch.load 将首先将模块加载到CPU，然后将每个参数复制到保存位置，这将导致同一台机器上的所有进程使用相同的设备集

def demo_checkpoint(rank, world_size):
    setup(rank, world_size)

    # setup devices for this process, rank 1 uses GPUs [0, 1, 2, 3] and
    # rank 2 uses GPUs [4, 5, 6, 7].
    n = torch.cuda.device_count() // world_size
    device_ids = list(range(rank * n, (rank + 1) * n))

    model = ToyModel().to(device_ids[0])
    # output_device defaults to device_ids[0]
    ddp_model = DDP(model, device_ids=device_ids)

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

    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()
    # configure map_location properly
    rank0_devices = [x - rank * len(device_ids) for x in device_ids]
    device_pairs = zip(rank0_devices, device_ids)
    map_location = {'cuda:%d' % x: 'cuda:%d' % y for x, y in device_pairs}
    ddp_model.load_state_dict(
        torch.load(CHECKPOINT_PATH, map_location=map_location))

    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_ids[0])
    loss_fn = nn.MSELoss()
    loss_fn(outputs, labels).backward()
    optimizer.step()

    # Use a barrier() to make sure that all processes have finished reading the
    # checkpoint
    dist.barrier()

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

    cleanup()

DDP 和 Model Parallelism 一起使用

DDP还可以与 Model Parallelism一起使用，但是不支持进程内的复制。您需要为每个module 副本创建一个进程，与每个进程的多个副本相比，通常可以提高性能。 这种训练方式在具有巨大的数据量较大的模型时特别有用。使用此功能时，需要小心地实现 multi-GPU model，以避免使用硬编码的设备，因为会将不同的模型副本放置到不同的设备上

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)

将multi-GPU model 传递给DDP时，不得设置device_ids和output_device，输入和输出数据将通过应用程序或模型forward() 方法放置在适当的设备中。

def demo_model_parallel(rank, world_size):
    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()


if __name__ == "__main__":
    run_demo(demo_basic, 2)
    run_demo(demo_checkpoint, 2)

    if torch.cuda.device_count() >= 8:
        run_demo(demo_model_parallel, 4)