MaxK-GNN: Extremely Fast GPU Kernel Design for Accelerating Graph Neural Networks Training
CoRR(2023)
摘要
In the acceleration of deep neural network training, the GPU has become the
mainstream platform. GPUs face substantial challenges on GNNs, such as workload
imbalance and memory access irregularities, leading to underutilized hardware.
Existing solutions such as PyG, DGL with cuSPARSE, and GNNAdvisor frameworks
partially address these challenges but memory traffic is still significant.
We argue that drastic performance improvements can only be achieved by the
vertical optimization of algorithm and system innovations, rather than treating
the speedup optimization as an "after-thought" (i.e., (i) given a GNN
algorithm, designing an accelerator, or (ii) given hardware, mainly optimizing
the GNN algorithm). In this paper, we present MaxK-GNN, an advanced
high-performance GPU training system integrating algorithm and system
innovation. (i) We introduce the MaxK nonlinearity and provide a theoretical
analysis of MaxK nonlinearity as a universal approximator, and present the
Compressed Balanced Sparse Row (CBSR) format, designed to store the data and
index of the feature matrix after nonlinearity; (ii) We design a coalescing
enhanced forward computation with row-wise product-based SpGEMM Kernel using
CBSR for input feature matrix fetching and strategic placement of a sparse
output accumulation buffer in shared memory; (iii) We develop an optimized
backward computation with outer product-based and SSpMM Kernel.
We conduct extensive evaluations of MaxK-GNN and report the end-to-end system
run-time. Experiments show that MaxK-GNN system could approach the theoretical
speedup limit according to Amdahl's law. We achieve comparable accuracy to SOTA
GNNs, but at a significantly increased speed: 3.22/4.24 times speedup (vs.
theoretical limits, 5.52/7.27 times) on Reddit compared to DGL and GNNAdvisor
implementations.
更多查看译文
AI 理解论文
溯源树
样例
生成溯源树,研究论文发展脉络
Chat Paper
正在生成论文摘要