Efficient micro-batch training method and device for large-scale graph neural network

Abstract

The invention provides a high-efficiency micro-batch training method and device for a large-scale graph neural network, introduces an auxiliary graph constructed based on message passing, designs a graph division method based on the auxiliary graph, and distributes weights for edges based on out-degrees of source nodes. The target nodes sharing the same source node can be effectively distributed into the same micro-batch to reduce node redundancy, and load imbalance is reduced by uniformly distributing the nodes with high degree to different micro-batches as much as possible; in addition, the out-degree center graph division process is optimized, the expenditure for re-division in each training period is eliminated, a data loading mechanism specific to micro-batches is introduced, the micro-batch data loading efficiency in the training process is improved, and the expenditure for data preparation is reduced. According to the method, the memory wall problem faced by large-scale graph neural network training can be overcome, and large-scale graph neural network training exceeding GPU memory capacity becomes feasible and efficient.

Description

Technical Field

[0001] The present invention belongs to the technical field of graph data analysis and processing, and specifically relates to an efficient micro-batch training method and device for large-scale graph neural networks. Background Art

[0002] GNNs have become a powerful tool for learning and analyzing graph-structured data, and have achieved remarkable predictive performance in a variety of graph-related tasks, such as link prediction, node classification, and graph classification. However, training complex GNN models on large graphs faces huge challenges, especially in terms of memory scalability and computational efficiency. For example, GNNs learn node vector representations by iteratively aggregating features from neighboring nodes. Increasing the depth of aggregation leads to an exponential increase in the number of feature vectors that must be loaded into memory for aggregation. Considering that device memory is a scarce resource on existing mainstream hardware accelerators (such as GPUs), methods such as deep aggregation to improve model quality are prone to out of memory (OOM). This memory capacity limitation hinders the ability to train GNNs on large graphs using more complex GNN architectures, such as using deeper aggregation depths, larger sampling rates, and more advanced aggregators (e.g., LSTM), which is critical for capturing subtle patterns in graph data.

[0003] In order to break the GPU memory capacity limit for large-scale GNN training, a lot of research efforts have emerged in recent years in both algorithm and system optimization. On the algorithm side, researchers have proposed a variety of sampling techniques (e.g., node-by-node sampling, layer-by-layer sampling, and subgraph-by-subgraph sampling) to downsample the original graph before applying the GNN model so that the data can fit into the device memory. Although sampling enables GNNs to be trained on larger graphs, studies have shown that this may come at the expense of model accuracy. On the system side, researchers have developed specialized GNN frameworks such as DGL (Minjie Wang, DaZheng, Zihao Ye, Quan Gan, Mufei Li, Xiang Song, Jinjing Zhou, Chao Ma, Lingfan Yu, Yu Gai, et al. Deep graph library: A graph-centric, highly-performant package for graph neural networks. arXiv preprint arXiv:1909.01315, 2019), PyG (Matthias Fey and Jan Eric Lenssen. Fast graph representation learning with pytorchgeometric. arXiv preprint arXiv:1903.02428, 2019), and NeuGraph (Lingxiao Ma, ZhiYang, Youshan Miao, Jilong Xue, Ming Wu, Lidong Zhou, and Yafei Dai. Neugraph: Parallel deep neural network computation on large graphs. In 2019USENIX AnnualTechnical Conference, pages 443-458, 2019), improve the efficiency of GNN training by implementing efficient graph operation primitives. However, these frameworks encounter memory scalability challenges when processing large graphs or complex GNN architectures.Another option is to use a distributed framework such as DistDGL (Da Zheng, Chao Ma, Minjie Wang, Jinjing Zhou, Qidong Su, Xiang Song, Quan Gan, Zheng Zhang, and George Karypis. Distdgl: distributed graph neural network training for billion-scale graphs. In2020IEEE / ACM 10th Workshop on Irregular Applications: Architectures and Algorithms, pages 36-44. IEEE, 2020) to partition and process large graphs in a distributed manner, leveraging the combined memory of multiple GPUs and / or computing nodes to scale model training. While this approach is effective, it typically involves significantly higher hardware costs, which makes it impractical for many model scientists with limited hardware budgets.

