本文针对BertLarge分布式并行训练所存在的问题,介绍Whale的并行化设计和方案。通过为模型并行和数据并行,搭配流水并行辅助并行策略,优化通信拓扑结构,以解决BertLarge分布式训练性能较差的问题。在Whale中,您可以通过模型划分、资源划分及映射三个步骤,实现大规模数据及模型的分布式训练。

背景信息

受益于深度神经网络的发展,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模型的分布式性能。详细方案如下:
  1. 模型并行

    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拆分的模型并行方式,使模型训练成为可能。
  2. 流水并行
    仅采用模型并行进行分布式训练,其中一张GPU卡计算时,其他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)。增大流水Micro-Batch的并行图
  3. 混合并行
    因为模型并行不可能将模型拆分为无限份,且切分数量过大会影响性能,所以仅采用模型并行和流水并行依然不能满足性能要求。针对业务场景产生的海量训练数据,需要提高大量计算资源并行处理,因此还需要采用数据并行解决大规模数据问题。在流水并行的基础上增加数据并行后的计算图如下所示。混合并行上图中以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分布式编程范式):模型划分、资源划分及映射。

  1. 模型划分。
    将模型各部分按照显存占用和算力需求划分为不同部分,即Scopes划分,详情请参见whale.scopes。上述的实现方案中将模型划分为以下三部分:
    • Stage 0:包括Embedding和Encoder Layer 1~8层
    • Stage 1:包括Encoder Layer 9~16层
    • Stage 2:包括Encoder Layer17~24层和Pooler层
    针对BertLarge模型的原始代码,在Whale中您可以通过如下步骤实现流水并行。
    1. 通过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)
    2. 在模型并行的代码基础上,通过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)
    3. 在流水并行的基础上,通过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)
  2. 资源划分。
    由于将模型划分为三个部分,因此需要将硬件资源进划分为三个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})
  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倍。