[pytorch00]模型并行在单机上的最佳实践

模型并行在单机上的最佳实践

​ 模型并行在分布式训练技术中被广泛使用。之前的文章已经解释了如何使用DataParallel训练多个gpu上的神经网络；该方法将相同的模型复制到所有GPU，其中每个GPU处理splited后的输入数据。虽然它可以显著地加速训练过程，但它在某些情况下无法工作，因为模型太大，无法装入一个GPU中（如数个G甚至十几个G时）。这篇文章展示了如何通过使用**模型并行(Model parallel)来解决这个问题，与数据并行(DataParallel)**不同的是，它将一个模型分割到不同的GPU上，而不是在每个GPU上复制整个模型。

​ 具体来说，假设模型$M$包含10个层，当使用DataParallel时，每个GPU将拥有这10个层的副本，而当在两个GPU上并行使用模型时，每个GPU保存5个层。

​ 模型并行的思想是将模型的不同子网络放到不同的设备上，并实现相应的前向(forward)方法，使中间输出在设备之间传输。由于模型只有一部分操作在单个设备上，因此一组设备可以共同服务于更大的模型。在这篇文章中，我们不会试图构建太大的模型，并将它们挤进有限数量的gpu中。相反，这篇文章只着重展示了模型并行的思想，具体如何使用，取决于读者们。

基本用法

​ 让我们从一个包含两个线性层的玩具模型开始。要在两个GPU上运行此模型，只需将每个线性层放在不同的GPU上，并相应地移动输入和中间输出以匹配层设备。

import torch
import torch.nn as nn
import torch.optim as optim


class ToyModel(nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Linear(10, 10).to('cuda:0')
        self.relu = torch.nn.ReLU()
        self.net2 = torch.nn.Linear(10, 5).to('cuda:1')

    def forward(self, x):
        x = self.relu(self.net1(x.to('cuda:0')))
        return self.net2(x.to('cuda:1'))

​ 可以看到，上面的玩具模型看起来非常类似于在单一GPU上实现它，除了在适当的设备上放置线性层和张量的to(device)调用。这是模型中唯一需要更改的地方。backward()和torch.optim会自动处理梯度的更新，就好像模型是在一个GPU上一样。另外，在调用loss函数时，需确保标签与输出在同一设备上。

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

optimizer.zero_grad()
outputs = model(torch.randn(20, 10))
labels = torch.randn(20, 5).to('cuda:1')
loss_fn(outputs, labels).backward()
optimizer.step()

对现有的模型进行并行

​ 仅需几行更改，就可以在多个GPU上运行现有的单GPU模块。下面的代码展示了如何将torchvision.models.reset50()分解到两个gpu上。其思想是继承现有的ResNet模块，并在构建期间将层拆分到两个gpu。然后，重写forward方法，通过相应地移动中间输出来“缝合”两个子网络。

from torchvision.models.resnet import ResNet, Bottleneck

num_classes = 1000


class ModelParallelResNet50(ResNet):
    def __init__(self, *args, **kwargs):
        super(ModelParallelResNet50, self).__init__(
            Bottleneck, [3, 4, 6, 3], num_classes=num_classes, *args, **kwargs)

        self.seq1 = nn.Sequential(
            self.conv1,
            self.bn1,
            self.relu,
            self.maxpool,

            self.layer1,
            self.layer2
        ).to('cuda:0')

        self.seq2 = nn.Sequential(
            self.layer3,
            self.layer4,
            self.avgpool,
        ).to('cuda:1')

        self.fc.to('cuda:1')

    def forward(self, x):
        x = self.seq2(self.seq1(x).to('cuda:1'))
        return self.fc(x.view(x.size(0), -1))

​ 上面的实现解决了模型太大而无法装入一个GPU的问题。然而，你可能已经注意到，它会比在单一GPU上运行要慢。这是因为，在任何时间点，两个gpu只有一个在工作，而另一个坐在那里什么也不做。 由于需要在layer2和layer3之间将中间输出从cuda:0复制到cuda:1，性能会进一步降低。

​ 让我们运行一个实验来获得更定量的执行时间对比。在这个实验中，我们通过随机的输入和标签分别训练ModelParallelResNet50()和现有的torchvision.models.reset50()。在训练之后，模型不会产生任何有用的预测，但是我们可以对执行时间有一个合理的理解。

import torchvision.models as models

num_batches = 3
batch_size = 120
image_w = 128
image_h = 128


def train(model):
    model.train(True)
    loss_fn = nn.MSELoss()
    optimizer = optim.SGD(model.parameters(), lr=0.001)

    one_hot_indices = torch.LongTensor(batch_size) \
                           .random_(0, num_classes) \
                           .view(batch_size, 1)

    for _ in range(num_batches):
        # generate random inputs and labels
        inputs = torch.randn(batch_size, 3, image_w, image_h)
        labels = torch.zeros(batch_size, num_classes) \
                      .scatter_(1, one_hot_indices, 1)

        # run forward pass
        optimizer.zero_grad()
        outputs = model(inputs.to('cuda:0'))

        # run backward pass
        labels = labels.to(outputs.device)
        loss_fn(outputs, labels).backward()
        optimizer.step()

MSELoss为损失函数，optim为损失函数。SGD作为优化器。训练128 X 128幅图像，num_batches=3，每批包含120幅图像。我们使用timeit运行train方法10次，并用标准差绘制执行时间。

import matplotlib.pyplot as plt
plt.switch_backend('Agg')
import numpy as np
import timeit

num_repeat = 10

stmt = "train(model)"

setup = "model = ModelParallelResNet50()"
# globals arg is only available in Python 3. In Python 2, use the following
# import __builtin__
# __builtin__.__dict__.update(locals())
mp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
mp_mean, mp_std = np.mean(mp_run_times), np.std(mp_run_times)

setup = "import torchvision.models as models;" + \
        "model = models.resnet50(num_classes=num_classes).to('cuda:0')"
rn_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
rn_mean, rn_std = np.mean(rn_run_times), np.std(rn_run_times)


def plot(means, stds, labels, fig_name):
    fig, ax = plt.subplots()
    ax.bar(np.arange(len(means)), means, yerr=stds,
           align='center', alpha=0.5, ecolor='red', capsize=10, width=0.6)
    ax.set_ylabel('ResNet50 Execution Time (Second)')
    ax.set_xticks(np.arange(len(means)))
    ax.set_xticklabels(labels)
    ax.yaxis.grid(True)
    plt.tight_layout()
    plt.savefig(fig_name)
    plt.close(fig)


plot([mp_mean, rn_mean],
     [mp_std, rn_std],
     ['Model Parallel', 'Single GPU'],
     'mp_vs_rn.png')

结果表明，模型并行实现的执行时间比现有的单GPU实现长4.02 / 3.75-1 = 7％。因此我们可以得出结论，在gpu之间来回复制张量的开销大约有7%。还有改进的空间！因为我们知道两个gpu中的一个在执行过程中处于空闲状态。一种选择是将每个批处理进一步分割，这样当一个分割到达第二个子网络时，剩下的分割可以被送入第一个子网络。通过这种方式，可以在两个gpu上并发地运行两个（多个）连续的分割。

通过pipelining inputs进行加速

在接下来的实验中，我们将每批120幅图像进一步划分为6组，每组20张图片。由于PyTorch异步启动CUDA操作，实现不需要派生多个线程来实现并发。

class PipelineParallelResNet50(ModelParallelResNet50):
    def __init__(self, split_size=20, *args, **kwargs):
        super(PipelineParallelResNet50, self).__init__(*args, **kwargs)
        self.split_size = split_size

    def forward(self, x):
        splits = iter(x.split(self.split_size, dim=0))
        s_next = next(splits)
        s_prev = self.seq1(s_next).to('cuda:1')
        ret = []

        for s_next in splits:
            # A. s_prev runs on cuda:1
            s_prev = self.seq2(s_prev)
            ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

            # B. s_next runs on cuda:0, which can run concurrently with A
            s_prev = self.seq1(s_next).to('cuda:1')

        s_prev = self.seq2(s_prev)
        ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))

        return torch.cat(ret)


