Accelerating Mini-batch HGNN Training by Reducing CUDA Kernels

A heterogeneous graph is a type of graph that contains multiple types of vertices and edges, each representing different kinds of entities and relations [27]. Unlike homogeneous graphs, which have a single type of vertex and edge, heterogeneous graphs capture the complexity and diversity of real-world data. By incorporating rich semantic information through diverse vertex and edge types, heterogeneous graphs enable more nuanced analysis and learning [9, 7, 8, 6].

Metapath is a crucial concept in heterogeneous graphs, defining a sequence of vertex types and edge types that outline a path through the graph. It captures complex structural patterns and semantic relationships between different types of entities, aiding in the understanding and definition of connectivity and interaction patterns within heterogeneous graphs.

Semantic graphs are subsets of the original graph that focus on specific types of vertices and relations, often derived based on metapaths. They are used to enhance the understanding of semantic relations and to facilitate more efficient and targeted computations in heterogeneous graphs.

An example of a heterogeneous graph is the ACM (Association for Computing Machinery) graph, as shown in Fig. 1. In this graph, different types of vertices represent various entities such as papers (P), authors (A), and subjects (S). The edges between these vertices denote different types of relations, such as an author writing a paper or a paper related to a subject. This rich structure allows the ACM graph to model the multifaceted interactions within the academic research community, capturing the complexity of real-world academic networks. By leveraging the diverse vertex and edge types, the ACM graph enables sophisticated analyses and insights, such as identifying influential authors and discovering emerging research subjects. Metapaths can be used to define paths through the ACM graph, such as “Author - Paper - Subject”, which captures the sequence where an author writes a paper that is related to a subject.

HGNNs are a class of GNN models specifically designed to handle heterogeneous graphs [6, 9, 8]. They offer a powerful ability for processing and learning from these graphs by effectively capturing and integrating the complex relations among different types of entities. This capability is crucial for accurately modeling the multifaceted nature of real-world networks, such as academic networks, knowledge graphs, biological networks, and others [6, 9, 8]. For example, in academic networks like the ACM graph, HGNNs can effectively model the relations between papers, authors, subjects, and other entities. They can distinguish between various types of connections, such as co-authorship, providing a more nuanced understanding of academic collaborations and research trends.

An HGNN layer typically involves four stages [18, 17, 19], as shown in Fig. 1.

• $\bullet$

  Semantic Graph Build: This stage splits the heterogeneous graph into multiple semantic graphs, each focusing on different types of vertices and edges. By isolating specific types of relations (metapaths), the model can better capture the unique interactions and dependencies within each subset of the graph.

• $\bullet$

  Feature Projection: In this stage, the feature vectors of vertices in each semantic graph are transformed using an MLP or a similar neural network. This transformation is essential for mapping the raw features into a latent space where they are more suitable for subsequent processing. It enhances the model’s ability to learn from complex and high-dimensional data.

• $\bullet$

  Neighbor Aggregation: This stage involves aggregating features from neighboring vertices for each vertex within each semantic graph. By doing so, the model captures the relational information inherent in the graph structure. This aggregation process is crucial for understanding how information propagates through the network and how the local structure influences each vertex.

• $\bullet$

  Semantic Fusion: The final stage integrates semantic information across all semantic graphs by combining the results from the previous stage. This fusion process provides holistic vertex representations that incorporate information from all types of vertices and edges. These enriched representations are essential for accurately capturing the diverse relations in the heterogeneous graph.

The final vertex representations generated by HGNNs can be used for various downstream tasks, such as vertex classification, where each vertex is assigned a label based on its features and relations; link prediction, which involves predicting the likelihood of a link between two vertices; and graph clustering, which groups vertices into clusters based on their similarities.

Full-batch and Mini-batch HGNN Training. HGNN training can be divided into two approaches: full-batch and mini-batch [28].

In full-batch training, the entire graph dataset is fed into the model as a single complete batch. This means that all vertices and edges participate in every forward and backward pass of the model. This approach is commonly used for small datasets or when maximizing the utilization of training data is crucial. However, it may encounter memory limitations and computational inefficiencies when dealing with large graph data.

On the other hand, mini-batch training divides the original graph dataset into multiple smaller batches, with each batch containing a subset of vertices and edges. These mini-batches are used for feature collection and model computation. This approach effectively handles large graph data by reducing the amount of data processed in each training iteration, thereby reducing memory consumption and improving computational efficiency.