[0004] Among many methods, mini-batch training has emerged as a promising technique in deep learning. It involves dividing a mini-batch into K smaller micro-batches and accumulating the gradients of these micro-batches to update the model weights, such as Figure 1 As shown. This approach is able to reduce memory usage while maintaining model convergence and quality. Compared with traditional neural network (NN) models, GNNs exhibit more complex dependencies (e.g., N:M) between output / labels and input features within a batch, rather than a simple 1:1 mapping. Applying micro-batch training to GNN models presents challenges such as node redundancy, load imbalance, and data preparation overhead. Although recent work such as Betty (Shuangyan Yang, Minjia Zhang, Wenqian Dong, and Dong Li. Betty: Enabling large-scale gnn training with batch-level graph partitioning. In Proceedings of the 28th ACM International Conference on Architectural Support for Programming Languages ​​and Operating Systems, Volume 2, pages 103-117, 2023) has explored GNN-based micro-batch training, it still has the following defects:

[0005] The shortcomings of (1) insufficient redundancy reduction, (2) computationally unfriendly load balancing strategies, and (3) high data preparation overhead have resulted in suboptimal performance in terms of memory consumption reduction and training efficiency, limiting the practicality and performance potential of micro-batch training methods for large-scale GNN model training, and thus limiting the applicability of large-scale CNN models in computer equipment applications. Summary of the invention

[0006] In order to solve the above problems existing in the prior art, the present invention provides an efficient micro-batch training method for large-scale graph neural networks. The technical problem to be solved by the present invention is achieved through the following technical solutions:

[0007] In a first aspect, the present invention provides an efficient micro-batch training method for large-scale graph neural networks, comprising:

[0008] S100, using a graph neural network GNN to perform feature aggregation on an input graph, and determining source nodes, target nodes, and edges between target nodes involved in a message passing process of the GNN;

[0009] S200, constructing a global message passing graph using the source node, the target node and the edges between the target nodes;

[0010] S300, dividing the global message passing graph into subgraphs corresponding to each micro-batch based on an out-degree partitioning method, and forming a batch training node list from nodes in each subgraph;

[0011] S400, traversing the batch training node list to create a corresponding micro-batch loader MBL for each micro-batch for data training of the micro-batch.

[0012] In a second aspect, the present invention provides an efficient micro-batch training device for large-scale graph neural networks, comprising:

[0013] A determination module is configured to perform feature aggregation on an input graph using a graph neural network (GNN) and determine source nodes, target nodes, and edges between target nodes involved in a message passing process of the GNN;

[0014] A construction module configured to construct a global message passing graph using the source node, the target node, and the edges between the target nodes;

[0015] A partitioning module is configured to partition the global message passing graph into subgraphs corresponding to each micro-batch based on an out-degree partitioning method, and form a batch training node list from nodes in each subgraph;

[0016] The traversal module is configured to traverse the batch training node list to create a corresponding micro-batch loader MBL for each micro-batch for data training of the micro-batch.

[0017] In a third aspect, the present invention provides a computer device comprising: a memory and a processor, wherein the memory and the processor are coupled;

[0018] The memory is used to store one or more computer instructions;

[0019] The processor is used to execute the one or more computer instructions to implement the efficient micro-batch training method for large-scale graph neural networks as described in the first aspect.

[0020] The present invention has at least the following beneficial effects:

[0021] 1. The present invention introduces a message passing flow graph to assist in realizing graph partitioning. By assigning weights to edges and based on the out-degree of source nodes, Wenger can effectively divide target nodes that share the same source node into the same micro-batch, while maintaining load balance by distributing nodes with high out-degree to different micro-batches as evenly as possible.

[0022] 2. The subgraph partitioning method of the present invention can reduce redundant nodes in micro-batches and reduce the load imbalance problem during training by intelligently allocating nodes with high connectivity.

[0023] 3. The present invention optimizes the out-degree central graph partitioning process by using a global message passing flow graph, thereby avoiding the overhead of re-partitioning in each training cycle.

[0024] 4. The present invention introduces a data loading mechanism for micro-batches, which improves the efficiency of data loading and reduces the overhead of data preparation.

[0025] 5. The present invention combines the message passing flow graph, the out-degree central graph partitioning method and the micro-batch data loading mechanism, so that Wenger can efficiently train GNN on large graphs outside the memory capacity limit, making the training more feasible and efficient. Although the present invention adopts micro-batch training, it still maintains the model convergence and accuracy comparable to full batch training. Compared with the existing state-of-the-art methods, the training method of the present invention achieves significant reductions in both memory consumption and training time.

