使用 RPC 的分布式管道并行化

原文：https://pytorch.org/tutorials/intermediate/dist_pipeline_parallel_tutorial.html

作者：Shen Li

先决条件：

本教程使用 Resnet50 模型来演示如何使用torch.distributed.rpc API 实现分布式管道并行性。 可以将其视为单机模型并行最佳实践中讨论的多 GPU 管道并行性的分布式对应物。

注意

本教程需要 PyTorch v1.6.0 或更高版本。

注意

本教程的完整源代码可以在pytorch/examples中找到。

基础知识

上一教程分布式 RPC 框架入门显示了如何使用torch.distributed.rpc为 RNN 模型实现分布式模型并行性。 该教程使用一个 GPU 来托管EmbeddingTable，并且提供的代码可以正常工作。 但是，如果模型驻留在多个 GPU 上，则将需要一些额外的步骤来增加所有 GPU 的摊销利用率。 管道并行性是在这种情况下可以提供帮助的一种范例。

在本教程中，我们使用ResNet50作为示例模型，单机模型并行最佳实践教程也使用了该模型。 类似地，ResNet50模型被分为两个碎片，输入批量被划分为多个拆分，并以流水线方式馈入两个模型碎片。 区别在于，本教程将调用异步 RPC，而不是使用 CUDA 流来并行执行。 因此，本教程中介绍的解决方案也可以跨计算机边界使用。 本教程的其余部分分四个步骤介绍了实现。

第 1 步：对 ResNet50 模型进行分片

这是在两个模型分片中实现ResNet50的准备步骤。 以下代码是从torchvision中的 ResNet 实现中借用的。 ResNetBase模块包含两个 ResNet 碎片的通用构件和属性。

  1. import threading
  3. import torch
  4. import torch.nn as nn
  6. from torchvision.models.resnet import Bottleneck
  8. num_classes = 1000
  10. def conv1x1(in_planes, out_planes, stride=1):
  11.     return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
  13. class ResNetBase(nn.Module):
  14.     def __init__(self, block, inplanes, num_classes=1000,
  15.                 groups=1, width_per_group=64, norm_layer=None):
  16.         super(ResNetBase, self).__init__()
  18.         self._lock = threading.Lock()
  19.         self._block = block
  20.         self._norm_layer = nn.BatchNorm2d
  21.         self.inplanes = inplanes
  22.         self.dilation = 1
  23.         self.groups = groups
  24.         self.base_width = width_per_group
  26.     def _make_layer(self, planes, blocks, stride=1):
  27.         norm_layer = self._norm_layer
  28.         downsample = None
  29.         previous_dilation = self.dilation
  30.         if stride != 1 or self.inplanes != planes * self._block.expansion:
  31.             downsample = nn.Sequential(
  32.                 conv1x1(self.inplanes, planes * self._block.expansion, stride),
  33.                 norm_layer(planes * self._block.expansion),
  34.             )
  36.         layers = []
  37.         layers.append(self._block(self.inplanes, planes, stride, downsample, self.groups,
  38.                                 self.base_width, previous_dilation, norm_layer))
  39.         self.inplanes = planes * self._block.expansion
  40.         for _ in range(1, blocks):
  41.             layers.append(self._block(self.inplanes, planes, groups=self.groups,
  42.                                     base_width=self.base_width, dilation=self.dilation,
  43.                                     norm_layer=norm_layer))
  45.         return nn.Sequential(*layers)
  47.     def parameter_rrefs(self):
  48.         return [RRef(p) for p in self.parameters()]

