网站开发 数据库,装修公司免费装修,平安保险网站,seo网站推广 沈阳1. 背景介绍
如果你拿到了两台8卡A100的机器#xff08;做梦#xff09;#xff0c;你的导师让你学习部署并且训练不同尺寸的大模型#xff0c;并且写一个说明文档。你意识到#xff0c;你最需要学习的就是关于分布式训练的知识#xff0c;因为你可是第一次接触这么多卡…1. 背景介绍
如果你拿到了两台8卡A100的机器做梦你的导师让你学习部署并且训练不同尺寸的大模型并且写一个说明文档。你意识到你最需要学习的就是关于分布式训练的知识因为你可是第一次接触这么多卡但你并不想深入地死磕那些看起来就头大的底层原理你只想要不求甚解地理解分布式的基本运行逻辑和具体的实现方法。那么我来帮你梳理关于大模型的分布式训练需要了解的知识。
1.1 分布式定义
分布式就是把模型或者数据分散分布到不同的GPU去。为什么要分散到不同的GPU当然是因为一个GPU的显存太小了不管从为了训练加速还是模型太大塞不进去这两个角度来看本质就是单个GPU显存不够。为什么GPU显存这么小是因为GPU对带宽的要求很高能达到这样高带宽的内存都很贵也就是说从成本的角度上导致单个GPU显存有限。
1.2 分布式方法分类
从需要加载的模型大小来分类一般可以这样对分布式的方法进行分类。 单卡能够容纳训练的小模型分布式就是为了训练加速一般使用的是数据并行的方法也就是每一块GPU复制一份模型然后将不同的数据放到不同的GPU训练。 模型变得越来越大单卡都无法支持一个模型训练的时候就会使用模型并行的方法模型并行又分为流水线并行Pipeline Parallelism和张量并行Tensor Parallelism其中流水线并行指的是将模型的每一层拆开分布到不同GPU。当模型大到单层模型都无法部署在单个gpu上的时候我们就会用到张量并行将单层模型拆开训练。 Deepspeed则是用了Zero零冗余优化的方法进一步压缩训练时显存的大小以支持更大规模的模型训练。
前排提示文末有大模型AGI-CSDN独家资料包哦 2. 必要知识补充
2.1 模型是怎么训练的
我们想了解模型训练时分布式是如何进行优化的那么知道模型是如何训练的就非常重要。我们以目前最广泛使用的混合精度训练为例按照训练运行的逻辑来讲 Step1:优化器会先备份一份FP32精度的模型权重初始化好FP32精度的一阶和二阶动量用于更新权重。 Step2:开辟一块新的存储空间将FP32精度的模型权重转换为FP16精度的模型权重用于前向和计算梯度。 Step3:运行forward和backward产生的梯度和激活值都用FP16精度存储。 Step4:优化器利用FP16的梯度和FP32精度的一阶和二阶动量去更新备份的FP32的模型权重。 Step5:重复Step2到Step4训练直到模型收敛。
我们可以看到训练过程中显存主要被用在四个模块上 模型权重本身(FP32FP16) 梯度(FP16) 优化器(FP32) 激活值(FP16)
具体训练的时候会消耗多少显存怎么计算呢多大的模型在训练的时候会超出我们的最大显存呢关于大模型训练时的显存占用分析可以看一看这篇文章“一文讲明白大模型显存占用只考虑单卡”讲得非常详细自卖自夸有助于后续我们理解怎么优化显存占用。 https://zhuanlan.zhihu.com/p/713256008 2.2 GPU是怎么通讯的
我们常常听说单机8卡、多机多卡、多节点这里的“单机”和“节点”也就是指的是单个服务器8卡自然指的是8张显卡。在同一个服务器内的8张GPU共享同一个主机系统的资源如CPU资源、系统内存等同一节点内的不同GPU之间通过PCIe总线或者NVLink仅限NVIDIA显卡进行通讯不同节点之间的GPU一般是通过InfiniBand网卡通讯NVIDIA的DGX系统就是一张GPU卡配一个InfiniBand网卡下图可以看出多节点之间的通讯带宽和单节点多GPU之间的通讯带宽是可以达到差不多速度的。你问为什么节点之间不用NVLink因为 NVLink设计用于非常短距离的通信通常在同一服务器机箱内 不方便扩展。除了上面介绍的硬件还有软件部分就是大名鼎鼎的NCCL 英伟达集合通信库专用于多个 GPU 乃至多个节点间通信的实现。AllReduceBroadcastReduceAllGatherReduceScatter就是NCCL主要实现的通信原语。
其实NCCL就是提供了一个接口我们也不需要了解底层如何实现的只需调用接口就可以实现GPU间的通信。
2.2.1 Broadcast广播
将一个GPU上的数据同时发送到所有其他GPU这种操作方式就叫做广播broadcast)。
2.2.2 Scatter散射
把一个GPU上的数据切分为多块按照顺序向其他GPU发送数据块就叫做散射。
2.2.3 Reduce规约
reduce在这里的作用是求和即将多个GPU的数据发送到指定的GPU然后求和。
2.2.4 Gather聚合
Gather是Scatter散射的逆操作将分散在各个GPU的数据块拼接在一起。
2.2.5 AllReduce全规约
每一个GPU把数据发送给其余GPU每个GPU对于接收到的数据进行求和。
2.2.6 AllGather全聚合
每一个GPU把数据发送到其余所有GPU每个GPU对于接受到的数据拼接为完整数据。
2.2.7 ReduceScatter散射规约
散射规约指的就是先做Scatter再做Reduce先将GPU的数据分割为多个数据块然后每一块数据按GPU序列进行散射Scatter。对所有的GPU都进行散射之后当前GPU所接受到的所有数据块进行求和也就是Reduce规约。
看完了感觉很混乱也不要紧先有个大概的印象就好你只要知道这些GPU间的基本通讯的组合拳能够能够实现我们分布式所需要的各种数据传输就好。
3. 数据并行
数据并行就是一开始大模型很小的时候单卡装下绰绰有余这时候我们想加速训练除了batchsize开大还有什么办法呢一个人刷盘子太慢那就多招几个人嘛多招几个人指的就是多复制几个模型每个模型占领一个GPU去干活嘛猴子猴孙身外身法启动
3.1 参数服务器
一个非常简单的方法就是从八个猴子里选一个作为一个主管这个主管的作用就是收集所有猴子的梯度然后经过九九八十一算得到平均梯度后再分别传回去。 https://zhuanlan.zhihu.com/p/713256008 这里我们引入一个通讯量的概念就是用来衡量单个GPU或者整个系统的通讯参数量我们设一份完整的参数为 Φ 在这里就是指的一份完整的梯度参数假设有N个GPU那么通讯量如下表有没有感觉这个方法很熟悉没错就是上面提到过的ALL_Reduce方法。这也是李沐大佬的参数服务器这一论文提出的方法想要深入了解的可以去看论文 Scaling Distributed Machine Learning with the Parameter Server 。
3.2 Ring AllReduce
上面的参数服务器方法虽然简单但是也存在很致命的两个缺陷 1冗余太多每一个GPU都要复制一个模型权重这些一样的权重就是冗余。 2通讯负载不均衡Server需要和每一个Worker进行通讯而Worker们却只用发送和接受各一次。
那么Ring AllReduce就是一种用来解决负载不均的方式我们把上面的方法理解为中心化方法那么该方法就是去中心化的方法不再有Server和Worker的区别做到了每一个GPU通讯量的均衡。整个过程其实就是之前提到的通讯原语里的ReduceScatter然后AllGather大家可以回去对照着看一下。对于Ring AllReduce的通讯量计算因为是一个并行度很高的通讯所以只需要看单个GPU的通讯量然后乘数量就是总通讯量了。需要注意的是参数服务器和Ring AllReduce在通讯总量上是几乎一样的但是他们的通讯时间却不一样就是因为前者在sever端通讯负载过大堵车了导致时间更长。现在我们torch常用的DDP Distributed Data Parallel 多机多卡训练就是用的Ring AllReduce进行通讯。
4. 模型并行
上面讲到的数据并行是在模型足够小能被单张GPU容纳的情况下的并行方法但是现在大模型时代模型动不动70B、几百B的单张卡都容纳不下该怎么办呢那就只能把模型肢解了呗。这样的分布式方法跟数据并行不一样的最深层次的原因就是数据并行的数据都是相互独立的数据并行是为了提高训练效率的并行而模型并行的数据不一定是独立的我把一个模型拆为两半那么我这两半的模型都得用同一批数据去跑才行模型并行是模型太大不得已的对模型进行肢解的一种并行。
4.1 流水线并行
流水线并行也就是横切模型不都一层一层的吗那我们就正好一层一层地剥开~~就是按照这样的思路拆开模型。
4.1.1 朴素流水线
朴素流水线如下图所示也是最简单的一种流水线方法按照时间顺序先一层一层地进行前向然后再一层一层地反向传播也非常好理解。需要注意的是我们在每一层之间传递的是当前层的输出张量而不是梯度信息故通讯量相对较小。但是通讯量小不代表总的利用效率越高因为通讯和计算的并行程度还有内存冗余也是很需要考虑的因素。下面这张图是更加具体一些的两张GPU实现朴素流水线的展示我们发现除了中间多了一步前向和后向的信息传输实际运行起来跟放在一张GPU里是没有任何区别的。朴素流水线非常的简单这导致的问题也非常的明显我们指出最主要的两个问题 GPU利用率过低 从我们的图4-2也能看出来图中所有的空白部分都是被我们浪费掉的GPU空间这些GPU空间被叫做Bubble也就是空泡。朴素流水线的空泡率可以计算为 ( \text{空泡率} \frac{G - 1}{G} ) 空泡率指的是时间上浪费的时间占比( G ) 指的就是GPU数量可以看到GPU越多空泡率就越接近1这会造成GPU利用率非常低。 通讯和计算没有交错 GPU在层之间进行通讯的时候所有GPU都被闲置了通讯和计算并没有实现并行模式这也会导致训练的效率大大降低。
4.1.2 微批次流水线并行
这个方法的核心思路就是用数据并行的方式来解决朴素流水线的GPU并行效率太低的问题将一个大batch拆为多个小batch进行流水线。我们举个最简单例子来理解张三和李四用朴素流水线处理一筐鱼张三负责洗鱼李四负责煮鱼李四必须得等张三洗完一整筐鱼才能煮鱼。现在他们改进为微批次流水线把一大筐鱼分为N小框鱼这时张三洗完一小筐鱼就给李四煮鱼然后张三又接着洗下一小筐鱼张三和李四之间就产生了并行处理的部分。
我们现在来看下面的图就很好理解了这种方法叫做F-then-B 模式的流水线策略在微批次的情况下先进行前向计算再进行反向计算。F-then-B 模式的空泡率可以计算为 空泡率 是GPU数量是minibatch数可以看到空泡率明显下降了而且随着相比越大的影响就会越小。什么你说还不够高效那么我们看看1F1BOne Forward pass followed by One Backward pass模式的流水线策略这是一种前向计算和反向计算交叉进行的方法。在 1F1B 模式下前向计算和反向计算交叉进行可以及时释放不必要的中间变量。如下图所示绿色的箭头就表示前向和反向之间的距离他们之间的距离比F-then-B里的距离要更近这就说明这种方法可以在反向后更快地释放不需要的激活值。1F1B 模式的空泡率可以计算为空泡率诶这里空泡率怎么跟F-then-B的模式一模一样啊他们的效率应该是一样的才对吧为什么还说优化了。这是因为1F1B 的交叉进行快速释放的策略模式节省了显存在设备显存一定的情况下就可以通过增大的值增大micro-batch的个数来降低Bubble率了。
4.2 张量并行
上面讲了横着切的流水线并行我们知道横着切指的是按模型分层的维度来切分那么我们的竖着切指的是什么呢指的就是按照模型每一层的层内参数张量进行切分。
张量并行从数学原理上来看就是对于linear层就是把矩阵分块进行计算然后把结果合并对于非linear层则不做额外设计。
如下图所示X是输入A是权重矩阵Y则是输出我们对权重矩阵A有横切和竖切两种切法。竖切的最后两个向量就是合并横切的两个向量就是相加这是由矩阵乘法性质决定的。其实这就是张量并行的最核心原理后面提到的方法无论怎么包装都离不开这两个矩阵乘法的切分。正如之前提到的只有在线性层中会出现张量并行下图就是关于线性层矩阵乘法的具体实现包括前向和反向会存在的具体计算和通讯。
5. 混合并行
上面几节介绍了关于数据并行和模型并行的基础实现思路那么这一小节我们就来讲一讲实际应用中我们常听到的Megatron和deepspeed具体是如何在数据并行和模型并行上混合优化并且玩出花来的。我们这篇文章还是主要讲思路逻辑帮助你去理解整体而不纠结于太细枝末节的技术毕竟如果要扣细节和深度Megatron和deepspeed随便一个单独拿出来讲解都是需要大量篇幅的。
5.1 Megatron
Megatron的核心思路就是混合了模型并行MP包括了张量并行TP、流水线并行PP和数据并行DP这三种并行方式也被称为3D并行的策略。其中张量并行是Megatron的重点也是我们要先学习的一个知识点你说上面不是已经讲过张量并行是基于线性层进行横切或者竖切吗那你知道具体在transformer架构里面是怎么切割的吗什么情况下横切什么情况下竖切呢transformer里面有哪些线性层是可以被这样切割的呢这些问题也就是Megatron首先要回答的问题。
5.1.1 复习一下 Transformer 模型结构
我自己画了一个简易版的 Transformer 结构旨在突出结构中的 linear 线性层有哪些部分可以很清晰地看到我们可以将模型分为三个板块一个是 Embedding 输入和输出层一个是多头注意力层一个是前馈层。接下来我们就会分别介绍这三个板块的张量并行是如何具体实现的。
5.1.2 FFN 前馈层的张量并行
我们了解前馈层是由两层 Linear 组成的也就是下图的 A 和 B 权重矩阵。下图的张量并行方式是很简单清晰的就是对 A 进行了竖切对 B 进行了横切。那么这时候就要问了为什么我们先来看两个式子如果
如果 其中冒号表示 concat
你觉得哪个式子是正确的这里就不卖关子了第一个加法的式子是不对的第二个合并的式子是正确的。
因为 Relu 或者其他任意的非线性函数是不具有类似乘法分配律 的性质的。
我们举一个简单的例子就很好理解了。
Relu (x) 就是 max (0, x) 嘛
可以很清楚地看到对于非线性激活函数是不能先分开经过激活函数然后求和的。合并的例子更简单这里就不写了。经过激活函数后再合并的结果和直接合并再经过激活函数的结果是相同的。
这里我们应该也明白为什么上图中的 要对 进行竖切了吧因为这样子可以不需要在计算非线性激活函数前做加法的通讯而仅仅在最后合并结果进行一次通讯。对 进行竖切后因为矩阵乘法的原因 的 就只能进行横切了。
具体来分析一下 f 的 forward 计算把输入 拷贝到两块 GPU 上每块 GPU 即可独立做 forward 计算。 g 的 forward 计算每块 GPU 上的 forward 的计算完毕取得 和 后GPU 间做一次 AllReduce相加结果产生 。 g 的 backward 计算只需要把 拷贝到两块 GPU 上两块 GPU 就能各自独立做梯度计算。 f 的 backward 计算当当前层的梯度计算完毕需要传递到下一层继续做梯度计算时我们需要求得 。则此时两块 GPU 做一次 AllReduce把各自的梯度 和 相加即可。
因此此时两块 GPU 做一次 AllReduce把各自的梯度 和 相加即可。
5.1.3 MHA 多头注意力层的张量并行
对于多头注意力来讲进行张量分割的方式则比较特殊如下图所示是按照多头的头数量来进行张量并行分割的这是因为每个头上都可以独立计算 attention。
具体来讲我们对 Q、K、V 矩阵做列切割对 O 矩阵做行切割。这样切割的原因和 FFN 层一样都是为了减少通讯量。forward 和 backward 计算也都差不多都是进行两次 AllReduce这里就不再赘述了。
5.1.4 Embedding 嵌入层的张量并行
我们知道对于 Embedding 层来讲矩阵维度要么是V,H输入矩阵要么是H,V输出矩阵。
我们要对其进行并行分割的最主要原因就是这个 V 词表维度动不动就上十几万相对于一般的 H 维度几千维是完全不在一个数量级。所以我们也只需要牢牢记住对于 Embedding 层就是对 V 词表维度进行分割。
对于输入 Embedding 层
那么我们要从 这个维度拆分这个 就会横着切每一个 GPU 上的权重维度就是
同理也可以分析得出对于输出 Embedding 层 会竖着切每一个 GPU 上的权重维度就是
5.1.5 3D 并行举例分析
这里我会给出一个例子来帮助大家理解 Megatron 的 3D 并行是如何具体应用的。
假设我们有两台机器每台机器 8 张 GPU我们的一个模型需要 8 张 GPU 才能完整装下假设模型一共 4 层 Linear 层构成。
首先我们先用流水线并行的思路将模型的四层拆开但是我们发现模型的一层也放不下一个 GPU所以我们再用到张量并行的思路将每一层的权重一分为二这时候我们将一个模型拆为八块正好每一块可以放在一个 GPU 里。一般而言通讯量排名是TP张量并行 DP数据并行 PP流水线并行。通讯量大的部分尽量放入一台机器内因为 node 内比 node 间带宽高。基于这个考虑我们进行如下的分割。
可以从下图看到红色和蓝色分别代表两个独立的完整模型被分割为 8 块L 和 R 对应的就是一层 Linear 层的张量分割模式。可以看到我们所有成对的 L 和 R 张量并行都是放在一个 Node 里的要优先保证张量并行的通讯。然后是考虑数据并行和流水线并行我们将 0、1 两层和 2、3 两层分别置于不同的 Node因为流水线并行的通讯量相对较少可以将其分别放在两个 Node。
5.2 DeepSpeed
想要很好地理解 DeepSpeed请一定复习一下之前提到的模型训练的流程和显存占用2.1 小结和 GPU 通讯原语2.2 小结这样会有更加深刻的理解。
假设我们有一个无限大显存的 GPU就完全没有数据并行的必要了这里可以自己思考一下模型并行为什么是有必要的即使是有无限大的 GPU也就是说我们的数据并行本质上就是利用通讯换取显存的思想DeepSpeed 大名鼎鼎的 Zero 机制也是基于这样的思想。
可能有了解过 Zero 机制的同学会觉得它的方法更像是模型并行因为它涉及对优化器参数、模型权重、模型梯度的分割。但 Zero 机制其实是一个数据并行因为在做 forward 和 backward 的时候仍然需要完整的模型权重。
模型并行的本质是将一个大的模型权重拆分为多个小的部分每一个小部分都可以独立地在单张 GPU 上运行数据并行则是需要一个完整的模型权重这里的完整模型权重可以是一整个模型也可以是被拆分过后的一层模型只要这个模型是独立的就可以在单张 GPU 上运行而我们只是想办法塞进更多数据。Zero 机制尽管进行了很多分割但在 forward 和 backward 时仍然需要完整的模型权重。现在看不太懂没关系继续看下去到后面就会慢慢理解。
模型训练的步骤我们先一起简单复习一下 模型权重进行 forward 和 backward得到梯度。 利用梯度去更新优化器adam一阶二阶动量。 优化器更新模型权重。
下面这张图是我之前一直非常纠结的地方每一个蓝框里的内容都是我以为当前步骤下一定需要的显存。 第一张图backward 的时候显存一定是有模型权重和模型梯度的优化器参数其实用不到。 第二张图在更新优化器一阶二阶动量的时候其实不需要加载模型权重了。 第三张图用优化器的一阶二阶动量去更新模型权重时梯度应该是不需要的可以释放掉。
但是实际情况是模型权重、模型梯度、优化器参数一直都占用显存。 优化器一阶二阶动量参数的意义是存储历史的梯度信息、梯度变化信息所以它不能被释放掉要一直占用等待更新。 模型权重是我们要更新的参数当然也需要一直保存在显存里等待更新。 梯度每一次都不一样这总可以先释放掉吧其实也不行因为第二步和第三步是合在一起的也就是 optimizer.step() 这个函数需要等更新完权重之后模型梯度才可以释放但这时候又进入下一次的 forward 和 backward 了梯度占用也就一直保存了其实这里我现在也感觉怪怪的不知道有没有大佬能解答一下感觉这个梯度理论上应该是可以释放的。
总之就是我们实际训练的时候会多出很多的冗余参数明明当前阶段用不到但却在占用我们的显存。那么该怎么优化这些冗余呢答案就是切分布式切分。切开之后用通讯换取显存需要的时候就传过来不需要的时候就释放掉。我们可以切优化器参数、切梯度、切模型权重这也是所谓的 Zero 123我们下面将具体介绍。
5.2.1 切优化器Optimizer StatesZero1
我们以 Ring AllReduce 为基准进行对比不考虑模型并行进行分析可以更加清晰地理解 Zero 的思路。
设模型的参数量是 GPU 总数为 模型权重的显存是 模型梯度的显存是 模型优化器的显存是 因为优化器可以用不同的Adam 占用 占用非常多这也是为什么要先切优化器。我们通过下面的动图和步骤来理解 Zero1 切优化器之后是如何训练的 forward 和 backward 得到梯度。 对梯度做 Ring AllReduce得到平均梯度单个 GPU 通讯量 。 根据优化器参数和对应的模型梯度更新对应的模型权重。 对模型权重的更新部分做 AllGather单个 GPU 通讯量 。最后我们来补充下面这个表格很好理解优化器部分的显存除以 被分布到各个 GPU 上单卡通讯量在刚才的步骤分析里可以得出是增加到 。有细心的人会发现这个阶段有个明显的可优化的点注意看第三步在每一个 GPU 上我们其实只用到了一部分的模型梯度也就是虚线框里的梯度其他的就是冗余的部分。这其实就是接下来 Zero2 要进行优化的点对梯度进行拆分。
5.2.2 切梯度GradientsZero2
我们直接上步骤和动图大家可以结合前面的通讯原语言慢慢理解还是挺有意思的 forward 和 backward 得到完整梯度。 对梯度进行 ReduceScatter单个 GPU 通讯量 。 根据部分梯度以及对应的优化器参数更新对应模型的部分权重。 对模型权重进行 AllGather 操作单个 GPU 通讯量 。再来补充一下显存和通讯量的表格有的人会注意到在第一步的时候在 ReduceScatter 之前我们不是得到了完整的梯度吗那这个时候我们的显存应该是 才对啊那么显存的上限不应该还是这么多吗我认为是用到了 CPU 内存去作为一个暂时缓存的 buffer缓存马上就要被 ReduceScatter 的梯度。因为在后面我们会给出的官方视频里也可以看到模型在做 AllGather 之后会有一个 temporary buffer 去暂存。
很神奇的发现 Zero2 相比于 Ring AllReduce用的通讯量是一样的显存占用却省了很多这就是所谓的去冗余的效果。
5.2.3 切模型权重ParametersZero3
这一部分本来也想做一个 GIF 图去展示但画图的时候遇到了一些疑惑然后就发现了下面这个官方的演示视频实在是太过清晰明了了如果你认真看懂了前面的内容那么这个视频会非常容易地看懂非常棒的视频。
个人建议一定要看一下我之前看很多博客介绍都没看明白一看这个视频就懂了。仔细想想这个动态实现的方式确实很难用语言去没有歧义地描述。相信我看这个视频绝对比你去慢慢琢磨那些博客更快。
总结一下
从整个流程来看模型权重进行了两次 AllGather对梯度做了一次 ReduceScatter。但这里的 AllGather 和 ReduceScatter 跟之前不太一样它是和 forward、backward 整个流程交叉着进行的这一点真的要结合视频细细体会你会恍然大悟。终于我们得到了《 ZeRO: Memory Optimizations Toward Training Trillion Parameter Models 》论文里这张总结的大图是否感觉一切明了了。
6. 总结
本来只想简单地勾勒出整个技术链路的逻辑没想到写了这么多的内容累鼠我了。感觉也没啥好总结的了就先这样吧有什么漏洞和问题的后面再补。毕竟要先做一坨屎出来然后再慢慢修改是最好的。
