本文针对BertLarge分布式并行训练所存在的问题,介绍Whale的并行化设计和方案。通过为模型并行和数据并行,搭配流水并行辅助并行策略,优化通信拓扑结构,以解决BertLarge分布式训练性能较差的问题。在Whale中,您可以通过模型划分、资源划分及映射三个步骤,实现大规模数据及模型的分布式训练。
背景信息
受益于深度神经网络的发展,NLP和CV领域的模型效果得到大幅度提升。同时,模型参数量也大幅度增加。以Imagenet分类任务为例,优胜算法从2014年的GoogleNet到2018年的Squeeze-and-Excitation Networks,参数量增长约36倍(从4百万增长至1.458亿参数)。

以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小又导致通信占比高,因此使用传统数据并行进行训练的加速效果极差。
问题描述

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加重训练通信占比,导致分布式扩展效果差。
- 模型并行
- 将模型按照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模型的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卡的利用率,具体操作如下:
- 将一个Mini-Batch拆分为多个Micro-Batch。
- 每张卡训练完一个Micro-Batch数据后,将Activation传给下一个GPU,然后立刻训练下一个Micro-Batch数据。
上图的示例中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实现流水并行
在Whale中实现流水并行需要三个步骤(标准的Whale分布式编程范式):模型划分、资源划分及映射。
性能
测试环境项 | 描述 |
---|---|
GPU型号 | ecs.gn6v-c10g1.20xlarge(V100 * 8) |
网络 | VPC-35 GB |
NCCL_MAX_NRINGS | NVIDIA官方参数,测试时取值为4。 |
NCCL_MIN_NRINGS | NVIDIA官方参数,测试时取值为4。 |
- Horovod数据并行
- Whale数据并行
- 结合Whale模型并行和流水并行

- Whale在数据并行场景的性能优于Horovod。尤其在GPU 64卡环境下,Whale的加速比是Horovod的1.74倍。
- 在GPU小于8卡的场景,Whale数据并行与结合Whale模型并行和流水并行的性能相差不大。因为数据并行的效果主要得益于服务器内GPU通过高速NVLink相连,通信效率很高。当GPU卡数逐渐增多,出现跨服务器通信时,流水并行的优势才会显现。
- 在GPU 64卡场景,结合Whale模型并行和流水并行的吞吐性能是Whale数据并行的1.34倍。
在文档使用中是否遇到以下问题
更多建议
匿名提交