作者：英特尔

概述

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

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

消息传递范式

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

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

Figure 1: The message passing paradigm

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

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

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

分散-归约

分散-归约 (scatter-reduce) 的模式本质上是并行的，它使用 src 张量的值更新 self 张量中由 index 指定的条目。理想情况下，在外部维度上并行化性能最佳。然而，直接并行化会导致写冲突，因为不同的线程可能同时尝试更新同一条目。

Figure 2: Scatter-reduce and its optimization scheme

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

为了优化此内核，我们使用排序后进行归约

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

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

SpMM-归约

稀疏矩阵-矩阵归约 (SpMM-Reduce) 是 GNN 中的一个基本操作符，其中 A 是 CSR 格式的稀疏邻接矩阵，B 是稠密特征矩阵，归约类型可以是求和 (sum)、平均值 (mean) 或最大值 (max)。

Figure 3: SpMM optimization scheme

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

优化此内核的最大挑战是如何在沿稀疏矩阵 A 的行进行并行化时平衡线程负载。A 中的每一行对应一个节点，其连接数可能相差很大；这会导致线程负载不平衡。解决此类问题的一种技术是在线程划分之前进行负载扫描 (payload scanning)。除此之外，还引入了其他技术来进一步挖掘 CPU 性能，例如向量化 (vectorization)、循环展开 (unrolling) 和分块 (blocking)。

这些优化通过 torch.sparse.mm 并使用 amax、amin、mean、sum 等归约标志来实现。

性能提升：最高加速 4.1 倍

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

模型 – 数据集 选项 加速比
GCN-Reddit (推理) 512-2-64-dense 1.22 倍
1024-3-128-dense 1.25 倍
512-2-64-sparse 1.31 倍
1024-3-128-sparse 1.68 倍
GraphSage-ogbn-products (推理) 1024-3-128-dense 1.15 倍
512-2-64-sparse 1.20 倍
1024-3-128-sparse 1.33 倍
full-batch-sparse 4.07 倍
GCN-PROTEINS (训练) 3-32 1.67 倍
GCN-REDDIT-BINARY (训练) 3-32 1.67 倍
GCN-Reddit (训练) 512-2-64-dense 1.20 倍
1024-3-128-dense 1.12 倍

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

从基准测试结果可以看出，我们在 PyTorch 和 PyG 中的优化为推理和训练带来了 1.1 倍至 4.1 倍的加速。

用于 PyG 的 torch.compile

PyTorch 2.0 的旗舰功能 torch.compile 与 PyG 2.3 版本完全兼容，借助面向 CPU 的 TorchInductor C++/OpenMP 后端，在 PyG 模型推理/训练方面比命令式模式带来了额外的加速。具体而言，在英特尔至强铂金 8380 处理器上测试 基本 GNN 模型的训练性能时，实现了 3.0 倍至 5.4 倍的性能加速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 生态系统的道路上又向前迈进了一步。

参考资料

脚注

产品和性能信息

1铂金 8380：1 节点，2 个英特尔至强铂金 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，英特尔于 2023 年 3 月 2 日测试。

2铂金 8380：1 节点，2 个英特尔至强铂金 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，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。