自学内容网 自学内容网
自学内容网 > 正文

图解大模型分布式训练：ZeRO系列方法

🕗 发布于 2024-12-12 16:06 分布式  

AI因你而升温，记得加个星标哦！

演进路线

  • 数据并行Data Parallelism：一台机器可以装下模型，所以将同一个模型同时部署在多台机器，用多份数据分开训练
  • 流水线并行Pipeline Parallelism：一台机器装不下模型，但模型的一层或多层一台设备装得下，所以同一个模型按层拆分到不同机器进行训练
  • 张量并行Tensor Parallelism：模型的一层都装不下了，所以同一个模型层内拆分开训练

不了解分布式训练的同学建议先阅读这几篇文章：

MapReduce：大模型分布式训练必备知识

图解大模型分布式训练：数据并行

图解大模型分布式训练：流水线并行

图解大模型分布式训练：张量并行Megatron-LM方法

目前大模型训练目前常见的解决方案是数据并行（Data Parallelisms, DP）和模型并行（Model Parallelisms, MP），但随着GPT-3把LLM的参数量推到了175B，训练所需参数大小更是到达了万亿规模，Megatron-LM面对如此大规模的参数也变得无能为力，而DeepSpeed的ZeRO方法问世则很好的解决了这个问题。

ZeRO

ZeRO 有两组优化：ZeRO-DP 旨在减少模型状态的内存占用，ZeRO-R 旨在减少剩余内存消耗。进行优化前到一些前置理解：

  • MP 的工作原理是垂直地切分模型，将网络和参数切分到多个设备，需要每层之间进行大量的通信。所以整体来看DP比MP具有更好的计算效率
  • DP 在内存方面效率低下，因为Model States在所有设备中都被冗余存储。相反，MP 分割模型状态以获得更高的内存效率
  • DP 和 MP 都保留了整个训练过程中需要的所有模型状态，但并非所有时间都需要所有内容。例如，每层对应的参数只在该层的前向传播和反向传播期间需要

ZeRO (Zero Redundancy Optimizer) 的开发便是用于解决 DP 和 MP 存在的问题，可以说是集百家之精华，并去其糟粕。ZeRO 消除了数据和模型并行训练中的内存冗余，同时保留了低通信量和高计算粒度，使我们能够根据设备数量按比例缩放模型大小，并保持持续的高效率。

通过实验发现，训练期间大部分的内存被模型状态（Model States），也就是：优化器参数（Optimizer states），梯度（Gradients）和模型（Model Parameters）所消耗。除此之外，残余状态（Residual States）消耗了剩余的内存，残余状态包括：前向传播时得到的 Activations，进行计算通信的临时缓冲区还有没有被妥善管理的内存碎片。各种数据的占比如下图所示：

可以看到模型参数仅占模型训练过程中所有数据的一部分，当进行混合精度运算时，其中模型状态参数（优化器状态+梯度+模型参数）占到了一大半以上。因此，我们需要想办法去除模型训练过程中的冗余数据。

针对模型状态的存储优化（去除冗余），DeepSpeed 提出了 ZeRO，ZeRO 使用的方法是分片，即每张卡只存 1/N 的模型状态量，这样系统内只维护一份模型状态参数。ZeRO对模型状态（Model States）参数进行不同程度的分割，主要有三个不同级别：

ZeRO-1 : 优化器状态分片（ Optimizer States Sharding）

设有 N d N_d Nd​台GPU，ZeRO-1将内存占用比例最高的 Optimizer States划分为 N d N_d Nd​份存储在不同的 GPU 上，并将该优化命名为$P_{os}
，如上图所示，当 ，如上图所示，当 ，如上图所示，当N_d=64 时，通过 时，通过 时，通过P_{os}$可以将内存开销降低了接近4倍。

ZeRO-1 非常适合使用类似Adam进行优化的模型训练，因为Adam拥有额外的参数m（momentum）与v（variance），特别是FP16混合精度训练。ZeRO-1 不适合使用SGD类似的优化器进行模型训练，因为SGD只有较少的参数内存，并且由于需要更新模型参数，导致额外的通讯成本。

ZeRO-2 : 优化器状态与梯度分片（Optimizer States & Gradients Sharding）

相比于ZeRO-1，ZeRO-2除了对optimizer state进行切分，还对Gradient进行了切分，该优化命名为$P_{os+g}
，如上图所示，通过 ，如上图所示，通过 ，如上图所示，通过P_{os+g}$可以将内存开销降低了接近8倍。

ZeRO-3 : 优化器状态、梯度和模型权重参数分片（Optimizer States & Gradients & Parameters Sharding）

为了进一步节省更多的内存，ZeRO-3提出进行模型参数的分片，让每个进程只存储与其分区对应的参数，当前向和反向传播需要其分区之外的参数时通过通信操作（broadcast）从对应的 DP 进程中接收它们。虽然乍一看这似乎会产生大量的通信开销，但论文表明这种方法仅将基础 DP 系统的总通信量增加到 1.5 倍而其重要的性质便是此时 N d N_d Nd​在分母处（如上图），那么就意味着只要GPU的数量够多，ZeRO 就能适应任意规模的模型。

以上优化模型状态的内存的方法称为ZeRO-DP。 在通信方面，使用$P_{os} 和 和 和P_{os+g} 不会产生额外的通信，使用 不会产生额外的通信，使用 不会产生额外的通信，使用P_{os+g+p} 时， Z e R O − D P 最多会产生 1.5 倍的通信量，但是可以进一步减少 时，ZeRO-DP 最多会产生 1.5 倍的通信量，但是可以进一步减少 时，ZeRO−DP最多会产生1.5倍的通信量，但是可以进一步减少N_d$倍的内存占用。

