[论文笔记] PiPAD: Pipelined and Parallel Dynamic GNN Training on GPUs

PiPAD: Pipelined and Parallel Dynamic GNN Training on GPUs

PiPAD: 基于 GPU 的流水线并行动态 GNN 训练 [Paper]
PPoPP’23

摘要

提出了 PiPAD, 一个流水线并行的 DGNN 训练框架, 用于 GPU 上的端到端性能优化. 从算法和运行时层面重构了整体训练范式.

1 介绍

动态图表示分类: 连续时间动态图(Continuous Time Dynamic Graphs, CTDGs) 和离散时间动态图(Discrete Time Dynamic Graphs, DTDGs).
DTDG: 使用静态 GNN(如 GCN) 在所有时间步上对各个快照进行空间图学习, 同时使用循环神经网络 (RNN) 以获得不同快照之间的时间特征; 采用多个连续快照的滑动窗口机制.
本文重点关注基于 DTDG 的 DGNN, 并将滑动窗口简称为帧(frame).

DGNN 的两个主要性能问题:

1. 沿着时间线不断更新图快照而导致数据传输开销主导了整个训练时间
2. 时间无关(GNN)与时间相关(RNN)组件的组合产生了大量并行机会, 相邻帧之间的快照重叠提供了大量的数据重用机会, 这些迫切需要得到充分利用.

两个关键发现:

1. 现实世界的动态图通常以缓慢的速度变化, 导致相邻快照之间存在大量拓扑重叠
2. 除了类 SpMM 聚合的不规则访问模式之外, GNN 还存在因结点特征维度不同导致的其他类型的内存低效问题.

一次一个快照的训练方式导致大量冗余数据传输和低效内存访问.
提取重叠拓扑并对多个快照进行单个时间无关聚合操作不仅可以减少通信量而且有更多机会启用合并内存访问并解决低带宽利用率问题.

提出了 PiPAD:

1. 提出了一种基于切片的图表示格式和重叠感知的多快照传输方法, 减少从 CPU 到 GPU 的数据加载时间.
2. 设计了帧内并行(具有线程感知切片合并的维度感知并行 GNN, 支持同时处理多个快照)和帧间重用机制.
3. 实现了管道执行框架来协调 CPU 端操作、PCIe 上的数据传输和 GPU 计算; 并集成了运行时调优器动态调整 DGNN 训练的并行性和重用级别.

文本贡献:

• 对三种典型 DGNN 模型的执行和性能特征进行了全面分析. 揭示了 DGNN 训练的几个关键发现被激发了本文的设计.
• 提出了 PiPAD 来优化单 GPU DGNN 训练的端到端性能. 将重叠感知数据组织、帧内并行、流水线执行机制和帧间重用与动态调优技术相结合, 以多快照的方式全面实现流水线并行训练.
• PiPAD 比最先进的 DGNN 框架实现了显著的加速.

2 背景

2.1 动态图神经网络

DTDG 是快照 $\{{\mathcal{G}}^1,{\mathcal{G}}^2...{\mathcal{G}}^t\}$ 的有序集合. ${\mathcal{G}}^t=\{{\mathcal{V}}^t,{\mathcal{E}}^t\}$ 为在时间步 $t$ 的具有结点 ${\mathcal{V}}^t$ 和边 ${\mathcal{E}}^t$ 的快照.
基于 DTDG 的 DGNN 遵循滑动窗口/帧的处理模式, 并且可以包括多个层. 每一层中, GNN 组件从每个快照中学习结构信息, 而 RNN 模块捕捉时间轴上的时间特征.

2.2 GNN 加速

静态 GNN
DGNN
管道/并行训练

3 动机

3.1 DGNN 训练的性能瓶颈

数据传输开销: DGNN在快照序列上操作, 需要不断地将新的图数据从 CPU 加载到 GPU.
拓扑重叠: 现实生活的动态图中相邻快照间的拓扑变化率通常有限.

3.2 GNN 中内存访问效率低下

现有解决方案忽略了由于稠密矩阵(输入或聚合特征向量)的维数而导致其他类型的内存访问低效.
两种类型的内存访问低效:

• 带宽不饱和: 特征维度小于 32/4 时, 有用数据的大小小于单个事务的最小访问粒度
• 请求突发: 特征维度大于 128/4 时, warp 访问一行需要向共享内存和全局内存发出多个请求

3.3 数据重用和并行的机会

数据重用: 连续帧之间存在大量重叠. GCN 第一层中的聚合操作对图拓扑的邻接矩阵和原始结点特征进行操作, 与沿时间轴的参数更新无关, 可以缓存相关的聚合结果并在下一帧或训练时期重用.
并行: GNN 对不同快照的相同操作可以并行. 设计一种新的图格式可以有效地提取重叠以消除冗余传输, 设计了一种可以同时处理多个图的"并行" GNN 计算模式.

4 PiPAD

4.1 重叠感知的数据组织

基于切片的图表示:
切片 CSR 格式:

• 原始的 $RowOffsets$ 数组改为 $RowIndices$(RI), 表示每个切片的行索引.
• 添加了一个新数组 $SliceOffsets$(SO) 存储每个切片的第一个元素在 $ColumnIndices$ 字段中的偏移
• 每个切片存储具有固定上限的 $nnz$

重叠感知的数据传输:
将每个帧划分为多个分区来执行并行训练, 每个分区包含连续的快照.
数据传输调整为分区粒度模式, 拓扑数据(邻接矩阵)被重组为一个所有快照的重叠部分和几个每个快照的独有部分.
利用单独的 GPU 流和固定内存(pinned memory)以按需方式实现异步数据传输.