setup = "model = PipelineParallelResNet50()"
pp_run_times = timeit.repeat(
    stmt, setup, number=1, repeat=num_repeat, globals=globals())
pp_mean, pp_std = np.mean(pp_run_times), np.std(pp_run_times)

plot([mp_mean, rn_mean, pp_mean],
     [mp_std, rn_std, pp_std],
     ['Model Parallel', 'Single GPU', 'Pipelining Model Parallel'],
     'mp_vs_rn_vs_pp.png')

​ 请注意，设备到设备张量拷贝操作在当前流(current streams)上是同步的。如果您创建多个流，您必须确保复制操作正确同步。在完成复制操作之前写入源张量或读取/写入目标张量可能导致不确定的行为。上面的实现只在源设备和目标设备上使用默认流，因此没有必要强制执行额外的同步。

​ 实验结果表明，对并行ResNet50进行建模的流水线输入可将训练过程大致加快3.75 / 2.51-1 = 49％。它离理想的100%加速还有相当大的距离。由于我们在流水线并行实现中引入了一个新的参数分割大小，所以还不清楚这个新参数如何影响整个训练时间。直观地说，使用小的分割大小会导致很多小的CUDA内核启动(GPU利用率下降)，而使用大的分割大小会导致在第一次和最后一次分割期间相对较长的空闲时间。对于这个特定的实验，可能有一个最佳的分割大小配置。让我们尝试通过使用几个不同的分割大小值运行实验来找到它。

means = []
stds = []
split_sizes = [1, 3, 5, 8, 10, 12, 20, 40, 60]

for split_size in split_sizes:
    setup = "model = PipelineParallelResNet50(split_size=%d)" % split_size
    pp_run_times = timeit.repeat(
        stmt, setup, number=1, repeat=num_repeat, globals=globals())
    means.append(np.mean(pp_run_times))
    stds.append(np.std(pp_run_times))

fig, ax = plt.subplots()
ax.plot(split_sizes, means)
ax.errorbar(split_sizes, means, yerr=stds, ecolor='red', fmt='ro')
ax.set_ylabel('ResNet50 Execution Time (Second)')
ax.set_xlabel('Pipeline Split Size')
ax.set_xticks(split_sizes)
ax.yaxis.grid(True)
plt.tight_layout()
plt.savefig("split_size_tradeoff.png")
plt.close(fig)

结果表明，将split_size设置为12可获得最快的训练速度，从而获得3.75 / 2.43-1 = 54％的加速比。还有机会进一步加快训练进程。例如，所有对cuda:0的操作都被放在它的默认流上。这意味着next分割的计算操作不能复制prev分割的。

但是，由于prev和next分块是不同的张量，所以没有重叠计算和复制的问题。

实现需要在两个gpu上都使用多个流，不同的子网结构需要不同的流管理策略。

由于没有一种通用的多流解决方案适用于所有的模型并行用例，我们将不在本教程中讨论它。