混合CPU-GPU训练新范式

GPU虽能高效执行神经网络张量运算，但内存有限。针对超出GPU内存容量的大规模图神经网络训练，传统方法需通过CPU随机选择节点和边组成小批量数据，并连同节点特征一并传输至GPU。这种递归邻居采样方式导致大量CPU-GPU数据传输，成为模型训练中最耗时的环节。

全局邻居采样(GNS)核心技术

  1. 特征缓存机制
  2. 在创建小批量数据前，先从全图中采样一组节点并将其特征预加载至GPU内存（称为缓存）。生成小批量时优先从缓存中检索邻居节点，仅当缓存不足时才从CPU补充获取。
  3. 高权重节点优先策略
  4. 采样概率与节点度数成正比，既确保高频连接节点优先缓存，又保留部分低频节点。通过数学证明，该方法能收敛至与完全随机采样相同的模型最优性能。
  5. 子图索引优化
  6. 在CPU端构建包含所有缓存节点及其直连邻居的子图，通过快速查询替代昂贵的集合运算，显著提升邻居匹配效率。

实验验证

  • 在5个数据集上的测试表明，GNS在混合CPU-GPU环境下比次优方案快2倍以上
    • 两种传统采样方法需耗费10倍时间才能达到相近精度
    • 当前支持单CPU-GPU架构，多GPU分布式训练的扩展研究正在进行中

性能突破关键

通过减少90%以上的CPU-GPU数据传输，GNS将训练瓶颈从数据搬运转向实际计算，为十亿级节点规模的图神经网络训练提供实用解决方案。该方法已发表于KDD 2021会议，相关代码实现同步公开。 更多精彩内容 请关注我的个人公众号 公众号（办公AI智能小助手）