研究摘要
在大规模图（具有超过十亿级边）上高效训练图神经网络（GNNs）面临巨大的挑战，尤其是在内存限制和数据传输瓶颈的制约下。
这些问题在基于GPU的采样中尤为突出。传统方法要么因为CPU-GPU之间的数据传输导致严重的性能瓶颈，要么在多GPU设置中因过多的数据重排和同步开销而受到限制。为了解决大规模图GNN训练中的这些问题，我们提出了HyDRA，一个开创性的框架，用于提升基于采样的小批次训练效率。HyDRA在多GPU内存共享和多节点特征检索方面引入了创新，通过将采样和数据传输无缝集成为单一内核操作，彻底变革了跨GPU采样的实现方式。框架采用了一种混合指针驱动的数据放置技术，以提高邻居检索效率；设计了一种针对高度节点的精确复制策略，减少通信开销；并利用动态跨批次数据编排和流水线机制来最小化冗余数据传输。
研究背景
图1： 基于采用的GNN训练流程
图结构数据在实际场景中极为常见，其顶点和边的表示可以用来描述复杂的实体关系，例如社交网络中的用户互动、生物网络中的分子结构。传统的机器学习方法无法充分利用图的拓扑信息，而 GNN 利用图卷积、注意力机制等技术，从结构与特征中提取深层次信息，实现节点分类、链路预测等任务的高效处理。
尽管 GNN 展现了强大的潜力，但其对大规模图数据的训练效率一直是研究难点。传统全批次方法需要将整个图数据加载到内存中进行训练。然而，当图数据达到十亿级别边的规模时，其内存需求轻松突破 60 GB，远超高端 GPU（如 40 GB A100）的容量限制。这使得基于全批次的训练方法不再现实。
为此，基于采样的 GNN 方法应运而生。它通过在训练中随机选择子图（即小批次的拓扑结构），既显著降低了内存开销，也在以下方面表现出优势：提高模型精度、降低计算复杂度、
增强模型泛化能力。
基于采样的 GNN 通常由以下两个关键步骤组成：图采样：从大型原始图中选择一个包含部分节点和边的子图，形成小批次的拓扑结构；特征检索：从存储介质中访问选定顶点及其邻居的属性信息。
这一流程解决了图数据规模带来的内存挑战，但随之也暴露了新的技术瓶颈。
研究现状
针对大规模图数据的 GNN 训练，国内外学者提出了一系列方法。从早期基于 CPU 内存的采样框架（如 DGL、PyG），到利用 GPU 虚拟内存技术的 DGL-UVA，再到更高效的多 GPU 分布式采样方法 DSP，技术手段逐步提升。然而，这些方法在以下几个方面仍面临瓶颈：
1. CPU-GPU 数据传输的低效率
大多数现有框架依赖 CPU 内存存储图数据，并在 CPU 上进行图采样。采样结果通过 PCIe 传输到 GPU 进行训练。然而，CPU 采样速度较慢，加之 PCIe 的有限带宽，导致 GPU 的利用率不足，无法充分发挥硬件的计算潜力。
2. GPU 内存的限制
近年来，多种基于 GPU 的缓存优化方法（如 PaGraph、BGL 和 GNNLab）试图通过缓存关键数据来优化 GPU 内存使用。但受限于单 GPU 的容量，这些方法仅能缓解部分数据传输开销，未被缓存的数据仍需通过 PCIe 传输，导致性能无法显著提升。
3. 分布式采样中的通信开销
作为多 GPU 采样领域的先驱，DSP 提出了将图拓扑分区到多 GPU 上的新方法，并通过集体通信实现跨 GPU 的邻居检索。然而，其设计依赖频繁的全互通信，并需对数据进行大量重排（如将目标顶点分发到不同 GPU 再重新聚合），进一步加剧了通信成本
研究方法
图2： HyDRA 的架构：内存共享图采样和分层跨批次特征检索
本研究提出了一个名为 HyDRA 的新型框架，旨在解决现有基于采样的图神经网络（GNN）训练方法在处理十亿级边规模图数据时面临的局限性。HyDRA 的设计核心在于通过单节点多 GPU 内存共享与多节点特征检索两大关键创新，有效优化图采样过程中的存储与通信效率。
首先，我们提出了一种新颖的多 GPU 图采样方法，通过 GPU 内存共享显著减少现有 DSP 框架中的数据重排和同步开销。在图拓扑数据管理方面，考虑到其数据量仅占整体图数据的约 10%，我们将其划分并分配到单节点内的多个 GPU 上，以避免昂贵的跨节点通信。为优化采样效率，我们将采样与远程邻居检索合并为单一内核操作，从而彻底消除大量的数据重排和频繁的跨 GPU 同步。进一步地，我们设计了一种基于混合指针驱动的数据管理机制，使图拓扑数据在多个 GPU 之间共享的同时，为每个 GPU 本地维护邻居列表的头尾地址映射。
然而，在 GPU 内存共享过程中，我们发现由于现有图划分方法效率不足，导致通信量大幅增加，最高可达原来的 8 倍。通过深入分析，我们将这一问题归因于高度邻接顶点（即具有大量邻居的顶点）。针对这一瓶颈，我们开发了一种基于高度邻接顶点复制的创新划分策略。具体来说，我们将这些顶点的邻居列表复制到每个 GPU 分区中，从而将远程访问转换为本地读取。该策略有效降低了通信开销，显著提升了多 GPU 采样的性能。
其次，为优化存储与计算效率，我们将占用大量存储空间的特征数据划分到集群内的多个节点和 GPU 上。在多节点特征检索中，我们设计了一种分层数据重用策略，解决传统方法中因数据频繁移动而导致的效率低下问题。在小批次训练过程中，尤其是较大批量情况下，不同批次之间的数据存在显著的重叠性。传统方法往往丢弃上一批次的数据而重新加载新数据，忽略了潜在的优化空间。为此，我们提出了一种动态跨批次数据调度机制，通过捕捉并重用这些重复数据，大幅提升了数据利用效率。
在多节点、多 GPU 系统中，由于通信带宽的不均衡性，我们进一步扩展了分层数据重用的概念。具体来说，这种分层方法分为三个层级：单个 GPU 内部、同节点多个 GPU 间以及跨节点之间。通过优先利用带宽更高的通信通道，我们的分层方法能够有效减少通信延迟。此外，这种系统化的数据重用方式不仅充分利用了数据的局部性与冗余性，还通过优化数据流与网络能力的匹配，显著降低了整体通信开销，为 GNN 的高效训练提供了有力保障。
研究成果
图3： 端到端训练时间对比
本研究提出的 HyDRA 框架在大规模图神经网络（GNN）训练领域取得了重要进展，成功解决了基于采样的 GNN 方法在处理十亿级别边图数据时的关键挑战。通过创新的单节点多 GPU 内存共享机制和多节点特征检索策略，HyDRA 显著提升了图采样与训练效率，并突破了现有框架在内存与通信方面的瓶颈。在多达64个A100 GPU的系统上进行评估时，HyDRA显著优于当前领先方法，与DSP和DGL-UVA相比，训练速度提高了1.4倍到5.3倍，同时在多GPU可扩展性方面实现了高达42倍的改进。HyDRA为大规模GNN高性能训练设立了新的标杆。
相关研究成果已收录于 CCF-A 类顶级会议 SC 24（The International Conference for High Performance Computing, Networking, Storage, and Analysis），论文题为 Scaling New Heights: Transformative Cross-GPU Sampling for Training Billion-Edge Graphs。