跳转至

5.1 ZeRO原理解析

数据并行原理

学习目标

  • 理解数据并行DP的底层原理

数据并行

  • 数据并行(Dapa Parallelism)的整体架构如下图所示:

  • 参考上面的架构图, 经典数据并行的流程如下:
    • 1: 若干块计算GPU, 如图中GPU0 ~ GPU2; 1块梯度收集GPU, 如图中AllReduce操作所在GPU.
    • 2: 在每块计算GPU上都拷贝一份完整的模型参数.
    • 3: 把一份数据X(例如一个batch)均匀分给不同的计算GPU.
    • 4: 每块计算GPU做一轮FWD和BWD后, 算得一份梯度.
    • 5: 每块计算GPU将自己的梯度push给梯度收集GPU, 做聚合操作. (这里的聚合操作一般指梯度累加)
    • 6: 梯度收集GPU聚合完毕后, 计算GPU从它那pull下完整的梯度结果, 用于更新模型参数W. 更新完毕后, 计算GPU上的模型参数依然保持一致.
    • 聚合再下发梯度的操作, 称为AllReduce.
  • 本章的重点内容将介绍由微软开发的ZeRO(零冗余优化), 它是DeepSpeed这一分布式训练框架的核心, 被用来解决大模型训练中的显存开销问题, ZeRO的思想就是用通讯换显存.

存储消耗分析

  • 首先分析一下大模型训练的过程中, GPU都需要存什么内容?
    • 模型状态参数(Model States Parameter), 和模型本身息息相关, 必须存储的内容
      • optimizer states: Adam优化器中的momentum和variance
      • gradients: 模型梯度
      • parameters: 模型参数
    • 冗余状态参数(Residual States Parameter), 并非模型必须, 但在训练中额外产生的内容
      • activation: 激活值, 在backward中会用到, 存储下来计算梯度会更快.
      • temporary buffers: 临时存储, 例如将梯度发送到某块GPU做聚合时产生的存储.
      • unusable fragment memory: 碎片化存储空间. 取不到连续空间会被Fail掉, 可以整理内存来解决.
  • 混合精度训练: 对于模型, 我们肯定希望其参数越精准越好, 即用fp32(单精度浮点数, 存储占4Byte)来表示参数W. 但是在forward和backward的过程中, fp32的计算开销也是庞大的. 那么能否在计算的过程中, 引入fp16或bf16(半精度浮点数, 存储占2Byte), 来减轻计算压力呢? 混合精度训练就产生了, 它的步骤如下图所示:
    • 1: 存储一份fp32的parameter, momentum和variance(统称model states).
    • 2: 在forward开始之前, 额外开辟一块存储空间, 将fp32 parameter减半到fp16 parameter.
    • 3: 正常做forward和backward, 在此之间产生的activation和gradients, 都用fp16进行存储.
    • 4: 用fp16 gradients去更新fp32下的model states.
    • 5: 当模型收敛后, fp32的parameter就是最终的参数输出.

  • 有了上述原理的理解, 接下来可以计算模型在训练时需要的存储大小了. 假设模型参数W大小以Byte为单位:

  • 这里暂时不统计activation, 原因是activation不仅与模型参数相关, 还与batch size相关. activation的存储不是必须的, 存储activation只是为了在backward的时候计算梯度更快一些. 但可以通过只保留输入X, 重新做一遍forward来得到每一层的activation.

ZeRO-DP

  • 知道了什么东西会占存储, 以及它们占了多大的存储之后, 我们就可以来思考如何优化存储了. 在整个训练中, 有很多states并不会每时每刻都用到:
    • 1: Adam优化下的optimizer states只在最终做update时才用到.
    • 2: 数据并行中, gradients只在最后做AllReduce和updates时才用到.
    • 3: 参数W只在做forward和backward的那一刻才用到.
  • 所以ZeRO想了一个简单粗暴的办法: 如果数据算完即废, 等需要的时候, 我再想办法从某个地方拿回来, 那不就省了一笔存储空间吗? 沿着这个思路, 我们逐一来看ZeRO是如何递进做存储优化的:
    • P_os: 优化状态
    • P_os + P_g: 优化状态, 梯度
    • P_os + P_g + P_p: 优化状态, 梯度, 参数

P_os

  • P_os: 优化状态, 将optimizer state分成若干份, 每块GPU上各自维护一份, 这样就减少了相当一部分的显存开销, 如下图所示:

  • 此时W=fp16, G=fp16, O=fp32, 整体数据并行的流程如下:
    • 1: 每块GPU上存一份完整的参数W, 将一个batch的数据分成3份, 每块GPU各吃一份, 做完一轮foward和backward后, 各得到一份梯度.
    • 2: 对梯度做一次AllReduce, 就能得到完整的梯度G, 产生单卡通讯量2Q.
    • 3: 得到了完整的梯度G, 就可以对W做更新. 我们知道W的更新由optimizer states和梯度共同决定, 由于每块GPU上只保存一部分optimizer states, 因此只能如下图中所示, 将相应的W(蓝色部分)进行更新.
    • 4: 此时每块GPU上都有部分W没有完成更新(白色部分). 所以我们需要对W做一次All-Gather, 从别的GPU上把更新好的部分W取回来, 产生单卡通讯量Q.

  • 做完P_os后, 设GPU个数为Nd, 则显存和通信量情况如下图所示:

  • 结论: P_os在增加1.5倍单卡通信量的开销基础上, 将单卡存储降低了4倍, 看起来是个不错的trade-off.

