分享免费的编程资源和教程

网站首页 > 技术教程 正文

学习笔记-模型训练加速 模型训练原理

goqiw 2024-09-30 19:07:12 技术教程 24 ℃ 0 评论

数据并行

数据并行

● 数据并行是最常见的并行形式, 因为它很简单

● 在数据并行训练中, 数据集被分割成几个碎片, 每个碎片被分配到一个设备上

● 每个设备将持有一个完整的模型副本, 并在分配的数据集碎片上进行训练

● 在反向传播之后, 模型的梯度将被Allreduce, 以便在不同设备上的模型参数能够保持同步

● 主要分为两个操作: 输入数据切分模型参数同步



输入数据切分

● 方式一:在每个训练 Epoch 开始前, 将整个训练数据集根据并行进程数划分, 每个进程只读取自身切分的数据。

● 方式二:数据的读取仅由具体某个进程负责 (假设为 rank0)。rank0 在数据读取后同样根据并行进程数将数据切分成 多块,再将不同数据块发送到对应进程上。

● 切分注意事项

要求所有进程每个训练 step 输入的 local batch size 大小相同

○ 要保证所有进程上分配到相同的 batch 数量

模型参数同步

数据并行实现的关键问题在于如何保证训练过程中每个进程上模型的参数相同。需要保证如下2点:

● 每个进程上模型的初始化参数相同;

○ 方法一:所有进程在参数初始时使用相同的随机种子并以相同的顺序初始化所有参数

○ 方法二:通过个具体进程初始化全部模型参数,之后由该进程向其他所有进程广播模型参数

● 每个进程上每次更新的梯度相同。

○ 前向计算

○ 反向计算

○ 参数更新

前向计算

每个进程根据自身得到的输入数据独立前向计算, 因为输入数据不同每个进程会得到不同的 Loss。

● 使用Batch Normalization的问题

○ 数据并行训练中 global batch size 被切分到不同的进程之上, 每个进程上只有部分的输入数据, 这样批归一化在计算输入tensor batch 维度的平均值 (Mean) 和方差 (Variance) 时仅使用了部分的 batch 而非 global batch,会导致部分对 batch size 比较敏感的模型的精度下降

○ 在数据并行训练中用 SyncBatchNorm 策略来保证模型精度, 该策略在模型训练前向 BN 层计算 mean 和 variance 时加入额外的同步通信, 使用所有数据并行进程上的 tensors 而非自身进程上的 tensor 来 计算 tensor batch 维度的 mean 和 variance



反向计算

每个进程根据自身的前向计算独立进行反向计算

● 在后续的更新步骤之前,对所有进程上的梯度进行同步, 保证后续更新步骤中每个进程使用相同的全局梯度更新模型参数

● 梯度同步过程通过Allreduce sum 同步通信操作实现的, 对梯度使用 Allreduce sum 操作后每个进程上得到的梯度是相同的,等于所有进程上梯度对应位置相加的和



参数更新

每个进程经过上述2步后得到相同全局梯度,然后各自独立地完成参数更新。

因为更新前模型各进程间的参数是相同的,更新中所使用的梯度也是相同的,所以更新后各进程上的参数也是相同的。

Ring-AllReduce

● 假设有N块GPU,每块GPU上的数据也对应被切成N份。AllReduce的最终目标,就是让每块GPU上的数据都变成汇总的结果



● Ring-ALLReduce则分两大步骤实现该目标:Reduce-ScatterAll-Gather

Reduce-Scatter

每个GPU只和其相邻的两块GPU通讯。每次发送对应位置的数据进行累加,每一次累加更新都形成一个拓扑环,因此被称为Ring



● 首先,gpu将数组划分为N个更小的块(其中N是环中的gpu数)



● 接下来,GPU将进行N-1次 Scatter-Reduce 迭代;在每次迭代中,GPU将向其右邻居发送一个块,并从其左邻居接收一个块并累积到该块中

● 每个GPU发送和接收的块在每次迭代中都是不同的

● 第n个GPU从发送块N和接收块N - 1开始,然后从那里向后进行

● 每次迭代都发送它在前一次迭代中接收到的块



● 在第一次发送和接收完成之后,每个GPU将拥有一个块,该块由两个不同GPU上相同块的和组成



● 在下一次迭代中,该过程继续进行,到最后,每个GPU将有一个块,该块包含所有GPU中该块中所有值的总和。





