TensorFlow 中的 bfloat16 数据类型

火山方舟向量数据库大模型

picture.image


今天是农历八月十五,先祝各位新老朋友中秋快乐,阖家幸福~ 加班的童鞋代码无 bug 一遍就过~

picture.image

今天在看代码时遇到 bfloat16,于是深入研究了下,跟各位分享。

我们代码中常用 32 位单精度浮点格式(float32)表示权重和各层响应,一个值占据 4 个字节存储空间。

picture.image

如上图为一般 float(float64、float32、float16)位段表示。首位为符号位,占据 1 bit 空间。剩下的 bit 由指数位和尾数位构成,IEEE 754 std 规定了每种浮点格式指数、尾数位宽。

picture.image

一个 float32 的数值通过如下公式计算得到:

picture.image

其中指数部分减去 127 是为 bias。同样,将上述公式稍加修改(主要两处:尾数下面分母和 bias)就能得到 float64、float16 的计算方法。

从上表还看到半精度浮点格式(float16)占用存储空间为单精度的一半,在一些实验中显示使用半精度浮点格式对于训练神经网络几乎没有影响,因此越来越多硬件(Nvidia Tesla V100/P100,Nvidia TX2,ARMv8.2,……)开始原生支持使用 float16(又称 fp16、FP16)做计算。

由于 Exponent 位宽不同,float16 与 float32 转换时需要做 re-bias 运算,开销较大。

在 TensorFlow 中除了标准 float16 之外还创造了一种新的数据类型 bfloat16,同样只占用一半的存储空间,但与 float32 转换更方便,简单讲就是将 float32 的高 16 位截取下来即为 bfloat16。bfloat16 有 1 bit 符号位,8 bit 指数位,7 bit 尾数位,其动态范围与 float32 相当。

picture.image

bfloat16 与 float32 相互转换的函数声明位于 tensorflow/core/framework/bfloat16.h

void FloatToBFloat16(const float* src, bfloat16* dst, int64 size);

void BFloat16ToFloat(const bfloat16* src, float* dst, int64 size);

我们写一个简单的测试,将 1~65536 数值转为 bfloat16,然后与原数值对比,结果如下图所示:

picture.image

由于截断效应,原本连续的数值呈现了阶梯状,说明引入了截断误差。我们进一步查看误差数值变化规律。

picture.image

随着被表示的数值绝对值增加,误差也在逐步递增,在 32768~65535 区间内最大误差可达 255。

作为对照,我们将 float32 与 float16 也做如上实验,结果如下:

picture.image

picture.image

我们发现 float16 数值与 float32 数值误差更小,在 32768~65535区间内最大误差不超过 16,这是由于 float16 为尾数分配了更多 bit,可以表示更细粒度的数值。float16 为此付出的代价是数值范围受限,大于 65504 的数值直接变为 Inf,即上溢出。而 bfloat16 由于指数位宽与 float32 相同,可表示的范围更大,表示更粗粒度的数值。

在实际应用中需要根据模型精度需求和硬件特性,有针对性地选择数值类型,利用更少存储和网络带宽实现更高效的计算。

本公众号已开通赞赏功能,如果认为文章对自己有帮助,欢迎打赏~


picture.image

0
0
0
0
评论
未登录
看完啦,登录分享一下感受吧~
暂无评论