BertLarge分布式训练（流水并行） - 机器学习PAI

背景信息

受益于深度神经网络的发展，NLP和CV领域的模型效果得到大幅度提升。同时，模型参数量也大幅度增加。以Imagenet分类任务为例，优胜算法从2014年的GoogleNet到2018年的Squeeze-and-Excitation Networks，参数量增长约36倍（从4百万增长至1.458亿参数）。

NLP领域很多模型相较于图像领域的模型规模更大。如下图所示，近几年NLP领域的训练数据和模型规模都非常大，例如GPT-3的参数达到了1750亿。由于模型参数的增长速度远超GPU显存的增长速度（从P4的8 GB到A100的40 GB），导致训练支持的Batch Size变小，进而使训练过程中的通信占比变高，影响模型分布式扩展。 NLP数据及模型规模走势

以BertLarge模型为例，它在很多NLP场景里取得了非常好的结果。BERT的Idea与ELMo和GPT都非常接近，但是BERT可以在很大的非监督语料里进行预训练，速度远超ELMo，效果优于GPT，在真实业务场景中广泛应用。

BertLarge只有3.4亿参数规模，与T5和GPT-3相比，模型参数非常小。但是在Nvidia V100 16G GPU上，训练Batch Size仅达到2~8（具体值和Embedding大小、Sequence Length等有关）。模型参数规模大导致通信梯度大，Batch Size小又导致通信占比高，因此使用传统数据并行进行训练的加速效果极差。

问题描述

以BertLarge模型为例，其模型结构如下图所示。 BertLarge模型结构

BertLarge的原始模型代码如下所示。

embedding_output = embedding(inputs)
encoder_output = encoder_layer_1_24(embedding_output)
pooler_output = pooler(encoder_output)

在数据并行场景，以max_seq_len等于384为例，使用V100单卡（16 GB显存），模型最大Batch Size仅达到6。每轮迭代同步1.245 GB的梯度，在50 GB网络环境下，通信耗时为2*1.245 GB/50 GB=398.4 ms（读取训练数据也需要消耗网络带宽，因此实际通信耗时大于该值）。因此，大规模训练主要存在两个问题：

Batch size过小，导致模型波动较大，从而使得收敛效果差。
梯度通信量大，小Batch Size加重训练通信占比，导致分布式扩展效果差。

为解决上述问题，Whale提出模型并行、流水并行及数据并行的混合并行策略。同时支持高性能通信Backend及调度优化，从而为您提供高性能、简洁且易用的分布式训练框架。模型并行、流水并行及数据并行的混合并行策略如下：

模型并行
- 将模型按照Layer粒度切分为不同的Stage，并将其分配至不同的GPU中执行，以降低每张GPU卡中的模型显存占用，从而提高Batch Size。
- Stage之间使用Activation通信代替梯度同步，大幅度降低卡间通信量。例如在BertLlarge模型中，每轮迭代仅需要传输27 MB的Activation数据。
- 增大Batch Size，使训练更加稳定，从而收敛效果更好。
流水并行
如果仅采用模型并行，则GPU卡间任务执行有依赖，因此同一时间只有一个GPU执行，其他GPU空闲，导致GPU利用率低。流水并行可以将一个Mini-Batch拆分为多个Micro-Batch，不同GPU中可以同时执行流水的不同Stage，从而提高GPU利用率。
数据并行
受限于模型Layer总数和流水并行效率，模型并行和流水并行的混合并行策略不能无限地进行分布式扩展。在超大训练数据规模场景，还需要结合数据并行进行分布式扩展。

实现方案

通过模型并行、流水并行及数据并行的混合并行策略，优化BertLarge模型的分布式性能。详细方案如下：