4.2 帧内并行

维度感知并行 GNN:

• 将来自不同快照的特征矩阵合并以执行一次聚合.
  根据感知重叠的数据组织, 将原始聚合分为两部分: 一个快照组中所有特征重叠部分的邻接矩阵的并行聚合(称为合并特征), 以及来自每个快照各自特征的独有拓扑部分上的非并行聚合.
• 利用向量内存指令支持加载一个请求中 32/64/128 位宽的浮点数作为解决请求突发问题的可选项.

线程感知的切片合并:
将邻接矩阵中的相邻切片合并为一组, 并将其设置为每个 warp 的基本处理单元. 同一 warp 内的线程被均分为几个线程组(thread group, TG), 并且一个 TG 操作一个切片.
共享内存上以交错模式实现切片分组的数据布局.
在每次计算迭代中, 每个 TG 对来自相应切片的一个元素和来自合并特征矩阵的一行进行操作.

• $slice\_size$: 每个切片拥有的最大非零元数
• $coalesce\_num$: 切片分组的最大大小, 即每个 warp 中的线程分组(TG)数
• $load\_index$: 分配给当前线程加载的切片索引 (相邻线程分配不同的切片)
• $comp\_index$: 分配给当前线程计算的切片索引 (相邻线程分配相同的切片)
• $com{p}_{dim}$: 分配给当前线程计算的切片对应结点的特征维度索引 (相邻线程分配同一切片的不同特征维度)
• $loa{d}_{size}$: 索引为 $load\_index$ 的切片的非零元数(切片对应结点的邻结点数), 即该切片需加载 $nnz$ 数
• $com{p}_{size}$: 计算所需的迭代次数 (每次迭代读取一组/TG 个非零元, 计算所有非零元所需的迭代次数)
• line 8: $comp\_index<coalesce\_num$: 计算的切片数量不能超过 TG 数, 从而确定实际工作的线程数 ($lane\_id<(featurens.dim\ast coalesce\_num)$)
• line 12~20: 由全局内存加载切片到共享内存. 一个 warp 每次迭代加载 $coalesce\_num$ 个切片, 每个切片至多加载 $featurens.dim$ 个非零元.
• line 22~39: 加载切片和特征进行计算. 一个 warp 每次迭代加载 $coalesce\_num$ 个元素, 和切片对应的特征矩阵.

局部性优化的权重重用:
将一个权重分片保留在共享内存中来连续计算所有快照的特征, 然后移动到下一个权重分片进行相同的迭代.

对于 DGNN 中的其他内核, 按照计算和数据依赖的顺序依次执行, 并参考 OOB 并利用 CUDA Graph API 进行启动.

4.3 管道执行框架

算法层面: 设计了切片 CSR 以便提取图重叠, 并重构 GNN 以实现帧内多个快照的并行计算.

运行时层面:
前几个 epoch (准备 epoch) 期间:

• 在线图分析器 ❶: 将所有快照从原始 CSR 格式转换为切片 CSR 格式
• 数据准备模块 ❷: 提取相邻快照之间的重叠部分, 并为分区的计算方式做准备.

后续 epoch:

• CPU 和 GPU 端缓冲区 ❸: 缓存聚合结果以支持不同帧之间的数据重用
• 动态调优器 ❹: 在线调整并行级别(帧内每个分区的快照数)和重用级别(GPU 端缓存的聚合结果的大小)
• 管道控制器 ❺: 调整整个训练过程

准备阶段: 单快照方式, 异步数据传输

• 主机端在线收集必要的统计信息, 用于指导动态调优器选择优化并行选项
• 一次性执行图切片和重叠提取
• 存储每个帧内快照间的图演化率, 以供之后在线调整决策

后续阶段: 分区方式

1. 确定当前帧的并行级别
2. 以分区方式传输训练数据

4.4 帧间重用和动态调整

帧间重用:
维护一个 GPU 端缓冲区来根据下一帧中的使用顺序缓存一些聚合结果.
根据需求为每个帧动态分配不同大小的缓冲区.

动态调整:
将一个帧中的快照均匀分布到每个分区.
确定每个帧的两个关键参数: 每个分区的快照数量($S_{per}$), 以及缓存在 GPU 中以供下一帧重用的聚合结果的大小.

$S_{per}$ 影响因素及选择:

• 内存消耗: 利用在线分析统计数据确定不会触发 OOM(out-of-memory) 的上限 $U$
• 并行 GNN 的计算加速: 相同重叠率(overlap rate, OR)或特征维度设置下, 首选较大的 $S_{per}$
• 计算和数据传输间的重叠程度: 去除导致管道停顿的方案后选择计算加速最高的方案

4.5 讨论

使用 C++、CUDA C 和 Python 在 PyTorch 之上实现 PiPAD.

局限性: 多快照并行训练会增加基于阶段(phase-based)的内存消耗, 并导致大图固有的扩展性问题.

5 评估

性能: Figure 10

GPU 利用率: Table 2
并行优化的性能提升: Figure 11
切片 CSR 性能提升: Figure 12

---

笔者总结

本文的核心在于针对 DTDG 类型的动态图神经网络, 提出了切片 CSR 的图表示格式和重叠感知的多快照传输方法, 帧内并行和帧间重用机制, 以及管道执行框架的优化方法进行加速.