如何训练千亿/万亿大模型

acecase

AI模型的规模最近几年快速提升，尤其是在Transformer结构出现以后，不断有新的模型刷新记录，当前各大土豪厂商已经将模型规模提升到了千亿&万亿级别，比如OpenAI的GPT-3、Microsoft&NVIDIA的Megatron-Turing、Google的T5/Switch-XXL/Switch-C，以及国内的盘古、紫极太初、M6、悟道等等。模型规模的提升，带来了模型精度的提升，但因为模型规模过大，也对模型的训练带来的极大的挑战，包括内存需求大、网络通信量大、训练/推理性能低等问题，需要通过分布式部署、内存优化、自动部署等技术来解决这些问题。
分布式技术

数据并行：Data Parallel

在一个存在有N个计算节点的系统，每个节点上都有完成的计算程序（模型参数、模型结构）。将每一步要训练的数据（global-batch）拆分成N份（mini-batch），每个节点负责完成1/N数据的计算。这种计算方式就是数据并行。
数据并行是一种非常常见的并行模式。如果将一个GPU看作一个系统，一个CUDA Core看作一个计算节点，那么CUDA Kernel就是以数据并行的方式在进行计算。在分布式训练中，我们通常将一个加速器（比如GPU）看作是一个计算节点。
因为训练的目的是完成对模型参数的更新，因此在采用数据并行时，最后需要在每个节点完成梯度计算后，共同进行模型参数的更新，因此数据并行的通信也就发生在此位置。
数据并行的参数更新支持两种模式

• 同步更新：所有节点都计算完成后，将梯度进行平均，然后以统一的梯度对参数进行更新；然后再同步启动下一个迭代的训练
  

在采用Parameter-Server模式进行训练时，需要在参数服务器对每个Node的梯度进行累积，并在所有（或一定个数）的节点的梯度累积完成后，计算平均梯度，再使用平均梯度更新参数。
在采用All-Reduce模式进行训练时，所有的节点在准备后各自的梯度后，通过All-Reduce算法求得平均梯度，然后在每个节点上各自更新自己本地的参数。因为All-Reduce模式具有更好的扩展性和性能，当前在同步更新时主要采用All-Reduce模式。

• 异步更新：每个节点计算完成后，独立进行参数更新，并立即启动下一个迭代的训练
  

异步更新仅在Parameter-Server模式下存在。在采用Parameter-Server模式进行训练时，参数服务器在每一个Node的梯度时，直接使用该梯度进行更新参数。在一次训练迭代中参数会进行多次更新。每个节点获取到的上参数服务器接收到请求时的参数值。
同步更新的好处时所有节点使用的参数值相同，模型训练过程中更容易收敛；但缺点时引入了全局同步点，训练时间受限于系统中执行最慢的节点；同时同步更新要求所有节点都可以正常工作，一量出现节点故障，就需要从checkpoint重训。
异步更新的好处则刚好相反，每个节点的执行性能不受其它节点影响，且可通过N+M的模式容忍少量节点的故障；但有时会因为参数不一致引入收敛问题。
数据并行的通信仅发生成参数更新阶段，且通信时间通常可以和计算overlap，因此数据并行具有非常好的扩展性，可以获得非常好的加速比。
当前基本所有的AI框架，都具有原生的数据并行支持，比如通过TensorFlow的ParameterStrategy/MirroredStrategy/MultiWorkerMirroredStrategy，Pytorch的nn.torch.DataParallel/DistributedDataParallel，可以比较简单的进行数据并行的实现。另外也可以借助第三方库如horovod等非常容易的实现数据并行。
层内模型并行：Tensor Parallel

数据并行要求每个节点上都部署完整的模型，但对于千亿/万亿级别的模型，完整模型所需的内存将远超加速器节点的内存容量。比如一个完整的GPT-3模型需要3TB以上的内存，此时就必须对模型进行拆分。
层内模型并行是每一个算子的计算，切分到多个不同的Device上，从而降低每个Device的内存需求。下图是Megatron-LM在Transformer结构的网络中使用的一种层内模型并行方式，对MLP的两个GEMM计算的右矩阵A和B，分别按K轴和N轴进行切分，并在Dropout前插入All-Reduce算子。



Transformer模型并行示例

