Pytorch并行分布式训练三 （模型并行）
Single-Machine Model Parallel Best Practices
https://pytorch.org/tutorials/intermediate/model_parallel_tutorial.html#pytorch 多GPU训练一般采用的是“数据并行”的方法，但它同样也支持“模型并行”。“模型并行”需要解决的首要问题就是模型太大，导致单张GPU无法完整的加载整个模型。由于数据并行方法模型会被复制（广播），面临这种情况时数据并行也无济于事，模型并行能够将模型拆分到各个GPU当中，以解决上述问题。
1.序言
上一章讲到的DataParallel已经能够实现单机多卡训练，但它不适合模型很大的情况，原因在于它会将模型复制成多份。接下来将介绍的方法会将模型进行分割，这也使得每张GPU中需要运行的网络层数将减少。此方法的关键思想在于：将不同的“子网络”放在不同的GPU当中forward，再跨GPU处理中间的输出（GPU之间进行通信）。
本章介绍的方法属于“模型并行”，与DP,DDP等“数据并行”的方法在实现原理上有着本质性的区别，但值得注意的是，DDP能够与“模型并行”方法结合使用。
2.基础用法
我们将使用一个toy模型来实现同一模型运行在两个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'))可以看到，在forward当中只有一些to(device)与使用单GPU时不同，backward()和torch.optim会自动处理跨GPU的情况，我们不用做特别的代码修改。
你还需要再注意的一件事情是，保证输入和网络第一层、输出和label都在同一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()3.将模型并行运用在已有的模型上
已有的单GPU模型也可以通过简单的改造变成GPU并行的模型，主要思想是：将已有模型进行分块，并放入不同的GPU当中。下面以已有的resnet50为例（注意继承关系）
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))如上，通过简单的nn.Sequential()包装，就能将已有的单GPU模型拆分到不同的GPU当中去。然而，这样的模型拆分只解决了“模型太大，单GPU无法训练“的问题。由于中间输出（layer2和layer3之间）必须进行跨GPU的同步（cuda:0 到 cuda:1），这样的多GPU训练并不能提升速度。实际上它将耗费更多的时间（如下图）（由于多出来的同步将增加耗时）。
4.使用管道化输入（Pipelining Inputs）进行加速
为了实现加速效果，我们可以让多个GPU设备进行并行处理，修改上面的代码可得（注意继承自ModelParallelResNet50）:
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') #compute on cuda:0
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) #compute on cuda:1
ret.append(self.fc(s_prev.view(s_prev.size(0), -1)))
return torch.cat(ret)通过将输入数据在BS的维度上进行拆分，可实现两个GPU的并行计算（GPU计算本身就是并行的，A.,B.操作都是GPU操作，不需要写额外的代码来支持并行）。
需要注意的是，这里不能使用多个GPU流(Stream)进行GPU内的并行，因为在做to(device)操作之前，必须要完成数据同步。此处的两个GPU都是用单个Stream，所以也不需要做额外的数据同步操作。
下面是实验时间，管道的方法相对于单GPU实现了3.75/2.51-1=49% 的加速效果：
另外，此处的split_size是一个超参数，将会影响性能。当split_size太小时，会产生许多小的CUDA kernel从而影响速度；当split_size太大是，开头和结尾的splits将消耗更多的时间（非并行的），影响总体速度。下面是在作者的实验环境中进行的split_size调参实验（不同的硬件及环境可能会得到不同的结果）：
最后，需要说明的是，此方法还存在一些缺点。需要使用其他方法来实现各个GPU中的多Stream运算以实现更快的加速