评价此页
★ ★ ★ ★ ★
recipes/amx

利用 Intel® Advanced Matrix Extensions#

创建日期：2023 年 6 月 14 日 | 最后更新：2023 年 6 月 14 日 | 最后验证：未验证

简介#

Advanced Matrix Extensions (AMX)，也称为 Intel® Advanced Matrix Extensions (Intel® AMX)，是一种 x86 扩展，引入了两个新组件：一个称为“tiles”的二维寄存器文件和一个能够对这些 tiles 进行操作的 Tile Matrix Multiplication (TMUL) 加速器。AMX 旨在处理矩阵以加速 CPU 上的深度学习训练和推理，非常适合自然语言处理、推荐系统和图像识别等工作负载。

Intel 通过 4 代 Intel® Xeon® 可扩展处理器和 Intel® AMX 提升 AI 能力，与上一代相比，提供 3 倍至 10 倍的推理和训练性能，请参阅 使用 Intel® AMX 加速 AI 工作负载。与运行 Intel® Advanced Vector Extensions 512 Neural Network Instructions (Intel® AVX-512 VNNI) 的 3 代 Intel Xeon 可扩展处理器相比，运行 Intel AMX 的 4 代 Intel Xeon 可扩展处理器每周期可以执行 2,048 次 INT8 操作，而不是 256 次 INT8 操作。它们还可以每周期执行 1,024 次 BF16 操作，而之前是 64 次 FP32 操作，请参阅 使用 Intel® AMX 加速 AI 工作负载 的第 4 页。有关 AMX 的更详细信息，请参阅 Intel® AMX 概述。

PyTorch 中的 AMX#

PyTorch 利用 AMX 通过其 oneDNN 后端对 BFloat16 进行计算密集型运算符以及进行 INT8 量化，从而在支持 AMX 的 x86 CPU 上开箱即用地获得更高的性能。有关 oneDNN 的更详细信息，请参阅 oneDNN。

操作完全由 oneDNN 根据生成的执行代码路径处理。例如，当支持的操作在具有 AMX 支持的硬件平台上执行到 oneDNN 实现时，AMX 指令将在 oneDNN 内部自动调用。由于 oneDNN 是 PyTorch CPU 的默认加速库，因此无需手动操作即可启用 AMX 支持。

利用 AMX 进行工作负载的指南#

本节提供了如何利用 AMX 进行各种工作负载的指南。

  • BFloat16 数据类型

    • 使用 torch.cpu.amp 或 torch.autocast("cpu") 将为支持的运算符利用 AMX 加速。

    model = model.to(memory_format=torch.channels_last)
with torch.cpu.amp.autocast():
   output = model(input)

注意

使用 torch.channels_last 内存格式以获得更好的性能。

  • 量化

    • 应用量化将为支持的运算符利用 AMX 加速。

  • torch.compile

    • 当生成的图模型进入支持的运算符的 oneDNN 实现时，AMX 加速将被激活。

注意

当在支持 AMX 的 CPU 上使用 PyTorch 时，框架将默认自动启用 AMX。这意味着 PyTorch 将在可能的情况下尝试利用 AMX 功能来加速矩阵乘法运算。但是，值得注意的是，将运算分派到 AMX 内核的决定最终取决于 oneDNN 库和量化后端的内部优化策略，PyTorch 依赖它们来增强性能。PyTorch 和 oneDNN 库内部如何处理 AMX 利用的具体细节可能会随着框架的更新和改进而发生变化。

可以利用 AMX 的 CPU 运算符：#

可以利用 AMX 的 BF16 CPU 运算符

  • conv1d

  • conv2d

  • conv3d

  • conv_transpose1d

  • conv_transpose2d

  • conv_transpose3d

  • bmm

  • mm

  • baddbmm

  • addmm

  • addbmm

  • linear

  • matmul

可以利用 AMX 的量化 CPU 运算符

  • conv1d

  • conv2d

  • conv3d

  • conv_transpose1d

  • conv_transpose2d

  • conv_transpose3d

  • linear

确认 AMX 正在被使用#

设置环境变量 export ONEDNN_VERBOSE=1，或使用 torch.backends.mkldnn.verbose 来启用 oneDNN 转储详细消息。

with torch.backends.mkldnn.verbose(torch.backends.mkldnn.VERBOSE_ON):
    with torch.cpu.amp.autocast():
        model(input)

例如，获取 oneDNN 详细信息

onednn_verbose,info,oneDNN v2.7.3 (commit 6dbeffbae1f23cbbeae17adb7b5b13f1f37c080e)
onednn_verbose,info,cpu,runtime:OpenMP,nthr:128
onednn_verbose,info,cpu,isa:Intel AVX-512 with float16, Intel DL Boost and bfloat16 support and Intel AMX with bfloat16 and 8-bit integer support
onednn_verbose,info,gpu,runtime:none
onednn_verbose,info,prim_template:operation,engine,primitive,implementation,prop_kind,memory_descriptors,attributes,auxiliary,problem_desc,exec_time
onednn_verbose,exec,cpu,reorder,simple:any,undef,src_f32::blocked:a:f0 dst_f32::blocked:a:f0,attr-scratchpad:user ,,2,5.2561
...
onednn_verbose,exec,cpu,convolution,jit:avx512_core_amx_bf16,forward_training,src_bf16::blocked:acdb:f0 wei_bf16:p:blocked:ABcd16b16a2b:f0 bia_f32::blocked:a:f0 dst_bf16::blocked:acdb:f0,attr-scratchpad:user ,alg:convolution_direct,mb7_ic2oc1_ih224oh111kh3sh2dh1ph1_iw224ow111kw3sw2dw1pw1,0.628906
...
onednn_verbose,exec,cpu,matmul,brg:avx512_core_amx_int8,undef,src_s8::blocked:ab:f0 wei_s8:p:blocked:BA16a64b4a:f0 dst_s8::blocked:ab:f0,attr-scratchpad:user ,,1x30522:30522x768:1x768,7.66382
...

如果为 BFloat16 获取 avx512_core_amx_bf16 的详细信息，或为 INT8 量化获取 avx512_core_amx_int8 的详细信息，则表示 AMX 已激活。

结论#

在本教程中，我们简要介绍了 AMX，如何在 PyTorch 中利用 AMX 加速工作负载，以及如何确认 AMX 正在被使用。

随着 PyTorch 和 oneDNN 的改进和更新，AMX 的利用可能会随之变化。

一如既往，如果您遇到任何问题或有任何疑问，可以使用 论坛 或 GitHub issues 来联系我们。

在此页面上
PyTorch 库