全网最全 - 混合精度训练原理

通常我们训练神经网络模型的时候默认使用的数据类型为单精度 FP32。近年来,为了加快训练时间、减少网络训练时候所占用的内存,并且保存训练出来的模型精度持平的条件下,业界提出越来越多的混合精度训练的方法。这里的混合精度训练是指在训练的过程中,同时使用单精度(FP32)和半精度(FP16)。

1、浮点数据类型

浮点数据类型主要分为双精度(Fp64)、单精度(Fp32)、半精度(FP16)。在神经网络模型的训练过程中,一般默认采用单精度(FP32)浮点数据类型,来表示网络模型权重和其他参数。在了解混合精度训练之前,这里简单了解浮点数据类型。

根据 IEEE 二进制浮点数算术标准(IEEE 754)的定义,浮点数据类型分为双精度(Fp64)、单精度(Fp32)、半精度(FP16)三种,其中每一种都有三个不同的位来表示。FP64 表示采用 8 个字节共 64 位,来进行的编码存储的一种数据类型;同理,FP32 表示采用 4 个字节共 32 位来表示;FP16 则是采用 2 字节共 16 位来表示。如图所示:

微信图片_20211206114518.jpg

从图中可以看出,与 FP32 相比,FP16 的存储空间是 FP32 的一半,FP32 则是 FP16 的一半。主要分为三个部分:

最高位表示符号位 sign bit。

中间表示指数位 exponent bit。

低位表示分数位 fraction bit。

以 FP16 为例子,第一位符号位 sign 表示正负符号,接着 5 位表示指数 exponent,最后 10 位表示分数 fraction。公式为:

01.png

同理,一个规则化的 FP32 的真值为:

02.png

一个规格化的 FP64 的真值为:

03.png

FP16 可以表示的最大值为 0 11110 1111111111,计算方法为:

04.png

FP16 可以表示的最小值为 0 00001 0000000000,计算方法为:

05.png

1.png

2、使用 FP16 训练问题

首先来看看为什么需要混合精度。使用 FP16 训练神经网络,相对比使用 FP32 带来的优点有:

  1.减少内存占用:FP16的位宽是FP32的一半,因此权重等参数所占用的内存也是原来的一半,节省下来的内存可以放更大的网络模型或者使用更多的数据进行训练。
  2.加快通讯效率:针对分布式训练,特别是在大模型训练的过程中,通讯的开销制约了网络模型训练的整体性能,通讯的位宽少了意味着可以提升通讯性能,减少等待时间,加快数据的流通。
  3.计算效率更高:在特殊的AI加速芯片如华为Ascend 910和310系列,或者NVIDIA VOTAL架构的Titan V and Tesla V100的GPU上,使用FP16的执行运算性能比FP32更加快。

但是使用 FP16 同样会带来一些问题,其中最重要的是 1)精度溢出和 2)舍入误差。

2.png

  2.舍入误差:Rounding Error指示是当网络模型的反向梯度很小,一般FP32能够表示,但是转换到FP16会小于当前区间内的最小间隔,会导致数据溢出。如0.00006666666在FP32中能正常表示,转换到FP16后会表示成为0.000067,不满足FP16最小间隔的数会强制舍入。

3、混合精度相关技术

为了想让深度学习训练可以使用 FP16 的好处,又要避免精度溢出和舍入误差。于是可以通过 FP16 和 FP32 的混合精度训练(Mixed-Precision),混合精度训练过程中可以引入权重备份(Weight Backup)、损失放大(Loss Scaling)、精度累加(Precision Accumulated)三种相关的技术。

3.1、权重备份(Weight Backup)

权重备份主要用于解决舍入误差的问题。其主要思路是把神经网络训练过程中产生的激活 activations、梯度 gradients、中间变量等数据,在训练中都利用 FP16 来存储,同时复制一份 FP32 的权重参数 weights,用于训练时候的更新。具体如下图所示。

v2-723c1d3de5f3730e94301735252ac581_1440w.jpg

从图中可以了解,在计算过程中所产生的权重 weights,激活 activations,梯度 gradients 等均使用 FP16 来进行存储和计算,其中权重使用 FP32 额外进行备份。由于在更新权重公式为:

1.png

深度模型中,lr x gradent 的参数值可能会非常小,利用 FP16 来进行相加的话,则很可能会出现舍入误差问题,导致更新无效。因此通过将权重 weights 拷贝成 FP32 格式,并且确保整个更新过程是在 fp32 格式下进行的。即:

2.png

权重用 FP32 格式备份一次,那岂不是使得内存占用反而更高了呢?是的,额外拷贝一份 weight 的确增加了训练时候内存的占用。 但是实际上,在训练过程中内存中分为动态内存和静态内容,其中动态内存是静态内存的 3-4 倍,主要是中间变量值和激活 activations 的值。而这里备份的权重增加的主要是静态内存。只要动态内存的值基本都是使用 FP16 来进行存储,则最终模型与整网使用 FP32 进行训练相比起来, 内存占用也基本能够减半。

3.2、损失缩放(Loss Scaling)

如图所示,如果仅仅使用 FP32 训练,模型收敛得比较好,但是如果用了混合精度训练,会存在网络模型无法收敛的情况。原因是梯度的值太小,使用 FP16 表示会造成了数据下溢出(Underflow)的问题,导致模型不收敛,如图中灰色的部分。于是需要引入损失缩放(Loss Scaling)技术。

v2-f52eb5731d9437fb6340f8a1ef22dadb_1440w.jpg

下面是在网络模型训练阶段, 某一层的激活函数梯度分布式中,其中有 68%的网络模型激活参数位 0,另外有 4%的精度在 2^-32～2^-20 这个区间内,直接使用 FP16 对这里面的数据进行表示,会截断下溢的数据,所有的梯度值都会变为 0。