现在，我们准备定义两个模型碎片。 对于构造器，我们只需将所有 ResNet50 层分为两部分，然后将每个部分移至提供的设备中。 两个分片的forward函数获取输入数据的RRef，在本地获取数据，然后将其移至所需的设备。 将所有层应用于输入后，它将输出移至 CPU 并返回。 这是因为当调用方和被调用方中的设备数量不匹配时，RPC API 要求张量驻留在 CPU 上，以避免无效的设备错误。

  1. class ResNetShard1(ResNetBase):
  2.     def __init__(self, device, *args, **kwargs):
  3.         super(ResNetShard1, self).__init__(
  4.             Bottleneck, 64, num_classes=num_classes, *args, **kwargs)
  6.         self.device = device
  7.         self.seq = nn.Sequential(
  8.             nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False),
  9.             self._norm_layer(self.inplanes),
  10.             nn.ReLU(inplace=True),
  11.             nn.MaxPool2d(kernel_size=3, stride=2, padding=1),
  12.             self._make_layer(64, 3),
  13.             self._make_layer(128, 4, stride=2)
  14.         ).to(self.device)
  16.         for m in self.modules():
  17.             if isinstance(m, nn.Conv2d):
  18.                 nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
  19.             elif isinstance(m, nn.BatchNorm2d):
  20.                 nn.init.constant_(m.weight, 1)
  21.                 nn.init.constant_(m.bias, 0)
  23.     def forward(self, x_rref):
  24.         x = x_rref.to_here().to(self.device)
  25.         with self._lock:
  26.             out =  self.seq(x)
  27.         return out.cpu()
  29. class ResNetShard2(ResNetBase):
  30.     def __init__(self, device, *args, **kwargs):
  31.         super(ResNetShard2, self).__init__(
  32.             Bottleneck, 512, num_classes=num_classes, *args, **kwargs)
  34.         self.device = device
  35.         self.seq = nn.Sequential(
  36.             self._make_layer(256, 6, stride=2),
  37.             self._make_layer(512, 3, stride=2),
  38.             nn.AdaptiveAvgPool2d((1, 1)),
  39.         ).to(self.device)
  41.         self.fc =  nn.Linear(512 * self._block.expansion, num_classes).to(self.device)
  43.     def forward(self, x_rref):
  44.         x = x_rref.to_here().to(self.device)
  45.         with self._lock:
  46.             out = self.fc(torch.flatten(self.seq(x), 1))
  47.         return out.cpu()

第 2 步：将 ResNet50 模型片段拼接到一个模块中

然后，我们创建一个DistResNet50模块来组装两个分片并实现流水线并行逻辑。 在构造器中，我们使用两个rpc.remote调用分别将两个分片放在两个不同的 RPC 工作器上，并保持RRef到两个模型部分，以便可以在正向传播中引用它们。 forward函数将输入批量分为多个微批量，并将这些微批量以流水线方式馈送到两个模型部件。 它首先使用rpc.remote调用将第一个分片应用于微批量，然后将返回的中间输出RRef转发到第二个模型分片。 之后，它将收集所有微输出的Future，并在循环后等待所有它们。 请注意，remote()和rpc_async()都立即返回并异步运行。 因此，整个循环是非阻塞的，并将同时启动多个 RPC。 中间输出y_rref保留了两个模型零件上一个微批量的执行顺序。 微批量的执行顺序无关紧要。 最后，正向函数将所有微批量的输出连接到一个单一的输出张量中并返回。 parameter_rrefs函数是简化分布式优化器构造的助手，将在以后使用。

  1. class DistResNet50(nn.Module):
  2.     def __init__(self, num_split, workers, *args, **kwargs):
  3.         super(DistResNet50, self).__init__()
  5.         self.num_split = num_split
  7.         # Put the first part of the ResNet50 on workers[0]
  8.         self.p1_rref = rpc.remote(
  9.             workers[0],
  10.             ResNetShard1,
  11.             args = ("cuda:0",) + args,
  12.             kwargs = kwargs
  13.         )
  15.         # Put the second part of the ResNet50 on workers[1]
  16.         self.p2_rref = rpc.remote(
  17.             workers[1],
  18.             ResNetShard2,
  19.             args = ("cuda:1",) + args,
  20.             kwargs = kwargs
  21.         )
  23.     def forward(self, xs):
  24.         out_futures = []
  25.         for x in iter(xs.split(self.split_size, dim=0)):
  26.             x_rref = RRef(x)
  27.             y_rref = self.p1_rref.remote().forward(x_rref)
  28.             z_fut = self.p2_rref.rpc_async().forward(y_rref)
  29.             out_futures.append(z_fut)
  31.         return torch.cat(torch.futures.wait_all(out_futures))
  33.     def parameter_rrefs(self):
  34.         remote_params = []
  35.         remote_params.extend(self.p1_rref.remote().parameter_rrefs().to_here())
  36.         remote_params.extend(self.p2_rref.remote().parameter_rrefs().to_here())
  37.         return remote_params

