查看原文
其他

揭秘!MLPerf Training v2.0飞桨何以力压NGC PyTorch,实现同等GPU配置BERT模型训练性能第一

将AI进行到底的 飞桨PaddlePaddle 2023-07-25

在6月30日最新发布的MLPerf Training v2.0里,百度使用飞桨框架( PaddlePaddle )和百度智能云百舸计算平台提交的BERT Large模型GPU训练性能结果,在同等GPU配置下的所有提交结果里排名第一,超越了高度定制优化且长期处于领先位置的NGC PyTorch框架,向全世界展现了飞桨框架的性能优势。

图1. MLPerf Training v2.0 BERT模型在8卡NVIDIA GPU A100( 400W功耗,80G显存 )下前五名的训练性能结果,百度飞桨方案比其他提交结果快5%-11%不等[1]

飞桨端到端

全流程优化概述



BERT模型是MLPerf Training v2.0自然语言处理领域的唯一基准模型。飞桨在BERT模型8卡GPU训练上创造了世界最优的训练性能,这来自于飞桨框架基础性能与分布式技术的领先性,以及飞桨与NVIDIA GPU的深度协同优化。对于深度学习模型训练任务,从数据读取到模型计算,从底层算子到上层分布式策略,从多设备负载均衡到全流程调度机制,都会影响最终训练性能。

飞桨基于领先的架构设计和长期的实践打磨,在高性能训练方面做出了系统性的优化工作,主要体现在如下几个方面:
  • 变长序列输入模型的计算加速
针对变长序列输入模型大多采用的padding填充对齐方式[2-4]带来的冗余计算问题,提供对变长输入和对应模型结构的高效支持,让GPU算力资源专注于有效计算,尤其是对Transformer类模型计算效率提升明显。
  • 数据读取和模型训练的负载均衡
针对分布式训练经常出现的负载不均衡问题,将模型训练和数据读取、预处理分配到不同设备上进行,确保异构算力物尽其用,实现数据IO和计算的平衡。
  • 高性能算子和融合优化技术
针对框架基础性能优化的极致需求,研发了高性能的高可复用算子库PHI ,充分优化GPU内核实现,提升算子内部计算的并行度,并通过算子融合降低仿存开销,发挥GPU的极限性能。
  • 高加速比的混合并行训练策略
针对传统数据并行性能、显存瓶颈受限的问题,实现了融合数据并行、模型并行、分组参数切片并行等策略的混合并行分布式训练策略,部分场景下可实现超线性加速的分布式训练性能。
  • 全流程异步执行调度
针对模型训练过程各环节存在的同步频率高、时间重叠度低等问题,设计异步调度机制,保证模型收敛的同时去除大部分同步操作,实现数据处理、训练和集合通信等各环节近乎全异步调度,提升端到端极致性能。

图2. 飞桨端到端训练全流程优化路径
以下以MLPerf BERT模型为例,逐一介绍飞桨框架的各项优化内容,更详细的优化细节请参见http://arxiv.org/abs/2208.08124 。

变长序列

输入模型的计算加速



自然语言处理领域的输入数据大部分都是变长的,为了组网和计算方便,往往会通过padding的方式将每个序列填充为等长的长度,但带来了冗余数据和冗余计算,如图3所示。为了提升MLPerf BERT模型性能,需要去除这些冗余数据带来的冗余计算量。

图3. 序列数据 padding 的示意图


支持变长Encoder计算

BERT Encoder是BERT的主要组成部分,其中的Multi-Head Attention部分的优化是提升BERT训练性能的关键。Multi-Head Attention的主体结构如图4所示。

图4. Multi-Head Attention的主体结构

在参考NVIDIA Apex Library实现的基础上[5],飞桨实现了支持变长的融合 Multi-Head Attention算子,即unpad FMHA(Fused Multi-Head Attention)算子,我们将Encoder模块使用unpad计算方法进行实现。相比padding的方法,unpad的实现传入的输入有一些变化。

首先,如图5所示,输入张量的维度由原来的[batch_size, max_seq_len]被压缩(不同样本的有效 token 被线性平铺)存储为[ntokens],其中batch_size为批大小,max_seq_len为padding的最大序列长度, ntokens代表当前mini-batch有效token总个数。

其次,要求传入能够反映变长序列实际长度的辅助信息,如batch_offset和max_seqlen_cur_bs等,其中batch_offset代表当前mini-batch中样本有效序列长度的前缀和信息,max_seqlen_cur_bs代表当前mini-batch样本的最大序列长度。这些信息会传入到unpad FMHA算子进行辅助计算。

