[论文解读] Helios: An Efficient Out-of-core GNN Training System on Terabyte-scale Graphs with In-memory Performance

Jie Sun, Mo Sun+7 more|arXiv (Cornell University)|Oct 2, 2023

Advanced Graph Neural NetworksComputer Science被引用 3

一句话总结

Helios 是一种新颖的离外存图神经网络（GNN）训练系统，通过利用 GPU 启动的异步 I/O、跨 CPU 和 GPU 内存的 GPU 管理异构缓存，以及深度 GNN 意识流水线，实现在 TB 级图上的内存内训练吞吐量。它在 TB 级图上相比最先进的 GPU 管理基线系统最高提升 6.43 倍，相比 CPU 管理系统提升超过 182 倍。

ABSTRACT

Training graph neural networks (GNNs) on large-scale graph data holds immense promise for numerous real-world applications but remains a great challenge. Several disk-based GNN systems have been built to train large-scale graphs in a single machine. However, they often fall short in terms of performance, especially when training on terabyte-scale graphs. This is because existing disk-based systems either overly focus on minimizing the number of SSD accesses or do not fully overlap SSD accesses with GNN training, thus resulting in substantial unnecessary overhead on the CPU side and then low GPU utilization. To this end, we propose Helios, a system that can train GNN on terabyte graphs in a single machine while achieving throughput comparable with in-memory systems. To achieve this, we first present a GPU-initiated asynchronous disk IO stack, allowing the GPU to directly access graph data on SSD. This design only requires about 30% GPU cores to reach the almost maximal disk IO throughput and wastes no GPU cores between IO submission and IO completion such that the majority of GPU cores are left for other GNN kernels. Second, we design a GPU-managed heterogeneous cache that extends the cache hierarchy to heterogeneous CPU and GPU memory and thus enhances cache lookup throughput significantly by GPU parallelism. Finally, we build a deep GNN-aware pipeline that seamlessly integrates the computation and communication phases of the entire GNN training process, maximizing the utility of GPU computation cycles. Experimental results demonstrate that Helios can match the training throughput of in-memory GNN systems, even for terabyte-scale graphs. Remarkably, Helios surpasses the state-of-the-art GPU-managed baselines by up to 6.43x and exceeds CPU-managed baselines by over 182x on all terabyte-scale graphs.

研究动机与目标

• 解决基于磁盘的 GNN 训练系统在处理 TB 级图时因 CPU 瓶颈和 GPU 利用率低而导致的关键性能瓶颈。
• 克服现有 CPU 管理系统存在的局限性，这些系统因 CPU-GPU 数据传输串行化和小批量准备阶段的高 CPU 开销而浪费 GPU 计算周期。
• 消除如 GIDS 等 GPU 管理系统所面临的吞吐量上限，这些系统因缺乏并行性和缓存容量不足，导致磁盘 I/O 利用率低且 GPU 资源耗尽。
• 在单机离外存架构下，无需内存存储，实现 TB 级图的类内存级训练吞吐量。
• 设计一种系统，通过 GPU 原生控制和智能数据层次管理，重叠计算、I/O 和缓存管理，最大化 GPU 利用率。

提出的方法

• 引入一种 GPU 启动的异步 I/O 堆栈，使 GPU 能够直接访问存储在 SSD 上的图数据，减少 CPU 参与，使约 30% 的 GPU 核心可维持接近最大值的磁盘吞吐量。
• 采用一种跨 CPU 和 GPU 内存的 GPU 管理异构缓存，利用 GPU 并行性加速缓存查找，降低 I/O 延迟。
• 设计一种深度 GNN 意识流水线，无缝集成计算和通信阶段，最小化 GPU 空闲周期，最大化硬件利用率。
• 通过重叠数据加载与 GNN 训练内核，优化磁盘 I/O，确保 GPU 工作负载持续，消除 I/O 提交与完成之间的空闲时间。
• 采用超批次和小批量生成策略，将 GPU 驱动的采样和特征提取调度至最小化冗余数据移动并最大化 I/O 效率。
• 利用 GPUDirect 技术实现 GPU 与 SSD 之间的直接数据传输，相比 CPU 中介的数据移动，降低延迟和内存复制开销。

实验结果

研究问题

• RQ1 通过最小化 CPU 瓶颈和最大化 GPU 利用率，基于磁盘的 GNN 训练系统是否能在 TB 级图上实现类内存级吞吐量？
• RQ2 GPU 启动的异步 I/O 在多大程度上能提升 GNN 训练工作负载中的磁盘 I/O 吞吐量并减少 GPU 空闲时间？
• RQ3 跨 CPU 和 GPU 内存的 GPU 管理异构缓存在多大程度上能减少磁盘 I/O 并加速大规模 GNN 的数据访问？
• RQ4 一种集成计算和通信阶段的深度、GNN 意识流水线是否能显著提升 GPU 利用率和训练吞吐量？
• RQ5 在 TB 级图上，基于磁盘的 CPU 管理与 GPU 管理系统之间存在多大性能差距？新系统能否缩小这一差距以匹配内存内性能？

主要发现

• 通过消除 CPU 瓶颈并最大化 GPU 利用率，Helios 即使在 TB 级图上也能实现与内存内 GNN 系统相当的训练吞吐量。
• 在所有 TB 级图上，Helios 相比最先进的 GPU 管理基线系统 GIDS，训练吞吐量最高提升 6.43 倍。
• 相比 Ginex 和 MariusGNN 等 CPU 管理基线系统，Helios 的训练吞吐量提升超过 182 倍，展现出显著的性能优势。
• GPU 启动的异步 I/O 堆栈仅使用约 30% 的 GPU 核心，即可实现接近最大值的磁盘 I/O 吞吐量，使大部分 GPU 核心可用于 GNN 内核计算。
• GPU 管理的异构缓存通过在 CPU 和 GPU 内存中缓存频繁访问的数据，显著减少了磁盘 I/O，缓解了外存存储带来的性能惩罚。
• Helios 在训练工作负载中实现了超过 90% 的 GPU 利用率，而 Ginex 等 CPU 管理系统仅观察到 0.6% 的 GPU 利用率，这得益于 I/O 与计算的有效重叠。