使用 bfloat16 提高模型性能

默认情况下，TPU 使用 bfloat16 值执行矩阵乘法运算，使用 IEEE float32 值执行累加运算。使用降低精确度的浮点数可以缩短收敛时间，而不会损失准确性。

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

显式使用 bfloat16

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

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

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

首先，我们建议您使用一个 Cloud TPU 参考模型获得一些经验。之后，性能剖析工具指南和问题排查指南提供了深入的技术信息，可帮助您自行创建和优化机器学习模型。

格式转换详情

XLA 编译器会自动插入从 float32 到 bfloat16 的格式转换。在 TPU 上，转换中的舍入方案为向最近的偶数舍入，并溢出到 inf。此外，Cloud TPU 上的 bfloat16 不支持小数，因此所有小数都会在转换期间清零。特殊值（例如 NaN 和 inf）会在转换中保留。

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

从 Cloud TPU 上训练的模型获取的检查点可以部署到其他硬件平台（例如，在 CPU 或 GPU 上推断或微调），而无需进行大量手动转换。