图 5. unpad方法输入张量的存储变化

对BERT Encoder模块使用unpad优化后,端到端获得了大约2.3倍的性能提升。


变长Attention计算优化

在变长输入下,很难实现出一个适合所有序列长度分布的、统一的最优CUDA Kernel。举例来说,我们编写了一个在max_seq_len_cur_bs=512的情况下性能最好的CUDA Kernel ,但很难保证在max_seq_len_cur_bs=128的情况下性能也是最优的。不同序列长度分布有对应的不同的最优CUDA Kernel实现,Kernel 的分块大小、shared memory和寄存器使用情况也会有所不同。
在实际训练过程中,unpad FMHA算子会根据当前mini-batch样本的max_seqlen_cur_bs的值来选择最合适的CUDA Kernel执行,其选取Kernel 的原则如下所示,其中FMHA_128 、FMHA_256、FMHA_384和FMHA_512为根据不同max_seq_len_cur_bs的值所实现的最优CUDA Kernel。
if (0 < max_seqlen_cur_bs <= 128)
    FMHA_128(...);
else if (128 < max_seqlen_cur_bs <= 256)
    FMHA_256(...);
else if (256 < max_seqlen_cur_bs <= 384)
    FMHA_384(...);
else if (384 < max_seqlen_cur_bs <= 512)
    FMHA_512(...);

由上述选取CUDA Kernel的原则可以发现,根据max_seqlen_cur_bs选取的CUDA Kernel,往往偏向于数值更大的CUDA Kernel。例如,如果一个mini-batch中只有少量序列的长度为512,其他序列的长度都小于128,那么上述原则选取的CUDA Kernel则是FMHA_512,而不是FMHA_128。我们在MLPerf BERT模型实际测试中发现,训练过程中几乎全部选择的都是FMHA_512的 CUDA Kernel,这使得序列长度较小的样本的计算无法达到更好的性能。

为了解决该问题,我们对mini-batch内样本进行分组,让不同小组的样本选择最合适的FMHA Kernel,从而提升整体性能。首先,对mini-batch内样本按照序列长度进行分组,例如根据上述CUDA Kernel选取原则分为 (0, 128], (128, 256], (256, 384], (384, 512] 这4组。然后根据分组情况,每个小组launch相应的 FMHA Kernel。每个小组在调用FMHA Kernel时,需要根据组内的样本个数和token信息重新设置batch_size大小和相关输入的地址偏移。如图6所示,同一个mini-batch的7个样本,根据样本的长度划分成了四个不同的小组(具有相同颜色标识的矩形块属于同一个小组),并分别launch了对应的CUDA Kernel。对于256这个分组(图中黄色的序列),其有2个样本,会调用FMHA_256 Kernel,调用时需设置batch_size为2,max_seqlen_cur_bs为256。

图6. 分组示意图

默认情况下,不同分组的FMHA Kernel在同一个stream中执行。同一个stream中的Kernel会按照Kernel launch的顺序依次来执行。只有当前一个launch的Kernel执行完成,后面launch的Kernel才开始调度。事实上,由于不同分组的FMHA Kernel是独立且没有数据依赖的,可以使用多stream技术进一步提升性能。多stream情况下,调度器会根据当前GPU资源使用情况,同时调度来自不同stream的FMHA Kernel,从而提高GPU的资源利用率,提升Multi-Head Attention模块的整体性能,如图7所示。

图7. FMHA kernels overlap示意图

我们需要使用CUDA event控制不同stream 上 Kernel的执行顺序,从而保证计算结果的正确性。图8展示了Kernel的依赖关系。我们将Multi-Head Attention模块之前的最后一个算子和之后的第一个算子对应的CUDA Kernel分别称为A Kernel和C Kernel。为了保证计算结果的正确性,需要保证FMHA模块的所有 Kernel 的执行必须在A Kernel完成之后才开始执行,且C Kernel必须等 FMHA 中所有Kernel执行完成后再开始执行。如图8所示,使用多stream优化后,Multi-Head Attention 模块的多个FMHA Kernel的执行获得了很好的overlap,GPU计算资源可以得到更充分的利用。

图8. FMHA Kernel间依赖关系

实验结果显示,对序列数据分组和多stream优化给Multi-Head Attention模块带来了大约20%的性能提升,端到端大约提升了3.5%~3.7%。


