概述
近年来，AI模型日益增长的复杂性对硬件提出了越来越高的计算能力要求。为了解决这个问题，人们提出了降低精度数字格式。Bfloat16是一种定制的16位浮点格式，专为AI设计，由一个符号位、八个指数位和七个尾数位组成。Bfloat16具有与float32相同的动态范围，因此无需特殊的处理，例如损失缩放。因此，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的Intel CPU上实现最佳性能。其他运算符，例如张量运算符和神经网络运算符，在PyTorch原生级别进行优化。我们已将bfloat16内核级优化扩展到稠密张量的大多数运算符，适用于推理和训练（稀疏张量bfloat16支持将在未来的工作中涵盖），具体包括：
- Bfloat16矢量化：Bfloat16以无符号16位整数存储，这要求它在进行加法、乘法等算术运算时转换为float32。具体来说，每个bfloat16向量将被转换为两个float32向量，进行相应的处理，然后再转换回来。而对于非算术运算，如cat、copy等，则是直接进行内存复制，不涉及数据类型转换。
- Bfloat16归约：Bfloat16数据的归约使用float32作为累加类型，以保证数值稳定性，例如求和、BatchNorm2d、MaxPool2d等。
- Channels Last优化：对于视觉模型，从性能角度来看，Channels Last是优于Channels First的首选内存格式。我们为所有常用的CV模块在Channels Last内存格式下实现了完全优化的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模型的推理性能进行了基准测试，每个socket单实例（批处理大小 = 2 x 物理核心数）。结果显示bfloat16比float32有1.4倍到2.2倍的性能提升。
bfloat16相对于float32的性能提升主要来自以下三个方面：
- 计算密集型运算符利用了新的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和Intel PyTorch团队的共同努力。特别感谢来自Meta的Vitaly Fedyunin和Wei Wei，他们付出了宝贵的时间并提供了实质性的帮助！我们共同为改善PyTorch CPU生态系统又迈出了一步。