由英特尔提供

概述

英特尔 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. 收集：收集相邻节点和边的边缘级信息。
  2. 应用：使用用户定义函数（UDFs）更新收集到的信息。
  3. 分散：聚合到节点级信息，例如，通过特定的聚合函数，如总和、平均值或最大值。

Figure 1: The message passing paradigm

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

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

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

Scatter-Reduce

散射-减少的模式本质上是并行的，它使用索引指定的 src 张量的值来更新自张量的值。理想情况下，在外部维度上并行化将是最高效的。然而，直接并行化会导致写冲突，因为不同的线程可能会同时尝试更新相同的条目。

Figure 2: Scatter-reduce and its optimization scheme

图 2：散射-减少及其优化方案（来源：马明飞）

为了优化这个内核，我们使用排序和减少：

  • 排序：使用并行基数排序按升序排序索引张量，使得指向自张量中相同条目的索引由同一线程管理。
  • 减少操作：在自身的外部维度并行化，并对每个索引的源条目进行向量化减少。

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

SpMM-Reduce

稀疏矩阵-矩阵减少是 GNN 中的一个基本算子，其中 A 是 CSR 格式的稀疏邻接矩阵，B 是密集特征矩阵，减少类型可以是求和、平均值或最大值。

Figure 3: SpMM optimization scheme

图 3：SpMM 优化方案（来源：马明飞）

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

这些优化是通过 torch.sparse.mm 使用 amax、amin、mean、sum 的 reduce 标志来完成的。

性能提升：最高可达 4.1 倍加速

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

模型 – 数据集 选项 加速比
GCN-Reddit（推理） 512-2-64-密集 1.22 倍
1024-3-128-密集 1.25 倍
512-2-64-稀疏 1.31 倍
1024-3-128-稀疏 1.68 倍
GraphSage-ogbn-products（推理） 1024-3-128-密集 1.15 倍
512-2-64-稀疏 1.20 倍
1024-3-128-稀疏 1.33 倍
全批稀疏 4.07 倍
GCN-蛋白质（训练） 3-32 1.67 倍
GCN-REDDIT-BINARY (训练) 3-32 1.67 倍
GCN-Reddit (训练) 512-2-64-密集 1.20 倍
1024-3-128-密集 1.12 倍

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

从基准测试结果来看，我们在 PyTorch 和 PyG 中的优化实现了推理和训练的 1.1 倍-4.1 倍速度提升。

PyG 的 torch.compile

PyTorch2.0 的旗舰功能 torch.compile 与 PyG 2.3 版本完全兼容，得益于 TorchInductor C++/OpenMP 后端对 CPU 的支持，在 PyG 模型推理/训练中带来了额外的速度提升。特别是，在基本 GNN 模型上，使用 Intel Xeon Platinum 8380 处理器进行模型训练时，性能速度提升了 3.0 倍-5.4 倍。

Figure 4: Performance Speedup with Torch Compile

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

Torch.compile 可以将消息传递的多个阶段融合成一个内核，从而由于节省内存带宽而提供显著的加速。请参阅这篇 PyTorch 几何教程以获取更多支持。

请注意，PyG 中的 torch.compile 处于测试版，并且正在积极开发中。目前，一些功能还不能无缝协同工作，例如 torch.compile(model, dynamic=True)，但来自 Intel 的修复正在路上。

结论与未来工作

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

致谢

本博客中展示的结果是英特尔 PyTorch 团队和 Kumo 的共同努力。特别感谢 Matthias Fey（Kumo）、Pearu Peterson（Quansight）和 Christian Puhrsch（Meta），他们花费宝贵的时间并给予了大量帮助！我们一起在改善 PyTorch CPU 生态系统的道路上迈出了新的一步。

参考文献列表

脚注

产品和性能信息

1 Platinum 8380：1 节点，2x Intel Xeon Platinum 8380 处理器，总计 256GB（16 插槽/ 16GB/3200）DDR4 内存，uCode 0xd000389，HT 开启，Turbo 开启，Ubuntu 20.04.5 LTS，5.4.0-146-generic，INTEL 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，由 Intel 于 2023 年 3 月 2 日测试。

2 铂金 8380：1 节点，2x 英特尔至强铂金 8380 处理器，总计 256GB（16 插槽/16GB/3200MHz）DDR4 内存，uCode 0xd000389，超线程开启，睿频开启，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。