「Pytorch」BF16 Mixed Precision Training

在深度学习领域，神经网络的训练性能瓶颈常常出现在 GPU显存的使用上。主要表现为两方面：

1. 单卡上可容纳的模型和数据量有限；
2. 显存与计算单元之间的带宽和延迟限制了运算速度；

为了解决显卡瓶颈的问题，涌现了不同的解决方法。

模型参数量估计

为了更好地估算模型所要占用的显存，首先需要分析模型训练过程中有哪些部分需要消耗存储空间。在 “ZeRO: Memory Optimizations Toward Training Trillion Parameter Model” 中提出，模型在训练时，主要有两大部分的空间占用。

1. 对于大模型来说，主要的空间占用是模型状态，包括优化器状态（eg：Adam优化器的动量和方差）、模型参数和模型参数的梯度；
2. 剩余的空间主要被模型训练中间激活值、临时缓冲区和不可用的内存碎片占用，统称为剩余状态。

以 3.5B 大模型为例。①35亿个参数，如果使用 FP16进行存储的话，即 70亿个字节，约7GB左右；②前向传播的激活值和反向传播的梯度大小跟模型参数保持一致，约 7GB；③以 Adam优化器为例，包括三部分，分别为 FP32格式模型参数的备份，FP32的动量和方差，加起来约28GB；因此，从理论上，要微调此模型的话，至少需要 49GB的空间。

1、BF16 半精度浮点数

双精度浮点 Float64；单浮点精度 Float32；半浮点精度 Float16，被广泛应用于model推理
为了进一步降低计算量和显存占用，可以考虑整数int4和int8量化推理

之前深度学习模型的训练通常都采用 Float32(FP32)的精度，而作者发现，使用较低的精度来进行模型的训练也是可行的，并且能够显著提升速度。通过采用混合精度训练，一般可以实现 2~3 倍的速度提升，极大的优化了模型的训练流程。

• FP32 整体长度为4个字节，即32位，其中有8位的指数位宽，23位的尾数精度和1位的符号位，能够表示的数值范围是 $1\times 2^{-126}\sim (2-ϵ)\times 2^{127}$ ；
• 在一些不太需要高精度计算的应用中，eg：图像处理和神经网络中，32位的空间其实有一些浪费，因此又出现了新的数据类型，半精度浮点数 FP16，使用16位（2个字节）来存储浮点值，有5位的指数位宽，10位的尾数精度和1位的符号位，能够表示的数值范围是 $1\times 2^{-14}\sim (2-ϵ)\times 2^{15}$；

格式	位数/位	指数位宽/位	尾数精度/位	符号位/位	数值范围
FP32	32	8	23	1	$1\times 2^{-126}\sim (2-ϵ)\times 2^{127}$
FP16	16	5	10	1	$1\times 2^{-14}\sim (2-ϵ)\times 2^{15}$
BP32	16	8	7	1	$1\times 2^{-126}\sim (2-ϵ)\times 2^{127}$

混合精度训练，即在模型训练时同时采用 FP32 和 FP16 两种精度。在实践过程中，研究人员发现在大语言模型的训练中直接使用 FP16会有一些问题，在训练过程中 loss 会非常不稳定，因此使用 FP16 训练大模型非常困难。问题在于 FP16的指数位宽只有 5位，能表示的最大整数为 65504，一旦权重超过这个值就会发生溢出，因此只能进行较小数的乘法，eg：可以计算 $250\times 250=62500$，但如果计算 $255\times 255=65025$ 就会溢出，这是导致训练出现问题的主要原因。这也意味着模型权重必须保持很小。一种成为损失缩放的技术可以缓解这个问题，但是当模型变得非常大时，FP16 较小的数值范围依旧是一个问题。

• 为了更好地解决 FP16的问题，谷歌开发了一种新的浮点数格式 BF16（Brain Floating Point, 2个字节），用于降低存储需求，提高机器学习算法的计算速度。BF16 的指数位宽为8位，于 FP32相同，尾数精度采用7位，因此当使用 BF16时，精度非常差。然而，在训练模型时一般采用随机梯度下降法及其变体，其过程像蹒跚而行，即使某一步没有找到最优方向也没关系，模型会在后续调整纠正。

