中国科大提出新型GNN训练框架Capsule，性能提升12倍
2月23日消息，2月11日，中国科学技术大学苏州高等研究院医学影像智能与机器人研究中心的DDL实验室宣布，其在大规模图神经网络（GNN）训练系统领域的研究成果已被国际知名学术会议ACM International Conference on Management of Data（SIGMOD）2025正式接收。
当前，主流的图神经网络训练框架如DGL和PyG通过利用GPU的并行计算能力，能够从图数据中高效提取结构信息，在推荐系统、自然语言处理、计算化学以及生物信息学等多个领域表现出优异性能。然而，尽管GPU为GNN训练提供了强大的计算支持，但其显存容量有限，难以满足超大规模图数据的存储需求，这使得现有GNN系统在扩展性方面仍存在显著挑战。
为解决这一问题，DDL实验室团队提出了一种名为Capsule的新型核外（Out-of-Core）GNN训练框架，为大规模图神经网络训练提供了高效的解决方案。与现有的核外GNN框架不同，Capsule通过图划分和图裁剪策略，将训练子图的结构及特征完全加载到GPU显存中，从而彻底消除了反向传播过程中CPU与GPU之间的I/O开销，显著提升了系统性能。
此外，Capsule采用基于最短哈密顿回路的子图加载方式，并引入流水线并行策略，进一步优化了系统性能。该框架还具备即插即用的特性，能够无缝集成到主流开源GNN训练框架中。在针对大规模真实图数据集的实验中，Capsule在仅使用22.24%内存的情况下，相比目前最优的系统性能提升了最高12.02倍，并首次提供了关于训练所得嵌入方差的理论上界。
这一突破标志着我国在图计算系统领域取得了重要进展，为社交网络分析、生物医学图谱构建等需要处理超大规模图数据的应用场景提供了全新的技术手段。相关论文已发表，详细内容可通过以下链接查阅：https://doi.org/10.1145/3709669。
