首页 A1训练模型 学习笔记-模型训练加速

学习笔记-模型训练加速

标题:学习笔记-模型训练加速

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

数据并行训练时,数据集被分成若干分片,每个分片分配给一个设备

每个设备将保存模型的完整副本并在分配的数据集片段上进行训练

反向传播后,模型的梯度会进行Allreduce,以便不同设备上的模型参数能够同步

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

输入数据分割 方法一:在每个训练Epoch开始之前,根据并行进程的数量对整个训练数据集进行分割,每个进程只读取自己的分割数据。

方法二:数据读取只负责特定进程(假设为rank0)。读取数据后,rank0还根据并行进程的数量将数据分成多个块,然后将不同的数据块发送给相应的进程。

分段注意事项

所有进程都需要为每个训练步骤输入相同的本地批量大小。

确保所有进程分配相同的批号

数据并行实现模型参数同步的关键问题是如何保证模型在训练过程中各进程的参数相同。您需要确保以下两点:

各进程上模型的初始化参数相同;

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

方法二:通过特定进程初始化所有模型参数,然后该进程将模型参数广播给所有其他进程

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

远期计算

逆算

参数更新

前向计算 每个进程根据其获取的输入数据独立进行前向计算。由于输入数据不同,每个过程会得到不同的Loss。

使用批量归一化的问题

在数据并行训练中,全局batch size被分成不同的进程,每个进程只有部分输入数据。这样,批量归一化就计算了输入张量批量维度的均值(Mean)和方差(Variance)。仅使用部分batch而不是全局batch,这会导致一些对batch size敏感的模型的准确率下降。

在数据并行训练中使用SyncBatchNorm策略来保证模型的准确性。该策略在模型训练之前向BN 层计算均值和方差时添加了额外的同步通信,并在所有数据并行进程上使用张量而不是在其自身进程上使用张量。计算张量批量维度的均值和方差

逆向计算 各进程在各自的正向计算的基础上独立进行逆向计算。

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

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

参数更新 每个进程经过上述两步后获得相同的全局梯度,然后独立完成参数更新。

由于更新前模型各过程的参数相同,更新时使用的梯度也相同,因此更新后各过程的参数也相同。

Ring-AllReduce 假设有N个GPU,每个GPU上的数据也被切分成N部分。 AllReduce的最终目标是将每个GPU上的数据变成汇总结果

Ring-ALLReduce 通过两个主要步骤实现这一目标:Reduce-Scatter 和All-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 阶段结束。进入全员聚集阶段。目标是将完整的聚合数据广播到其他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 数量无关。

Pipeline并行背景 当训练较大的模型时,每个GPU需要存储模型参数和中间结果,需要较大的BatchSize,这会增加GPU的使用率。

当模型达到一定大小时,单个GPU 将无法存储整个模型。

直接的解决方案是将模型分成不同的层,并将每个层放在GPU 上

简单流水线并行流程介绍

Naive pipeline 并行是实现pipeline 并行训练最直接的方法

将模型按照层数划分为多个部分(Stage),并将每个部分(Stage)分配给一个GPU

定期对小批量数据进行训练,并在模型分割的边界处进行通信

下标表示批次。这里只有1个batch,所以下标都是0。每一行代表一个GPU,每一列是一个时间步。

1个batch的数据依次在4个GPU上进行前向计算,然后反向进行后向计算。

缺点

1.GPU利用率不足

一个。下图中每行的空白区域表示GPU处于空闲状态。

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

我。假设每个GPU前进+后退的时间为t

二.假设总共有K个GPU

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

d.每行已使用面积为1 t,可用面积为(K-1) * t

e.可以计算出每个GPU的空闲率为:(K-1)/K

f.当K越大,即GPU数量越多时,空闲率接近1,即GPU资源被浪费。

2、中间结果占用大量内存

一个。在向后计算梯度的过程中,我们需要使用每一层的中间结果z。