[0026] The present invention will be further described in detail below with reference to the accompanying drawings and embodiments. BRIEF DESCRIPTION OF THE DRAWINGS

[0027] Figure 1 This is a comparative diagram of micro-batch training and small batch training;

[0028] Figure 2 It is a flow chart of an efficient micro-batch training method for large-scale graph neural networks of the present invention;

[0029] Figure 3 It is a workflow diagram of the efficient micro-batch training method for large-scale graph neural networks of the present invention;

[0030] Figure 4 It is a comparative schematic diagram of a local message passing graph and a global message passing graph provided by the present invention;

[0031] Figure 5 It is a schematic diagram comparing the out-degree-based graph partitioning provided by the present invention with other graph partitioning results. DETAILED DESCRIPTION

[0032] The present invention is further described in detail below with reference to specific embodiments, but the embodiments of the present invention are not limited thereto.

[0033] Aiming at the challenges of applying micro-batch methods to GNN training in the prior art and the GPU memory capacity limitation and computational efficiency problems of existing GNN micro-batch training methods, the present invention provides an efficient micro-batch training method designed specifically for large-scale GNN training. The key to the training method of the present invention is the combination of graph partitioning technology assisted by message passing flow graph and a specific data loading mechanism for micro-batches.

[0034] In order to reduce redundancy and load imbalance, the training method of the present invention adopts an out-degree-based auxiliary graph partitioning method, which uses message passing flow partitioning to construct a global message passing graph. This method assigns weights to edges according to the out-degree of the source node, embeds the out-degree information in the global message passing graph, and can effectively group the target nodes that share the same source node into the same micro-batch, while maintaining load balance by distributing nodes with high out-degrees in different micro-batches as evenly as possible.

[0035] In order to reduce the data preparation overhead, the present invention adapts the graph partitioning method process of the out-degree center to the global message passing flow graph to avoid repeated partitioning overhead. In addition, the training of the present invention introduces a data loading mechanism for micro-batches to improve the data loading efficiency of micro-batch GNN training. These schemes jointly reduce the redundancy and load imbalance within different micro-batches, and reduce the data preparation overhead in micro-batch training, making large-scale GNN training beyond the GPU memory capacity feasible and efficient.

[0036] The technical solution details of the present invention are described in detail below.

[0037] like Figure 2 and Figure 3As shown, the present invention provides an efficient micro-batch training method for large-scale graph neural networks. Figure 2 It is a flow chart of the efficient micro-batch training method for large-scale graph neural networks of the present invention. Figure 3 It is a workflow diagram of the efficient micro-batch training method for large-scale graph neural networks of the present invention. The present invention utilizes an auxiliary graph called a message passing flow graph and combines the following key technologies: (1) an out-degree-based partitioning strategy to solve the redundancy and load imbalance problems of different micro-batches in micro-batch GNN training, and (2) a duplicate partitioning elimination mechanism and a specific data loader for micro-batches to reduce the partitioning and data loading overhead in micro-batch GNN training. The present invention constructs a global message passing flow graph (GMFG) by tracking the nodes and edges involved in the GNN message passing flow. Then, using the out-degree-based partitioning method, the GMFG is partitioned into subgraphs representing each micro-batch according to a predetermined number of micro-batches. This partitioning minimizes redundancy and load imbalance within different micro-batches and eliminates the need for repeated graph partitioning across training cycles. Similar to Betty, Wenger uses a memory estimation-based method to determine the optimal number of micro-batches before partitioning. Finally, Wenger uses the partitioned input node list to instantiate a Micro-Batch Loader (MBL) to independently load the corresponding graph data for each micro-batch, which helps improve the data loading efficiency of subsequent micro-batch training.

[0038] Combination Figure 2 and Figure 3 , an efficient micro-batch training method for large-scale graph neural networks of the present invention includes:

[0039] S100, using a graph neural network GNN to perform feature aggregation on an input graph, and determining source nodes, target nodes, and edges between target nodes involved in a message passing process of the GNN;

[0040] It is worth noting that a message passing flow graph (MFG) is a graph that represents the direction of data flow, where the target node receives messages from a subset of the source nodes, and the edges represent the computational dependencies involved in the aggregation of messages during neighbor sampling. The MFG is referred to as "local" here, meaning that there is an independent message passing flow graph for each training cycle. LMFG dynamically captures the directional relationship between nodes during the message aggregation process. Specifically, for each layer of the GNN, the target node aggregates features from its neighboring nodes, and these source nodes of one layer become the target nodes of the next layer. By tracking this process, a LMFG that reflects the directionality used during training can be constructed, such as Figure 4As shown, in each training cycle, LMFG changes as the data is shuffled.

