使用 bfloat16 提升模型的性能

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

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

明确使用 bfloat16

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

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

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

要开始使用,我们建议您先体验一下 Cloud TPU 参考模型。更新后 性能剖析工具指南 问题排查指南 深入的技术信息,可帮助您打造和优化机器学习 构建自己的机器学习模型。

格式转换详情

float32bfloat16 的格式转换由 XLA 编译器。在 TPU 上,转换中的舍入方案为 四舍五入为最接近的偶数 并溢出至 inf。此外,Cloud TPU 上的 bfloat16 不支持小数,因此所有小数都会在转换期间清零。特殊值(如 NaNinf)在转换中会保留。

XLA 编译器还会自动插入从 bfloat16float32 的格式转换。由于 float32 可以表示 bfloat16 中的所有确切值, 转换会在尾数位中填充 16 个零。特殊值包括 都会保留下来

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