A distributed graph neural network training method based on heterogeneous devices

Abstract

This paper discloses a distributed graph neural network training method based on heterogeneous devices. This method processes large graph data and uses a graph cut algorithm to obtain several subgraphs with balanced nodes and edges. A dynamic programming algorithm is then used to determine an optimal device placement plan. The training workload for each device is planned based on the training capabilities of the heterogeneous devices, allowing all devices to complete training tasks as simultaneously as possible within the same iteration, thereby improving edge device utilization. A fault-tolerance mechanism is also added to ensure that the network maintains a certain level of functional performance even when network nodes fail or are lost.

Description

Technical Field

[0001] The present invention relates to the field of heterogeneous device training, and in particular to a distributed graph neural network training method based on heterogeneous devices. Background Art

[0002] Graphs, composed of nodes and the relationships (edges) between them, possess powerful expressive capabilities. With the maturity of machine learning (ML) technology, researchers have combined convolutional neural networks (GCNs) with graph embedding algorithms (GEAs) to process complex graph data structures, proposing the concept of graph neural networks (GNNs). GNNs extend machine learning algorithms from feature extraction in Euclidean space to non-Euclidean space, achieving state-of-the-art results in a variety of graph-related tasks, such as node classification, graph classification, and link prediction.

[0003] With the advent of the big data era, the scale of graphs has continued to expand. The memory and training resources of a single machine are insufficient to support the training of large graphs, posing a new challenge for graph neural networks. Unlike previous approaches, this time, training is based on sampled subgraphs (mini-batch) rather than full-graph training. Full-graph training involves all nodes participating in training and gradient updates. This approach offers the advantage of no accuracy loss, but suffers from high memory usage and the inability to train large graphs. Sampled subgraph training effectively addresses these drawbacks. Typically, a group of central nodes is randomly selected, and subgraphs are generated based on node sampling. Graph neural network computations and gradient updates are then performed on these subgraphs. This training model is called a distributed graph neural network (DGNN). Distributed graph neural networks have the advantages of low memory usage and ease of parallelization. However, the random sampling algorithm has a high degree of randomness, which can compromise model accuracy. This distributed model also introduces many new challenges. Inter-machine communication overhead and parallel training time are both significant components of distributed graph neural network training.

[0004] Current mainstream distributed training methods are based on DistDGL (referenced paper: DistDGL: Distributed Graph Neural Network Training for Billion-Scale Graphs) and use GPUs with the same training capabilities for training. However, the resources of many edge devices are not fully utilized, and training resources are extremely valuable. Improving the utilization of edge devices is a major challenge in the field of graph neural networks.

[0005] In practice, a system typically involves multiple devices with varying training capabilities (i.e., heterogeneous devices). Therefore, this paper proposes a distributed graph neural network training method based on heterogeneous devices. This method is not only applicable to heterogeneous device applications but also effectively improves edge device utilization. Summary of the Invention

[0006] In response to the defects of the existing technology, the present invention proposes a distributed graph neural network training method based on heterogeneous devices. The training task volume of each device is planned according to the training capabilities of the heterogeneous devices, so that all devices can complete the training tasks as simultaneously as possible in the same iteration, thereby improving the utilization rate of edge devices.

[0007] The present invention specifically comprises the following steps:

[0008] Step A: Evaluate the training capability of each heterogeneous device participating in the training, obtain the parameters of each device's training capability, and rank the training capabilities of each device, specifically:

[0009] The central node sends the same subgraph training task to each device using the same training model and training parameters. It records the forward propagation and backward propagation time of each device, thereby obtaining the training capacity of each device and ranking them to prepare for the subsequent maximization of each training device.

[0010] During the evaluation process, in order to reduce the evaluation error caused by the performance evaluation or testing of the device, each device is tested repeatedly, and after obtaining multiple test data, the average is calculated.

[0011] Step B: Process the large graph data and use a graph cutting algorithm to obtain several subgraphs with balanced nodes and edges (the subgraphs must be able to be trained on the device with the weakest training capability); divide the large graph data that would overflow the memory when directly trained into subgraphs suitable for edge device training, so that edge devices can participate in sharing the training tasks, improve the utilization rate of edge devices, and effectively break through the limitation that large graphs cannot be trained on edge devices.

[0012] Step C: Use a dynamic programming algorithm to derive an optimal device placement plan. Devices with strong training capabilities will receive more training tasks, while devices with weak training capabilities will receive fewer training tasks. Based on the resulting plan strategy, each subgraph is placed on the corresponding device for training, ensuring that each device completes the training tasks assigned by the algorithm as much as possible within the same timeframe. Edge devices share the training tasks of the neural network without increasing the overall training time, further improving the utilization of edge devices.

