大家好，我是Tim。
在进入AI 2.0时代后，大模型发展已是人工智能的核心，但训练大模型确实一项比较艰巨的工作和挑战，因为它需要大量的 GPU 内存和较长的训练时间。
然而，单个 GPU 工作线程的内存有限，并且许多大型模型的大小已经超出了单个 GPU 的范围。
目前已有多种并行范例可以实现跨多个 GPU 的模型训练，以及各种模型架构和内存节省设计，有助于训练非常大的神经网络。
今天我们就来简单总结下。
1. 数据并行性
数据并行性(DP) 就是指在不同的 GPU 上运行批次的不同数据子集。
它会将将相同的参数复制到多个 GPU（通常称为“workers”），并为每个 GPU 分配不同的example数据以同时处理。
仅数据并行性仍然需要模型大小能够适配单个 GPU 的内存。数据并行可以让您利用多个 GPU 的计算，但代价是存储参数的许多重复副本。
这也意味着，随着并行的GPU卡数的增加，模型训练所需的内存资源也会进行成倍增加。
话虽如此，有一些策略可以增加 GPU 可用的有效 RAM，例如在使用时临时将参数拷贝到 CPU 内存，即实现双层的存储策略GPU作为缓存，CPU进行全量参数存储。
当每个数据并行工作线程更新其参数副本时，它们需要进行协调以确保每个工作线程继续拥有相似的参数。
最简单的方法是在worker之间引入阻塞通信：
(1) 独立计算每个worker上的梯度；
(2) 平均各个Worker的梯度；
(3) 在每个Worker上独立计算相同的新参数。
步骤（2）就是引入阻塞通信的方法，在每个小批量结束时它需要同步传输大量数据，数据量与工作线程*参数大小成正比。
它可以防止模型权重过时和良好的学习效率，但每台机器都必须停止并等待其他机器发送梯度，所以这可能会损害模型的训练吞吐量。horovod就是一种深度学习训练中常用的分布式框架。
其次，还有有各种 异步同步方案可以消除这种开销，每个 GPU 工作线程异步处理数据，无需等待或停止。
然而，它很容易导致使用过时的权重，从而降低统计学习效率。尽管它增加了计算时间，但可能不会加快收敛的训练时间。
所以在实践中，我们通常还坚持使用同步方法。
2. 模型并行性
模型并行（MP）旨在解决模型权重无法适应单个节点的情况。
模型并行的计算和模型参数分布在多台机器上，它与每个worker处理整个模型的完整副本的数据并行不同，模型并行仅在一个worker上分配一小部分模型参数，因此内存使用量和计算量都减少了。
由于深度神经网络通常包含一堆垂直层，因此逐层分割大型模型感觉很简单，其中一小部分连续层被分组为一个Worker的一个分区。
然而，通过具有顺序依赖关系的多个此类工作线程来运行每个数据批次的简单实现会导致等待时间的巨大泡沫和计算资源的严重利用不足。
3. 管道并行性
管道并行（PP）将模型并行与数据并行相结合，以减少低效的时间“泡沫”。
核心思想是将一个批次拆分为多个微批次，每个微批次的处理速度应该成比例地更快，并且每个Worker一旦可用就开始处理下一个微批次，从而加快管道执行速度。
如果有足够的微批次，则可以充分利用Worker(GPU卡)，并在步骤开始和结束时将气泡降至最低。其中，梯度在微批次之间取平均值，并且仅在所有微批次完成后才会更新参数。
模型的分区（worker）数量也称为管道深度。
在前向传递过程中，Worker只需将其层块的输出（称为激活）发送给下一个工作人员；在向后传递期间，它仅将这些激活的梯度发送给前一个Worker。
如何安排这些传递以及如何跨微批次聚合梯度有很大的设计空间，其中比较出名的就是GPipe与PipeDream。
GPipe 让每个Worker连续向前和向后传递，然后在最后同步聚合来自多个微批次的梯度，如上图所示。
这种方式，无论Worker数量如何，同步梯度下降都能保证学习的一致性和效率，不过可以看出的是其气泡仍然存在。
GPipe 论文观察到，如果微批次数量超过分区数量的 4 倍，气泡开销几乎可以忽略不计。
GPipe 的吞吐量随着设备数量的增加而几乎实现线性加速。
相反， PipeDream 安排每个工作进程交替处理前向和后向传递。
PipeDream 将每个模型分区命名为“阶段”，每个阶段工作线程可以有多个副本来运行数据并行性。
在此过程中，PipeDream 使用确定性循环负载平衡策略在阶段的多个副本之间分配工作，以确保同一小批量的前向和后向传递发生在同一副本上。
由于 PipeDream 没有跨所有工作进程的批次结束全局梯度同步，因此一次前向一次反向的本机实现很容易导致使用不同版本的模型权重的一个微批次的前向和后向传递，从而降低学习效率。
4. 张量并行性
管道并行性按层“垂直”分割模型。我们还可以“水平”分割层内的某些操作，这通常称为 张量并行 训练。
对于许多现代模型（例如 Transformer），计算瓶颈是将激活批矩阵与大权重矩阵相乘。
矩阵乘法 可以被认为是行和列对之间的点积；可以在不同的 GPU 上计算独立的点积，或者在不同的 GPU 上计算每个点积的一部分并对结果求和。
无论采用哪种策略，我们都可以将权重矩阵分割成大小均匀的“分片”，将每个分片托管在不同的 GPU 上，并使用该分片计算整个矩阵乘积的相关部分，然后再进行通信以组合结果。
Megatron-LM就是一个例子 ，它在 Transformer 的自注意力层和 MLP 层中并行化矩阵乘法。
PTD-P 使用张量、数据和管道并行性；它的管道调度为每个设备分配多个非连续层，以更多网络通信为代价减少气泡开销。
有时，网络的输入可以跨维度并行化，并具有相对于交叉通信的高度并行计算。
序列并行 就是这样一种想法，其中输入序列在时间上被分割成多个子示例，通过允许计算以更细粒度的示例进行，按比例减少峰值内存消耗。
5. 混合专家 (MoE)
随着研究人员试图突破模型大小的限制，专家混合 (MoE)方法最近引起了广泛关注。该思想的核心是集成学习：多个弱学习器的组合给你一个强学习器！
使用 专家混合 (MoE) 方法，仅使用网络的一小部分来计算任何一个输入的输出。
一种示例方法是拥有多组权重，网络可以在推理时通过门控机制选择使用哪一组权重。
这可以在不增加计算成本的情况下启用更多参数。每组权重都被称为“专家”，希望网络能够学会为每个专家分配专门的计算和技能。
不同的专家可以托管在不同的 GPU 上，从而提供了一种清晰的方法来扩展模型所使用的 GPU 数量。
恰好一层 MoE 包含
作为专家的前馈网络{E_i}^n_{i=1}
可训练的门控网络G学习概率分布n”专家”，以便将流量路由到少数选定的”专家”。
根据门控输出，并非每个”专家”都必须进行评估。当”专家”数量过多时，我们可以考虑使用两级分层MoE。
GShard（Lepikhin 等人，2020）通过分片将 MoE 变压器模型扩展至 6000 亿个参数。MoE 变压器用 MoE 层替换所有其他前馈层。分片 MoE 变压器仅具有跨多台机器分片的 MoE 层，而其他层只是简单地复制。
Switch Transformer（Fedus 等人，2021 ）通过用稀疏开关 FFN层替换密集前馈层（其中每个输入仅路由到一个专家网络），将模型大小扩展到数万亿个参数，并且具有更高的稀疏性。
6. 其他节省内存的设计
Mixed Precision Training
混合精度AMP是指训练时在模型中同时使用 16 位和 32 位浮点类型，从而加快运行速度，减少内存使用的一种训练方法。
其加速的原理是在NVIDIA GPU 使用 float16 执行运算的速度比使用 float32 快一倍多，大大提升了算力的天花板。
但是简单的将模型的运算变成FP16并不能完全work，这是因为FP16所表示的数值范围，远小于FP32和TF32的数值范围，这使得模型的运算大大受限，因此需要额外的一些trick来保证模型能够收敛到跟FP32一样的结果。
避免以半精度丢失关键信息的三种技术：
1. 权重备份（Weight Backup）
训练时，权重、激活值和梯度都使用FP16进行计算，但是会额外保存TF32的权重值，在进行梯度更新时对TF32的权重进行更新。在下一步训练时将TF32的权重值转换为FP16再进行前向和反向的计算。
这里主要是因为使用FP16进行梯度更新的话，有一些梯度过小就会变成0，从而没有更新。还有就是权重值比梯度值大太多的时候，相减也会导致梯度消失。
2. 损失缩放（Loss Scaling）
一般我们在训练模型的时候，梯度的量级都非常小，由于使用了FP16，就会导致一些小的梯度直接变成了0。
如下面这张图展示了激活函数的梯度值的分布情况。可以看到非0梯度中有一大半都不在FP16的表示范围内。
我们注意到其实FP16的右半部分其实没有用到，我们可以把梯度乘以一个较大的值，从而让整个梯度分布向右移动，从而能够落在FP16的表示范围内。
一个非常简单的方法就是在梯度计算前对loss乘以一个很大的值，这样根据链式求导法则，计算到的梯度都会被放大，当我们真正更新梯度的时候需要再将梯度缩小回来原来的值，用TF32进行更新。
3. 精度累加（Precision Accumulated）
在FP16的模型中，一些算术运算比如矩阵乘法需要用TF32来累加乘积的结果，然后再转换为FP16，这样的效果会更好一些。Tensor cores已经提供这种支持。
例如在Nvidia GPU设备中已经带有Tensor Core，可以利用FP16混合精度来进行加速，还能保持精度。Tensor Core主要用于实现FP16的矩阵相乘，在利用FP16或者TF32进行累加和存储。在累加阶段能够使用TF32大幅减少混合精度训练的精度损失。
Gradient Accumulation
梯度累积是一种训练神经网络的数据Sample样本按Batch拆分为几个小Batch的方式，然后按顺序计算。
按顺序执行Mini-Batch，同时对梯度进行累积，累积的结果在最后一个Mini-Batch计算后求平均更新模型变量，它是一种用时间换存储的技术。
在进一步讨论梯度累积之前，我们来看看神经网络的计算过程。
深度学习模型由许多相互连接的神经网络单元所组成，在所有神经网络层中，样本数据会不断向前传播。
在通过所有层后，网络模型会输出样本的预测值，通过损失函数然后计算每个样本的损失值（误差）。
神经网络通过反向传播，去计算损失值相对于模型参数的梯度。最后这些梯度信息用于对网络模型中的参数进行更新。
而梯度累积就是，每次获取1个batch的数据，计算1次梯度（前向），此时梯度不清空，不断累积，累积一定次数后，根据累积的梯度更新网络参数，然后清空所有梯度信息，进行下一次循环。
CPU Offloading
将未使用的数据暂时下载到CPU或不同设备之间，并在需要时将其读回。
由于CPU的存储相比与GPU，其空间大且便宜，实现一种双层的存储，将大大扩展训练时的存储空间。
但是如果简单的实现将会大大减慢训练速度，复杂的实现需要实现预取数据，以便设备永远不需要等待。
这一想法的一种实现是 ZeRO ，它将参数、梯度和优化器状态分割到所有可用的硬件上，并根据需要具体化它们。
Activation Recomputation
Recompute是把在前向计算中将tensor释放掉，在反向传播时要用时，在重新计算该tensor，占据内存大，并且重计算量小的tensor适合用这种方式。
重计算有很多种方式：
speed centric在重新计算出一个tensor后，会保留这个tensor，以备后续使用；
memory centric在计算出一个tensor后还会把它释放掉，再次使用时再重新计算。
Cost aware再重新计算出一个tensor后有两种选择，当继续保存这个tensor可能导致大于内存峰值时，就释放掉，否则就保留。
可以将swap和recompute结合在一起使用，对于conv的tensor，采用swap的方式，而对于BN的tesnor，采用recompute的方式，这种就是对于特定的op采用不同的方式。
此外，还有一种结合方式是预先迭代几次，收集必要的内存和运行时间信息，然后判断哪些tesnor该swap，哪些该recompute
Compression
模型压缩是将大模型采用裁剪、权重共享等方式进行处理，以减少模型参数量。通常这种方式，比较少用，因为非常容易使精度降低。
模型压缩方法通常有
修剪：可采用对连接、kernel、channel进行裁剪的方式
权重共享
低秩分解
二值化权重，权重由32bits换成8 bits、16bits，混合精度训练
知识蒸馏：使用训练好的teacher模型，来指导student模型训练
Memory Efficient Optimizer
优化器在模型训练中会有一定程度的内存消耗。
以流行的 Adam 优化器为例，它内部需要保持动量和方差，两者都与梯度和模型参数保持相同的规模。突然之间，我们需要4 倍的模型权重内存。
已经提出了几种优化器来减少内存占用。例如，Adafactor ( Shazeer et al. 2018 ) 没有像 Adam 那样存储完整的动量和变化，而是仅跟踪移动平均值的每行和每列总和，然后根据这些总和估计二阶矩。SM3（Anil et al. 2019）描述了一种不同的自适应优化方法，也导致内存大幅减少。
ZeRO（零冗余优化器；Rajbhandari et al. 2019）根据对大型模型训练的两个主要内存消耗的观察，优化了用于训练大型模型的内存：
大部分由模型状态占据，包括优化器状态（例如 Adam 动量和方差）、梯度和参数。混合精度训练需要大量内存，因为除了 FP16 版本之外，优化器还需要保留 FP32 参数和其他优化器状态的副本。
剩余的被激活、临时缓冲区和不可用的碎片内存（本文中称为残留状态）消耗。
ZeRO 结合了两种方法：ZeRO-DP和ZeRO-R。ZeRO-DP 是一种增强的数据并行性，以避免模型状态上的简单冗余。
它通过动态通信调度跨多个数据并行进程划分优化器状态、梯度和参数，以最大限度地减少通信量。ZeRO-R 使用分区激活重新计算、恒定缓冲区大小和动态内存碎片整理来优化残留状态的内存消耗。
7. 总结
训练神经网络是一个迭代过程。
在每次迭代中，我们都会向前传递模型的各层，以计算一批数据中每个训练示例的输出。然后另一遍向后穿过各层，通过计算每个参数的梯度来传播每个参数对最终输出的影响程度 。
批次的平均梯度、参数和一些每个参数的优化状态被传递给优化算法，例如 Adam，它将计算下一次迭代的参数和新的每个参数优化状态。
随着训练对批量数据进行迭代，模型不断发展以产生越来越准确的输出。
然而随着大模型的到来，使得单机难以完成模型的训练，这时产生很多的并行技术。
各种并行技术将训练过程划分为不同的维度，包括：
数据并行性——在不同的 GPU 上运行批次的不同子集；
管道并行性——在不同的 GPU 上运行模型的不同层；
张量并行性——分解单个运算的数学运算，例如将矩阵乘法拆分到多个 GPU 上；
混合专家——仅通过每层的一小部分来处理每个示例。
除此以外，在机器资源和内存资源不足的情况下，也衍生了很多有趣的模型训练策略，包括：
混合精度训练——训练时在模型中同时使用 16 位和 32 位浮点类型，从而加快运行速度，减少内存使用的一种训练方法。
梯度累积——每次获取1个mini-batch的数据，计算1次梯度（前向），此时梯度不清空，不断累积，累积一定次数后，根据累积的梯度更新网络参数，然后清空所有梯度信息，进行下一次循环。
模型卸载CPU——将未使用的数据暂时下载到CPU或不同设备之间，并在需要时将其读回。
重算——把在前向计算中将tensor释放掉，在反向传播时要用时，在重新计算该tensor。
模型压缩——模型压缩是将大模型采用裁剪、权重共享等方式进行处理，以减少模型参数量。
内存优化版优化器——使用内存优化版本的优化器进行训练。