博客

使用 Bfloat16 增强 PyTorch 在 Intel® Xeon® 可扩展处理器上的性能

概述

近年来，AI 模型日益增长的复杂性对硬件的计算能力提出了更高的要求。为了解决这一问题，人们提出了降低精度的数值格式。Bfloat16 是一种专为 AI 设计的 16 位浮点格式，包含 1 个符号位、8 个指数位和 7 个尾数位。由于与 float32 具有相同的动态范围，bfloat16 不需要特殊的处理（如损失缩放）。因此，在运行深度神经网络的推理和训练任务时，bfloat16 可以直接替代 float32。

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

英特尔和 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 的英特尔 CPU 上实现最优性能。其他算子（如张量算子和神经网络算子）则在 PyTorch 原生级别进行优化。我们将 bfloat16 内核级优化扩展到了大多数稠密张量算子上，同时适用于推理和训练（稀疏张量 bfloat16 支持将在未来工作中涵盖），具体包括：

  • Bfloat16 向量化：Bfloat16 存储为无符号 16 位整数，在进行加法、乘法等算术运算时需要将其转换为 float32。具体来说，每个 bfloat16 向量将被转换为两个 float32 向量，进行相应处理后再转换回 bfloat16。而对于 cat、copy 等非算术运算，则直接进行内存拷贝，不涉及数据类型转换。
  • Bfloat16 规约 (Reduction)：Bfloat16 数据上的规约操作使用 float32 作为累加类型以保证数值稳定性，例如 sum、BatchNorm2d、MaxPool2d 等。
  • 通道最后 (Channels Last) 优化：对于视觉模型，从性能角度看，Channels Last 是优于 Channels First 的内存格式。我们针对所有常用的计算机视觉 (CV) 模块，在 Channels Last 内存格式上实现了全面优化的 CPU 内核，同时兼顾了 float32 和 bfloat16。

使用自动混合精度 (AMP) 运行 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。激活值和参数的混合数据类型有助于在保持精度的同时提高性能。

性能提升

我们在英特尔® 至强® Platinum 8380H CPU @ 2.90GHz（代号 Cooper Lake）上对 TorchVision 模型进行了推理性能基准测试（单插槽单实例，batch size = 2 x 物理核心数）。结果显示，bfloat16 相比 float32 有 1.4 倍至 2.2 倍的性能提升。

bfloat16 相比 float32 的性能提升主要源于以下 3 个方面：

  • 计算密集型算子利用了新的 bfloat16 原生指令 VDPBF16PS，该指令使硬件计算吞吐量翻倍。
  • Bfloat16 的内存占用仅为 float32 的一半，因此理论上内存带宽密集型算子的速度可以提高一倍。
  • 在 Channels Last 格式下，我们有意为所有识别内存格式的算子保持相同的并行化方案（在 Channels First 下无法做到），这增加了将每一层的输出传递到下一层时的数据局部性。本质上，它使数据更靠近 CPU 核心，而数据无论如何都会驻留在缓存中。在这种情况下，由于内存占用较小，bfloat16 相比 float32 将具有更高的缓存命中率。

结论与未来工作

在本篇博客中，我们介绍了 PyTorch 1.12 中引入的 bfloat16 最新软件优化。在第三代英特尔® 至强® 可扩展处理器上的结果显示，bfloat16 在 TorchVision 模型上相比 float32 有 1.4 倍至 2.2 倍的性能提升。预计在支持 AMX 指令的下一代英特尔® 至强® 可扩展处理器上会有进一步的提升。尽管本篇博客的性能数据是基于 TorchVision 模型收集的，但其收益覆盖了所有拓扑结构。未来，我们将继续致力于将 bfloat16 优化扩展到更广的范围！

致谢

本博客中展示的结果是 Meta 和英特尔 PyTorch 团队共同努力的成果。特别感谢来自 Meta 的 Vitaly Fedyunin 和 Wei Wei，他们贡献了宝贵的时间并提供了巨大的帮助！我们共同朝着改善 PyTorch CPU 生态系统的目标又迈进了一步。

参考资料