NeutronAscend: Optimizing GNN Training with Ascend AI Processors

Graph Neural Networks (GNNs) have achieved remarkable success in graph data analysis. However, its application requires extensive computing resources. Recently, a new generation of Neural Processing Units (NPUs), such as the Ascend AI processor, has emerged as an essential component of computing power in data centers and technology companies due to its outstanding performance and low threshold for acquisition. Deploying GNN models on these accelerators for performance optimization is a natural choice; however, it remains a challenging task. The Ascend AI processor is typically architected on multiple AI Cores that are physically decoupled and designed for dense matrix computation. When processing graph data with inherent sparsity and power-law distribution, the Ascend AI processors suffer from the inter-core workload imbalance and inefficient intra-core resource utilization. In this article, we present NeutronAscend, an efficient GNN training framework tailored for the Ascend AI processor. NeutronAscend employs two critical designs for both inter-core and intra-core performance optimization. At the inter-core level, we employ GNN tensor-data hybrid parallelism that ensures load balance by partitioning vertex features evenly across AI Cores and partitioning graph data across small dimensions to maximize computing resource utilization. At the intra-core level, computation-aware task scheduling aligns the GNN process with computational units, leveraging locality-aware graph compression to reduce irregular memory access and redundant computation, while utilizing inter-unit pipelining to further enhance overall performance. Experimental results on the Ascend 910B processor demonstrate that NeutronAscend achieves an average 4.71× speedup compared to the current publicly available baselines on Ascend NPUs. While our work is tailored to Ascend, its design principles are broadly applicable to other NPUs that adopt physically decoupled, matrix-centric compute architectures.