层内模型并行需要插入通信，并且因为后面的计算依赖于前面的输出，所以通信和计算必须串行执行，而不能像数据并行一样进行计算和通信的overlap。因此，层内模型并行通常要求采用高带宽的网络，比如NVIDIA的Megatron-LM在训练GPT-3大模型时，仅在部署了NVLink的Node内使用模型并行；即使如此，层内模型并行的通信时长仍然远远高于数据并行以及流水并行的时长。参考《DIVE DEEP INTO THE PERFORMANCE MODEL OF GPT-3 TRAINING ON MEGATRON LM: STORAGE, COMPUTATION,AND COMMUNICATION》。
流水并行：Pipeline Parallel

流水并行是一种符合人类直觉的并行模式。此并行模式是将模型拆分为多个stage，将每一个stage（或多个）部署在不同的Device上串接起来，前一个stage的输出作为后一个stage的输入的一种方式。
在训练场景下，每个stage的前向和后向需要部署在相同的Device上，以避免跨Device的weights传输。如果采用了参数同步更新的方式训练，一个Device在执行完正向的stage后，需要等待该stage的反向也执行完成，才能开始下一个迭代的训练。因此每个Device会存在等待时间，即“Bubble”（参考GPipe）。



Bubble Time

为了减小Bubble以提升效率，流水并行需要将一个mini-batch切分为多个micro-batch，每个Device在执行完一个micro-batch的计算后，可以立即执行下一个micro-batch的计算。采用了micro-batch后，Bubble的时长为(K - 1) / (M + K - 1)，K是stage数，M是micro-batch数。要减少Bubble就需要尽可能增加micro-batch数。比如Megatron-LM中的M = 192。



mini-batch拆分micro-batch

在切分了多个micro-batch后，每个Device上需要缓存多个micro-batch的activation，因此需要占用大量的内存。而完全的优先调度反向又会导致只有反向完成后才能调度前向，使得Device出现空闲，因此PipeDream采用了一种1F1B的流水方案，每个Device以交替的方式执行每一个micro-batch的正向和反向计算，如下图：



1F1B

1F1B使得每个stage的反向计算可以尽早得到调度执行，因此可以及时释放对应的内存，减少内存占用。但此方式存在两个问题：

• 同一个mini-batch的前向和后向weight不一致（如上图的Machine 1的mini-batch 5的正反向）
  

要解决此问题，需要为每个mini-batch引入weight副本，同一个mini-batch的前向和反向使用相同的weight副本。此方式解决了weight不一致的问题，但引入了多份副本，又增加了内存占用。

• 同一个mini-batch的不同stage的参数不一致（mini-batch 5的正向在Machine 1和2上）；
  

要解决此问题，需要引入Vertical Sync机制，每个mini-batch在进入Pipeline开始使用最新的weight，并且将信息传递给后续stage，所有stage使用相同的weight副本。
《Memory-Efficient Pipeline-Parallel DNN Training》则采用了PipeDream-2BW的方案来降低PipeDream方案的内存需求。PipeDream-2BW采用double-buffer的方式，在每一个Device上保存2个weights副本，当所有stage都完成了新版本的weights的更新后，新的micro-batch就可以采用新的weights进行正向和反向的计算。



PipeDream-2BW方案

在该文中还提出了一个PipeDream-Flush的方案，该方案和GPipe一样只保留一份weight，但又采用了1F1B的调度方式，通过这种方式以执行性能为代价降低了GPipe的峰值内存。



PipeDream-Flush

Megatron-LM采用了PipeDream-Flush的方案，并进一步增加了interleaved Pipeline优化来优化Bubble Time。此方案中，每一个Device上可以处理v（下图中v=2）个不相邻的stage。以下图为例，原来的4个stage被拆分为了8个stage，Device 1负责stage 0和4，在micro-batch 1经过Device 1/2/3/4处理完成后，又回到了Device 1进行stage 4的计算。在warm-up完成后采用1F1B的方式调度。使用此方案，要求micro-batch数必须是stage数的整数倍。



Megatron-LM的pipeline方案

interleaved pipeline，深/浅蓝表示正向的2个stage，深/浅绿表示反向的2个stage。interleaved Pipeline将Bubble Time降低了v倍，但代价时引入了v倍的通信量。
流水并行模式下，通信仅发生在stage之间，通信量相对较小。但因为存在Bubble，会导致设备利用率受限。并且stage的切分要同时兼顾内存、通信与计算的均衡。
专家并行：Expert Parallel