All-Gather

● 在scatter-reduce步骤完成之后,每块GPU上都有一块数据拥有了对应位置完整的聚合。此时,Reduce-Scatter阶段结束。进入All-Gather阶段。目标是把完整聚合的数据广播到其余GPU对应的位置上。

● All-Gather过程与scatter-reduce是相同的(发送和接收的N-1次迭代),只是gpu接收的值没有累加,而是简单地覆盖块。

● 第n个GPU首先发送第n+1个块并接收第n个块,然后在以后的迭代中总是发送它刚刚接收到的块。



● 第一次迭代完成后,每个GPU将拥有最终数组的两个块。

● 下一个迭代中,该过程将继续,到最后,每个GPU将拥有整个数组的完整累积值






Ring-AllReduce通讯量分析

● 假设GPU数量为N,每块GPU存储的数据量为K,那么每个梯度块为K / N

● 每个GPU要进行N-1次reducer-scatter操作和N-1次all-gather操作,总传输量为


● 随着N增大,近似为2K,单卡传输量为恒定值,与GPU数量无关

流水线并行

背景

● 训练更大的模型时,每块GPU里需要存放模型参数、中间结果,同时需要更大的BatchSize,这加大GPU的使用

● 当模型大到一定的程度,单张GPU会无法存放整个模型

● 一个直接的解决办法是,把模型隔成不同的层,每一层都放到一块GPU上



朴素的流水线并行

流程介绍

● 朴素流水线并行是实现流水线并行训练的最直接的方法

● 将模型按照层间切分成多个部分(Stage),并将每个部分(Stage)分配给一个 GPU

● 对小批量数据进行常规的训练,在模型切分成多个部分的边界处进行通信

○ 下标表示batch,此处只有1个batch,因此下标都是0。每一行表示一个gpu,每一列是一个时间步

○ 1个batch的数据,依次在4块GPU上完成前向计算,再反向分别进行backward计算


弊端

1. GPU利用率不够

a.下图中每一行空白区域,表示该gpu处于空闲状态

b.可计算下图黑色矩阵中空白区域的占比, 来计算利用率

i. 假设每块gpu一次forward + backward的时间为t

ii. 假设共有K块GPU

c.整个矩阵的长为: K * t

d.每一行上,利用区域为 1 t, 空闲区域为 (K-1) * t

e.可计算出,每块GPU的空闲率: (K-1)/ K

f.当K越大,即GPU的数量越多时,空置的比例接近1,即GPU的资源都被浪费掉了



2. 中间结果占据大量内存

a.在做backward计算梯度的过程中,我们需要用到每一层的中间结果z。

b.假设我们的模型有L层,每一层的宽度为d,则对于每块GPU,不考虑其参数本身的存储,额外的空间复杂度为O(B * L*d / K)。

c.从这个复杂度可以看出,随着模型的增大,N,L,d三者的增加可能会平滑掉K增加带来的GPU内存收益。因此,这也是需要优化的地方

优化后的流水线并行

切分micro-batch

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

● 未划分前的数据,叫mini-batch。在mini-batch上再划分的数据,叫micro-batch

计算流程如下

○ 下图中,第一个下标表示GPU编号,第二个下标表示micro-batch编号

○ 空闲率为:

■ M的micro-batch,截止最后一个被送进第0块GPU,耗时M个时间步

■ 最后一个micro-batch过完所有gpu,需要K-1个时间步

■ 因此,过完全部数据,需要K + M-1个时间步(分母)

■ 每块GPU需要完成M个数据的计算,需要M个时间步,即,有效利用为M个时间步

■ 故,空闲时间步为 K - 1 (分子)

○ 当 M >= 4K, 空闲率可忽略不记


在micro-batch的划分下,计算Batch Normalization时会有影响。

Gpipe的方法是,在训练时计算和运用的是micro-batch里的均值和方差

同时持续追踪全部mini-batch的移动平均和方差,以便在测试阶段进行使用

Layer Normalization则不受影响。

re-materialization(active checkpoint)

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

● backward的时候再由保存下来的z重新进行forward来算出每层的中间值



张量并行

基本思想就是把模型的参数纵向切开,放到不同的GPU上进行独立计算,然后再做聚合。用 GEMM 为例,来看看如何进行模型并行,这里要进行的是 XA = Y,对于模型来说,X 是输入,A是权重,Y是输出。从数学原理上来看,对于linear层就是把矩阵分块进行计算,然后把结果合并,对于非linear层则不做额外设计。