模型并行
BertLarge模型的Batch Size通常仅达到2~8（具体值与Embedding大小、Sequence Length等有关），导致模型波动大，进而使得收敛效果差。此时，可以将模型以Layer为单位拆分为多份，并将其放至不同GPU卡中进行分布式训练，即模型并行。
无论在何种带宽、硬件拓扑下，模型分片可能受Load Balance（包括显存均衡度、算力均衡度）影响。BertLarge模型的每一层Activation、显存及Flops计算几乎都一致，因此从均衡各个模型部分显存占用、算力需求的角度出发，将BertLarge中的Encoder Layer 1~8层、Encoder Layer 9~16层，Encoder Layer 17~24层分别放至不同的GPU卡中进行训练，并行化后的计算图如下图所示。如上图所示，将Embedding Layer和Encoder Layer 1~8层放在GPU 0中进行运算，将Encoder Layer 9~16层放在GPU 1中进行运算，Encoder Layer 17~24层和Pooler层放在GPU 2中进行运算。该处理方式具有以下优势：
- 每张GPU卡的模型大幅度减小，因此可以增大Batch Size，以提升收敛加速。
- 只需要在GPU之间通信Activation（每轮迭代约27 MB的Activation数据），省掉了梯度通信过程。
- 对于模型过大导致单卡显存无法存放的情况，通过Layer拆分的模型并行方式，使模型训练成为可能。
流水并行
仅采用模型并行进行分布式训练，其中一张GPU卡计算时，其他GPU卡空闲，导致GPU资源利用率低。如下图所示。上图中，在时间维度，每一个时间点上只有一张GPU卡在执行Forward或Backward，其他GPU卡都处于空闲等待状态。Whale设计了流水并行，以提高模型并行场景GPU卡的利用率，具体操作如下：
1. 将一个Mini-Batch拆分为多个Micro-Batch。
2. 每张卡训练完一个Micro-Batch数据后，将Activation传给下一个GPU，然后立刻训练下一个Micro-Batch数据。
按照上述操作，在多个GPU之间形成流水，执行逻辑如下图所示。上图的示例中Micro-Batch数量为4，流水并行优化后的时间轴显示，同一时间点上多张GPU卡可以并行计算。当4个Micro-Batch结束后，每张GPU卡将每个Micro-Batch的梯度进行本地累计（即上图中的Grads Acc）之后，进行Update权重操作。与模型并行相比，提升了GPU利用率。然而，上图中依然存在大量空白，即某些时间点上依然存在某些GPU处于空闲状态。针对该情况，可以增大流水并行的Micro-Batch数量，以降低空闲时间。此外，Whale采用Backward-Prefered（调度每个Micro-Batch间的前后向计算顺序，使GPU卡尽量一直处于计算状态）调度优化策略提升流水并行性能，从而降低GPU空闲时间，如下图所示（Micro-Batch配置为5）。
混合并行
因为模型并行不可能将模型拆分为无限份，且切分数量过大会影响性能，所以仅采用模型并行和流水并行依然不能满足性能要求。针对业务场景产生的海量训练数据，需要提高大量计算资源并行处理，因此还需要采用数据并行解决大规模数据问题。在流水并行的基础上增加数据并行后的计算图如下所示。上图中以3个Worker为例，每个Worker拥有3个GPU。Whale将每个Worke内的模型部分分片放至本Worker内的多张GPU卡中进行流水并行计算。每个Worker完成Micro-Batch数量个Batch Size的流水训练后，先累计本地的梯度，再在Worker之间进行梯度AllReduce同步。
为了更高效利用服务器内GPU间的NVLink（如果没有，则通过PCI-e）带宽，降低梯度同步的开销，Whale进行了如下优化：
- 将模型并行拆分放至不同服务器中进行计算。
- 将数据并行部分的梯度同步（通信量约为1.2 GB）放至服务器内进行。
- 将Layer间切分产生的Activation通信（约27 MB）通过TCP网络通信。
该方式可以提高通信效率，从而降低每轮迭代时间。在Whale中实现该优化后的并行化方案，无需修改模型定义部分代码，只需要修改硬件资源分组逻辑，即将参数layout从Column模式改成Row模式，详情请参见whale.cluster。该方案的计算图如下所示。

Whale实现流水并行

在Whale中实现流水并行需要三个步骤（标准的Whale分布式编程范式）：模型划分、资源划分及映射。

模型划分。

将模型各部分按照显存占用和算力需求划分为不同部分，即Scopes划分，详情请参见whale.scopes。上述的实现方案中将模型划分为以下三部分：

Stage 0：包括Embedding和Encoder Layer 1~8层
Stage 1：包括Encoder Layer 9~16层
Stage 2：包括Encoder Layer17~24层和Pooler层

针对BertLarge模型的原始代码，在Whale中您可以通过如下步骤实现流水并行。

通过whale.stage实现模型并行，代码示例如下。

import whale as wh

with wh.stage():
    embedding_output = embedding(inputs)
    encoder_output = encoder_layer_1_8(embedding_output)

with wh.stage():
    encoder_output = encoder_layer_9_16(embedding_output)

with wh.stage():
    encoder_output = encoder_layer_17_24(embedding_output)
    pooler_output = pooler(encoder_output)

在模型并行的代码基础上，通过whale.pipeline实现流水并行，示例代码如下。

import whale as wh

