跳转到主要内容
博客

如何加速 Intel® CPU 上的 PyTorch Geometric

作者： 2023 年 7 月 10 日2024 年 11 月 14 日暂无评论

概述

英特尔 PyTorch 团队一直与 PyTorch Geometric (PyG) 社区合作，为图神经网络 (GNN) 和 PyG 工作负载提供 CPU 性能优化。在 PyTorch 2.0 版本中，引入了几项关键优化，以提高 CPU 上的 GNN 训练和推理性能。开发人员和研究人员现在可以利用英特尔的 AI/ML 框架优化，显著加快模型训练和推理速度，从而能够直接使用 PyG 进行 GNN 工作流。

在本博客中，我们将深入探讨如何在使用 PyTorch 2.0 的旗舰功能 torch.compile 加速 PyG 模型的同时，优化 PyG 训练和推理性能。

消息传递范式

消息传递是指节点通过相互发送消息与其各自的邻居交换信息的過程。在 PyG 中，消息传递过程可以概括为三个步骤：

  1. Gather（收集）：收集相邻节点和边的边级信息。
  2. Apply（应用）：使用用户定义函数（UDF）更新收集到的信息。
  3. Scatter（分散）：聚合到节点级信息，例如，通过特定的归约函数（如求和、平均或最大值）。
Figure 1: The message passing paradigm

图 1：消息传递范式（来源：Matthias Fey）

消息传递性能与图的邻接矩阵的存储格式高度相关，邻接矩阵记录了节点对的连接方式。存储格式有两种方法：

  • COO（坐标格式）中的邻接矩阵：图数据物理存储为二维张量形状，为 [2, num_edges]，它映射源节点和目标节点的每个连接。性能热点是分散-归约（scatter-reduce）。
  • CSR（压缩稀疏行）中的邻接矩阵：与 COO 相似的格式，但在行索引上进行了压缩。这种格式允许更高效的行访问和更快的稀疏矩阵-矩阵乘法 (SpMM)。性能热点是稀疏矩阵相关的归约操作。

分散-归约

分散-归约模式本质上是并行的，它使用 `src` 张量中的值，通过 `index` 指定的条目来更新 `self` 张量的值。理想情况下，在外部维度上并行化将是最有效的。然而，直接并行化会导致写入冲突，因为不同的线程可能同时尝试更新相同的条目。

Figure 2: Scatter-reduce and its optimization scheme

图 2：分散-归约及其优化方案（来源：Mingfei Ma）

为了优化此内核，我们使用排序后归约的方法：

  • 排序：使用并行基数排序将 `index` 张量按升序排列，以便指向 `self` 张量中相同条目的索引由同一线程管理。
  • 归约：在 `self` 的外维度上并行，并对每个索引的 `src` 条目进行向量化归约。

在训练过程的反向路径（即收集）中，不需要排序，因为其内存访问模式不会导致任何写入冲突。

SpMM-归约

稀疏矩阵-矩阵归约是图神经网络中的一个基本操作，其中 **A** 是 CSR 格式的稀疏邻接矩阵，**B** 是密集特征矩阵，归约类型可以是 *sum*、*mean* 或 *max*。

Figure 3: SpMM optimization scheme

图 3：SpMM 优化方案（来源：Mingfei Ma）

优化此内核的最大挑战是如何在稀疏矩阵 **A** 的行上并行化时平衡线程负载。**A** 中的每一行对应一个节点，其连接数可能差异巨大；这导致线程负载不平衡。解决此类问题的一种技术是在线程分区之前进行负载扫描。除此之外，还引入了其他技术以进一步利用 CPU 性能，例如向量化、展开和分块。

这些优化通过 `torch.sparse.mm` 完成，并使用 `amax`、`amin`、`mean`、`sum` 等归约标志。

性能提升：高达 4.1 倍加速

我们收集了 pytorch_geometric/benchmark 和 Open Graph Benchmark (OGB) 中的推理和训练基准性能，以展示上述方法在英特尔® 至强® 铂金 8380 处理器上的性能提升。