步骤 3：定义训练循环

定义模型后，让我们实现训练循环。 我们使用专门的“主”工作器来准备随机输入和标签，并控制分布式反向传递和分布式优化器步骤。 它首先创建DistResNet50模块的实例。 它指定每个批量的微批数量，并提供两个 RPC 工作程序的名称（即worker1和worker2）。 然后，它定义损失函数，并使用parameter_rrefs()帮助器创建DistributedOptimizer以获取参数RRefs的列表。 然后，主训练循环与常规本地训练非常相似，除了它使用dist_autograd向后启动并为反向和优化器step()提供context_id之外。

  1. import torch.distributed.autograd as dist_autograd
  2. import torch.optim as optim
  3. from torch.distributed.optim import DistributedOptimizer
  5. num_batches = 3
  6. batch_size = 120
  7. image_w = 128
  8. image_h = 128
  10. def run_master(num_split):
  11.     # put the two model parts on worker1 and worker2 respectively
  12.     model = DistResNet50(num_split, ["worker1", "worker2"])
  13.     loss_fn = nn.MSELoss()
  14.     opt = DistributedOptimizer(
  15.         optim.SGD,
  16.         model.parameter_rrefs(),
  17.         lr=0.05,
  18.     )
  20.     one_hot_indices = torch.LongTensor(batch_size) \
  21.                         .random_(0, num_classes) \
  22.                         .view(batch_size, 1)
  24.     for i in range(num_batches):
  25.         print(f"Processing batch {i}")
  26.         # generate random inputs and labels
  27.         inputs = torch.randn(batch_size, 3, image_w, image_h)
  28.         labels = torch.zeros(batch_size, num_classes) \
  29.                     .scatter_(1, one_hot_indices, 1)
  31.         with dist_autograd.context() as context_id:
  32.             outputs = model(inputs)
  33.             dist_autograd.backward(context_id, [loss_fn(outputs, labels)])
  34.             opt.step(context_id)

第 4 步：启动 RPC 进程

最后，下面的代码显示了所有进程的目标函数。 主要逻辑在run_master中定义。 工作器被动地等待主服务器发出的命令，因此只需运行init_rpc和shutdown即可，其中默认情况下shutdown会阻塞，直到所有 RPC 参与者都完成。

  1. import os
  2. import time
  4. import torch.multiprocessing as mp
  6. def run_worker(rank, world_size, num_split):
  7.     os.environ['MASTER_ADDR'] = 'localhost'
  8.     os.environ['MASTER_PORT'] = '29500'
  9.     options = rpc.TensorPipeRpcBackendOptions(num_worker_threads=128)
  11.     if rank == 0:
  12.         rpc.init_rpc(
  13.             "master",
  14.             rank=rank,
  15.             world_size=world_size,
  16.             rpc_backend_options=options
  17.         )
  18.         run_master(num_split)
  19.     else:
  20.         rpc.init_rpc(
  21.             f"worker{rank}",
  22.             rank=rank,
  23.             world_size=world_size,
  24.             rpc_backend_options=options
  25.         )
  26.         pass
  28.     # block until all rpcs finish
  29.     rpc.shutdown()
  31. if __name__=="__main__":
  32.     world_size = 3
  33.     for num_split in [1, 2, 4, 8]:
  34.         tik = time.time()
  35.         mp.spawn(run_worker, args=(world_size, num_split), nprocs=world_size, join=True)
  36.         tok = time.time()
  37.         print(f"number of splits = {num_split}, execution time = {tok - tik}")

下面的输出显示通过增加每批中的拆分数量而获得的加速。

  1. $ python main.py
  2. Processing batch 0
  3. Processing batch 1
  4. Processing batch 2
  5. number of splits = 1, execution time = 16.45062756538391
  6. Processing batch 0
  7. Processing batch 1
  8. Processing batch 2
  9. number of splits = 2, execution time = 12.329529762268066
  10. Processing batch 0
  11. Processing batch 1
  12. Processing batch 2
  13. number of splits = 4, execution time = 10.164430618286133
  14. Processing batch 0
  15. Processing batch 1
  16. Processing batch 2
  17. number of splits = 8, execution time = 9.076049566268921