TF32和AMP训练为何可以保证训练精度收敛

目标

自NV的Ampere架构后，越来越多的AI研究者使用TF32或者AMP（混合精度）进行网络训练。在网络训练精度的同时，也极大的提升了训练速度。（其中更多的应用者选择TF32）

相对于原先深度学习领域使用FLOAT32进行训练，明明训练中降低了数据表示精度，为什么还能取得满意的训练和推理结果呢？本文主要基于这个基本问题展开讨论。

下面先简单看一下A100支持的浮点运算单元和其对应算力。

基本数据类型介绍：

各数据类型算力介绍：

问题解答

如果想要解答上面这个问题，最主要还是要探究：1）TF32或者AMP相对与FLOAT32训练到底有什么区别？ 2）训练过程中对于数据到底有什么要的要求？

下面我们通过对比TF32和FLOAT32训练来解答第一个问题。TF32和FLOAT32训练过程中，更多地集中数据类型上差异，而非机制上的差异。（相对AMP是通过机制上补强来解决数据类型上的短板。）

1 基本知识IEEE754

IEEE754描述了4中类型浮点数，即单精度（float32），双精度（double），延伸单精度（43bit以上），延伸双精度（73bit位以上）。

1.1 单精度的内存存储格式

这里主要看一下单精度float32。其包含1bit符号位，8bit指数位和23bit小数位。

 通过名称我们也能基本了解，浮点数的二进制存储其实是一种科学计数法的体现。

符号位为0表示正数，为1表示负数；

指数位存在正负，故其表示范围为：-127到128，但是实际存储中通过偏置，转化为正数存储。

（为了简化计算，正数在计算机中按照原码存储，而负数则是补码存储，故方便浮点数运算，向右偏移127，即加上127进行存储，将指数的符号位规避。）

小数位用来存储科学计数法的数值部分（其从十进制转化为二进制可以理解为一种级数的处理），其表示范围为：  (其中规格数存在，而非规格书不存在)

从上面的计算公式看，浮点的小数部分计算的复杂度相对更高。

1.2 二进制和十进制的转换计算公式

公式：

正如公式所示： 指数位的固定偏置是127，也就是，其中n为指数位bit位数。

 举例1：11.25转化位二进制存储

第一步： 11.25转换为二进制位 1011.01，使用科学计数法计数为：

第二步：符号位存储 0， 小数位为： 3+127=130（二进制表示：1000 0010），符号位为：1.01101, (二进制存储： 0110 1000 0000 0000 0000 000)

1.3 规格数（normal number）和非规格数（denormal number）

IEEE754中定义规格数和非规格数，避免数据下溢。其中非规格数可以理解为 指数位 全部为0。

通过如下图了解基本的数据取值范围：

（摘自：从科学记数法到浮点数标准IEEE 754）

放在坐标轴看：

（摘自：从科学记数法到浮点数标准IEEE 754）

浮点数表示范围脑图：

2 TF32为什么可以保证训练精度

TF32相对FLOAT32，其只是缩减了数据的精度，相对保证了数据的取值范围。那么为什么可以保证训练精度呢？

从NV的AMP官方文档和论文中，可以看到训练中grad数据的取值范围对于训练的影响更大。

数据类型	数据表示范围	数据标识精度
FLOAT32 (1位符号、8位指数和23位尾数)	(+或-)1.4e-45 ~  (+或-)3.4e+38	可以近似理解为7位
TF32 (1位符号、8位指数和10位尾数)	(+或-)1.4e-45 ~  (+或-)3.4e+38	可以近似理解为7位
FLOAT16 (1位符号、5位指数和10位尾数)	(+或-)65504 （ ） ~  (+或-)  5.96e+8	可以近似理解为4位

关于数据精度可以参看： IEEE754详解（最详细简单有趣味的介绍）_明月几时有666的博客-CSDN博客_ieee754

案例：（摘自https://arxiv.org/pdf/1710.03740.pdf）

如下图是通过SSD网络查看grad数据取值范围，如果单纯地使用FLOAT16进行运算，会导致很多grad被抹成0，导致训练不收敛。

3 训练中是如何使用TF32数据类型

目前Pytorch框架中涉及TF32训练的算子主要如下：`nn.Linear`, `nn.Conv*`, cdist, tensordot, affine grid and grid sample, adaptive log softmax, GRU and LSTM。

其在具体运算中是由硬件指令支持的。输入为FLOAT32，硬件指令读取输入数据，通过舍入方法将其转换为TF32进行矩阵乘运算，部分和在通过FLOAT32进行累加。此过程中只有在矩阵乘时对数据进行了舍入处理，存在部分精度丢失。

4 AMP训练

待补充

引用:

从科学记数法到浮点数标准IEEE 754

IEEE 754格式是什么? - 知乎

https://arxiv.org/pdf/1710.03740.pdf