支持变长Embedding计算

BERT模型的最开始部分是Embedding模块。该模块由三个Embedding层、两个Add层和一个Dropout和LayerNorm层组成。通过分析发现,如果对batch_size和max_seq_len维度合并,对该模块所有算子的计算逻辑没有影响。因此,我们对该模块也进行了unpad的计算优化。实验结果显示,BERT Embedding模块使用unpad优化后,端到端获得了0.4%~0.9%左右性能提升。

数据读取

和模型训练的负载均衡




Exchange padding优化

为了实现负载均衡、改善模型收敛性,NVIDIA提出了一种名为Exchange Padding的多卡输入数据编排方法[6],如图9所示。具体的步骤为(假设输入数据已经padding至最大序列长度):
  • 对所有卡上的padding后的输入数据进行AllGather操作,使得每个卡均能获得所有卡上的全量输入数据。
  • 每个卡根据实际序列长度( padding 前的实际序列长度)分别对AllGather后的数据进行排序。
  • 根据每个卡的rank id进行Interleaving Slice操作,得到每个卡上最终的输入数据,即rank id为i的GPU进程获取到的样本编号为i, i + num_devices, i + 2 * num_devices, ...,其中 num_devices为GPU总数。例如,在图9中,Worker 0分到了编号为0和2的样本数据,Worker 1则分到了编号为1和3的样本数据。

图9. Exchange Padding的流程

在NGC PyTorch的实现中,Exchange Padding的三个步骤均是在GPU上进行的,而且Exchange Padding和模型训练是串行进行的,每个mini-batch训练前必须先等待Exchange Padding操作完成。然而,在MLPerf极限性能优化的情况下,GPU计算负载是非常重的,利用率往往可达98%以上,在GPU上进行Exchange Padding操作必然会消耗一定的GPU计算资源,影响模型训练性能;此外,Exchange Padding和模型训练过程串行进行也会降低模型整体训练效率。
基于此,我们提出了一种在CPU端做Exchange Padding的方法,具体的Timeline如图10所示,整体端到端训练时间提升了1.5%。具体方式为:
  • Exchange Padding操作在CPU端完成,减少GPU端的计算负载。
  • 在模型训练过程中,预先对下一个mini-batch的输入数据进行Exchange Padding操作,使得CPU端的Exchange Padding操作和GPU端的模型训练操作完全overlap起来,消除现有方案每个mini-batch开始前必须串行等待Exchange Padding完成所耗费的时间。

图10. CPU Exchange Padding实现方案的Timeline


Padding移除过程与模型训练的Overlap

在上述Exchange Padding的操作中,AllGather要求输入数据必须是带padding信息的。因此在输入变长Encoder模块前,必须进行padding移除的操作来去除填充的无用token。在这个padding移除的过程中会产出许多耗时较短的算子,这些算子仅与输入数据相关,与模型参数无关。在NVIDIA的BERT实现中,这些小算子均是在GPU中计算的,且与训练过程串行进行,影响模型训练速度。借助我们提出的CPU端的Exchange Padding操作,我们将这些小算子的计算从GPU端全部挪到CPU端进行,所有计算均与Exchange Padding过程一同完成,使得这些小算子的计算完全与模型训练过程overlap起来,测试结果表明,此项优化策略带来了1.3%的端到端性能提升。

高性能算子

和融合优化技术



在深度学习训练中,高性能算子开发与优化是非常重要的一环。对于单个算子而言,飞桨框架通过提高算子并行度、使用更优的GPU指令、减少GPU仿存开销等方式对进行优化,提升算子性能。对于多个算子而言,飞桨通过算子融合的方式,减少仿存开销和CUDA Kernel Launch延时 ,以发挥GPU的极限性能。


Embedding算子优化

对于BERT模型来说,混合精度是提升训练速度的有效方法,除了sum和mean之外的几乎所有算子都采用FP16进行计算。然而FP16类型的Embedding反向性能存在一定的问题。Embedding反向计算时,不同线程可能会写入同一个结果,因此会存在写冲突问题。为此,飞桨采用原子指令来解决写冲突,当输入为FP16类型时,底层会调用CUDA的 atomicAdd(half* ) 指令进行计算。然而,由于硬件对该指令的支持不够好,使得Embedding算子的耗时比较异常,达到了模型总耗时的9.1%左右,严重影响了模型的吞吐。
针对此问题,飞桨使用了CUDA的atomicAdd(half2* ) 指令替换原先的 atomicAdd(half* ) ,对Embedding算子的反向计算进行了优化。相比于 atomicAdd(half* ) ,atomicAdd(half2* ) 能一次完成两个FP16数据的原子加法,而且在GPU底层有相关的硬件加速支持,性能更优。优化后,Embedding 算子总耗时占比从原来的9.1%降到了0.5%,模型端到端获得了10%左右的性能提升。


