使用 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)在转换中会保留。

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

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