在现代 CPU 上扩展类 BERT 模型推理 - 第 2 部分

TL;DR

Hugging Face 已经证明，优化软件组件——特别是内存分配器、并行化库以及使用贝叶斯优化——可以显著增强类 BERT 模型在 Intel Ice Lake Xeon CPU 上的推理性能。这些优化结合了 AVX-512 和 VNNI 等硬件特性，使得在基于 CPU 的基础设施上更高效地扩展 NLP 工作负载。

利用 Intel 软件实现 AI 高效性

为了最大限度地提高 Intel 硬件上 AI 工作负载的性能，开发人员可以利用 Intel oneAPI 框架下的软件优化栈。这些工具旨在弥合 PyTorch 和 TensorFlow 等高级框架与底层 CPU 架构之间的差距。

关键软件组件

  • Intel oneMKL (Math Kernel Library): 提供高效的线性代数例程。
  • Intel OpenMP 和 oneTBB (Threading Building Blocks): 用于计算的高级并行化框架。
  • Intel oneDNN: 一个深度神经网络原语库（例如 ReLU、全连接层），已原生集成到 PyTorch 和 TensorFlow 中（自 2.5.0 版本起）。
  • Intel PyTorch Extension (IPEX): 一个专门的框架，在优化方案被合并到 PyTorch 主库之前，充当优化的实验室。

CPU 推理的性能调优参数

推理性能主要由三个因素驱动：内存中的数据表示、数学算子的实现以及并行化的效率。

内存分配与管理

内存分配（向 OS 请求动态内存的过程）会影响速度和碎片化。虽然像 glibc 这样的默认分配器是通用型的，但专门的分配器可以减少多线程深度学习工作负载中的同步开销。

  • tcmalloc (Google): 通过为每个线程维护本地内存段，减少全局关键路径，通常在各种工作负载中提供最佳性能。
  • jemalloc (Facebook): 在特定的低并发情况下可能最快。
  • mimalloc (Microsoft): 另一种改进内存管理的替代方案。

计算的并行化

高效利用多个 CPU 核心需要不仅仅是增加核心数。CPU 缓存失效和并发数据访问等因素可能会阻碍扩展性。 Hugging Face 建议在 Intel 硬件上使用 Intel 实现的 OpenMP 规范 ("IOMP")，以优化线程调度和资源绑定。

优化的数学算子

现代 CPU 使用 SIMD (Single Instruction Multiple Data) 指令来在每个时钟周期内对多个项目进行操作。Intel CPU 支持 SSE2, AVX, AVX2, 和 AVX-512。像 Intel MKL 和 oneDNN 这样的库高效地实现了这些算子，从而为常见的模式（如 Linear + ReLU 或 Convolution）实现性能加速。

在 Intel Ice Lake Xeon CPU 上进行基准测试

基准测试是在 Ubuntu 20.04.2 LTS 上使用 Intel Ice Lake Xeon Platinum 8380 CPU 进行的，测试了 PyTorch 1.9.0 和 TensorFlow 2.5.0 在各种批量大小 (1 到 128) 和序列长度 (8 到 512) 下的表现。

Eager 模式 vs. 图模式

  • Eager 模式 (PyTorch, TensorFlow): 计算图是在执行期间发现的。这提供了灵活性，但引入了运行时开销，并使得算子融合（例如 Convolution + ReLU）变得更加困难。
  • Graph 模式 (TorchScript, TensorFlow Graph, Intel TensorFlow): 图是预先已知的，允许进行剪枝、算子融合和预先规划的内存分配。

扩展性关键发现

  • 核心扩展性: 增加核心数通常会降低延迟，但并非单调递增。在工作负载大小和分配的资源之间存在权衡。
  • 跨插槽开销 (Inter-Socket Overhead): 在具有多个 CPU（多个插槽）的系统上使用所有核心，由于跨插槽通信，通常会引入显著的延迟开销。
  • 分配器影响: 在需要动态管理内存的 Eager 模式下，分配器（例如 tcmalloc）的选择对性能的影响比在可以预先保留资源的 Graph 模式下更为显著。

使用 Intel SigOpt 进行自动性能调优

由于寻找最优设置的搜索空间非常庞大（结合了核心数、内存分配器、并行化库、Transparent Huge Pages 和 KMP block time），暴力搜索调优是低效的的。 Hugging Face 利用了 Intel SigOpt，它使用贝叶斯优化来更快地找到接近最优的配置。

SigOpt 结果

  • 效率: SigOpt 提供的性能非常接近暴力搜索的结果，最大的差距仅为 8.6%。
  • 参数重要性: 核心数始终是最关键的参数。然而，对于较大的序列长度（例如 512），并行化库的选择（OpenMP vs. Intel OpenMP）相对于内存分配器变得更加重要。
  • 资源优化: SigOpt 识别出使用较少的核心（例如 16 个核心）有时可以产生最佳延迟，从而允许并行运行多个模型实例以提高整体吞吐量。

结论

在 Intel Ice Lake Xeon CPU 上为生产环境优化 BERT-like 模型需要一种分层方法：从底层硬件特性开始，逐步过渡到框架特定的优化（如 oneDNN），最后调优系统级参数，如内存分配器和 OpenMP 实现。这些工作已集成到 Hugging Face Optimum 库和 Infinity 产品中，以提供高效的容器化推理解决方案。

Sources