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实现流水并行

在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):  # 按顺序依次处理num_micro_batch个Mini-Batch大小的数据，进行流水并行。每轮迭代每个副本训练num_micro_batch*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数据并行
• Whale数据并行
• 结合Whale模型并行和流水并行

对比结果如下图所示。

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

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

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