Our work focuses on mini-batch training because it significantly improves efficiency and speed by optimizing hardware utilization and reducing memory consumption [28]. Mini-batch training enhances convergence stability and supports scalable training on large datasets. It enables the effective use of vectorized operations and facilitates distributed computing.

Workflow of Mini-batch HGNN Training. Mini-batch HGNN training follows a structured workflow designed to efficiently process and learn from heterogeneous graph [21], as illustrated in Fig. 2. Sampling: mini-batches are sampled from the original graph on CPU. Feature Collection: features specific to each mini-batch are then collected on CPU according to its topology. Data Loading: the mini-batch and features are transferred from CPU to GPU. Model Computation: these features undergo a forward pass on GPU through a series of stages, including semantic graph build, feature projection, neighbor aggregation, and semantic fusion. Then, the model performs a backward pass on GPU for gradient computation and parameter update.

To address inefficiencies in previous GPU-based solutions for HGNN training, we introduce HiFuse, an enhancement for PyG aimed at accelerating mini-batch HGNN training on CPU-GPU systems. HiFuse employs a two-pronged strategy: one focusing on data reorganization and the other on workflow optimization.

From the data perspective, we propose a method for feature reorganization and merging during the neighbor aggregation stage to optimize data access and execution, thereby enhancing GPU processing efficiency. The reorganization technique adjusts the layout of vertex features by grouping features of the same vertex type together, improving memory access patterns, data locality, and overall memory efficiency. The merging technique combines smaller feature matrices into larger ones, allowing single kernels to process more data, reducing the number of kernel launches, and thus improving GPU utilization and throughput.

From the workflow perspective, we offload the edge index selection in the semantic graph build stage from GPU to CPU, thereby avoiding the overhead of numerous small CUDA kernel launches and the under-utilization of compute resources. To ensure this transition does not introduce bottlenecks, we employ two strategies: first, parallel processing on CPU handles multiple tasks concurrently using multi-threading to fully utilize CPU cores; second, an asynchronous pipeline overlaps execution stages across CPU and GPU, ensuring effective utilization of both and minimizing idle times.

Observations. After sampling, each mini-batch mainly includes the edge indices of the sampled subgraph. Each edge index consists of the IDs of the source vertex and target vertex of an edge. These edge indices are stored in a $2\times N$ tensor in coordinate format, representing the graph topology for all relations. Before processing each semantic graph, the edge indices of the same edge type must be selected for each semantic graph, a process known as edge index selection in the semantic graph build stage. This involves edge type matching and index selection, carried out by the compare and index-select operations.

Observation : The kernels used for compare and index-select operations exhibit high control intensity and a high proportion of integer instructions because they primarily involve integer comparison and addressing operations. Additionally, these kernels have low computational demands and short execution times, resulting in underutilization of GPU’s processing capabilities.

Observation : During each mini-batch training process, aside from data sampling on CPU, all other computations occur on GPU. Consequently, CPU remains mostly idle, leading to an imbalance in the utilization of CPU and GPU. As shown in Table 1, the CPU utilization is generally lower than that of GPU. The ratio of CPU time to GPU time is as low as 0.04.

Offloading. Based on our observations, we offload the kernels for edge index selection from GPU to CPU, driven by several key considerations as follows.

• $\bullet$

  Enhanced Resource Utilization: CPU is underutilized, with most workloads assigned to GPU, leaving it mostly idle. By offloading control-intensive tasks to CPU, we can balance the workload between CPU and GPU, making better use of available resources.

• $\bullet$

  Reduced Overhead: Kernels with short execution time do not significantly benefit from GPU acceleration and can incur overhead from frequent kernel launches. Executing these tasks on control-friendly CPU can eliminate this overhead, leading to more efficient processing.

Since each semantic graph’s computation in each layer requires edge indices, it is most efficient to perform edge index selection after sampling but before model computation. The edge indices for all relations are temporarily stored, allowing direct access during computation without further indexing. The specific implementation is detailed in Algorithm 2. Edge indices (EdgeIndex) and edge identifiers (EdgeID) are obtained from the mini-batch, representing the graph topology and the identifiers of edges within the sampled subgraph, respectively. Before the indexing computation for each layer, the EdgeID is used to index all edge types within the sampled subgraph for that layer, resulting in a tensor of edge types. During the indexing operation for each relation, edge type matching is first performed using the compare operation to select indices of the same edge type from the edge type tensor. Subsequently, the index-select operation is used to select the edge index tensor for the current relation based on these indices. Finally, the indices edge tensor for each relation is stored, completing the edge index selection.

