DNN之LNN：训练大型神经网络的核心技术(数据并行+管道并行+张量并行+专家混合MoE+内存优化策略【CheckPoint梯度检查点/AMP混合精度训练/Offloading数据卸载/优化器内存优化

DNN之LNN：训练大型神经网络的核心技术(数据并行+管道并行+张量并行+专家混合MoE+内存优化策略【CheckPoint梯度检查点/AMP混合精度训练/Offloading数据卸载/优化器内存优化/压缩技术)之详细攻略

导读：大型神经网络的训练是一项复杂的任务，需要利用多个GPU进行并行计算。数据并行、管道并行、张量并行和专家混合等并行化技术将训练过程切片处理，提高了训练效率。此外，还有其他内存节省设计，如梯度检查点、混合精度训练、数据卸载、内存高效优化器和压缩等，进一步优化了大型神经网络的训练过程。
这篇文章介绍了一些训练大型神经网络的核心技术：

>>数据并行(Data Parallelism)：将不同部分的数据集分配给不同GPU进行训练，需要集中计算平均梯度来同步参数更新。

>>管线并行(Pipeline Parallelism)：将神经网络分解为多个模块，不同模块分布在不同GPU上训练，需要处理模块间运行瓶颈问题。

>>张量并行(Tensor Parallelism)：将单层操作如矩阵乘法水平拆分到多个GPU上计算，然后汇总结果。

>>专家混合(Mixture of Experts)：输入样本只通过网络中的一小部分专家层进行处理，增加模型规模而不增加计算量。

>>混合策略：如PTD-P结合数据并行、管线并行和张量并行。Megatron-LM和Switch Transformer采用张量并行和专家混合策略。

>>内存优化策略：如CheckPoint重计算activations以节省内存、混合精度训练、参数与梯度的在设备与CPU之间交换等。

>>优化器内存优化：如Adafactor降低运行状态内存开销。

>>压缩技术内存优化：如中间结果压缩传输从而减少通信开销。

总之，这篇文章系统论述了目前主流的一些规模化训练模型的重要并行化和内存优化策略，为大规模神经网络训练提供了指导。

DNN之LNN：《Techniques for training large neural networks训练大型神经网络的技术》训练大型神经网络的核心技术(数据并行+管道并行+张量并行+专家混合MoE+内存优化策略【CheckPoint梯度检查点/AMP混合精度训练/Offloading数据卸载/优化器内存优化/压缩技术)之详细攻略

目录

《Techniques for training large neural networks》翻译与解读

No parallelism没有并行性：训练NN是一个迭代过程，基于并行技术的四个维度(数据/管道/张量/专家混合)

这是一个三层模型上各种并行策略的说明。每种颜色指代一层，虚线分隔不同的GPU。

Various Parallelism techniques各种并行技术

(1)、Data parallelism数据并行性【各个GPU独自处理】：拆分全量数据集(得到多个数据块)+将数据块分配到不同GPU上去执行→复制副本参数到多个GPU(优化策略【可临时转移到CPU内存】)+需协调确保每个GPU具有相似参数【每个GPU独立计算梯度→平均各个GPU的梯度{需传输大量参数故影响训练吞吐量}→独立计算相同的新参数】

(2)、Pipeline Parallelism管道并行性【“纵向”按层拆分】：大模型拆分为连续的层块【分配到不同的GPU】+每个GPU只持有部分参数(占内存更小)+但层间存在顺序依赖性(导致气泡时间块【读取数据+GPU间通信】)→降低气泡成本(每个GPU一次只处理数据元素的子集)+巧妙地将新计算与等待时间重叠(提出再次拆分为多个微批次的思想)+两种优化技巧(如GPipe【连续处理+最后同步聚合】和PipeDream【交替处理】)

这是一个朴素流水线并行设置的说明，其中模型按层垂直划分为4个分区。工作进程1托管网络第一层（最接近输入层）的模型参数，而工作进程4托管第四层（最接近输出层）。“F”、“B”和“U”分别代表前向、反向和更新操作。下标表示操作在哪个工作进程上运行。由于顺序依赖性，数据一次只能由一个工作进程处理，导致大量的空闲时间“气泡”。

比较GPipe和PipeDream流水线方案，每个批次使用4个微批次。微批次1-8对应于两个连续的数据批次。在图片中，“(数字)”表示在哪个微批次上执行操作，下标标记工作进程ID。请注意，PipeDream通过使用过时的参数执行某些计算来获得更高的效率。

(3)、Tensor Parallelism张量并行性【层内“横向”拆分】：Transformer的计算瓶颈(矩阵乘法)+将权重矩阵分片(到不同的GPU上)+各GPU计算部分乘积+最后通信组合结果，比如Megatron-LM

(4)、Mixture-of-Experts (MoE)专家混合：只使用网络一小部分权重+基于多组权重集(每组权重即一个专家并托管在不同的GPU上)+门控机制选择+，如GShard(基于MoE的600B参数)、Switch Transformer(基于MoE的数万亿参数)

混合专家层（MoE）的说明。由门控网络从n个专家中选择2个。（图像改编自：Shazeer等人，2017年）

Other memory saving designs其他节省内存的设计

(1)、CheckPoint(或激活重计算)技术(节省内存)：通过权衡计算成本和内存成本+采用重计算中间激活可节省大量内存(因计算成本相对较低)

(2)、MP混合精度训练技术(节省内存)：采用低精度(FP16)训练模型

(3)、Offloading技术(节省内存)：未使用的数据临时卸载到CPU+需要时再读取，比如ZeRO技术(参数/梯度/优化器状态)