专家并行是Google在GShard中提出的一种并行部署技术，它在MoE结构的模型中被广泛应用，本质也是一种模型并行技术。在MoE结构的模型中引入了Gating和Expert，通过将每个Expert部署在不同的Device上，Expert的个数可以随着集群规模的扩展，可以提供非常好的扩展性。


但在MoE的前后需要通过All2All算子进行全局通信，且这部分通信必须和计算串行，因此MoE模型的训练受All2All性能的影响非常大。而All2All的性能随集群规模变大而降低，这极大的限制了MoE模型的扩展性。
内存优化技术

重计算

重计算是一种以时间换空间的技术，它的原理可以参考Rematerialization，陈天奇在Training Deep Nets with Sublinear Memory Cost中将其应用在了NN训练中。
NN训练反向计算中，需要依赖于正向计算过程中生成的中间结果，因此这部分中间结果需要保存在内存中。重计算的思路是在正向计算中不再保存这些中间结果，而是在反向计算时进行重新计算，从而减少NN训练中的内存需求。
重计算的关键是在尽可能小的影响计算性能的情况下，降低内存需求，因此在网络中哪些算子的输出结果需要重新计算就至关重要。应该优先选择输出内存大、计算耗时小的算子进行重计算，而对于计算耗时大、输出内存小的，则不做重计算。
参数Sharding

模型的参数，包括权重参数和优化器参数，占用了大量的内存。如果可以采用分布式存储的方式将参数存储在多个Device上，则可以降低每个Device的内存需求。Automatic Cross-Replica Sharding of Weight Update in Data-Parallel Training提出了一种对模型参数进行切分部署的方案，Zero-DP也采用了相同的方案。



weight&opt parameter sharding

采用Sharding后，每个Device上仅保存了1/N的参数，包括权重和优化器；权重和优化器可以分别做Sharding。
如果权重进行了切分，在正向计算使用权重进行计算之前，通过AllGather获取到完整的权重进行计算。因为每个Device上仅保存了1/N的权重，因此在更新参数前，只需要使用1/N的梯度进行本Device上的参数更新，因此需要对梯度进行ReduceScatter。
如果优化器参数进行了切分，也只需要使用1/N的梯度进行优化器参数的计算和更新。如果权重没有切分，要在优化器计算后增加AllGather用于更新完整的权重。
引入参数Sharding会大幅降低参数的内存需求，但因为要获取到其它Device上的参数进行计算，会影响训练性能。因此需要引入参数的预取，在参数使用之前提前将其获取到本Device上，通过参数预取可以掩盖一定的通信时间。预取的时机受Device内存容量、计算算子的耗时、网络带宽等因素影响。
Swapping

Swapping方案是将Device内存中的数据交换到容量更大的Host内存中，以释放掉Device内存。Swap的数据可以是权重参数、优化器参数，也可以是activations、gradients等Tensor。
TFLMS提出了一种Tensor的Swapping方案，在计算过程中将算子产生的输出数据swap-out到Host侧，释放Device侧的内存；在数据要被使用之前，再通过增加swap-in操作将数据从Host拉回到Device。TFLMS通过对swap操作的融合，以及增加控制边的方式触发swap-in的提前释放来优化Swapping的性能。



tensor swapping

Zero-Offload则将权重、优化器和梯度进行Swapping。此方案将权重和优化器保存在Host侧。对于权重参数，在使用时传输给Device；如果使用了混合精度，在传输前进行FP32->FP16的转换以降低带宽需求。Device在梯度计算完成后，将FP16的梯度传输到Host侧进行优化器计算和参数更新。



ZeRO-offload

Zero-Offload的一个关键问题是优化器计算放在了Host侧计算。优化器计算主要是Memory Bound类型的矩阵计算，因为Host侧DDR内存的带宽远低于Device上的HBM带宽，同时很多CPU没有SVE矩阵计算的加速能力，因此Zero-Offload方案的Host性能将成为模型训练的性能瓶颈。
结合应用

以上的各种分布式并行和内存优化技术，单一应用时都无法满足千亿/万亿级别的模型训练需求，需要在部署时同时采用多种不同的技术。
在选择这些技术时，需要同时考虑模型结构、计算资源（Host/Device的算力和内存容量）、网络拓扑和带宽（Host-Device、Device-Device、Node-Node）等因素。下面举几个常见的部署方式：
3D并行（DP + TP + PP） + 重计算

