bfloat16 数值格式

使用降低精确率的浮点数是一种缩短收敛时间而不会损失准确率的常用方法。TPU 在执行矩阵运算时使用 bfloat16 数字格式。矩阵乘法运算是针对 bfloat16 值执行，而累积是针对 IEEE float32 值执行。

bfloat16 是用于机器学习的自定义 16 位浮点格式，由一个符号位、八个指数位和七个尾数位组成。下图显示了三种浮点格式的内部：float32：IEEE 单精度，float16：IEEE 半精度和 bfloat16。

图片

bfloat16 和 float32 的动态范围是等效的。不过，bfloat16 会占用一半的内存空间。如需详细了解 bfloat16 性能，请参阅《针对深度学习训练的 BFLOAT16 研究》。

选择 bfloat16

Google 硬件团队为 Cloud TPU 选择了 bfloat16，旨在提高硬件效率，同时保持准确训练深度学习模型的能力，并将 float32 的转换费用降至最低。硬件乘法器的物理大小根据尾数宽度的平方进行扩缩。如果尾数位比 FP16 少，则 bfloat16 乘法器的芯片大小是典型 FP16 乘法器的一半，并且比 float32 乘法器小 8 倍。

与尾数相比，神经网络对指数的大小更敏感。为确保下溢、上溢和 NaN 的行为相同，bfloat16 的指数大小与 float32 相同。bfloat16 处理非正规数的方式与 float32 不同，它会将它们刷新为零。与通常需要进行特殊处理（如损失扩缩）的 float16 不同，bfloat16 是在训练和运行深度神经网络时可以直接替代 float32。

混合精度训练

深度神经网络中的大多数计算都可以使用精度较低的值完成具有相同精确度的任务。某些模型甚至可以用更低的值实现更高的准确率。

在对 Cloud TPU 进行编程时，XLA 编译器会自动在 float32 和 bfloat16 之间转换值。

有关格式转换的详细信息

从 float32 到 bfloat16 的格式转换由 XLA 编译器自动插入。在 TPU 上，转换的舍入方案为四舍五入到最接近的偶数并溢出到 inf。此外，Cloud TPU 上的 bfloat16 不支持次正态数，因此所有超正态数在转换期间都会被刷新为零。转换中会保留特殊值，例如 NaN 和 inf。

XLA 编译器也会自动插入从 bfloat16 到 float32 的格式转换。由于 float32 可以表示 bfloat16 中的所有精确值，因此转换只会在尾数位中填充 16 个零。转化中将保留特殊值。

模型可移植性

模型中的参数值和激活可以以 32 位格式存储，因为 TPU 硬件可以将这些值自动转换为 bfloat16。从 Cloud TPU 上训练的模型获取的检查点可以部署到其他硬件平台（例如，在 CPU 或 GPU 上推断或微调），而无需进行大量手动转换。

使用 bfloat16 提高性能

虽然 TPU 中的自动格式转换可让您避免考虑数值精确率，但可以通过将值明确转换为 bfloat16 来实现进一步的性能提升。将值明确类型转换为 bfloat16 有两个原因：

  1. 以 bfloat16 格式存储值可节省片上内存，使 Cloud TPU 能够训练更大的模型或使用更大的批量大小。

  2. 某些操作受内存带宽限制，这意味着从内存加载数据所需的时间会减慢执行计算的总体时间。以 bfloat16 格式存储这些运算的操作数和输出可减少必须传输的数据量，从而提高整体速度。

如需开始使用，我们建议您通过一个已针对 Cloud TPU 进行优化的已启用 bfloat16 的参考模型亲自体验一下。之后，我们的性能指南、性能剖析工具指南和问题排查指南提供了深入的技术信息，可帮助您自行创建和优化机器学习模型。