Pytorch并行分布式训练三 （模型并行）

Single-Machine Model Parallel Best Practiceshttps://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运算以实现更快的加速