基于概率邻域扩展分析与缓存的通信高效图神经网络
机器学习
2023-05-08 v1 分布式、并行与集群计算
性能
摘要
自图神经网络（GNN）问世以来，由于在推荐系统和金融取证等应用中的广泛使用与成功，在大规模图上的 GNN 训练与推理得到了积极研究。本文关注分布式设置下采用节点级采样的 GNN 小批量训练与推理，其中顶点特征在分布式存储间的必要划分使得特征通信成为阻碍可扩展性的主要瓶颈。为在不损害预测精度的前提下显著减少通信量，我们提出了一种缓存远程分区中频繁访问顶点相关数据的策略。该策略基于对多跳邻域采样过程中顶点级包含概率（VIP）的分析，此类采样可能将邻域扩展至图的划分边界之外。VIP 分析不仅有助于消除通信瓶颈，还提供了一种通过优先将 GPU 存储分配给最频繁访问的顶点特征来组织内存中数据的方法。我们提出 SALIENT++，其扩展了先前的先进系统 SALIENT 以处理分区特征数据，并利用 VIP 驱动的缓存策略。SALIENT++ 通过使用深度流水线并大幅减少通信量，同时仅消耗 SALIENT 所需存储的一小部分，保留了 SALIENT 的本地训练效率与可扩展性。我们给出了基于 Open Graph Benchmark 数据集的实验结果，并证明在 8 台单 GPU 机器上使用 SALIENT++ 训练 3 层 GraphSAGE 模型比在 1 台单 GPU 机器上使用 SALIENT 快 7.1 倍，比在 8 台单 GPU 机器上使用 DistDGL 快 12.7 倍。
引用
@article{arxiv.2305.03152,
title = {Communication-Efficient Graph Neural Networks with Probabilistic Neighborhood Expansion Analysis and Caching},
author = {Tim Kaler and Alexandros-Stavros Iliopoulos and Philip Murzynowski and Tao B. Schardl and Charles E. Leiserson and Jie Chen},
journal= {arXiv preprint arXiv:2305.03152},
year = {2023}
}
备注
MLSys 2023. Code is available at https://github.com/MITIBMxGraph/SALIENT_plusplus