b.假设我们的模型有L层,每层的宽度为d,那么对于每个GPU,无论其参数本身的存储如何,额外的空间复杂度都是O(B * L*d/K)。

c.从这个复杂度可以看出,随着模型的增加,N、L、d的增加可能会平滑K增加带来的GPU显存增益。因此,这也是一个需要优化的区域

优化的流水线并行分割微批次 在模型并行的基础上,进一步引入数据并行方法,即将原始数据分成若干个批次,发送到GPU 进行训练。

分区前的数据称为mini-batch。在mini-batch上细分的数据称为微batch

计算过程如下

下图中,第一个下标表示GPU编号,第二个下标表示微批次编号。

闲置率为:

M 的微批次需要M 个时间步,直到最后一个被发送到第0 个GPU。

最后一个微批次通过所有GPU,需要K-1个时间步长

因此,需要K + M-1个时间步(分母)来遍历所有数据

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

因此,空闲时间步长为K – 1(分子)

当M=4K时,空闲率可以忽略

在micro-batch的划分下,会对Batch Normalization的计算产生影响。

Gpipe的方法是在训练时计算并使用micro-batch中的均值和方差。

继续同时跟踪所有mini-batch的移动平均值和方差,以供测试阶段使用

图层标准化不受影响。

重新物化(主动检查点) 在每个GPU 上,仅保存前一个块的最后一层输入z,其余中间结果在计算后被丢弃。

反向时,使用保存的z 再次进行前向,计算各层的中间值。

张量并行的基本思想是将模型的参数纵向切分,放到不同的GPU上独立计算,然后聚合。我们以GEMM为例,看看如何并行化模型。这里我们需要做的是XA=Y。对于模型来说,X是输入,A是权重,Y是输出。从数学的角度来看,对于线性层来说,将矩阵分成块进行计算,然后将结果进行合并。对于非线性层,不需要额外的设计。

行并行 将A 按行分成两部分。为了保证运行,我们还将X按照列分为两部分:

将X1 A1放在第一个GPU上计算,X2 * Y2放在第二个GPU上,然后合并

列并行度根据列划分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 的张量并行方案 对于self-attention 块,利用多头注意力操作中固有的并行性,对与键(K)、查询(Q) 和值(V) 相关的GEMM 进行列并行处理执行分区,以便与每个注意力头对应的矩阵乘法在一个GPU 上本地完成。这允许我们在GPU 之间分割每个注意力头参数和工作负载,每个GPU 获得一部分输出

对于后续的全连接层,由于每个GPU 上都有部分输出,因此将权重矩阵B 按行划分,直接与输入Y1、Y2 计算,然后通过g 中的all-reduce 操作和Dropout 得到得到最终结果Z

混合精度训练 FP16 训练的问题

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

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

权重备份 计算过程中产生的权重、激活值、梯度等均使用FP16存储和计算

使用FP32 额外备份权重,以便在训练期间进行更新。

由于在深度模型中,学习率梯度的参数值可能很小,如果使用FP16进行加法,很可能会出现舍入误差,导致更新无效。

因此,通过将权重复制为FP32格式并确保整个更新过程以FP32格式执行

损失缩放 当使用混合精度训练时,网络模型可能无法收敛。

由于梯度值太小,使用FP16表示会导致数据下溢(Underflow)

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

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

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

动态损耗缩放:

从相对较高的缩放因子(如2^24)开始,检查训练迭代过程中数字是否会溢出(Infs/Nans)

如果不存在梯度溢出,则不进行缩放,继续迭代;

如果检测到梯度溢出,则缩放因子减半,并重新确认梯度更新,直到数量在不会导致溢出的范围内;

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

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

精度累加 FP16 进行矩阵乘法,FP32 在矩阵乘法中间进行累加,然后将FP32 的值转换为FP16 存储

矩阵乘法用FP16,加法计算用FP32,以弥补精度损失。这样可以有效减少计算过程中的舍入误差,最大限度地减少精度损失的问题。

热门文章