作者：Mingfei Ma (Intel), Vitaly Fedyunin (Meta), Wei Wei (Meta)

概述

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

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

英特尔 (Intel) 和 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 作为累积类型，以保证数值稳定性，例如求和 (sum)、批量归一化 2D (BatchNorm2d)、最大池化 2D (MaxPool2d) 等。
  • Channels Last 优化：对于视觉模型，从性能角度来看，Channels Last 是比 Channels First 更优选的内存格式。我们已为所有常用 CV 模块在 Channels Last 内存格式上实现了完全优化的 CPU 内核，同时支持 float32 和 bfloat16。

使用自动混合精度 (Auto Mixed Precision) 运行 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 的性能提升主要来自以下 3 个方面

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

结论与未来工作

在这篇博客中，我们介绍了 PyTorch 1.12 中引入的关于 bfloat16 的最新软件优化。在第 3 代 Intel® Xeon® 可扩展处理器上的结果显示，bfloat16 在 TorchVision 模型上相对于 float32 具有 1.4 倍至 2.2 倍的性能提升。预计在下一代支持 AMX 指令的 Intel® Xeon® 可扩展处理器上，性能将进一步提升。尽管这篇博客的性能数据是使用 TorchVision 模型收集的，但其优势适用于所有拓扑结构。未来，我们将继续将 bfloat16 的优化工作扩展到更广泛的范围！

致谢

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

参考资料