将模型参数类型从 FP16换为 BF16，训练的大模型 loss值的下降也会变得更加稳定。

这种低精度和 混合精度训练的方法逐渐被广泛接受和应用，深度学习框架、GPU以及 神经网络加速器的设计也因此受到了深渊的影响。可以说，混合精度训练的提出，对深度学习领域起到了关键的推动作用，有效地解决了 GPU显存的使用问题，提升了模型训练的效率。

2、混合精度训练

paper：Mixed Precision Training

1. 维护一个权重的单精度副本，在每个优化器步骤后累计梯度(对于前向和反向传播，此副本四舍五入到半精度)；
2. 提出了损失缩放来保持小幅度的梯度值；
3. 使用半精度算法，该算法累积为单精度输出，在存储到内存之前将其转化为半精度；

FP32 为主副本权重，

在混合精度训练时，权重、激活函数、梯度被保存为 FP16，为了与 FP32模型的精度相匹配，在optimizer step时，维持 FP32权重为主线，并使用权重梯度进行更新。在每次迭代时，主权重的 FP16副本用于前向和反向传播，将 FP32训练所需的存储和带宽减半，如上图所示。

虽然对 FP32住权重的需求并不普遍，但许多模型还需要 FP32的两个可能原因是：

1. 权重更新变得太小，无法在 FP16中表示，任何大小小于 $2^{-24}$ 的值在 FP16中都将变为 零，当与学习率相乘时，这些小值梯度在优化器中都会变为零，并对模型的准确性产生不利影响。使用单精度进行更新可以解决这一问题；
2. 权重值 与 权重更新的比例非常大。在这种情况下，即使权重更新可以在 FP16中表示，当加法操作将其右移以使二进制点与权重对齐时，它仍然可能变为零。当归一化权重值的幅度比权重更新的新幅度大 至少2048倍，就会发生这种情况。由于 FP16有10位尾数，隐式位必须右移11位或更多位置，才能潜在地创建一个零。在比例大于2048时，隐式比特将右移12位或更多位。这将导致权重更新变得无法恢复的零。更大的比例将导致非标注化数字的效果。同样，可以通过计算 FP32中的更新来抵消这种影响。

图2-a所示，在 FP16前后传播更新 FP32权重主线时，匹配FP32训练结果，而更新FP16权重会导致 80%的相对精度损失。

由于更大的 batch-size 和每层的激活被保存以在反向传播过程中重复使用，因此训练内存消耗主要由激活决定。由于激活也以半精度格式存储，因此训练深度神经网络的整体内存消耗大约减半。

2.1 损失缩放

FP16指数偏差将归一化指数的范围集中到 [-14,15]，而实践中的梯度值往往由小幅度(负指数)主导，如图3所示，显示了Multibox SSD模型的 FP32训练期间在所有层上的急活梯度值的直方图，FP16 可表示范围的大部分未必使用，而许多值低于最小可表示范围 变为指数为0。放大梯度将使它们占据更多的可表示范围，并保留否则会丢失为0的值。当梯度未被缩放时，这个特定的网络会发散，但将其压缩8倍(指数为3)就足以匹配 FP32训练所达到的精度。这说明激活

2.2运算精度

神经网络模型分为三类：vector dot-products, reductions, point-wise operations。当涉及到降低精度的算法时，这些类别受益于不同的处理，为了保持模型的准确性，发现一些模型要求 FP16矢量点积将部分累加成 FP32，在写入内存之前将其转换为 FP16。如果 FP32中没有这种积累，一些 FP16模型与基线模型的精度不匹配。

之前的GPU只支持 FP16 乘/加法运算，而 NVIDIA Volta GPU引入了 Tensor Core，可以将 FP16输入矩阵相乘，并将乘积累加到 FP16或 FP32输出中。

FP32中应进行大幅缩减(向量各元素之和)。在累积统计数据和 softmax层时，这种建好主要出现在 batch-normalization 层中。在两种层类型种，仍然从内存中读取和写入FP16张量，在 FP32中执行算术运算，这并没有减缓训练过程，因为这些层的内存带宽有限，对算术速度不敏感。

逐点操作，如非线性和逐像素矩阵运算，是内存带宽有限，由于算术精度不影响这些运算的速度，因此可以使用FP16 或 FP32.