Heta: Distributed Training of Heterogeneous Graph Neural Networks

TL;DR

Heta is a novel distributed training framework for Heterogeneous Graph Neural Networks that reduces communication overhead and improves training efficiency by leveraging the graph structure and a new partitioning strategy.

Contribution

Heta introduces a relation-aggregation-first computation paradigm and a graph partitioning method tailored for HGNNs, addressing communication bottlenecks in distributed training.

Findings

Heta outperforms DGL and GraphLearn in training speed.

The framework effectively reduces communication overhead.

Evaluation on large datasets confirms significant efficiency gains.

Abstract

Heterogeneous Graph Neural Networks (HGNNs) leverage diverse semantic relationships in Heterogeneous Graphs (HetGs) and have demonstrated remarkable learning performance in various applications. However, current distributed GNN training systems often overlook unique characteristics of HetGs, such as varying feature dimensions and the prevalence of missing features among nodes, leading to suboptimal performance or even incompatibility with distributed HGNN training. We introduce Heta, a framework designed to address the communication bottleneck in distributed HGNN training. Heta leverages the inherent structure of HGNNs - independent relation-specific aggregations for each relation, followed by a cross-relation aggregation - and advocates for a novel Relation-Aggregation-First computation paradigm. It performs relation-specific aggregations within graph partitions and then exchanges partial aggregations. This design, coupled with a new graph partitioning method that divides a HetG based on its graph schema and HGNN computation dependency, substantially reduces communication overhead. Heta further incorporates an innovative GPU feature caching strategy that accounts for the different cache miss-penalties associated with diverse node types. Comprehensive evaluations of various HGNN models and large heterogeneous graph datasets demonstrate that Heta outperforms state-of-the-art systems like DGL and GraphLearn by up to 5.8x and 2.3x in end-to-end epoch time, respectively.