数据并行 - DP/DDP/ZeRO


Vincent
发布于 2024-08-03 / 25 阅读 / 0 评论 /
数据并行 - DP/DDP/ZeRO 数据并行DP DP - 梯度同步更新 数据并行的核心思想是:在各个GPU上都拷贝一份完整模型,各自吃一份数据,算一份梯度,最后对梯度进行累加来更新整体模型。理念不复杂,但到了大模型场景,巨大的存储和GPU间的通讯量,就是系统设计要考虑的重点了。在本文中,我们将递

数据并行 - DP/DDP/ZeRO

数据并行DP

DP - 梯度同步更新

数据并行的核心思想是:在各个GPU上都拷贝一份完整模型,各自吃一份数据,算一份梯度,最后对梯度进行累加来更新整体模型。理念不复杂,但到了大模型场景,巨大的存储和GPU间的通讯量,就是系统设计要考虑的重点了。在本文中,我们将递进介绍三种主流数据并行的实现方式:

  1. DP(Data Parallelism):最早的数据并行模式,一般采用参数服务器(Parameters Server)这一编程框架。实际中多用于单机多卡

  2. DDP(Distributed Data Parallelism):分布式数据并行,采用Ring AllReduce的通讯方式,实际中多用于多机场景

  3. ZeRO:零冗余优化器。由微软推出并应用于其DeepSpeed框架中。严格来讲ZeRO采用数据并行+张量并行的方式,旨在降低存储。

1)若干块计算GPU,如图中GPU0~GPU2;1块梯度收集GPU,如图中AllReduce操作所在GPU。

2)在每块计算GPU上都拷贝一份完整的模型参数。

3)把一份数据X(例如一个batch)均匀分给不同的计算GPU。

4)每块计算GPU做一轮FWD和BWD后,算得一份梯度G。

5)每块计算GPU将自己的梯度push给梯度收集GPU,做聚合操作。这里的聚合操作一般指梯度累加。当然也支持用户自定义。

6)梯度收集GPU聚合完毕后,计算GPU从它那pull下完整的梯度结果,用于更新模型参数W。更新完毕后,计算GPU上的模型参数依然保持一致。

7)聚合再下发梯度的操作,称为AllReduce。

总结一下:打散 – 收集 – 反向分发(更新)

实现DP的一种经典编程框架叫“参数服务器”,在这个框架里,计算GPU称为Worker,梯度聚合GPU称为Server。在实际应用中,为了尽量减少通讯量,一般可选择一个Worker同时作为Server。比如可把梯度全发到GPU0上做聚合。需要再额外说明几点:

  • 1个Worker或者Server下可以不止1块GPU。

  • Server可以只做梯度聚合,也可以梯度聚合+全量参数更新一起做

  • 在参数服务器的语言体系下,DP的过程又可以被描述下图

那么问题所在:

  • 存储开销大。每块GPU上都存了一份完整的模型,造成冗余。

  • 通讯开销大。Server需要和每一个Worker进行梯度传输。当Server和Worker不在一台机器上时,Server的带宽将会成为整个系统的计算效率瓶颈。

概括一下:每一个节点干完自己的活儿提交上去,等sever的反馈更新,这个等待的过程就是浪费时间,且sever的压力非常大。


DP - 梯度异步更新

所以,梯度异步更新的idea就出来了

在梯度异步更新的场景下,某个Worker的计算顺序为:

  • 在第10轮计算中,该Worker正常计算梯度,并向Server发送push&pull梯度请求。

  • 但是,该Worker并不会实际等到把聚合梯度拿回来,更新完参数W后再做计算。而是直接拿旧的W,吃新的数据,继续第11轮的计算。这样就保证在通讯的时间里,Worker也在马不停蹄做计算,提升计算通讯比。

  • 当然,异步也不能太过份。只计算梯度,不更新权重,那模型就无法收敛。图中刻画的是延迟为1的异步更新,也就是在开始第12轮对的计算时,必须保证W已经用第10、11轮的梯度做完2次更新了。

意思就是,work的参数阶段性更新,隔多久更新一次由延迟时间步决定

三种更新方式:

(a) 无延迟

(b) 延迟但不指定延迟步数。也即在迭代2时,用的可能是老权重,也可能是新权重,听天由命。

(c ) 延迟且指定延迟步数为1。例如做迭代3时,可以不拿回迭代2的梯度,但必须保证迭代0、1的梯度都已拿回且用于参数更新。

总结一下,异步很香,但对一个Worker来说,只是等于W不变,batch的数量增加了而已,在SGD下,会减慢模型的整体收敛速度。异步的整体思想是,比起让Worker闲着,倒不如让它多吃点数据,虽然反馈延迟了,但只要它在干活在学习就行。


分布式数据并行(DDP)

受通讯负载不均的影响,DP一般用于单机多卡场景。

因此,DDP作为一种更通用的解决方案出现了,既能多机,也能单机。

