SINGLE-MACHINE MODEL PARALLEL BEST PRACTICES

 模型并行在分布式训练技术中得到了广泛的应用。以前的帖子已经解释了如何使用DataParallel数据并行在多个GPU上训练神经网络；这个特性将相同的模型复制到所有GPU，其中每个GPU消耗输入数据的不同分区。虽然它可以显著加速训练过程，但对于模型太大而不能适应单个GPU的一些用例，它不起作用。这篇文章展示了如何通过使用model parallel模型并行来解决这个问题，与数据并行相比，它将单个模型分割到不同的GPU上，而不是在每个GPU上复制整个模型(具体来说，假设模型m包含10层：当使用数据并行时，每个GPU将拥有这10层中的每一层的副本，而当在两个GPU上使用模型并行时，每个GPU可以承载5层)。

模型并行的高级思想是将模型的不同子网络放置在不同的设备上，并相应地实现前向方法，以在设备之间移动中间输出。由于模型只对任何单个设备进行操作，一组设备可以共同为更大的模型服务。在这篇文章中，我们不会试图构建巨大的模型，并将它们压缩到有限数量的GPU中。相反，本文的重点是展示模型并行的思想。读者有责任将这些想法应用到现实世界的应用中。

 

NOTE

For distributed model parallel training where a model spans multiple servers, please refer to Getting Started With Distributed RPC Framework for examples and details.对于模型跨越多个服务器的分布式模型并行训练，请参阅分布式RPC框架的入门示例和详细信息。

 

Basic Usage

让我们从一个包含两个线性层的玩具模型开始。要在两个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上实现它，除了在适当的设备上放置线性层和张量的五到（设备）调用之外。这是模型中唯一需要更改的地方。向后()和torch.optim将自动处理梯度，就像模型在一个GPU上一样。在调用损失函数时，只需要确保标签与输出位于同一设备上。

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()

 

Apply Model Parallel to Existing Modules

还可以在多个GPU上运行现有的单GPU模块，只需几行更改。下面的代码显示了如何将torchvision.models.resnet50()分解为两个GPU。其思想是继承现有的ResNet模块，并在构建过程中将层拆分为两个GPU。然后，重写前向方法，通过相应地移动中间输出来缝合两个子网络。

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中只有一个在工作，而另一个坐在那里什么也不做。 性能进一步恶化，因为中间输出需要在第2层和第3层之间从cuda：0复制到cuda：1。

让我们运行一个实验，以获得更定量的执行时间视图。 在本实验中，我们通过运行随机输入和标签来训练模型并行ResNet50和现有的torchvision.models.resnet50()。 经过培训，模型不会产生任何有用的预测，但我们可以对执行时间有一个合理的理解。

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()

上面的训练（模型）方法使用n。 以MSELoss为损失函数，并进行优化。 作为优化器的SGD。 它模拟128X128图像的训练，这些图像被组织成3批，其中每批包含120幅图像。 然后，我们使用timeit运行列车（模型）方法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上运行。

 

Speed Up by Pipelining Inputs

在接下来的实验中，我们进一步将每个120个图像批次划分为20个图像分割。 由于Py Torch异步启动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')

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

 

 

实验结果表明，平行ResNet50模型的流水线输入使训练过程加快了约3.75/2.51-1=49%。 它仍然离理想的100%加速相当远。 由于我们在管道并行实现中引入了一个新的参数split_sizes，因此不清楚新参数如何影响整个训练时间。 直观地说，使用小split_size会导致许多微小的CUDA内核启动，而在第一次和最后一次分裂期间，使用大的split_size结果会导致相对较长的空闲时间。 两者都不是最佳选择。 对于这个特定的实验，可能有一个最优的split_size配置。 让我们尝试通过使用几个不同的split_size值运行实验来找到它。

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上的所有操作都放在其默认流上。 这意味着下一个拆分上的计算不能与prev拆分的复制操作重叠。 然而，由于prev和下一个分裂是不同的张量，没有问题将一个人的计算与另一个人的副本重叠。 实现需要在两个GPU上使用多个流，不同的子网结构需要不同的流管理策略。 由于没有一般的多流解决方案适用于所有模型并行用例，我们将不在本教程中讨论它。

注：

这篇文章显示了几个性能度量。 在自己的机器上运行相同的代码时，大家可能会看到不同的数字，因为结果取决于底层硬件和软件。为了获得环境的最佳性能，一个适当的方法是首先生成曲线来计算最佳分割大小，然后使用该分割大小来流水线输入。

​

接下来，给大家介绍一下租用GPU做实验的方法，我们是在智星云租用的GPU，使用体验很好。具体大家可以参考：智星云官网： http://www.ai-galaxy.cn/，淘宝店：https://shop36573300.taobao.com/公众号: 智星AI