面向十亿规模图的图神经网络分布式CPU与GPU混合训练
分布式、并行与集群计算
2022-03-15 v3
摘要
图神经网络（GNN）在图结构数据学习中取得了巨大成功。它们被广泛应用于推荐、欺诈检测和搜索等各种场景。在这些领域中，图通常规模庞大且异构，包含许多数百万或数十亿的不同类型顶点和边。为应对这一挑战，我们开发了DistDGLv2，该系统扩展自DistDGL，采用分布式CPU/GPU混合训练以迷你批次方式在大规模异构图上训练GNN。DistDGLv2将图数据置于分布式CPU内存中，并在GPU中执行迷你批次计算。为便于使用，DistDGLv2采用与Deep Graph Library（DGL）的迷你批次训练和异构图API兼容的接口，从而实现几乎无需修改代码的分布式训练。为保证模型精度，DistDGLv2采用同步训练方法，并允许构成迷你批次的自我网络包含非本地顶点。为确保数据局部性和负载均衡，DistDGLv2使用具有最小边割和多重均衡约束的多级划分算法划分异构图。DistDGLv2部署了异步迷你批次生成流水线，使计算与数据访问异步以充分利用所有硬件（CPU、GPU、网络、PCIe）。这种组合使DistDGLv2在训练高质量模型的同时实现高并行效率和内存可扩展性。我们在多种GNN工作负载上展示了DistDGLv2。结果表明，DistDGLv2相较DistDGL实现2-3倍加速，相较Euler实现18倍加速。在64块GPU的集群上，对含数亿顶点的图完成一个epoch仅需5-10秒。
引用
@article{arxiv.2112.15345,
title = {Distributed Hybrid CPU and GPU training for Graph Neural Networks on Billion-Scale Graphs},
author = {Da Zheng and Xiang Song and Chengru Yang and Dominique LaSalle and George Karypis},
journal= {arXiv preprint arXiv:2112.15345},
year = {2022}
}
备注
arXiv admin note: text overlap with arXiv:2010.05337