DDP首先要解决的就是通讯问题:将Server上的通讯压力均衡转到各个Worker上。实现这一点后,可以进一步去Server,留Worker。

聚合梯度 + 下发梯度这一轮操作,称为AllReduce。

接下来我们介绍目前最通用的AllReduce方法:Ring-AllReduce。它由百度最先提出,非常有效地解决了数据并行中通讯负载不均的问题,使得DDP得以实现。

太妙了,直接看图吧

通过分组累加到此,每一块GPU上都有了一个分块完整的梯度,即1/4的完整梯度。谓之,Reduce-Scatter

下面就是将每个分开的1/4完整的梯度同步到其他的GPU,也就是替换操作。谓之,ALLTOGETHER

再迭代两轮,就OK了

小结一下:

朴素数据并行(DP)与分布式数据并行(DDP)。两者的总通讯量虽然相同,但DP存在负载不均的情况,大部分的通讯压力集中在Server上,而Server的通讯量与GPU数量呈线性关系,导致DP一般适用于单机多卡场景。而DDP通过采用Ring-AllReduce这一NCCL操作,使得通讯量均衡分布到每块GPU上,且该通讯量为一固定常量,不受GPU个数影响,因此可实现跨机器的训练。

DDP做了通讯负载不均的优化,但还遗留了一个显存开销问题:数据并行中,每个GPU上都复制了一份完整模型,当模型变大时,很容易打爆GPU的显存


ZeRO

由微软开发的ZeRO(零冗余优化),它是DeepSpeed这一分布式训练框架的核心,被用来解决大模型训练中的显存开销问题。

ZeRO的思想就是用通讯换显存

存储主要分为两大块:Model States和Residual States

Model States指和模型本身息息相关的,必须存储的内容,具体包括:

optimizer states:Adam优化算法中的momentum和variance

gradients:模型梯度

parameters:模型参数W

Residual States指并非模型必须的,但在训练过程中会额外产生的内容,具体包括:

activation:激活值。在流水线并行中我们曾详细介绍过。在backward过程中使用链式法则计算梯度时会用到。有了它算梯度会更快,但它不是必须存储的,因为可以通过重新做Forward来算它。

temporary buffers: 临时存储。例如把梯度发送到某块GPU上做加总聚合时产生的存储。

unusable fragment memory:碎片化的存储空间。虽然总存储空间是够的,但是如果取不到连续的存储空间,相关的请求也会被fail掉。对这类空间浪费可以通过内存整理来解决。


混合精度运算

精度混合训练,对于模型,我们肯定希望其参数越精准越好,也即我们用fp32(单精度浮点数,存储占4byte)来表示参数W。但是在forward和backward的过程中,fp32的计算开销也是庞大的。那么能否在计算的过程中,引入fp16或bf16(半精度浮点数,存储占2byte),来减轻计算压力呢?

FP16是指使用16位浮点数表示的半精度浮点数,而BF16是指使用16位浮点数表示的bfloat16格式。它们之间的主要区别在于它们的存储格式和精度范围。

FP16使用1位符号位,5位指数位和10位尾数位,可以表示范围较广的数字,但精度较低。

BF16使用1位符号位,8位指数位和7位尾数位,可以提供更好的精度,并且能够保留更多有效数字,但范围较窄。

因此,FP16和BF16的主要差别在于它们的精度范围和存储格式,用户在使用时需要根据需求选择合适的格式。

于是,混合精度训练就产生了

  • 存储一份fp32的parameter,momentum和variance(统称model states)

  • 在forward开始之前,额外开辟一块存储空间,将fp32 parameter减半到fp16 parameter。

  • 正常做forward和backward,在此之间产生的activation和gradients,都用fp16进行存储。

  • 用fp16 gradients去更新fp32下的model states。

  • 当模型收敛后,fp32的parameter就是最终的参数输出。

也就是,模型参数存储时使用fp32,模型fw,bw计算时使用fp16

即设模型参数w为ϕ

因为采用了Adam优化,所以才会出现momentum和variance,adam好像很费内存

这里暂不将activation纳入统计范围,原因是:

  • activation不仅与模型参数相关,还与batch size相关

  • activation的存储不是必须的。存储activation只是为了在用链式法则做backward的过程中,计算梯度更快一些。但你永远可以通过只保留最初的输入X,重新做forward来得到每一层的activation(虽然实际中并不会这么极端)。

因为activation的这种灵活性,纳入它后不方便衡量系统性能随模型增大的真实变动情况。因此在这里不考虑它,在后面会单开一块说明对activation的优化。


知道了什么东西会占存储,以及它们占了多大的存储之后,我们就可以来谈如何优化存储了。

注意到,在整个训练中,有很多states并不会每时每刻都用到,举例来说;

  • Adam优化下的optimizer states只在最终做update时才用到

  • 数据并行中,gradients只在最后做AllReduce和updates时才用到

  • 参数W只在做forward和backward的那一刻才用到