[0013] Step D: Sample the training nodes of the subgraph to obtain a set of mini-batches, and load the mini-batches into the neural network for training.

[0014] During the sampling process, different sampling strategies are adopted according to the corresponding needs, such as neighbor sampling, partial neighbor sampling, random walk and other different sampling algorithms.

[0015] Step E: The subgraph is trained using the corresponding graph neural network model, and the parameters are synchronized to the central node for parameter aggregation and storage. The resulting parameters are then broadcast to each training device for a new round of training. The graph neural network aggregates the edge neighbor information of the target vertex to generate an intermediate embedding. The neural network then uses the aggregated features to convert them into output features, which are sent as input features to the next layer. The final embedding is obtained by stacking multiple graph neural network layers.

[0016] Specifically, the representation of graph neural network training vertices is as follows: ; ;

[0017] in Vertex Aggregated neighbor information, For the vertex connected vertex sets, Represents a vertex The number of neighbors of Vertex In the Layer embedding, represents a nonlinear activation function, For the The learnable parameters of the layer.

[0018] When each round of training is completed, The parameters and the number of training rounds are sent to the central node. When the central node receives the parameters of each device, it aggregates them to obtain the new parameters. , and broadcast it to all training devices. All devices will perform the next round of training after receiving the new parameters.

[0019] Step F: Enable the fault-tolerant detection mechanism to check whether the training equipment can train and transmit parameters normally in each round until the training is completed. During the training process, the training may not be completed normally due to problems such as memory overflow, bit misalignment, and network communication failure. Check whether the training equipment can send the training parameters to the central node normally and whether it can normally receive the parameters transmitted by the central node. If an abnormality occurs, return to step C to obtain a new placement strategy and broadcast the parameters of the previous round with the abnormality to all training equipment for parameter synchronization and retraining. This reduces the losses caused by the impact of the failure, further improves resource utilization, and also improves the robustness of the distributed model. Training is stopped when the graph data training reaches a convergence state.

[0020] Compared with the prior art, the present invention has the following beneficial effects:

[0021] 1. It breaks through the limitation that mainstream training methods can only use devices with the same training capabilities. Different placement strategies can be designed by combining and changing heterogeneous devices to train graph neural networks. Since directly training with large graphs will cause memory overflow, the use of a graph slicing algorithm reduces the memory occupied by the graph, allowing devices with weaker edge capabilities such as Raspberry Pi and development boards to be used for training. At the same time, copying the neighbor nodes of some edge points in the graph can reduce the loss rate of structural information, thereby greatly improving the utilization rate of edge devices.

[0022] 2. A fault-tolerant mechanism has been added. When a network node fails or is lost, the network can still maintain a certain level of functional performance. When an edge device fails to work properly, the dynamic programming algorithm is called again to plan a new placement strategy. At the same time, the central node retains the training parameters. When a training node is lost, other edge devices can still be used normally and efficiently, improving the stability and reliability of the network. BRIEF DESCRIPTION OF THE DRAWINGS

[0023] Figure 1 This is a flowchart of the distributed graph neural network training method based on heterogeneous devices;

[0024] Figure 2 A diagram of the structure of heterogeneous equipment, a central training node and other heterogeneous training equipment;

[0025] Figure 3 It is a sampling and aggregation method for mini-batch of graph neural networks. DETAILED DESCRIPTION

[0026] The following examples are provided to further illustrate the implementation steps and objectives of the present invention. The following examples are provided to illustrate the present invention, but are not intended to be exhaustive.

[0027] like Figure 1As shown, a distributed graph neural network training method based on heterogeneous devices specifically includes the following steps:

[0028] Step A: Evaluate the training capability of all equipment involved in the training, such as Figure 2 As shown, the central node sends the same subgraph training task to each device using the same training model and training parameters, records the time for each device to perform forward propagation and backward propagation, and thus obtains the training capacity of each device. It then sorts them and obtains the training capacity ranking of m devices in preparation for obtaining the placement plan next.

[0029] Step B: Use the METIS algorithm to process graph data. The METIS algorithm partitions the graph with minimal edge cutting. The algorithm assigns densely connected vertices to the same partition to reduce the number of edge cutting between partitions. After a portion of fixed points is divided into one partition, the neighbor nodes of these vertices are copied to the points in other partitions. The copied nodes are called HALO nodes. This ensures that when sampling nodes within the partition, there is no need to communicate with the equipment training the neighboring partitions to reduce training time. At this time, s subgraphs can be obtained.

[0030] Step C: At this point, there are m devices and s subgraphs, where the number of subgraphs trained on the last device is k. Define To define the minimum training time and placement plan for training s subgraphs on m devices, It represents the time required for the mth device to train k images, where , train the k subgraphs on the last device and the first sk subgraphs on the m-1 devices. Based on the dynamic programming algorithm, the state transition equation can be obtained:

[0031]

[0032] According to the state transition equation, the shortest training time for s subgraphs on m devices can be obtained, and the placement plan P with the shortest training time can be obtained.

[0033] Step D: Obtain a set of mini-batches by using neighbor sampling for the training nodes in the subgraph, such as Figure 3 As shown, the mini-batch is loaded into the neural network for training.

[0034] Step E: Load the mini-batch sampled in step D into the input of the graph neural network to obtain structural information, and then load the node information into the graph neural network for training. The graph neural network is trained according to ; ;in Vertex Aggregated neighbor information, For the vertex connected vertex sets, Represents a vertex The number of neighbors, Vertex In the Layer embedding, represents a nonlinear activation function, For the The learnable parameters of the layer.

[0035] Because training neural networks on graph data is far more accurate than traditional multi-layer perception, in this embodiment, the graph neural network aggregates the edge neighbor information of the target vertex to generate an intermediate embedding. The neural network then converts the aggregated features into output features, which are sent as input features to the next layer. The final embedding is obtained by stacking multiple graph neural network layers.

[0036] Step F: Enable the fault tolerance mechanism. The specific process is as follows: The parameters are sent to the central node together with the number of training rounds. The central node first saves these parameters and then aggregates the parameters of all devices to obtain a new The parameters are sent to all training devices. After receiving the parameters, the device proceeds to the next round of training. At the same time, it detects whether the device is abnormal. If the device is abnormal, it returns to step C to redetermine the placement plan. The number of devices becomes m-1, and dynamic planning is performed. The central node sends the previous round number and parameters of the abnormal round to the new corresponding device and continues training. The training ends when the training data converges.

Claims

1. A distributed graph neural network training method based on heterogeneous devices, characterized by: The specific steps include: Step A: Evaluate the training capability of each heterogeneous device participating in the training, obtain the parameters of each device's training capability, and rank the training capabilities of each device, specifically: The central node sends the same subgraph training task to each device using the same training model and training parameters. It records the forward propagation and backward propagation time of each device, thereby obtaining the training capacity of each device and ranking them to prepare for the subsequent maximization of each training device. Step B: Process the large graph data and use the graph cutting algorithm to obtain several subgraphs with balanced nodes and edges; Step C: Use a dynamic programming algorithm to determine an optimal device placement plan. Devices with strong training capabilities will receive more training tasks, while devices with weak training capabilities will receive fewer training tasks. Based on the resulting plan strategy, each subgraph is placed on the corresponding device for training, so that each device can complete the training tasks assigned by the algorithm at the same time. Step D: Sample the training nodes of the subgraph to obtain a set of mini-batches, and load the mini-batches into the neural network for training; Step E: Use the corresponding graph neural network model to train the subgraph, and synchronize the parameters to the central node for parameter aggregation and storage. The corresponding parameters obtained after aggregation are distributed to each training device through broadcasting for a new round of training; the graph neural network generates intermediate embeddings by aggregating the edge neighbor information of the target vertex, and then uses the neural network to convert the aggregated features into output features, which are sent as input features to the next layer; after stacking multiple graph neural network layers, the final embedding is obtained; Specifically, the representation of graph neural network training vertices is as follows: ; ; in Vertex Aggregated neighbor information, For the vertex connected vertex sets, Represents a vertex The number of neighbors of Vertex In the Layer embedding, represents a nonlinear activation function, For the Learnable parameters of the layer; When each round of training is completed, The parameters and the number of training rounds are sent to the central node. When the central node receives the parameters of each device, it aggregates them to obtain the new parameters. , and broadcast it to all training devices. All devices will perform the next round of training after receiving the new parameters; Step F: Enable the fault-tolerant detection mechanism to detect whether the equipment can train normally and transmit parameters normally in each round of training until the training is completed; During the training process, if an exception occurs, return to step C to obtain a new optimal placement strategy, and broadcast the parameters of the previous round in which the exception occurred to all training devices for parameter synchronization, retrain, and reduce the loss caused by the failure. The training is stopped when the graph data training reaches the convergence state.

2. The distributed graph neural network training method based on heterogeneous devices according to claim 1, characterized in that: In the evaluation process of step A, in order to reduce the evaluation error generated when the device is evaluated or tested for performance, each device is tested repeatedly, and after obtaining multiple test data, the average value is calculated.

3. The distributed graph neural network training method based on heterogeneous devices according to claim 1, characterized in that: The subgraph obtained in step B must be trainable on the device with the weakest training capability.

4. The distributed graph neural network training method based on heterogeneous devices according to claim 1, characterized in that: During the sampling process of step D, different sampling strategies are adopted according to corresponding requirements, and the sampling strategies include neighbor sampling, partial neighbor sampling or random walk.

Images