Mega: More Efficient Graph Attention for GNNs

Graph neural networks (GNNs) have demonstrated effectiveness across diverse application domains by leveraging graph information to uncover intrinsic correlations alongside feature representation. This enables GNNs to explore richer information compared to conventional neural networks, resulting in enhanced predictive performance. However, the integration of graphstructured data into the learning process poses two challenges. First, the sparsity and irregularity of the graph representation result in inefficient and expensive memory accesses on throughput-oriented accelerators, such as GPUs. Second, as GNN training involves interleaved graph operations to extract topological information and neural operations to update node or edge embeddings, the joint optimization of these two operations on accelerators is challenging due to their distinct resource requirements. Among GNN graph operations, graph attention which helps focus GNN training on highly correlated nodes is critical to training performance and model accuracy. However, our profiling of representative GNNs reveals that irregular memory access during graph attention accounts for the dominating overhead in GNN training. To address this issue, this paper proposes a more efficient graph attention method Megato accelerate GNN training. MegA converts the original graph representation into one that regularizes memory access patterns for graph attention. Specifically, during preprocessing, Megatraverses a graph to derive a schedule for graph attention and uses the schedule to reorganize the graph representation for optimized memory access. MegA explores several techniques to balance memory access efficiency and preserve the original graph properties to avoid the loss of model accuracy. Experimental results with representative GNNs and graph data sets show that Megaconsistently outperforms conventional graph attention methods with up to 3x speedup.