使用 bfloat16 提升模型的性能
默认情况下,TPU 执行
使用 bfloat16
进行矩阵乘法运算
通过 IEEE 计算float32
值。使用降低精确度的浮点数可以缩短收敛时间,而不会损失准确性。
bfloat16
和 float32
的动态范围是等效的。不过,bfloat16
使用一半的内存空间如需详细了解 bfloat16
性能,请参阅对深度学习训练的 BFLOAT16 进行的一项研究。
明确使用 bfloat16
虽然 TPU 中的自动格式转换可让您避免考虑数值精确率,但您可以通过将值明确转换为 bfloat16
来实现性能提升。将值明确类型转换为 bfloat16
有两个原因:
以
bfloat16
格式存储值可节省片上内存,使 Cloud TPU 能够训练更大的模型或使用更大的批量大小。某些操作受内存带宽限制,这意味着从内存加载数据所需的时间会减慢执行计算的总体时间。以
bfloat16
格式存储这些运算的操作数和输出可减少必须传输的数据量,从而提高整体速度。
要开始使用,我们建议您先体验一下 Cloud TPU 参考模型。更新后 性能剖析工具指南 问题排查指南 深入的技术信息,可帮助您打造和优化机器学习 构建自己的机器学习模型。
格式转换详情
从 float32
到 bfloat16
的格式转换由
XLA 编译器。在 TPU 上,转换中的舍入方案为
四舍五入为最接近的偶数
并溢出至 inf
。此外,Cloud TPU 上的 bfloat16
不支持小数,因此所有小数都会在转换期间清零。特殊值(如 NaN
和 inf
)在转换中会保留。
XLA 编译器还会自动插入从 bfloat16
到 float32
的格式转换。由于 float32
可以表示 bfloat16
中的所有确切值,
转换会在尾数位中填充 16 个零。特殊值包括
都会保留下来
从 Cloud TPU 上训练的模型获取的检查点可以部署到其他硬件平台(例如,在 CPU 或 GPU 上推断或微调),而无需进行大量手动转换。