所以,ZeRO想了一个简单粗暴的办法:如果数据算完即废,等需要的时候,我再想办法从个什么地方拿回来,那不就省了一笔存储空间吗?


Zero1 - 优化状态分割

优化参数在模型的W中,优化状态分割的意思是把W切开,每一个GPU单独自己更新属于自己的,然后在同步一下

显存下降的非常明显, 在增加1.5倍单卡通讯开销的基础上,将单卡存储降低了4倍。看起来是个还不错的trade-off

Zero2 - 优化梯度分割

接着切,同理可得,切梯度G

每块GPU用自己对应的O和G去更新相应的W。更新完毕后,每块GPU维持了一块更新完毕的W。同理,对W做一次All-Gather,将别的GPU算好的W同步到自己这来

Zero3 - 优化模型参数分割

全部都切开!!!

每块GPU置维持对应的optimizer states,gradients和parameters

最后数据并行的流程如下:

(1)每块GPU上只保存部分参数W。将一个batch的数据分成3份,每块GPU各吃一份。

(2)做forward时,对W做一次All-Gather,取回分布在别的GPU上的W,得到一份完整的W,单卡通讯量 。forward做完,立刻把不是自己维护的W抛弃。

(3)做backward时,对W做一次All-Gather,取回完整的W,单卡通讯量 。backward做完,立刻把不是自己维护的W抛弃。

(4)做完backward,算得一份完整的梯度G,对G做一次Reduce-Scatter,从别的GPU上聚合自己维护的那部分梯度,单卡通讯量 。聚合操作结束后,立刻把不是自己维护的G抛弃。

(5)用自己维护的O和G,更新W。由于只维护部分W,因此无需再对W做任何AllReduce操作。

ZeRO - Residual states

现在来看对residual states的优化

固定大小的内存buffer:

  • 提升带宽利用率。当GPU数量上升,GPU间的通讯次数也上升,每次的通讯量可能下降(但总通讯量不会变)。数据切片小了,就不能很好利用带宽了。所以这个buffer起到了积攒数据的作用:等数据积攒到一定大小,再进行通讯。

  • 使得存储大小可控。在每次通讯前,积攒的存储大小是常量,是已知可控的。更方便使用者对训练中的存储消耗和通讯时间进行预估。

设置机制,对碎片化的存储空间进行重新整合,整出连续的存储空间。

防止出现总存储足够,但连续存储不够而引起的存储请求fail

ZeRO - Offload

最后,简单介绍一下ZeRO-Offload。它的核心思想是:显存不够,内存来凑。如果我把要存储的大头卸载(offload)到CPU上,而把计算部分放到GPU上,这样比起跨机,是不是能既降显存,也能减少一些通讯压力呢?

forward和backward计算量高,因此和它们相关的部分,例如参数W(fp16),activation,就全放入GPU。

update的部分计算量低,因此和它相关的部分,全部放入CPU中。例如W(fp32),optimizer states(fp32)和gradients(fp16)等

核心思想是:显存不够,内存来凑。

模型并行

当你有一个单卡装不下的大模型时,一个直接的解决办法是,把模型隔成不同的层,每一层都放到一块GPU上(简而言之,切模型)

如下图:

F表示Forward, B表示Backward,4块GPU

这张图的含义是:我在GPU0上做完一次forward,然后将GPU0上最后一层的输入传给GPU1,继续做forward,直到四块GPU都做完forward后,我再依次做backward。等把四块GPU上的backward全部做完后,最后一个时刻我统一更新每一层的梯度。(是一个串行的过程)

两个问题:

(1)GPU利用度不够。

(2)中间结果占据大量内存

流水线并行

切分micro batch

流水线并行的核心思想是:

在模型并行的基础上,进一步引入数据并行的办法,即把原先的数据再划分成若干个batch,送入GPU进行训练。未划分前的数据,叫mini-batch。在mini-batch上再划分的数据,叫micro-batch。

切分完的数据,不用串行等待了,可以并行计算。

re-materialization(active checkpoint)

每块GPU上,我们只保存来自上一块的最后一层输入z,其余的中间结果我们算完就废。

等到backward的时候再由保存下来的z重新进行forward来算出(显存计算峰值)。

张量并行TP

详解笔记

目前基于Transformer做大模型预训练最基本的并行范式:来自NVIDIA的张量模型并行(TP)

实现的框架:MegatronLM

  • 它的基本思想就是把模型的参数纵向切开,放到不同的GPU上进行独立计算,然后再做聚合。

  • 关键问题:如何切分?

按行切

我们用N来表示GPU的数量。有几块GPU,就把W按行维度切成几份。下图展示了N=2时的切割方式:

X和W都要切开,前向传播过去相加,反向传播回来是concat

按列切分

X不用动,切W就行,前向传播的Y是concat,反向传播回来的梯度是相加