支持 cuBLASLt GEMM fusions

cuBLASLt提供了支持GEMM和其他小算子融合的API [7]。我们充分利用 cuBLASLt的融合API对BERT模型中广泛存在的Linear算子进行融合优化。如下所示,前向和反向过程中,分别对GEMM和bias (广播加法操作),以及GEMM和 dbias(规约操作) 进行融合。
  • 前向:GEMM + bias融合

  • 反向:GEMM + dbias融合
实验结果显示,使用该优化,单模块大约提升5%~20%,模型端到端大约获得5.5%性能提升。


支持inplace addto优化

BERT模型常见Residual Add等分支结构,同一个Tensor在前向阶段可能同时作为Linear和另一个算子的输入,从而导致反向阶段会出现一次梯度累加操作,即在反向阶段可能会出现如下算子列表:
dX_1 = ... # X的一部分梯度源自于另一个OP
dX_2 = linear_grad(dOut, Weight)
dX = dX_1 + dX_2
反向阶段需要计算linear_grad和add两个OP。根据矩阵乘法反向求导的数学公式,Linear的反向也是矩阵乘法操作,因此上述计算公式等价于:
dX_1 = ...
dX.ShareDataWith(dX_1) # dX与dX_1共享显存空间
dX = dX + linear_grad(dOut, Weight) = dX + matmul(dOut, transpose(Weight))

其中,dX=dX+matmul(dOut, transpose(Weight))可以利用cuBlasLt GEMM API设置beta参数等于1来实现,从而将原先的linear_grad和add两个算子融合为一个新的linear_grad算子(cublasLt GEMM API beta=1),减少框架调度开销和Kernel Launch Latency。实际测试结果表明,此优化带来了端到端1.1%的训练加速效果。



算子dropout

与算子residual add+layer_norm融合

在Encoder模块中,有两次dropout+residual add+layer_norm的调用。由于dropout算子和residual add算子都属于elementwise类型的算子,且 layer_norm属于批规约类型的算子,因此,这三个算子可以非常自然的按照 layer_norm的并行逻辑进行融合。相比融合前,前向和反向Kernel launch次数由原来的6次降为2次,GPU显存访问量大约由原来的12次降为4次。实验结果显示,使用该优化,单模块大约提升5%~30%,模型端到端大约获得2.3%性能提升。

高加速比的

混合并行训练策略



在MLPerf BERT模型分布式训练过程中,LAMB优化器需要进行L2-Norm等复杂计算,计算量往往较大。借鉴NVIDIA Apex libraryDistributedFusedLAMB的实现,我们在飞桨上实现了优化版本的分布式优化器,其底层原理是通过ZeRO Stage 2[2]的方式,将每个卡上需要更新的参数量减少至原来的1/num_devices (其中,num_devices 为GPU总数),以减少计算开销。

在此基础上,我们针对优化器内部的重点计算Kernel进行了深度优化。DistributedFusedLAMB 内部会将FP16参数梯度、FP16参数、FP32 Master Weight的显存空间融合为连续空间,以减少通信开销。DistributedFusedLAMB需要计算以下几个重要的CUDA Kernel:
  • FP16参数梯度的平方和,用于做global norm clipping。
  • 每个参数的L2-Norm。
  • 每个参数对应的Trust Ratio Tensor的L2-Norm。