By isolating edge index selection from each semantic graph’s computation, the forward pass retains only those operations demanding significant GPU resources. This approach improves the continuity and compactness of CUDA kernel scheduling, thus enhancing overall GPU utilization and computational efficiency during model computation.

Parallelizing. Offloading tasks to CPU is a relatively straightforward idea, but maximizing the CPU’s parallel processing capabilities is crucial to ensure timely data transfer to GPU. To achieve this, it is necessary to leverage idle CPU cores to parallelize the edge index selection, compensating for the CPU’s slower computation speed and improving overall efficiency.

The edge type matching and index selection for each relation are independent, making them ideal for parallel processing. Thus, we utilize the multi-threading capabilities of CPU to parallelize this process. This is implemented in PyTorch’s CPP backend environment, LibTorch, using the OpenMP multi-threading parallel programming framework. The specific implementation follows a similar approach to Algorithm 2. After obtaining the EdgeIndex and EdgeID from mini-batch graph data sampling, we use OpenMP precompiled directives to enable multi-threaded parallel acceleration in the critical code segments. By employing appropriate OpenMP directives based on the CPU’s multi-core resources, we enhance the performance of the multi-threaded program, ensuring efficient edge index selection and timely data transfer to GPU.

Pipelining. By combining parallel processing with an asynchronous pipeline, we can enhance the utilization of both CPU and GPU. This approach minimizes idle times, reduces latency, and maintains a steady flow of data, leading to significant performance improvements in complex computations.

Based on the above parallelization strategy, we propose an asynchronous execution method for CPU and GPU tasks to address the issue where CPU and GPU tasks in each mini-batch training session are executed sequentially, with GPU tasks taking significantly longer than CPU tasks. The optimized asynchronous computation flow is depicted in Fig. 6. Initially, the mini-batch sampling and edge index selection tasks are scheduled by a CPU thread. The processed mini-batch and edge index array are used for the model computation of the mini-batch. Another CPU thread transfers the mini-batch and edge index array from CPU to GPU. Data transfer is managed by a dedicated CUDA stream, which operates independently of the default CUDA stream used for model computation. Subsequently, the model computation begins on GPU. The computation tasks on GPU proceed asynchronously, including forward, backward, and parameter update for the current training mini-batch, managed by the default CUDA stream. After GPU computation tasks are completed, the CPU thread gathers and returns the results. This cycle continues until all mini-batch computations are finished.

In this CPU and GPU asynchronous mini-batch training pipeline, CPU and GPU computations can be executed in parallel. Each mini-batch’s sampling and edge index selection tasks are completed before the model computation tasks begin, ensuring seamless transitions between GPU computation tasks and avoiding GPU idle states. This approach also efficiently utilizes idle CPU computing resources. Overall, this method improves the efficiency of each mini-batch training, accelerating the entire training process of the HGNN model.

For our baseline comparison, we select PyG [21] for the following reasons. PyG is a widely used framework for training HGNNs and undergoes extensive performance optimizations in an open-source environment, resulting in significant kernel optimizations and recent performance improvements. Therefore, it represents the current state-of-the-art GPU-based solution for HGNN training.

Both HiFuse and PyG are executed on a Linux server equipped with an Intel Xeon Silver 4208 CPU and an NVIDIA T4 GPU, utilizing single precision floats. We utilize the latest versions of PyG (2.5.3) and ensure consistent runtime environments (CUDA 12.1 with driver version 530.30.02) across all solutions.

For evaluation, we employ two well-known HGNN models: RGCN [30], representing an HGNN model with a simple architecture, and RGAT [31], representing an HGNN model with a complex architecture. Additionally, we utilize four popular datasets (aifb, mutag, bgs, and am) [30], as shown in Table 2.

To assess performance, we utilize NVIDIA Nsight System and Nsight Compute. Each model, along with its respective dataset, is executed ten times, and outlier results are discarded for reliability. These tools offer detailed insights into the performance metrics of our models, facilitating accurate analysis and comparison of efficiency and optimization benefits.