v2-108fd200cff0b32d44478142734bb7ec_1440w.jpg

为了解决梯度过小数据下溢的问题,对前向计算出来的 Loss 值进行放大操作,也就是把 FP32 的参数乘以某一个因子系数后,把可能溢出的小数位数据往前移,平移到 FP16 能表示的数据范围内。根据链式求导法则,放大 Loss 后会作用在反向传播的每一层梯度,这样比在每一层梯度上进行放大更加高效。

v2-a50b999d275c849d974c9bced916cba0_1440w.jpg

损失放大是需要结合混合精度实现的,其主要的主要思路是:

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

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

动态损失缩放(Dynamic Loss Scaling):上面提到的损失缩放都是使用一个默认值对损失值进行缩放,为了充分利用 FP16 的动态范围,可以更好地缓解舍入误差,尽量使用比较大的放大倍数。总结动态损失缩放算法,就是每当梯度溢出时候减少损失缩放规模,并且间歇性地尝试增加损失规模,从而实现在不引起溢出的情况下使用最高损失缩放因子,更好地恢复精度。

动态损失缩放的算法如下:

动态损失缩放的算法会从比较高的缩放因子开始(如 2^24),然后开始进行训练迭代中检查数是否会溢出(Infs/Nans);

如果没有梯度溢出,则不进行缩放,继续进行迭代;如果检测到梯度溢出,则缩放因子会减半,重新确认梯度更新情况,直到数不产生溢出的范围内;

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

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

3.3、精度累加(Precision Accumulated)

在混合精度的模型训练过程中,使用 FP16 进行矩阵乘法运算,利用 FP32 来进行矩阵乘法中间的累加(accumulated),然后再将 FP32 的值转化为 FP16 进行存储。简单而言,就是利用 FP16 进行矩阵相乘,利用 FP32 来进行加法计算弥补丢失的精度。 这样可以有效减少计算过程中的舍入误差,尽量减缓精度损失的问题。

例如在 Nvidia Volta 结构中带有 Tensor Core,可以利用 FP16 混合精度来进行加速,还能保持精度。Tensor Core 主要用于实现 FP16 的矩阵相乘,在利用 FP16 或者 FP32 进行累加和存储。在累加阶段能够使用 FP32 大幅减少混合精度训练的精度损失。

v2-0abb630431816d5797d341b59a38d2d9_1440w.jpg

4、混合精度训练策略(Automatic Mixed Precision,AMP)

混合精度训练有很多有意思的地方,不仅仅是在深度学习,另外在 HPC 的迭代计算场景下,从迭代的开始、迭代中期和迭代后期,都可以使用不同的混合精度策略来提升训练性能的同时保证计算的精度。以动态的混合精度达到计算和内存的最高效率比也是一个较为前言的研究方向。

以 NVIDIA 的 APEX 混合精度库为例,里面提供了 4 种策略,分别是默认使用 FP32 进行训练的 O0,只优化前向计算部分 O1、除梯度更新部分以外都使用混合精度的 O2 和使用 FP16 进行训练的 O3。具体如图所示。

v2-7dcd15731a4086b8de82623b3b326e9a_1440w.jpg

这里面比较有意思的是 O1 和 O2 策略。

O1 策略中,会根据实际 Tensor 和 Ops 之间的关系建立黑白名单来使用 FP16。例如 GEMM 和 CNN 卷积操作对于 FP16 操作特别友好的计算,会把输入的数据和权重转换成 FP16 进行运算,而 softmax、batchnorm 等标量和向量在 FP32 操作好的计算,则是继续使用 FP32 进行运算,另外还提供了动态损失缩放(dynamic loss scaling)。

而 O2 策略中,模型权重参数会转化为 FP16,输入的网络模型参数也转换为 FP16,Batchnorms 使用 FP32,另外模型权重文件复制一份 FP32 用于跟优化器更新梯度保持一致都是 FP32,另外还提供动态损失缩放(dynamic loss scaling)。使用了权重备份来减少舍入误差和使用损失缩放来避免数据溢出。

当然上面提供的策略是跟硬件有关系,并不是所有的 AI 加速芯片都使用,这时候针对自研的 AI 芯片,需要找到适合得到混合精度策略。

5、实验结果

从下图的 Accuracy 结果可以看到,混合精度基本没有精度损失:

v2-3f3cf64ec7c1777343fd74a3db9efa70_1440w.jpg

Loss scale 的效果:

v2-bef81be5ed6dc2766f1c757f70f3f70b_1440w.png

题外话,前不久去 X 公司跟 X 总监聊下一代 AI 芯片架构的时候,他认为下一代芯片可以不需要加入 INT8 数据类型,因为 Transformer 结构目前有大一统 NLP 和 CV 等领域的趋势,从设计、流片到量产,2 年后预计 Transformer 会取代 CNN 成为最流行的架构。我倒是不同意这个观点,目前来看神经网络的 4 个主要的结构 MLP、CNN、RNN、Transformer 都有其对应的使用场景,并没有因为某一种结构的出现而推翻以前的结构。只能说根据使用场景的侧重点比例有所不同,我理解 Int8、fp16、fp32 的数据类型在 AI 芯片中仍然会长期存在,针对不同的应用场景和计算单元会有不同的比例。

参考文献:

Micikevicius, Paulius, et al. "Mixed precision training."arXiv preprint arXiv:1710.03740(2017).

Ott, Myle, et al. "Scaling neural machine translation."arXiv preprint arXiv:1806.00187(2018).

https://en.wikipedia.org/wiki/Half-precision_floating-point_format

apex.amp - Apex 0.1.0 documentation.

Automatic Mixed Precision for Deep Learning.

Training With Mixed Precision.

Dreaming.O:浅谈混合精度训练.