基于半精度的GNN训练优化
机器学习
2025-09-17 v2
摘要
近年来，低精度训练（如半精度浮点数）在深度学习中展现出提升系统性能和减少内存占用的同时保持精度的趋势。然而，现有GNN系统无法在GNN上实现这一目标，我们的分析显示它们在使用半精度时大幅表现不佳，同时出现异常精度。这些系统因精度降低导致数值溢出问题、硬件资源利用不足以及训练性能差。为缓解此问题，我们提出了HalfGNN，一种基于半精度的GNN系统。HalfGNN提出了新技术：面向半精度数据类型的新型向量操作，提升数据加载和归约性能；离散化SpMM，克服数值溢出并原生提供负载均衡。这些技术提高了硬件利用率，减少内存占用，并消除了原子写入。评估表明，HalfGNN在GAT、GCN和GIN上相对于基于浮点数的DGL，平均实现了2.30倍的训练时间加速，同时达到相似精度，并节省了2.67倍内存。
引用
@article{arxiv.2411.01109,
title = {Optimization of GNN Training Through Half-precision},
author = {Arnab Kanti Tarafder and Yidong Gong and Pradeep Kumar},
journal= {arXiv preprint arXiv:2411.01109},
year = {2025}
}