Speedup. To demonstrate the effectiveness of our optimizations, we compare HiFuse with the state-of-the-art framework, PyG. Both HiFuse and PyG pre-load vertex features into GPU device memory. For large heterogeneous graphs, they can be partitioned into several subgraphs, enabling the vertex features of these subgraphs to be pre-loaded into device memory. In the comparison, GM represents the geometric mean. As shown in Fig. 7, HiFuse achieves an average performance improvement of 2.38$\times$ compared to PyG, with a maximum improvement of up to 3.04$\times$. The improvement primarily comes from the following reasons: 1) The number of GPU kernels is significantly reduced, resulting in a substantial decrease in kernel launch time and idle time. 2) The pipeline and parallel execution between CPU and GPU ensures that CPU promptly provides input data to GPU. 3) The workload is balanced between CPU and GPU to some extent, enhancing overall efficiency.

Number of Kernels. To demonstrate the effect of reducing the number of kernels, we collect kernel counts for the training of one epoch. As shown in Fig. 8, HiFuse achieves a kernel reduction ratio of 43.6% to 73.2% compared to PyG across different datasets and models. This significant reduction is attributed to our offloading and merging optimizations. The reduction ratio in the RGAT model is smaller than in the RGCN model due to the additional kernels required for attention calculation.

Speedup. To evaluate the effect of each optimization on performance, we conduct an ablation study. The results, shown in Fig. 9, demonstrate that, compared to the baseline, using only the reorganization optimization results in a performance improvement of up to 1.17$\times$. This is due to improved data locality. Using both the reorganization and merging optimizations yields a performance increase of up to 1.83$\times$, benefiting from reduced kernel launch overhead and further improvements in data locality. Applying the reorganization, offloading, and parallelization optimizations results in a speedup of up to 1.7$\times$ due to reduced kernel launch time, improved data locality, and parallel execution. Additionally, incorporating the pipelining optimization achieves a speedup of up to 2.2$\times$.

An interesting observation is that the RGAT model benefits more from the merging optimization, while the RGCN model gains more from the offloading, parallelizing, and pipelining optimizations. This is because the RGAT model has a more computation-intensive model computation phase, making the GPU’s execution time more dominant. Therefore, optimizations that enhance the neighbor aggregation stage on GPU have a greater effect on reducing the execution time of the RGAT model. On the other hand, the RGCN model has a less computation-intensive model computation phase, meaning the GPU’s execution time is not as dominant. Consequently, optimizations targeting the CPU-executed phases are more effective for the RGCN model.

Ratio of CPU Time to GPU Time. To gain deeper insights into our optimizations, we analyze the ratio of CPU execution time to GPU execution time. The closer this ratio is to one, the more balanced the utilization of CPU and GPU, indicating more effective pipelined execution. As shown in Fig. 10, HiFuse achieves a ratio closer to one compared to PyG across different datasets and models. This indicates that our optimizations balance CPU and GPU utilization more effectively, resulting in more efficient pipelined execution. HiFuse achieves this by offloading edge index selection from GPU to CPU, parallelizing this process, and effectively bridging the dataflow between CPU and GPU with pipeline optimizations. Additionally, we observe that the impact differs between the RGAT and RGCN models. The RGAT model, which has a more computation-intensive model computation phase, benefits differently from these optimizations compared to the RGCN model.

Number of Kernel Reduction. To illustrate the effectiveness of the offloading and merging optimizations, we delve into the edge index selection and neighbor aggregation stages. As shown in Fig. 11, offloading edge index selection from GPU to CPU can reduce the number of kernels by up to 39.8%. Additionally, merging features to enable a single kernel for neighbor aggregation can reduce the number of kernels by up to 46.3%.

Compute and Memory Throughput. To demonstrate the optimization achieved through feature reorganization and merging, we investigate GPU compute and memory throughput. Compute throughput measures the efficiency of GPU computational resource utilization, while memory throughput indicates the efficiency of data transfer between the GPU’s memory and its computational units. As shown in Table 3, compared to PyG, the compute and memory throughput of the main kernel (‘scatter’) used in the neighbor aggregation stage are improved by up to 136$\times$ and 14$\times$, respectively. This improvement stems from feature reorganization and merging the inputs of multiple kernels into a single GPU kernel, offering several benefits: 1) Larger kernel sizes can better leverage the GPU’s parallel processing capabilities. 2) Fewer kernels significantly reduce kernel launch time overhead. 3) Enhanced data locality improves coalesced memory access and reduces memory access latency in the GPU pipeline, enabling faster data supply.