[0041] In a specific embodiment of the present invention, S100 includes:

[0042] S110, inputting the input graph into the graph neural network GNN, so that the target node of each layer in the GNN aggregates the features of the input graph from its neighbor nodes, and then transmits them to the nodes of the next layer in a message passing manner; wherein the node in the GNN used to aggregate the features is the target node, and the target node receives the message from the source node;

[0043] S120, determining the target nodes and edges between the target nodes involved in the message passing process in each layer of the GNN, wherein each edge represents the computational dependency involved in message aggregation during neighbor node sampling.

[0044] The purpose of acquiring the source node, the target node and the edge in the present invention is to construct a global message passing graph (GMFG). The process of the global message passing graph is consistent with the process of the local message passing graph.

[0045] S200, constructing a global message passing graph using the source node, the target node and the edges between the target nodes;

[0046] It is worth noting that LMFG is an auxiliary graph that captures the message passing process of GNN and plays a vital role in reducing redundancy and load imbalance during graph partitioning. With the help of LMFG, Wenger's partitioning method uses the out-degree of nodes in LMFG to transform the problem of minimizing redundancy and load imbalance into the problem of minimizing edge cut weights. This strategy enables Wenger to simultaneously minimize redundancy and load imbalance within different micro-batches.

[0047] Leveraging the directional relationships embedded in LMFG, Wenger's partitioning strategy aims to optimize the distribution of training nodes in micro-batches, promoting balanced memory usage and efficient computation throughout the training process. The key to this graph partitioning method is its ability to transform the redundancy and load imbalance minimization problem into the problem of minimizing edge cut weights. In out-degree-based graph partitioning, the edge weights in LMFG are determined by the out-degree of the source node. This choice is crucial because the out-degree provides valuable insights into the partitioning in two ways: 1) Higher edge weights mean that the source node is reused more frequently, which may lead to increased redundancy if the edge is cut. 2) Higher edge weights indicate larger node clusters around the source node, which helps to evenly distribute highly connected nodes in different subgraphs to improve load balancing.

[0048] With the help of the informative partitioning auxiliary property of LMFG, Wenger effectively groups the target nodes with common source nodes into the same micro-batch, reducing redundancy. In addition, the method alleviates the load imbalance by evenly distributing nodes with high out-degree in micro-batches, since these nodes can significantly affect the peak memory consumption during training.

[0049] Leveraging the directional relationships embedded in LMFG, Wenger's partitioning strategy aims to optimize the distribution of training nodes in micro-batches, promoting balanced memory usage and efficient computation throughout the training process. The key to this graph partitioning method is its ability to transform the redundancy and load imbalance minimization problem into the problem of minimizing edge cut weights. In out-degree-based graph partitioning, the edge weights in LMFG are determined by the out-degree of the source node. This choice is crucial because the out-degree provides valuable insights into the partitioning in two ways: 1) Higher edge weights mean that the source node is reused more frequently, which may lead to increased redundancy if the edge is cut. 2) Higher edge weights indicate larger node clusters around the source node, which helps to evenly distribute highly connected nodes in different subgraphs to improve load balancing.

[0050] With the help of the informative partitioning auxiliary property of LMFG, Wenger effectively groups the target nodes with common source nodes into the same micro-batch, reducing redundancy. In addition, the method alleviates the load imbalance by evenly distributing nodes with high out-degree in micro-batches, since these nodes can significantly affect the peak memory consumption during training.

[0051] S300, dividing the global message passing graph into subgraphs corresponding to each micro-batch based on an out-degree partitioning method, and forming a batch training node list from nodes in each subgraph;

[0052] There are high redundancy and load imbalance problems in current graph partitioning methods, which limit the effectiveness of micro-batch training in reducing GNN peak memory consumption and improving computational efficiency. Unlike previous methods that focus on minimizing the weight of edge cuts to balance the number of edges or nodes in subgraphs without solving the redundancy problem, the present invention introduces an out-degree-centric graph partitioning method to simultaneously minimize redundancy and load imbalance, thereby improving performance. Due to the use of micro-batch GNN training methods such as Local Message-Passing Graph (LMFG), repeated partitioning and inefficient data loading are required during the training cycle, resulting in high data preparation overhead. The key innovation of the out-degree-centric graph partitioning method proposed in the present invention is the use of an auxiliary graph, called the Global Message-Passing Graph (GMFG), and the corresponding GMFG partitioning scheme.