with wh.pipeline(num_micro_batch=5): # 按顺序依次处理5个Mini-Batch大小的数据，进行流水并行。每轮迭代每个副本训练5*Mini-Batch条样本数据。
    with wh.stage():
        embedding_output = embedding(inputs)
        encoder_output = encoder_layer_1_8(embedding_output)

    with wh.stage():
        encoder_output = encoder_layer_9_16(embedding_output)

    with wh.stage():
        encoder_output = encoder_layer_17_24(embedding_output)
        pooler_output = pooler(encoder_output)

在流水并行的基础上，通过whale.replica实现数据并行，示例代码如下所示。

import whale as wh

with wh.replica():  # 数据并行。
    with wh.pipeline(num_micro_batch=5):  # 按顺序依次处理5个Mini-Batch大小的数据，进行流水并行。每轮迭代每个副本训练5*Mini-Batch条样本数据。
        with wh.stage():
            embedding_output = embedding(inputs)
            encoder_output = encoder_layer_1_8(embedding_output)

        with wh.stage():
            encoder_output = encoder_layer_9_16(embedding_output)

        with wh.stage():
            encoder_output = encoder_layer_17_24(embedding_output)
            pooler_output = pooler(encoder_output)

资源划分。
由于将模型划分为三个部分，因此需要将硬件资源进划分为三个Virtual Devices。假设申请3*3的服务器规模，划分策略为layout="row"（详情请参见whale.cluster），则资源分为如下三组：
- Virtual Devices 0：[[/worker:0/gpu:0], [/worker:0/gpu:1], [/worker:0/gpu:2]]
- Virtual Devices 1：[[/worker:1/gpu:0], [/worker:1/gpu:1], [/worker:1/gpu:2]]
- Virtual Devices 2：[[/worker:2/gpu:0], [/worker:2/gpu:1], [/worker:2/gpu:2]]
Whale提供的Cluster工具可以对申请的Worker进行划分，示例代码如下。
```
cluster = wh.cluster(layout={"row":3})
```
映射。
根据模型划分和资源划分结果，将模型各部分分别映射至Virtual Devices：
- Stage 0部分采用数据并行放至Virtual Device 0，即[[/worker:0/gpu:0], [/worker:0/gpu:1], [/worker:0/gpu:2]]。每张GPU卡都有一个Stage 0的模型副本。
- Stage 1部分采用数据并行放至Virtual Device 1，即[[/worker:1/gpu:0], [/worker:1/gpu:1], [/worker:1/gpu:2]]。每张GPU卡都有一个Stage 1的模型副本。
- Stage 2部分采用数据并行放至Virtual Device 2，即[[/worker:2/gpu:0], [/worker:2/gpu:1], [/worker:2/gpu:2]]。每张GPU卡都有一个Stage 2的模型副本。
对于映射部分，在Whale中只需要对已生成的Cluster执行with语法，即可轻松完成模型到硬件资源的映射。通过Whale完成并行化的模型核心代码如下（包含流水并行的可执行代码请参见pipelined_bert_models_3_staged.py）。
```
import whale as wh
cluster = wh.cluster(layout={"row":3}) #资源划分。

with cluster:                         # 映射。
    with wh.replica():
        with wh.pipeline(num_micro_batch=5):
            with wh.stage():
                embedding_output = embedding(inputs)
                encoder_output = encoder_layer_1_8(embedding_output)

            with wh.stage():
                encoder_output = encoder_layer_9_16(embedding_output)

            with wh.stage():
                encoder_output = encoder_layer_17_24(embedding_output)
                pooler_output = pooler(encoder_output)
```

性能

以Horovod的数据并行性能数据为Baseline进行对比，测试环境如下。


测试环境项	描述
GPU型号	ecs.gn6v-c10g1.20xlarge（V100 * 8）
网络	VPC-35 GB
NCCL_MAX_NRINGS	NVIDIA官方参数，测试时取值为4。
NCCL_MIN_NRINGS	NVIDIA官方参数，测试时取值为4。

本文中的流水并行对比测试场景分别为：

Horovod数据并行
Whale数据并行
结合Whale模型并行和流水并行

对比结果如下图所示。流水并行和其他的性能对比

说明横坐标表示GPU卡数，纵坐标表示加速比。

从上图中，可以得到以下结论：

Whale在数据并行场景的性能优于Horovod。尤其在GPU 64卡环境下，Whale的加速比是Horovod的1.74倍。
在GPU小于8卡的场景，Whale数据并行与结合Whale模型并行和流水并行的性能相差不大。因为数据并行的效果主要得益于服务器内GPU通过高速NVLink相连，通信效率很高。当GPU卡数逐渐增多，出现跨服务器通信时，流水并行的优势才会显现。
在GPU 64卡场景，结合Whale模型并行和流水并行的吞吐性能是Whale数据并行的1.34倍。