NeutronAscend：利用Ascend AI处理器优化图神经网络（GNN）训练

《ACM Transactions on Architecture and Code Optimization》：NeutronAscend: Optimizing GNN Training with Ascend AI Processors

图神经网络在Ascend AI处理器上的异构计算优化研究，提出NeutronAscend框架通过跨核特征分区与负载均衡、内核内存访问优化与流水线调度，有效解决稀疏图数据处理中的计算资源利用不足和核间负载不均问题，实验表明在Ascend 910B上平均加速比达4.71倍，设计理念适用于其他矩阵中心架构的NPUs。

　　

摘要

图神经网络（GNN）在图数据分析领域取得了显著的成功。然而，其应用需要大量的计算资源。最近，新一代的神经处理单元（NPUs），如Ascend AI处理器，由于其出色的性能和较低的采购门槛，已成为数据中心和技术公司计算能力的重要组成部分。在这些加速器上部署GNN模型以优化性能是一个自然的选择；然而，这仍然是一项具有挑战性的任务。Ascend AI处理器通常由多个物理上解耦的AI核心组成，这些核心专为密集矩阵计算而设计。在处理具有固有稀疏性和幂律分布的图数据时，Ascend AI处理器会面临核心间工作负载不平衡以及核心内资源利用效率低的问题。

在本文中，我们提出了NeutronAscend，这是一个专为Ascend AI处理器设计的高效GNN训练框架。NeutronAscend采用了两种关键的设计方法来优化核心间和核心内的性能。在核心间层面，我们利用GNN张量数据混合并行性，通过将顶点特征均匀分配到各个AI核心上，并将图数据分割成较小的维度来确保负载平衡，从而最大化计算资源的利用率。在核心内层面，计算感知的任务调度使GNN过程与计算单元相匹配，通过利用基于局部性的图压缩来减少不规则的内存访问和冗余计算，同时利用单元间的流水线技术进一步提高整体性能。在Ascend 910B处理器上的实验结果表明，与当前公开的Ascend NPU基线相比，NeutronAscend的平均加速比为4.71倍。虽然我们的工作是针对Ascend处理器设计的，但其设计原则可以广泛应用于其他采用物理解耦、以矩阵为中心的计算架构的NPUs。