[0053] Refer to the following algorithm 1:

[0054]

[0055] Algorithm 1 shows the graph partitioning process of out-degree center. Initially, the algorithm takes LMFG and the specified number of mini-batches as input. It first extracts source nodes and target nodes from LMFG (line 1), and then creates a graph structure using these nodes (line 2). Next, the out-degree of each source node is calculated and set as the weight of the edge connected to the source node in the graph (lines 3-4), which reflects the influence of each source node on the training node. Nodes with higher out-degrees receive larger weights, indicating that they have a greater impact on peak memory consumption during training. The algorithm then optimizes the graph by eliminating non-output nodes and self-loops (lines 5-6), ensuring that the partitioning process is focused and efficient.

[0056] In Algorithm 1, an existing partitioning algorithm is applied to partition the weighted LMFG into subgraphs corresponding to each micro-batch (line 7). The present invention uses the partitioning algorithm of Metis (George Karypis. Metis: Unstructured graph partitioning and sparse matrix ordering system. Technical report, 1997) because it is scalable when handling large-scale graph partitioning tasks. This step ensures that nodes with high out-degree are fairly distributed among micro-batches, promotes load balancing, and thus minimizes peak memory consumption. Finally, the partitioning algorithm traverses the subgraph, extracts the target node, and appends the target node to `batched_train_nodes_list` (lines 8-12) as the output of the partitioning process. This list clearly defines the nodes assigned to each micro-batch, which helps to achieve micro-batch training that reduces peak memory consumption and balances computational load.