行并行(Row Parallelism)

● 把 A 按照行分割成两部分。为了保证运算,同时我们也把 X 按照列来分割为两部分:


● 把X1 A1 放到第一块GPU计算,X2 * Y2放到第二块GPU,然后在合并



列并行(Column Parallelism)

把 A按照列来分割


Transformers-MLP张量并行方案




第一个mlp + gelu的切分

行切分

● 括号之中的两项,每一个都可以在一个独立的GPU之上完成,然后通过 all-reduce 操作完成求和操纵

● 因为Gelu非线性,所以需要一个同步节点,同步两块gpu上的计算结果

列切分


● 无需同步点,直接把两个 GeLU 的输出拼接在一起就行

因此,对第一个mlp层,通常使用列切分

第二个mlp层的切分

● 使用行切分,把权重矩阵切分到两个 GPU 之上,得到 B1,B2

● 前面输出 Y1和 Y2 正好满足需求,直接可以和 B 的相关部分(B1,B2)做相关计算,不需要通信或者其他操作,就得到了 Z1,Z2。分别位于两个GPU之上

● Z1,Z2 通过 g 做 all-reduce(这是一个同步点),再通过 dropout 得到了最终的输出 Z

Transformer-SelfAttention的张量并行方案

● 对于自我注意力块, 利用多头注意力操作中固有的并行性,以列并行方式对与键(K)、查询(Q)和值(V)相关联的GEMM进行分区,从而在一个GPU上本地完成与每个注意力头对应的矩阵乘法。这使我们能够在GPU中分割每个attention head参数和工作负载,每个GPU得到了部分输出

● 对于后续的全连接层,因为每个GPU之上有了部分输出,所以对于权重矩阵B就按行切分,与输入的 Y1,Y2进行直接计算,然后通过 g 之中的 all-reduce 操作和Dropout 得到最终结果 Z

混合精度训练

● FP16训练的问题

数据溢出:数据溢出很容易理解,就是FP16能表示的数据范围比FP32小很多

○ 舍入误差:舍入误差是指当网络模型中有一些很小的反向梯度,在FP32中可以正常表示,但是转换到FP16后会小于当前区间内的最小间隔,导致数据丢失

权重备份

在计算过程中所产生的权重,激活值,梯度等均使用 FP16 来进行存储和计算

● 权重使用FP32额外进行备份,用于训练时候的更新


○ 由于深度模型中,学习率×梯度的参数值可能会非常小,如果利用FP16来进行相加的话,则很可能会出现舍入误差问题,导致更新无效。

○ 因此通过将权重拷贝成FP32格式,并且确保整个更新过程是在FP32格式下进行的

loss缩放

● 使用混合精度训练,会出现网络模型无法收敛的情况。

○ 因为梯度的值太小,使用FP16表示会导致数据下溢出(Underflow)

● 需要引入损失放大(Loss Scaling)技术

Scale up阶段,网络模型前向计算后在反响传播前,将得到的损失变化值DLoss增大2^K倍

Scale down阶段,反向传播后,将权重梯度缩2^K倍,恢复FP32值进行存储

动态损失缩放:

○ 从比较高的缩放因子开始(如2^24),进行训练迭代中检查数是否会溢出(Infs/Nans)

■ 如果没有梯度溢出,则不进行缩放,继续进行迭代;

■ 如果检测到梯度溢出,则缩放因子会减半,重新确认梯度更新情况,直到数不产生溢出的范围内;

○ 在训练的后期,loss已经趋近收敛稳定,梯度更新的幅度小,可以允许更高的损失缩放因子来再次防止数据下溢。

○ 动态损失缩放算法会尝试在每N(N=2000)次迭代将损失缩放增加F倍数,然后执行步骤2检查是否溢出。

精度累加

● FP16进行矩阵乘法运算,FP32来进行矩阵乘法中间的累加,然后再将FP32的值转化为FP16进行存储

● 就是利用FP16进行矩阵相乘,利用FP32来进行加法计算弥补丢失的精度。 这样可以有效减少计算过程中的舍入误差,尽量减缓精度损失的问题。

本文暂时没有评论,来添加一个吧(●'◡'●)

欢迎 发表评论:

最近发表
标签列表