DAHA: Accelerating GNN Training with Data and Hardware Aware Execution Planning

Abstract

Graph neural networks (GNNs) have been gaining a reputation for effective modeling of graph data. Yet, it is challenging to train GNNs efficiently. Many frameworks have been proposed but most of them suffer from high batch preparation cost and data transfer cost for mini-batch training. In addition, existing works have limitations on the device utilization pattern, which results in fewer opportunities for pipeline parallelism. In this paper, we present DAHA, a GNN training framework with data and hardware aware execution planning to accelerate end-to-end GNN training. We first propose a data and hardware aware cost model that is lightweight and gives accurate estimates on per-operation time cost for arbitrary input and hardware settings. Based on the cost model, we further explore the optimal execution plan for the data and hardware with three optimization strategies with pipeline parallelism: (1) group-based in-turn pipelining of batch preparation neural training to explore more optimization opportunities and prevent batch preparation bottlenecks; (2) data and hardware aware rewriting for intra-batch execution planning to improve computation efficiency and create more opportunities for pipeline parallelism; and (3) inter-batch scheduling to further boost the training efficiency. Extensive experiments demonstrate that DAHA can consistently and significantly accelerate end-to-end GNN training and generalize to different message-passing GNN models.