在ZeRO-R减少剩余内存消耗方面，主要也有三种优化方法：

P α P_\alpha Pα​：Partitioned Activation Checkpointing

ZeRO-R同样通过Activations分区来降低内存冗余，在在模型非常大和设备内存非常有限的情况下，这些分区的Activation Checkpoints也可以卸载到 CPU（Offload），以额外的通信成本将 Activation 的内存开销减少到接近零，这个优化称为 P α + c p u P_{\alpha+cpu} Pα+cpu​。

$C_B $: Constant Size Buffers

ZeRO-R 使用恒定大小的缓冲区来避免临时缓冲区随着模型大小的增加而爆炸，同时使缓冲区足够大以保持效率。

M D M_D MD​：Memory Defragmentation

ZeRO-R 通过为Activation Checkpoints 和 Gradient 预先分配连续的内存块，并在生成时将它们复制到预先分配的内存中，即时进行内存碎片整理，并将该优化命名为 M D M_D MD​。

ZeRO-Offload

设计等等并行的根因就是显存不足，ZeRO-Offload的想法很简单：显存不足，内存来补。相比于昂贵的显存，内存廉价多了，能不能在模型训练过程中结合内存呢？其实已经有很多工作了，但是他们几乎只聚焦在内存上面，没有用到CPU计算，更没有考虑多卡的场景。ZeRO-Offload则将训练阶段的某些模型状态下放（offload）到内存以及CPU计算。

ZeRO-Offload并不希望为了最小化显存占用而让系统的计算效率下降，否则的话，我们只用CPU和内存就可以了。但是将部分GPU的计算和存储下放到CPU和内存，必然需要满足以下两天内容的情况下，再考虑最小化显存的占用：

  • CPU和GPU之间的通信必然增加，所以不能让通信成为瓶颈
  • GPU的计算效率远远强于CPU，所以也不能让CPU参与过多计算，避免成为计算瓶颈

为了找到最优的offload策略，作者将模型训练过程看作数据流图（data-flow graph）。

  • 圆形节点表示模型状态，比如参数、梯度和优化器状态
  • 矩形节点表示计算操作，比如前向计算、后向计算和参数更新
  • 边表示数据流向

下图是用Adam优化器进行一次参数更新的流程（iteration/step），假设模型参数量是$M $，使用混合精度训练，边的权重要么是2M（fp16），要么是4M（fp32），前向计算（FWD）需要用到上一次的激活值（activation）和本层的参数（parameter），反向传播（BWD）也需要用到激活值和参数计算梯度：

我们现在要做的就是沿着边把数据流图切分为两部分，分别对应GPU和CPU，计算节点（矩形节点）落在哪个设备，哪个设备就执行计算，数据节点（圆形）落在哪个设备，哪个设备就负责存储，将被切分的边权重加起来，就是CPU和GPU的通信数据量。

ZeRO-Offload的切分思路是：图中有FWD、BWD、Param update和float2half四个计算类节点，前两个计算复杂度大致是 O ( M B ) O(MB) O(MB)， B B B是batch size，后两个计算复杂度是 O ( M ) O(M) O(M)，为了不降低计算效率，将前两个节点放在GPU，后两个节点不但计算量小还需要和Adam状态打交道，所以放在CPU上，Adam状态自然也放在内存中，为了简化数据图，将前两个节点融合成一个节点FWD-BWD Super Node，将后两个节点融合成一个节点Update Super Node。

以上是单卡情况，拓展到多卡，ZeRO-Offload利用了ZeRO-2（切片Adam状态和梯度到 1 N \frac{1}{N} N1​），而ZeRO-Offload做的同样是将这 1 N \frac{1}{N} N1​Adam状态和梯度都offload到内存，在CPU上进行参数更新。并且CPU和GPU的通信量和 N N N无关，因为传输的是fp16 gradient和fp16 parameter，总的传输量是固定的，由于利用多核并行计算，每个CPU进程只负责 1 N \frac{1}{N} N1​的计算，反而随着卡数增加节省了CPU计算时间。

但是有一个问题，当batch size很小时，GPU上每个micro-batch计算很快，此时CPU计算时长会成为训练瓶颈，一种方法是让CPU在某个节点更新参数时延迟一步，前N-1步，不进行延迟，避免早期训练不稳定，模型无法收敛。在第N步，CPU拿到GPU计算的梯度后，不更新参数（因为没计算完），但第N步GPU继续计算，也就是GPU空算了一步，到N+1步时，CPU才将第N步的梯度计算完成，此时GPU拿到第N步更新后的参数开始算N+1步的梯度。

实际应用

到这里为止，我们基本把分布式的计算架构介绍完了。在实际应用中，对Transformer类的模型，采用最经典方法是张量并行 + 数据并行，并在数据并行中引入ZeRO做显存优化。

其中，node表示一台机器，一般我们在同一台机器的GPU间做张量并行，在不同的机器上做数据并行。

在这里插入图片描述


原文地址：https://blog.csdn.net/Antai_ZHU/article/details/144409197

免责声明：本站文章内容转载自网络资源，如本站内容侵犯了原著者的合法权益，可联系本站删除。更多内容请关注自学内容网（zxcms.com）！

相关文章

自学内容网
Copyright © 2015-2024