P_os + P_g

  • P_os + P_g: 优化状态和梯度, 现在, 更近一步, 我们把梯度也拆开, 每个GPU各自维护一部分梯度, 如下图所示:

  • 数据并行的整体流程如下:
    • 1: 每块GPU上存一份完整的参数W, 将一个batch的数据分成3份, 每块GPU各吃一份, 做完一轮foward和backward后, 算得一份完整的梯度(下图中绿色+白色).
    • 2: 对梯度做一次Reduce-Scatter, 保证每个GPU上所维持的那块梯度是聚合梯度. 例如对GPU1, 它负责维护G1, 因此其他的GPU只需要把G1对应位置的梯度发给GPU1做加总就可. 汇总完毕后, 白色块对GPU无用, 可以从显存中移除. 单卡通讯量Q.
    • 3: 每块GPU用自己对应的O和G去更新相应的W, 更新完毕后, 每块GPU维持了一块更新完毕的W. 同理, 对W做一次All-Gather, 将别的GPU算好的W同步到自己这来. 单卡通讯量Q.

  • 再次比对下显存和通讯量:

  • 结论: 和朴素DP相比, 存储降了8倍; 相比于只优化P_os, 存储下降了2倍; 单卡通讯量持平, 效果更好了!

P_os + P_g + P_p

  • P_os + P_g + P_p: 优化状态, 梯度, 参数. 现在, 我们把参数也进行切分, 每块GPU置维持对应的optimizer states, gradients和parameters(即W), 如下图所示:

  • 数据并行的流程如下:
    • 1: 每块GPU上只保存部分参数W, 将一个batch的数据分成3份, 每块GPU各吃一份.
    • 2: 做forward时, 对W做一次All-Gather, 取回分布在别的GPU上的W, 得到一份完整的W, 单卡通讯量Q, forward做完, 立刻把不是自己维护的W抛弃.
    • 3: 做backward时, 对W做一次All-Gather, 取回完整的W, 单卡通讯量Q, backward做完, 立刻把不是自己维护的W抛弃.
    • 4: 做完backward, 算得一份完整的梯度G, 对G做一次Reduce-Scatter, 从别的GPU上聚合自己维护的那部分梯度, 单卡通讯量Q, 聚合操作结束后, 立刻把不是自己维护的G抛弃.
    • 5: 用自己维护的O和G, 更新W. 由于只维护部分W, 因此无需再对W做任何AllReduce操作.
  • 显存和通讯量如下图所示:

  • 结论: 我们用1.5倍的通讯开销, 换回60多倍的显存. 只要梯度计算和异步更新做的好, 通讯时间大部分可以被计算时间隐藏, 因此这样的额外通讯开销是很划算的!
  • 下图中展示了原始论文中的说明图. 经过以上分析, 这张说明图就很好懂了! 虽然ZeRO的设计不复杂, 但原始论文直接看的话很难懂.

  • 小领悟: 仔细想想, ZeRO其实掌握了降本增效(开猿节流)的精髓: 用完即弃, 需要再补! 反正我补一个和你差不多的, 也不会花费很多通(找)讯(人)时间, 还大大降低了我的成本. 模型的每一层多算(造)几(轮)遍(子)有啥关系呢, 反正在我的预算里每个人都一刻不停地干活, 就行啦!!!

ZeRO与模型并行

  • ZeRO是模型并行的形式, 数据并行的实质:
    • 模型并行, 是指在forward和backward的过程中, 我只需要用自己维护的那一部分W来计算就行. 即同样的输入X, 每块GPU上各自算模型的一部分, 最后通过某些方式聚合结果.
    • 但对ZeRO来说, 它做forward和backward的时候, 是需要把各GPU上维护的W聚合起来的, 即本质上还是用完整的W进行计算. 它是不同的输入X, 完整的参数W, 最终再做聚合.

ZeRO-R

  • 前面的ZeRO-DP是针对model states的显存优化, 现在的ZeRO-R是针对residual states的优化.

三个优化方向

  • Pa: Partitioned Activation Checkpointing, 对activation的存储是灵活的. 不像optimizer states, gradients和parameters对模型更新是必须的, activation只是起到加速梯度计算的作用. 因此, 在哪几层保存activation, 保存哪些activation都是可以灵活设置的. 同样, 我们也可以仿照以上切分方式, 每块GPU上只维护部分的activation, 需要时再从别的地方聚合过来就行. 需要注意的是, activation对显存的占用一般会远高于模型本身, 通讯量也是巨大的, 所以这块要灵活, 有效地实验设计.
  • C_B: Constant Size Buffer, 固定大小的内存buffer, 它的目的在于:
    • 提升带宽利用率. 当GPU数量上升, GPU间的通讯次数也上升, 每次的通讯量可能下降(但总通讯量不会变). 数据切片小了, 就不能很好利用带宽了. 所以这个buffer起到了积攒数据的作用: 等数据积攒到一定大小, 再进行通讯.
    • 使得存储大小可控. 在每次通讯前, 积攒的存储大小是常量, 是已知可控的, 更方便使用者对训练中的存储消耗和通讯时间进行预估.
  • M_D: Memory Defragmentation, 设置机制对碎片化的存储空间进行重新整合, 整理出连续的存储空间. 防止出现总存储足够, 但连续子存储不够而引起的存储请求fail.

ZeRO-Offload

  • ZeRO-Offload的核心思想是: 显存不够, 内存来凑! 如果我把要存储的大头卸载(offload)到CPU上, 而把计算部分放到GPU上, 这样比起跨机是不是能既降显存, 也能减少一些通讯压力呢?
  • ZeRO-Offload的具体做法如下图所示:
    • forward和backward计算量高, 因此和它们相关的部分, 例如参数W(fp16), activation, 就全放在GPU上.
    • update的部分计算量低, 因此和它相关的部分, 全部放入CPU中, 例如W(fp32), optimizer states(fp32)和gradients(fp16).