SINGLE-MACHINE MODEL PARALLEL BEST PRACTICES
本文译自PYTORCH并行处理:Author: Shen Li
模型并行在分布式训练中很常用。pytorch本身就用 DataParallel 做并行训练,使用非常简单。思想也比较直观:将模型复制到多个GPU上,然后每个gpu计算输入的一部分。尽管这个方法可以加速训练,但是当模型太大以至于放不下一个单独gpu时就不太好用了。本文利用另一种方法:model parallel(并非像DataParallel是一个函数,而是一种操作)。相比于DataParallel,该方法将一个模型分解后放到不同GPU上,而不是复制模型到不同GPU上。(栗子:一个总共10层的模型,当利用DataParallel,每个GPU将会有这个模型的全部10层,而利用model parallel在两个GPU上时,每个GPU仅仅有5层。)
model parallel是将一个模型的不同子网络放到不同设备上,在实现forward函数时手动移动或合并不同device上的数据。由于每个设备上仅有一部分模型,所以全部设备足以撑起一个比较大的模型。本文不会构建一个大型模型然后将其压缩至有限的GPU上,而是关注弄清model parallel的原理。具体应用取决于读者的具体应用。
对于多机model parallel训练,参考: Getting Started With Distributed RPC Framework
Basic Usage
从包含两个线性层的简易模型开始。要在两个GPU上运行此模型,只需将每个线性层放在不同的GPU上,然后相应地移动输入和中间输出以匹配devices。
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'))
注意to(devide)的用法,放在不同设备上。
注意到ToyModel实现看起来与在单一GPU上实现非常相似,除了一些to(device)的操作,将线性层个tensor置于合适的devices。这是唯一需要的一些改动。backward()和torch.optim将自动整理梯度。你只需要确保在计算loss时labels和输出在同一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模型run在多个GPU上。以下代码展示了如何将torchvision.models.reset50()
分解到两个GPU上。思想史将其layers分解为两部分,然后重载forwaed函数来移动中间的输出。
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上却会稍微变慢。这是因为在任何时候,仅仅只有一个GPU在工作,另一个只是在等待啥也不干。对于后者在反复的移动不同device上的中间结果时会导致额外的时间开销,以至于多个比一个还慢点。
为此做个实验验证一下,分别用随机数来训练这两个模型: 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()
train(model)函数在3个批量上训练,每个批量120张图,利用timeie来run10次后统计时间。
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')
从上图可以看到model parallel的实现比但GPU模型开销增加了4.02/3.75-1=7%
。所以可以认为这7%的开销花在了copy不同GPU上的tensor。当然这里肯定有空间提升多GPU的表现:上面的model parallel方法在实现时将模型划分成分多个部分,这多个部分工作室串行的!!!即一个批量走完GPU-0再走GPU-1。试想这个批量在走GPU-0的时候其他GPU在休息,在走GPU-1的时候其他GPU也在休息。如何让所有GPU同时工作起来?一个解决方案就是划分批量!例如将批量大小均匀化分为两部分,第一部分数据先走GPU-0,走完之后,第一部分数据接着走GPU-1,同时第二部分开始走GPU-0...以此循环,没有GPU会空闲着。
Speed Up by Pipelining Inputs
在下面将会将120的批量大小划分为每份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')
注意到,设备到设备tensor复制操作在源设备和目标设备上的当前流上同步。如果创建多个流,则必须确保复制操作正确同步。在完成复制操作之前写入源tensor或读/写目标张量可能会导致未定义的行为。上述实现仅在源设备和目标设备上使用默认流,因此不必强制执行额外的同步。
可以看到,这种方法不比model parallel ResNet50 快了3.75/2.51-1=49% !!!但是距离100%的理想加速还很远。为此我们可以去试着找出一种最优的split_size,上面我们的这个参数为20,但这不一定是最优参数。所以可以画图找出最优参数看看最多能加速多少。
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和next拆分是不同的张量,因此不存在将一个计算与另一个计算重叠的问题。实现需要在两个gpu上使用多个流,不同的子网结构需要不同的流管理策略。由于没有一个通用的多流解决方案适用于所有模型并行用例,我们将在本教程中不讨论它。