由 马明飞（英特尔），维塔利·费季努宁（Meta），魏巍（Meta） 撰写

概述

近年来，人工智能模型日益复杂，对硬件的计算能力提出了更高的要求。为了解决这个问题，提出了降低精度的数值格式。Bfloat16 是一种针对人工智能的定制 16 位浮点格式，它包含一个符号位、八个指数位和七个尾数位。与 float32 具有相同的动态范围，bfloat16 在运行深度神经网络进行推理和训练时不需要特殊处理，如损失缩放。因此，bfloat16 可以直接替换 float32。

第三代英特尔® 至强® 可扩展处理器（代号 Cooper Lake），是首款支持原生 bfloat16 的通用 x86 CPU。在英特尔® 高级矢量扩展 512（英特尔® AVX-512）中引入了三个新的 bfloat16 指令：VCVTNE2PS2BF16、VCVTNEPS2BF16 和 VDPBF16PS。前两个指令执行从 float32 到 bfloat16 的转换，最后一个指令执行 bfloat16 对的点积。在 Cooper Lake 上，bfloat16 的理论计算吞吐量是 float32 的两倍。在下一代英特尔® 至强® 可扩展处理器中，将通过高级矩阵扩展（英特尔® AMX）指令集扩展进一步增强 bfloat16 的计算吞吐量。

英特尔和 Meta 之前合作，使 PyTorch 支持 bfloat16，相关工作在 Cooper Lake 发布期间的一篇早期博客中发布。在那篇博客中，我们介绍了对原生 bfloat16 支持的硬件进步，并展示了 DLRM、ResNet-50 和 ResNext-101-32x4d 在 bfloat16 上相对于 float32 的性能提升 1.4 倍至 1.6 倍。

在这篇博客中，我们将介绍 PyTorch 1.12 中 bfloat16 的最新软件增强，这将适用于更广泛的用户场景，并展示更高的性能提升。

Bfloat16 的本地级优化

在 PyTorch CPU bfloat16 路径上，计算密集型算子，例如卷积、线性运算和 bmm，使用 oneDNN（oneAPI 深度神经网络库）在支持 AVX512_BF16 或 AMX 的 Intel CPU 上实现最佳性能。其他算子，如张量算子和神经网络算子，在 PyTorch 本地级别进行优化。我们已经将 bfloat16 内核级优化扩展到大多数密集张量算子，适用于推理和训练（稀疏张量 bfloat16 支持将在未来的工作中涵盖），具体包括：

  • Bfloat16 向量化：Bfloat16 以无符号 16 位整数的形式存储，在进行加、乘等算术运算时需要将其转换为 float32。具体来说，每个 bfloat16 向量将被转换为两个 float32 向量，进行相应处理后再转换回来。而对于 cat、copy 等非算术运算，则是直接进行内存复制，不涉及数据类型转换。
  • Bfloat16 降维：对 bfloat16 数据的降维操作使用 float32 作为累加类型以保证数值稳定性，例如求和、BatchNorm2d、MaxPool2d 等。
  • 通道最后优化：从性能角度来看，对于视觉模型，通道最后（Channels Last）的内存格式比通道第一（Channels First）更优。我们已经为所有常用的 CV 模块在通道最后内存格式上实现了完全优化的 CPU 内核，同时考虑了 float32 和 bfloat16。

使用自动混合精度运行 Bfloat16

要在 bfloat16 上运行模型，通常用户可以选择显式地将数据和模型转换为 bfloat16，例如：

# with explicit conversion
input = input.to(dtype=torch.bfloat16)
model = model.to(dtype=torch.bfloat16)

或者使用 torch.amp（自动混合精度）包。autocast 实例作为上下文管理器或装饰器，允许脚本的部分区域以混合精度运行，例如：

# with AMP
with torch.autocast(device_type="cpu", dtype=torch.bfloat16):
    output = model(input)

通常，显式转换方法和 AMP 方法具有相似的性能。尽管如此，我们建议使用 AMP 运行 bfloat16 模型，因为：

  • 更好的用户体验和自动回退：如果您的脚本中包含没有 bfloat16 支持的运算符，autocast 会隐式地将它们转换回 float32，而显式转换的模型将给出运行时错误。

  • 激活和参数的混合数据类型：与将所有模型参数显式转换为 bfloat16 不同，AMP 模式将在混合数据类型下运行。具体来说，输入/输出将保持为 bfloat16，而参数（例如，权重/偏差）将保持为 float32。激活和参数的混合数据类型将有助于提高性能，同时保持准确性。

性能提升

我们在 Intel® Xeon® Platinum 8380H CPU @ 2.90GHz（代号 Cooper Lake）上对 TorchVision 模型的推理性能进行了基准测试，每个插槽实例一个（批大小=2 x 物理核心数）。结果显示，bfloat16 相对于 float32 具有 1.4 倍到 2.2 倍的性能提升。

bfloat16 相对于 float32 的性能提升主要来自三个方面：

  • 计算密集型算子利用了新的 bfloat16 原生指令 VDPBF16PS，将硬件计算吞吐量翻倍。
  • Bfloat16 的内存占用只有 float32 的一半，因此理论上内存带宽密集型的操作将快两倍。
  • 在通道最后，我们故意为所有内存格式感知操作保持相同的并行化方案（在通道第一中无法这样做），这增加了将每一层的输出传递到下一层时的数据局部性。基本上，它将数据保持得更接近 CPU 核心，而数据本来就会驻留在缓存中。在这种情况下，由于内存占用更小，bfloat16 的缓存命中率将比 float32 更高。

结论与未来工作

在这篇博客中，我们介绍了 PyTorch 1.12 中引入的 bfloat16 的最新软件优化。在第三代英特尔 Xeon 可扩展处理器上的结果表明，bfloat16 在 TorchVision 模型上比 float32 有 1.4 倍到 2.2 倍的性能提升。预计在支持 AMX 指令的下一代英特尔 Xeon 可扩展处理器上会有进一步的改进。尽管本博客的性能数据是通过 TorchVision 模型收集的，但这一好处适用于所有拓扑结构。我们还将继续扩大 bfloat16 优化工作的范围，在未来涵盖更广泛的领域！

致谢

本博客中展示的结果是 Meta 和 Intel PyTorch 团队共同努力的结果。特别感谢 Meta 的 Vitaly Fedyunin 和 Wei Wei，他们花费宝贵的时间并给予了实质性的帮助！我们一起在改善 PyTorch CPU 生态系统之路上迈出了新的一步！

参考文献