模型 – 数据集选项加速比
GCN-Reddit（推理）512-2-64-dense1.22x
1024-3-128-dense1.25x
512-2-64-sparse1.31x
1024-3-128-sparse1.68x
GraphSage-ogbn-products（推理）1024-3-128-dense1.15x
512-2-64-sparse1.20x
1024-3-128-sparse1.33x
full-batch-sparse4.07x
GCN-PROTEINS（训练）3-321.67x
GCN-REDDIT-BINARY（训练）3-321.67x
GCN-Reddit（训练）512-2-64-dense1.20x
1024-3-128-dense1.12x

表 1：PyG 基准测试的性能加速1

从基准测试结果可以看出，我们在 PyTorch 和 PyG 中的优化在推理和训练方面实现了 **1.1x-4.1x 的加速**。

PyG 的 torch.compile

PyTorch 2.0 的旗舰功能 torch.compile 与 PyG 2.3 版本完全兼容，凭借适用于 CPU 的 TorchInductor C++/OpenMP 后端，它在 PyG 模型推理/训练中比命令式模式带来了额外的加速。特别地，在英特尔® 至强® 铂金 8380 处理器上，针对基本 GNN 模型的训练测量到了 **3.0x – 5.4x 的性能加速**2。

Figure 4: Performance Speedup with Torch Compile

图 4：使用 Torch Compile 获得的性能加速

Torch.compile 可以将消息传递的多个阶段融合成一个内核，这大大节省了内存带宽，从而显著提高了速度。有关更多支持，请参阅此pytorch geometric 教程。

请注意，PyG 中的 torch.compile 处于 Beta 模式，并正在积极开发中。目前，某些功能（例如 torch.compile(model, dynamic=True)）尚未无缝协同工作，但英特尔正在积极修复中。

结论与未来工作

在本博客中，我们介绍了 PyTorch 2.0 在 CPU 上包含的 GNN 性能优化。我们将与 PyG 社区紧密合作，进行未来的优化工作，重点将放在 torch.compile 的深入优化、稀疏优化以及分布式训练上。

致谢

本博客中呈现的成果是英特尔 PyTorch 团队和 Kumo 共同努力的结晶。特别感谢 Matthias Fey (Kumo)、Pearu Peterson (Quansight) 和 Christian Puhrsch (Meta) 付出的宝贵时间和提供的实质性帮助！我们共同在改进 PyTorch CPU 生态系统的道路上又迈进了一步。

参考文献

脚注

产品与性能信息

1Platinum 8380：1 节点，2x 英特尔至强铂金 8380 处理器，256GB (16 插槽/16GB/3200) DDR4 总内存，uCode 0xd000389，HT 开启，Turbo 开启，Ubuntu 20.04.5 LTS，5.4.0-146-generic，英特尔 SSDPE2KE016T8 1.5T；GCN + Reddit FP32 推理，GCN+Reddit FP32 训练，GraphSAGE + ogbn-products FP32 推理，GCN-PROTAIN，GCN-REDDIT-BINARY FP32 训练；软件：PyTorch 2.1.0.dev20230302+cpu，pytorch_geometric 2.3.0，torch-scatter 2.1.0，torch-sparse 0.6.16，英特尔于 2023 年 3 月 2 日测试。

2铂金 8380：1 节点，2x 英特尔至强铂金 8380 处理器，256GB (16 插槽/16GB/3200) 总 DDR4 内存，uCode 0xd000389，HT 开启，Turbo 开启，Ubuntu 20.04.5 LTS，5.4.0-146-generic，英特尔 SSDPE2KE016T8 1.5T；GCN、GraphSAGE、GIN 和 EdgeCNN，FP32；软件：PyTorch 2.1.0.dev20230411+cpu，pytorch_geometric 2.4.0，torch-scatter 2.1.1+pt20cpu，torch-sparse 0.6.17+pt20cpu，英特尔于 2023 年 4 月 11 日测试。

3性能因使用、配置和其他因素而异。详情请访问 www.Intel.com/PerformanceIndex。