NVIDIA的Megatron-LM，以及Microsoft&NVIDIA在DeepSpeed&Megatron，通过3D并行的方式，实现了Transformer类模型的部署。在每一个数据并行的Device集合内，模型被拆分为多个stage，在每个stage内有若干个Transformer Layer，Tranformer Layer内的Attention和MLP通过模型并行方式进行训练。同时，对每一个（或多个）Transformer Layer开启重计算。



Megatron-LM GPT-3 3D并行

在一个1K GPU的集群上，Megatron-LM采用如下参数完成了GPT-3 175B模型的部署：


DP + TP + EP

Google的Switch-Transformer中，对MoE结构的模型，采用DP + TP + EP的方式进行部署。MoE结构的模型通常深度不深，因此在进行了专家并行的拆分后，模型不再需要采用Pipeline的方式进行部署，以避免引入Bubble Time。



DP + TP + EP

设计原则

• 在模型可以部署下的前提下，优先选择数据并行；
  
• 模型并行因为带宽需求高且无法流水，需要将其控制在高带宽的网络内；
  
• 流水并行要减小stage数，增加macro-batch数，以减小Bubble Time；
  
• 流水并行的Stage切分，要将切分点放在通信需求少的位置，并尽可能使各节点的计算均衡；
  
• 重计算会降低计算性能，优先选Memory Bound的算子做重计算，避免选Compute Bound算子做重计算；
  
• 如果Host-Device带宽较高，可以考虑采用Swapping方案代替重计算方案来提升性能；
  
• CPU的内存带宽如果比较低且不支持向量加速，要避免将计算放在CPU侧；
  

自动部署方案

要将一个千亿/万亿级别的模型部署在集群中，从前面的技术分析可以看出，要考虑非常多的因素。要想要实现一个最优的部署是非常困难的，最好能有一些辅助工具自动或半自动的帮助我们完成部署，当前学术界和业界已经有了一些方案；但因为搜索空间过大，目前仍然没有看到一个成熟的解决方案，能够综合以上所有的技术来完成模型的最优部署。大多的方案都是在一定限定条件下，或是仅考虑部分技术的部署方案。下面是一些已有的方案简略介绍：
NVIDIA Megatron-LM

针对Transformer类的模型提供半自动的分布式部署。用户只需要提供DP、TP、PP的并行度，以及重计算的参数，Megatron-LM可以自动完成网络的分布式部署。下面是在Megatron-LM上启动GPT-3训练的脚本参数示例：
options=" \
--tensor-model-parallel-size 8 \
--pipeline-model-parallel-size 16 \
...... \
--micro-batch-size 1 \
--global-batch-size 1536 \
--activations-checkpoint-method uniform "
MindSpore AutoParallel
MindSpore支持自动并行部署，支持数据并行、模型并行与混合并行。
使用方式参考：https://www.mindspore.cn/tutorial/en/0.3.0-alpha/advanced_use/distributed_training.html#
FlexFlow
提出了一种基于仿真 + CostModel的自动并行搜索方案，以尽可能的找到最优的模型切分策略。该方案可以实现模型并行的自动搜索，但无法实现带Pipeline的混合并行搜索。
GSPMD
提出了一种模型定义和分布式策略分离的分布式训练的表达和切分方式。算法工程师在定义网络时将整个分布式环境当成一个逻辑设备定义，然后再通过单独的device_mesh标识网络拓扑（通常使用2维，以表示server内的高速互联和设备间的互联网络），通过mesh_split在Tensor上增加annotation表示切分策略。GSPMD可以通过少量的用户annotation推导出整个网络的annotation，然后再根据annotation进行自动切分。
Alpa
Alpa支持将JAX表达的计算图，lower为XLA HLO IR，通过inter-op pass和intra-op pass完成分布式部署的编译优化，并实现了runtime实现流水并行的Stage调度，然后通过GSPMD完 成每个stage内的并行部署。
在通过JAX开发模型时，只需要指定并行策略，并在模型前通过@alph.parallelize的decorator装饰，就可以完成自动部署，示例如下：
method = alpa.PipeshardParallel(num_micro_batches=16,
layer_option=alpa.AutoLayerOption(layer_num=2),
stage_option="auto")
# Define the training step. The function body is the same as the above one.
@alpa.parallelize(method=method)
def auto_pipeline_train_step(state, batch):
# ......