使用 bfloat16 提高模型的性能

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

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

明确使用 bfloat16

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

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

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

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

格式转换详情

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

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

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