对齐 NVIDIA BF16 算术模块的尝试

整形算术单元容易预测实现硬件行为，而浮点单元由于 （1）不遵守结合律（2）rounding 模式和特殊情况处理（subnormal、nan、-0、+inf、-inf） 往往更难预测硬件计算结果。神经网络中运算 MAC 运算累加超长数组同时涉及 (1) 和 (2) 问题，不满足交换律使得遍历保证 100% 和算法（GPU 结果）一致近乎不可能，且常见深度学习库中为了减少累加误差往往在更高精度进行累加，涉及层层 GPU 软件栈选项使得对齐更加困难^[1]；而 element-wise 运算只涉及数值计算，理论只最多需要遍历 2**32 种组合存在对齐可能性。

根据 CUDA 文档^[2]浮点运算特性可通过编译器选项控制，默认编译选项遵循 IEEE-754 round to nearest。GPU 算术支持配置选项如下：

• FTZ, Flush to Zero, 将 subnormal 舍入到 0，编译器选项控制
• Fast Div，快速近似除法，编译器选项控制
• Fast Sqrt，快速近似开方，编译器选项控制
• rounding mode ，代码层面控制

根据选项猜测 GPU 算术单元输入端口包含一个表示舍入模式，一个控制除法/开方迭代次数上限，而 FTZ 可能是在算术模块之外通过识别为 0 更改 sparsity 进而调度算术模块的输入数据。

虽说 IEEE-754 是一个参数化的定义，但毕竟 bf16 并不属于 IEEE-754 预定义的浮点格式，上周对 GPU element-wise bf16 浮点运算做了点小实验观察 GPU 算术行为，实验环境是 RTX Mobile 4060 + pytorch 2.2/CUDA 12.1。中科院开源fudian模块可以参数化生成浮点算术模块，但仅在 float32/double64 上进行测试满足 IEEE-754 标准^[3]。其中除法模块中间迭代位宽在 bf16 设置下会报错无法编译，乘法、加法模块可以正常编译通过。在 BF16/RNE 配置下遍历测试乘法计算 RTL 和 GPU 结果，得到结果十分有趣：

• 负 0 处理：Fudian 在 bf16 模式配置无法处理 -0 计算，-0 * 0 会输出 nan，-0 * inf 会输出 0 （因为浮点中的 0 可能是舍入后的 0 而非真 0，所以 IEEE-754 规定 0 乘 inf 应当是 inf），而 GPU 表现符合 IEEE-754 标准；
• NaN：即使同时 NaN 的输出，由于 NaN 有多种 bit 表示，二者数值也不相同。这部分在后文分析时假设 NaN 分析假设相同了，且 NaN 在测试集中肯定属于少数；
• 舍入：Fudian 部分计算和 GPU 能保持每一个bit 相同即 bit-accuracy，有意思的是保持一致的部分呈现高度的规律性，假设固定两个 bf16 输入的其中一个输入 A，遍历另一个输入 B，那么保持 bit-accuracy 的数值范围大致在 B 一个对称的区间，这个 B 的起点和大小随着 A 的值变化，大致范围在 40~60% 左右。也就是说只有大致一半的结果 Fudian 能和 GPU BF16 乘法保持一致。

具体挑了一个对不上的例子，结果如下：

A	B	GPU	Fudian
0x8001	0xe401	0x2181	0x2180

可见二者仅在 mantissa 相差 1 的舍入误差，遵循 IEEE-754 BF16 究竟是 0x2181 还是 0x2180 呢？A、B十进制数值分别是 -9.1835e-41 和 -9.51852e+20，GPU 计算结果是 8.74138e-20，Fudian 计算结果是 8.6736174e-20，而理论值应是 8.7413328420e-20。GPU 的计算结果与理论值误差更小，似乎此时 GPU 遵循了 IEEE-754 而 Fudian 和 IEEE-754 存在出入。

---

1. https://docs.pytorch.org/docs/stable/notes/numerical_accuracy.html ↩︎

2. https://docs.nvidia.com/cuda/floating-point/index.html ↩︎

3. https://github.com/OpenXiangShan/fudian ↩︎