在NVIDIA Apex library的DistributedFusedLAMB的实现中,以上3 个CUDA Kernel都是通过MultiTensorApply的方式来实现的,其原理是将每个Tensor按给定的chunk_size切分为一系列的chunk,并每隔一定数量的Tensor或chunk launch一次CUDA Kernel,每个CUDA Kernel的同一个block对同一个chunk内的数据进行处理。在NVIDIA Apex library的实现中,launch Kernel的时候需要记录每个chunk的起始显存指针,每个CUDA Kernel的入参必须包含N 个chunk 的起始显存指针( N 为 chunk 的数量),因此限制了每个CUDA Kernel处理的Tensor或chunk数量( CUDA Kernel 入参总大小一般不能超过4KB),即CUDA Kernel的gridDim大小,并行度较低。
然而,在 DistributedFusedLAMB的实现中,FP16参数梯度、FP16参数、FP32 Master Weight的显存空间是连续的。利用显存连续的特点,我们对CUDA Kernel的入参进行优化,仅传入每个chunk所属Tensor的指针首地址和每个Tensor的长度即可。由于Tensor个数往往远小于chunk的数量,因此此方案下CUDA Kernel的gridDim大于原有方案,一方面可以提高CUDA Kernel的计算并行度,另一方面可以减少CUDA Kernel launch的个数。特别地,在计算FP16参数梯度的平方和时,由于参数梯度空间是连续的,我们可以直接调用cub::DeviceReduce来直接计算,而不需要通过MultiTensorApply的方式来完成。
利用参数、参数梯度的显存连续的特点,我们对DistributedFusedLamb中的主要CUDA Kernel进行了优化。实际测试结果表明,在MLPerf BERT模型训练中,飞桨的DistributedFusedLAMB性能较NVIDIA Apex library提升22% ,端到端带来了18.3%的性能提升,8卡加速比达到8.1至8.2 ,实现了超线性加速。

全流程

异步执行调度



在以上讲述数据读取和模型训练负载均衡的时候,我们已经提到了飞桨的数据读取和模型训练异步执行的方案。除此以外,飞桨也默认实现了异步通信的功能。在这些基础上,我们还对模型训练和预测过程中存在的CPU与GPU间的内存/显存拷贝问题进行了优化,这些拷贝操作会阻塞GPU的异步执行,影响端到端的训练效果。

例如,学习率的值是在CPU上计算得到的,却需要在GPU上进行使用( Host to device memory copy )。为了减少CPU、GPU之间的交互,我们将学习率计算挪到GPU上。这可以借助飞桨自定义外部算子的能力,新编写一个计算学习率的算子进行实现。

又如,在每个mini-batch训练完成后,我们常常会打印出当前mini-batch的loss和accuracy等指标,但这会带来额外的Device to host memory copy。在满足MLPerf规则的前提下,我们将原先每个mini-batch均打印一次loss和accuracy的方式改为了每隔335步进行,大幅度地减少了CPU和GPU间的同步开销。

再如,在预测阶段,MLPerf需要统计全量数据集下的Masked Language Model Accuracy指标。常规的实现是,获取每个mini-batch下的Accuracy指标,然后在CPU端进行累加统计,但这会带来每个mini-batch预测均有Device to host memory copy的同步开销。我们通过新增GPU 上的 Accuracy累加统计算子,当且仅当最后一个mini-batch预测时才取出GPU上累计的Accuracy指标的值,减少同步开销。

测试结果显示,上述训练阶段的优化手段在端到端中带来了约0.4%的性能提升,预测阶段的优化手段在端到端中大约提升了0.17% 。


总结



飞桨在MLPerf Training v2.0中获得了BERT模型训练性能世界第一的瞩目成绩。这不仅得益于飞桨框架在性能优化领域的长期耕耘,更离不开硬件生态的助力。近年来,飞桨的技术实力深受广大硬件厂商认可,合作日趋紧密,软硬一体协同发展,生态共创硕果累累。前不久(5月26日),NVIDIA与飞桨合作推出的NGC-Paddle正式上线。同时在本次MLPerf中,Graphcore也通过使用飞桨框架取得了优异成绩。未来,飞桨将继续打造性能优势,在软硬协同性能优化和大规模分布式训练方面持续技术创新,为广大用户提供更加便捷、易用、性能优异的深度学习框架。

本文封面图由AI作画神器文心·一格创作

(传送门:yige.baidu.com)

参考文献

[1]The MLPerf™ Training v2.0 results. 

https://github.com/mlcommons/training_results_v2.0

[2]DeepSpeed library.

https://github.com/microsoft/DeepSpeed

[3]DeepLearningExamples library. 

https://github.com/NVIDIA/DeepLearningExamples

[4]PaddleNLP library.

https://github.com/PaddlePaddle/PaddleNLP

[5]Apex unpad fmha. https://github.com/NVIDIA/apex

[6]The MLPerf™ Training v1.1 results. 

https://github.com/mlcommons/training_results_v1.1

[7]cuBLASLt GEMM fusions. 

https://docs.nvidia.com/cuda/cublas/index.html#cublasLtEpilogue_t


关注【飞桨PaddlePaddle】公众号

获取更多技术内容~


您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存