[0057] refer to Figure 5 , Figure 5 A concrete example is shown, demonstrating how to partition a particular LMFG into two micro-batches while minimizing redundant nodes. The figure shows the actual partitioning results achieved by Wenger ( Figure 5 For comparative analysis, the figure also shows two additional partitioning results, which correspond to alternative partitioning methods that do not prioritize minimizing edge cut weights ( Figure 5(a) and (c) in the figure). Wenger's out-degree-based partitioning technique focuses on minimizing the total weight of the cut edges and successfully cuts the edges with a total weight of 6. This process produces two different subsets of training nodes, denoted as set AC and set B, with only 2 redundant nodes. In contrast, the total weights of the cut edges produced by the two alternative partitioning methods are 7 and 9, respectively, resulting in 3 and 4 redundant nodes, respectively. This example demonstrates the efficiency of the graph partitioning method provided by the present invention in reducing redundancy.

[0058] It is worth noting that the micro-batch GNN training methods in the prior art, such as Betty, require repeated partitioning and inefficient data loading during the training cycle, resulting in high data preparation overhead. In order to solve these two problems, the present invention introduces the global message passing flow graph GMFG as an extended version of the local message passing flow graph. GMFG provides a comprehensive graph structure before training begins, covering all potential message passing paths. Using GMFG, the present invention proposes two solutions: (1) Adapting the out-degree-based graph partitioning so that the one-time partitioning results can be reused in each cycle, eliminating the need for repeated partitioning; (2) Based on the partitioning results, a data loading mechanism for micro-batches is designed to improve the data loading efficiency in micro-batch GNN training.

[0059] In prior art training schemes like Betty, the dynamic changes in the topology of the input graph require new partitioning at each training cycle. This requirement stems from the cross-cycle data shuffling, which aims to improve the generalization ability of the model by changing the neighboring nodes and edges associated with each training node. It is worth noting that Wenger faces a similar requirement. Although the basic Wenger with the auxiliary graph LMFG effectively alleviates redundancy and load imbalance, it still requires a new partitioning step with an updated LMFG at each training cycle.

[0060] The necessity of repartitioning stems from the inherent limitation of LMFG in capturing comprehensive graph structural information. LMFG only captures a subset of the message passing flow in the original graph by sampling neighboring nodes at each cycle. Therefore, although LMFG is effective in partitioning, it only has a partial representation of the graph structure, which requires repeated partitioning to adapt to the evolving graph dynamics.

[0061] To address this limitation, the present invention utilizes the inherent structural consistency between various LMFGs by introducing a global message flow graph GMFG. Figure 3As shown, constructing GMFG follows a message passing process similar to LMFG. The notable difference is that GMFG, as a comprehensive graph structure, includes all potential neighboring nodes and edges involved in passing messages to training nodes. This overall representation ensures that each LMFG is a subgraph of GMFG, and promotes more extensive extraction of redundant information through the additional nodes and edges in GMFG. Therefore, the present invention adjusts the partitioning method of the out-degree central graph based on GMFG.

[0062] In a specific implementation of the present invention, S300 includes:

[0063] S310, extracting a source node and a target node in the global message passing graph;

[0064] S320, calculating the out-degree of each source node;

[0065] S330, according to the out-degree of each source node, setting corresponding weights for the edges connected to the source nodes to obtain a weighted global message passing graph;

[0066] S340: using a partitioning method, partition the weighted global message passing graph into subgraphs corresponding to each micro-batch.

[0067] In a specific implementation of the present invention, S340 includes:

[0068] S341, determining the optimal number of micro-batches using a memory estimation-based method;

[0069] S342 , dividing the weighted global message passing graph into subgraphs corresponding to the respective micro-batches by using a partitioning method according to the number of subgraphs being the same as the optimal number of micro-batches.

[0070] In this step, the number of subgraphs is set to be the same as the optimal number of micro-batches; and the weighted global message passing graph is divided into subgraphs corresponding to each micro-batch using a partitioning method according to the principle of minimizing the total weight of the cut edges.

[0071] In this improved partitioning method of the present invention, instead of using a new LMFG for each training cycle, the proposed graph partitioning algorithm is applied directly on the GMFG to create a batch list of training nodes, including reusable micro-batch training nodes. By adopting this strategy, the one-time partitioning results can be reused in consecutive training cycles, eliminating the need for repeated partitioning during micro-batch GNN training. Importantly, the transition from LMFG to GMFG does not weaken the effectiveness of Wenger's graph partitioning technique in reducing redundancy and load imbalance, because GMFG retains the same basic partitioning auxiliary properties as LMFG that are critical to solving these challenges.

[0072] S400, traversing the batch training node list to create a corresponding micro-batch loader MBL for each micro-batch for data training of the micro-batch.

[0073] In a specific implementation of the present invention, S400 includes:

[0074] S410, traversing the batch training node list to create a corresponding loader MBL for each micro-batch for data training of the micro-batch;

[0075] S420, aggregate all MBLs into a list of micro-batch data loaders.

[0076] As mentioned above, the data loading efficiency of Betty recorded in the prior art is low, which is caused by the cumbersome data extraction and index mapping process. To avoid this inefficiency, the present invention introduces a micro-batch specific data loading mechanism to simplify the data loading process. Utilizing highly optimized data loaders in GNN frameworks like DGL, the present invention creates a dedicated micro-batch loader MBL for each micro-batch. According to the requirements of DGL, building a data loader using DGL's DataLoader class requires fixed training nodes. This can be easily met by the results of the out-degree center partitioning based on GMFG.

[0077] As shown in Algorithm 2 below:

[0078] Algorithm 2 shows the generation process of MBL. First, GMFG is constructed using training nodes and input graph (line 1). After GMFG is constructed, the present invention performs out-degree-based partitioning to obtain a list of batch training nodes (line 2). Then, the algorithm traverses this list to instantiate a dedicated MBL for each micro-batch (lines 3-6), and finally returns a list of MBLs (line 7). These MBLs can be used for data loading across micro-batches during the training cycle.

[0079] In the proposed scheme, each MBL runs as an iterable entity and performs data shuffle independently, thereby separating the data loading process of different micro-batches. This separation ensures that Wenger's data loading process is independent of the full batch data loader, eliminating the need for cumbersome data extraction and index mapping procedures seen in Betty. All MBLs are stored in the host memory and share access to the same input graph as the full batch data loader. It has been observed that the use of MBLs only adds tens of megabytes of additional memory overhead, which is usually acceptable in practical scenarios because host memory resources are usually abundant.

[0080] It is worth noting that when using MBLs, the same training node may aggregate features from different neighboring nodes compared to using the full batch loader due to the different ways they operate during the shuffling of training nodes and neighboring nodes, with each MBL shuffling its own nodes independently, while the full batch loader shuffles all nodes at once. Despite this difference, given that MBLs and full batch loaders operate on the same input graph, the set of neighboring nodes and associated features for each training node are consistent in both loading methods, which ensures similar model convergence between them.

[0081] The method of the present invention can be implemented based on DGL 2.0.0 and PyTorch 2.0.1. It uses the API of DGL to build a global message passing flow graph (GMFG), out-degree-based graph partitioning, and micro-batch data loading. The training method of the present invention is set in a Python module compatible with DGL to facilitate large-scale GNN training that breaks through memory capacity limitations.

[0082] Specifically, the implementation process of the global GMFG construction is as follows: Wenger constructs the GMFG by tracing the message passing process of all training nodes. It creates a dgl.dataloader, covers all neighbors by sampling nodes, so that the neighbor nodes used in the training node message passing can be retrieved layer by layer through the dataloader iteration. This process forms a dgl.graph representing the GMFG. The implementation process of the out-degree graph partitioning is as follows: Wenger uses dgl.graph.edges()[0] to get the index of all source nodes, and uses dgl.graph.out_degrees() to calculate the out-degree of these nodes. These out-degrees are used as the weights of the edges after removing the self-loops through dgl.rm_self_loop(). Subsequently, the graph is partitioned through dgl.metis(), and the partition results are iterated to extract the training nodes of the micro-batch. The implementation process of the micro-batch data loader is as follows: Wenger uses the training nodes and the input graph to create a dgl.dataloader for each micro-batch, and then aggregates them into a list of micro-batch data loaders, which is convenient for efficient retrieval of each micro-batch data layer by layer during micro-batch training.

[0083] In an optional implementation of the present invention, after S400, the efficient micro-batch training method for large-scale graph neural networks further includes:

[0084] For each micro-batch of training data, all MBLs of the micro-batch are selected from the list of the micro-batch data loader to synchronously train the training data of the micro-batch.

[0085] In a second aspect, the present invention provides an efficient micro-batch training device for large-scale graph neural networks, including:

[0086] A determination module is configured to perform feature aggregation on an input graph using a graph neural network (GNN) and determine source nodes, target nodes, and edges between target nodes involved in a message passing process of the GNN;

[0087] A construction module configured to construct a global message passing graph using the source node, the target node, and the edges between the target nodes;

[0088] A partitioning module is configured to partition the global message passing graph into subgraphs corresponding to each micro-batch based on an out-degree partitioning method, and form a batch training node list from nodes in each subgraph;

[0089] The traversal module is configured to traverse the batch training node list to create a corresponding micro-batch loader MBL for each micro-batch for data training of the micro-batch.

[0090] In a third aspect, the present invention provides a computer device, comprising: a memory and a processor, wherein the memory and the processor are coupled;

[0091] The memory is used to store one or more computer instructions;

[0092] The processor is used to execute the one or more computer instructions to implement an efficient micro-batch training method for large-scale graph neural networks.

[0093] The present invention provides an efficient micro-batch training method and device for large-scale graph neural networks, introduces an auxiliary graph built based on message passing, and designs a graph partitioning method based on the auxiliary graph. The method assigns weights to edges based on the out-degree of the source node, and can effectively divide the target nodes sharing the same source node into the same micro-batch to reduce node redundancy, and reduce load imbalance by allocating nodes with high out-degree to different micro-batches as evenly as possible; in addition, the present invention optimizes the out-degree center graph partitioning process, eliminates the overhead of re-dividing each training cycle, and introduces a data loading mechanism specific to micro-batches, which improves the efficiency of micro-batch data loading during training and reduces the overhead of data preparation. The present invention can overcome the memory wall problem faced by large-scale graph neural network training and make large-scale graph neural network training that exceeds the GPU memory capacity feasible and efficient.

[0094] It is worth noting that the terms "first" and "second" in the present invention are only used for descriptive purposes and cannot be understood as indicating or implying relative importance or implicitly indicating the number of the indicated technical features. Therefore, the features defined as "first" and "second" may explicitly or implicitly include one or more of the features. In the description of the present invention, the meaning of "plurality" is two or more, unless otherwise clearly and specifically defined.

[0095] Although the present application is described herein in conjunction with various embodiments, in the process of implementing the claimed application, those skilled in the art may understand and implement other variations of the disclosed embodiments by reviewing the drawings, the disclosure, and the appended claims. In the claims, the word "comprising" does not exclude other components or steps, and "a" or "an" does not exclude a plurality of components or steps.

[0096] The above contents are further detailed descriptions of the present invention in combination with specific preferred embodiments, and it cannot be determined that the specific implementation of the present invention is limited to these descriptions. For ordinary technicians in the technical field to which the present invention belongs, several simple deductions or substitutions can be made without departing from the concept of the present invention, which should be regarded as falling within the protection scope of the present invention.

Claims

1. An efficient micro-batch training method for large-scale graph neural networks, characterized in that: include: S100, using a graph neural network GNN to perform feature aggregation on an input graph, and determining source nodes, target nodes, and edges between target nodes involved in a message passing process of the GNN; S200, constructing a global message passing graph using the source node, the target node and the edges between the target nodes; S300, dividing the global message passing graph into subgraphs corresponding to each micro-batch based on an out-degree partitioning method, and forming a batch training node list from nodes in each subgraph; S400, traversing the batch training node list to create a corresponding micro-batch loader MBL for each micro-batch for data training of the micro-batch.

2. The efficient micro-batch training method for large-scale graph neural networks according to claim 1 is characterized in that: S100 includes: S110, inputting the input graph into the graph neural network GNN, so that the target node of each layer in the GNN aggregates the features of the input graph from its neighbor nodes, and then transmits them to the nodes of the next layer in a message passing manner; wherein the node in the GNN used to aggregate the features is the target node, and the target node receives the message from the source node; S120, determining the target nodes and edges between the target nodes involved in the message passing process in each layer of the GNN, wherein each edge represents the computational dependency involved in message aggregation during neighbor node sampling.

3. The efficient micro-batch training method for large-scale graph neural networks according to claim 1 is characterized in that: S300 includes: S310, extracting a source node and a target node in the global message passing graph; S320, calculating the out-degree of each source node; S330, according to the out-degree of each source node, setting corresponding weights for the edges connected to the source nodes to obtain a weighted global message passing graph; S340: using a partitioning method, partition the weighted global message passing graph into subgraphs corresponding to each micro-batch.

4. The efficient micro-batch training method for large-scale graph neural networks according to claim 3 is characterized in that: S340 includes: S341, determining the optimal number of micro-batches using a memory estimation-based method; S342 , dividing the weighted global message passing graph into subgraphs corresponding to the respective micro-batches by using a partitioning method according to the number of subgraphs being the same as the optimal number of micro-batches.

5. The efficient micro-batch training method for large-scale graph neural networks according to claim 4 is characterized in that: S342 includes: Setting the number of subgraphs to be equal to the optimal number of micro-batches; According to the principle of minimizing the total weight of the cut edges, the weighted global message passing graph is divided into subgraphs corresponding to each micro-batch using a partitioning method.

6. The efficient micro-batch training method for large-scale graph neural networks according to claim 1 is characterized in that: S400 includes: Traversing the batch training node list to create a corresponding loader MBL for each micro-batch for data training of the micro-batch; Aggregate all MBLs into a list of micro-batch data loaders.

7. The efficient micro-batch training method for large-scale graph neural networks according to claim 6 is characterized in that: After S400, the efficient micro-batch training method for large-scale graph neural networks further includes: For each micro-batch of training data, all MBLs of the micro-batch are selected from the list of the micro-batch data loader to synchronously train the training data of the micro-batch.

8. The efficient micro-batch training method for large-scale graph neural networks according to claim 1 is characterized in that: The efficient micro-batch training method for large-scale graph neural networks is implemented based on DGL 2.0.0 and PyTorch 2.0.

1.

9. An efficient micro-batch training device for large-scale graph neural networks, characterized in that: include: A determination module is configured to perform feature aggregation on an input graph using a graph neural network (GNN) and determine source nodes, target nodes, and edges between target nodes involved in a message passing process of the GNN; A construction module configured to construct a global message passing graph using the source node, the target node, and the edges between the target nodes; A partitioning module is configured to partition the global message passing graph into subgraphs corresponding to each micro-batch based on an out-degree partitioning method, and form a batch training node list from nodes in each subgraph; The traversal module is configured to traverse the batch training node list to create a corresponding micro-batch loader MBL for each micro-batch for data training of the micro-batch.

10. A computer device, characterized in that: include: A memory and a processor, wherein the memory and the processor are coupled; The memory is used to store one or more computer instructions; The processor is used to execute the one or more computer instructions to implement the efficient micro-batch training method for large-scale graph neural networks as described in any one of claims 1-7.

Images