Graph neural network partitioning method based on heterogeneous GPU load balancing

Abstract

The present invention discloses a graph neural network partitioning method based on heterogeneous GPU load balancing, which involves GNN technology. This solution is proposed to solve the problem that the communication impact between partitions is not considered in the prior art. S1. Obtain the original graph G and the target number of partitions N; S2. Use the graph partitioning algorithm to perform coarse partitioning; S3. Randomly pre-assign each subgraph to each GPU; S4. Model the computational overhead of each GPU; S5. Model the communication overhead of each GPU; S6. Set a heuristic partitioning function and memory constraints based on the computational overhead and communication overhead; S7. Adjust the coarse partitioning based on the heuristic partitioning function and memory constraints; S8. Assign the partitioned subgraphs to the corresponding GPUs, execute the graph computing tasks and summarize the outputs. The advantage is that the divided subgraphs can not only adapt to the performance of the GPU with little impact on the accuracy, but also reduce the communication requirements as much as possible and prevent memory overflow, which effectively reduces the training time and improves the training efficiency.

Description

Technical Field

[0001] The present invention relates to GNN technology, and in particular to a graph neural network partitioning method based on heterogeneous GPU load balancing. Background Art

[0002] Unlike highly structured data such as images and text, graphs are data structures composed of vertices and edges, which can effectively capture the complex relationships between irregular data in non-Euclidean space. Graph Neural Network (GNN) is a deep learning model designed specifically for graph data. Compared with traditional deep learning methods such as CNN, it can directly process graph structured data. In real life, any entity and its connection relationship can be represented as graph structured data. Due to the special properties of graph neural networks, they have a very wide range of application scenarios.

[0003] However, as the scale of data continues to increase, traditional GNN processing becomes increasingly inefficient. In addition, since GNNs also have complex node relationships, dual modes of irregular and regular memory access and computation, high feature dimensions, and power-law distribution, the computing and storage requirements for large-scale graph data may exceed the capacity of a single machine. In addition, in some scenarios (such as cloud computing), users cannot fully control the type of GPU used. Therefore, in recent years, more and more studies have begun to explore GNN training methods on multiple GPUs. How to achieve efficient training of graph neural networks has become a key issue that urgently needs to be solved in practical applications.

[0004] In multi-GPU training, data needs to be distributed to different GPUs for training and gradient update. Due to the limitations of the synchronization strategy, the training time will be affected by the slowest GPU. Most existing methods pursue similar partition sizes because they assume the same GPU type. For example, CN 107122248A directly divides the partitions evenly according to the total number of nodes. The METIS method used in CN118504609A is to make the partition sizes consistent. Although it prunes the redundant nodes in the partitions, it is not for the purpose of adapting the GPU computing power. However, GPUs are developing rapidly, and a server usually has multiple GPUs of different types. Therefore, this partitioning method will lead to complex imbalances between GPUs, which will greatly increase the training time. Figure 1 Although CN 115658975A considers the load balancing problem of partitions, it is not related to the computing power of the GPU and does not consider the communication impact between partitions, and it is still difficult to meet the various changing situations of actual applications. Summary of the invention

[0005] The purpose of the present invention is to provide a graph neural network partitioning method based on heterogeneous GPU load balancing to solve the problems existing in the above-mentioned prior art.

[0006] The graph neural network partitioning method based on heterogeneous GPU load balancing described in the present invention includes the following steps.

[0007] S1. Obtain the original graph G and the number of partition targets N;

[0008] S2. Use graph partitioning algorithm to perform coarse partitioning;

[0009] S3. Randomly pre-assign each subgraph to each GPU;

[0010] S4. Model the computational overhead of each GPU;

[0011] S5. Model the communication overhead of each GPU;

[0012] S6. Set a heuristic partition function and memory constraints based on computational overhead and communication overhead;

[0013] S7. adjusting the coarsened partition according to the heuristic partition function and the memory constraint;

[0014] S8. Assign the partitioned subgraphs to the corresponding GPUs, execute graph computing tasks, and summarize the outputs.

[0015] The graph neural network partitioning method based on heterogeneous GPU load balancing described in the present invention has the advantage that the divided subgraphs can not only adapt to the performance of the GPU with little impact on the accuracy, but also minimize communication requirements and prevent memory overflow, ultimately effectively reducing training time and improving training efficiency. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] Figure 1 This is a schematic diagram of the difference in computing power between different GPUs.

[0017] Figure 2 It is a flow chart of the graph neural network partitioning method described in the present invention.

[0018] Figure 3 It is a schematic diagram of the communication interaction between GPUs in the graph neural network partitioning method described in the present invention.

[0019] Figure 4 It is a partitioning schematic diagram of the graph neural network partitioning method described in the present invention.

[0020] Figure 5 It is a schematic diagram of three communication modes of the graph neural network partitioning method described in the present invention.

[0021] Figure 6This is a comparison chart of the effects before and after sub-graph partitioning of the graph neural network partitioning method described in the present invention. DETAILED DESCRIPTION

[0022] like Figures 2 to 6 As shown, the graph neural network partitioning method based on heterogeneous GPU load balancing described in the present invention includes the following steps.

[0023] S1. Obtain the original graph G and the target number of partitions N.

[0024] The original graph G and the number of partition targets N are input by the user, and the number of partition targets N is the same as the number of GPUs. The GPU set P is represented as: P = {P1, P2, ... P i ,...,P N}, where P i Represents the i-th GPU.

[0025] S2. Use graph partitioning algorithm to perform coarse partitioning.

[0026] There is no special requirement for the graph partitioning algorithm, and it can be the existing METIS or Random method. The original graph G is divided into N partition subgraphs, and the set of N partition subgraphs is represented as S = {S1, S2, ... S i ,...,S N},S i Represents the i-th partition subgraph.

[0027] The partitioned subgraph contains two types of nodes: one is the inner node, which is the node that actually belongs to the current partition; the other is the outer node, which only contains the structural information of the outer node in the current partition, and the outer node actually belongs to the neighboring partition of the current partition. However, in order to work efficiently, it is also divided into the current partition, so a large part of the communication overhead is to obtain the features of these outer nodes, such as Figure 3 The partitioning process is shown in the figure Figure 4 shown.

[0028] S3. Randomly pre-assign each subgraph to each GPU.

[0029] There is no special requirement for the allocation method, and it can be allocated at will. The current allocation work is just mapping, and does not actually move the partition subgraph to the corresponding GPU memory.

[0030] The partition subgraph S i Assigned to GPU P i The mapping function is expressed as:

[0031]

[0032] S4. Model the computational cost of each GPU, including the following sub-steps:

[0033] S41. Evaluate the computing power of each GPU: Since in GNN, the calculation mainly involves sparse matrix multiplication SpMM and dense matrix multiplication MM. Therefore, the computing power of each GPU is evaluated for SpMM and MM separately.

[0034] Evaluate SpMM: Use a sparse matrix with a sparsity equivalent to most graph data, about 0.9. Specify the size of the sparse matrix as 5000x5000, repeat the calculation for 100 rounds, calculate the average value, and use the result as the computing power of the GPU for SpMM. The smaller the average value, the stronger the computing power of SpMM.

[0035] Evaluate MM: Specify the size of the dense matrix as 5000x5000, repeat the calculation for 100 rounds, calculate the average value, and use the result as the computing power of the GPU for MM. The smaller the average value, the stronger the computing power of MM.

[0036] S42. Modeling computational overhead based on computational power: Considering that SpMM is related to the total number of edges and MM is related to the number of inner points, the computational overhead model is expressed as:

[0037]

[0038] in, Represents GPU P i In the partition subgraph S i The computational overhead on GPU P is smaller, and the smaller the overhead is, the i The stronger the computing power is. α is used to adjust the weight of focusing on SpMM or MM. Represents the partition subgraph S i The number of all edges on . Represents GPU P i The SpMM computing capability. Indicates the maximum value of SpMM computing capability of all GPUs. Represents the partition subgraph S i The number of interior points on . Represents GPU P i MM computing power. Indicates the maximum value of the MM computing capability of all GPUs.

[0039] S5. Model the communication overhead of each GPU, including the following sub-steps:

[0040] S51. Evaluate the communication capability of each GPU: In GNN, when the GPUs do not support point-to-point transmission mode, there are three main communication modes involved: CPU to GPU, i.e. H2D; GPU to CPU, i.e. D2H; GPU internal, i.e. D2D; Figure 5 As shown, the communication overhead is therefore modeled based on these three evaluation results.

[0041] Evaluate the CPU to GPU: Specify the amount of data to be transmitted as 10MB, count the time it takes to send data from the CPU to the GPU, repeat the calculation 100 times, and calculate the average value. The result is used as the communication capability of the GPU from the CPU to the GPU. The smaller the value, the stronger the communication capability.

[0042] Evaluate the GPU to CPU communication: Specify the amount of data to be transmitted as 10MB, count the time it takes to send data from the GPU to the CPU, repeat the calculation 100 times, and calculate the average value. The result is used as the communication capability of the GPU from the GPU to the CPU. The smaller the value, the stronger the communication capability.

[0043] Evaluate the GPU internally: Specify the amount of data to be transferred as 10MB, apply for two memory spaces in the GPU, and count the time it takes to send data from one memory space in the GPU to the other memory space. Repeat the calculation for 100 rounds and calculate the average value. The result is used as the communication capability of the GPU within the GPU. The smaller the value, the stronger the communication capability.

[0044] S52. Modeling the communication overhead based on communication capability: Considering that as the number of GPUs increases, the impact of H2D or D2H will be greater than D2D. In addition, only nodes that have edge connections with other partitions need to communicate. Therefore, the model of communication overhead is expressed as:

[0045]

[0046] in, Represents GPU P i In the partition subgraph S i The communication overhead on GPU P is smaller, and the smaller the overhead is, the i The stronger the communication capability. Represents the partition subgraph S i The number of edges with source nodes outside the above. Represents GPU P i H2D capabilities. Indicates the maximum H2D capability of all GPUs. Represents GPU P i D2H capabilities. Indicates the maximum D2H capability of all GPUs. Represents GPU P iD2D capability. Indicates the maximum D2D capability of all GPUs.

[0047] S6. Set the heuristic partition function and memory constraints according to the computational overhead and communication overhead.

[0048] In order to minimize the total cost and make the costs of different subgraphs similar, set the heuristic partition function:

[0049] min{λ+Std(λ i )};

[0050] The function constraints are:

[0051] λ=max{λ i};

[0052] Std({λ i})≤ε;

[0053] λ i =T comp +T comm ;

[0054]

[0055] Among them, λ represents the maximum value of all partition overheads. i Represents the partition subgraph S i The total cost is determined by the corresponding communication cost and computational overhead Composition. Std({λ i}) represents all λ i The variance of the calculation. ε represents the variance threshold, which can be set to 1% of the total overhead average. M i Represents the partition subgraph S i The maximum available memory.

[0056] To prevent memory overflow, set memory constraints:

[0057] M i ≥|V i |×M node +|E i |×M edge +|d i |×M feat +β;

[0058] Among them, M node Indicates the amount of memory occupied by each node, M edge Indicates the amount of memory occupied by each edge, M feat Indicates the amount of memory used by each feature. |V i | represents the partition subgraph Si The number of nodes, |E i | represents the partition subgraph S i The number of edges, |d i | represents the partition subgraph S i β represents the reserved memory, which can be used for gradients, etc. and can be set by the user according to the model size.

[0059] S7. Adjust the coarsened partition according to the heuristic partition function and the memory constraint, including the following sub-steps:

[0060] S71. Initialize sequence number i=1;

[0061] S72. Select partition subgraph S i , calculate the total cost and corresponding cost variance of the current partition subgraph; when the cost variance is greater than the variance threshold ε, proceed to sub-step S73, otherwise proceed to sub-step S74.

[0062] S73. Sort the external points in the current partition subgraph by the number of edges, remove the external point with the smallest number of edges, and jump back to substep S72; if the current partition subgraph does not contain an external point, directly proceed to substep S74.

[0063] S74. Set i+1 and return to sub-step S72; if i=N, end the process.

[0064] Finally, the partitioning effect using the above heuristic partitioning is as follows Figure 6 As shown in Figure 2, it can be seen that the proposed subgraph partitioning method can effectively adjust the number of nodes and edges and make the final score meet the requirements of the objective function.

[0065] S8. Assign the partitioned subgraphs to the corresponding GPUs, execute graph computing tasks, and summarize the outputs.

[0066] The adjusted partitioned subgraphs are actually assigned to the corresponding GPUs, that is, directly assigned to the memory of the corresponding GPUs. Each GPU starts to execute the corresponding graph computing task locally. Finally, the calculation results on each GPU are aggregated to form the output result of the graph computing task.

[0067] The graph neural network partitioning method based on heterogeneous GPU load balancing described in the present invention has the following advantages over the prior art:

[0068] 1. Automatic adaptation of GPU performance: The size of the divided subgraphs corresponds to the computing power of the GPU. That is, in the combination of each subgraph and the corresponding GPU, a GPU with strong computing power can be allocated a larger subgraph size, and a GPU with weak computing power can be allocated a smaller subgraph size. Therefore, the computing time of the two is the same or very close, avoiding the GPU that has completed the calculation first waiting for the GPU that has not yet completed the calculation.

[0069] 2. Minimize communication: Minimize the number of edge nodes that need to communicate but have little impact. This will reduce the amount of communication and the communication time. And because the impact of these nodes is small, the impact on the final accuracy is also small, which can ensure the overall accuracy control of the solution.

[0070] 3. Prevent memory overflow: Taking into account the situation where the graph information is large but the GPU memory capacity is small, it is necessary to further reduce the nodes that communicate but have little impact according to the memory situation. This will avoid the problem of memory overflow caused by the GPU not being able to fit the partitioned subgraph.

[0071] For those skilled in the art, various other corresponding changes and deformations can be made according to the technical solutions and concepts described above, and all of these changes and deformations should fall within the protection scope of the claims of the present invention.

Claims

1. A graph neural network partitioning method based on heterogeneous GPU load balancing, characterized in that: The following steps are involved: S1. Obtain the original graph G and the number of partition targets N; S2. Use graph partitioning algorithm to perform coarse partitioning; S3. Randomly pre-assign each subgraph to each GPU; S4. Model the computational overhead of each GPU; S5. Model the communication overhead of each GPU; S6. Set a heuristic partition function and memory constraints based on computational overhead and communication overhead; S7. adjusting the coarsened partition according to the heuristic partition function and the memory constraint; S8. Assign the partitioned subgraphs to the corresponding GPUs, execute graph computing tasks, and summarize the outputs; The step S4 comprises the following sub-steps: S41. Evaluate the computing capability of each GPU for sparse matrix multiplication SpMM and dense matrix multiplication MM; S42. Modeling computational overhead based on computational power: in, Represents GPU P i In the partition subgraph S i The computational overhead on GPU P is smaller, and the smaller the overhead is, the smaller the GPU P is. i The stronger the computing power is; α is used to adjust the weight that focuses on SpMM or MM; Represents the partition subgraph S i The number of all edges on Represents GPU P i SpMM computing capabilities; Indicates the maximum value of SpMM computing capability of all GPUs; Represents the partition subgraph S i The number of interior points on ; Represents GPU P i MM computing capabilities; Indicates the maximum value of MM computing capability of all GPUs; The step S5 comprises the following sub-steps: S51. Evaluate the communication capabilities of each GPU: including the H2D capability from CPU to GPU, the D2H capability from GPU to CPU, and the D2D capability within the GPU; S52. Modeling communication overhead based on communication capabilities: in, Represents GPUP i In the partition subgraph S i The communication overhead on the GPU is smaller, the smaller the overhead is, the higher the GPUP i The stronger the communication capability; Represents the partition subgraph S i The number of edges with source nodes outside the above; Represents GPUP i H2D capabilities; Indicates the maximum H2D capability of all GPUs; Represents GPU P i D2H capabilities; Indicates the maximum D2H capability of all GPUs; Represents GPU P i D2D capability; Indicates the maximum D2D capability of all GPUs; In step S6, the heuristic partitioning function is: min{λ+std(λ i )); The constraints on the heuristic partition function are: λ=max{λ i }; Std((l i })≤e; λ i =T i comp +T i comm ; Among them, λ represents the maximum value of all partition overheads; λ i Represents the partition subgraph S i The total cost is determined by the corresponding communication cost T i comm and computational overhead T i comp Composition; Std({λ i }) represents all λ i The variance of the calculation; ε represents the variance threshold, which is set to 1% of the average value of the total cost; M i Represents the partition subgraph S i Maximum available memory; To prevent memory overflow, set memory constraints: M i ≥|You | ×M node +|E i |×M edge +|d i |×M feat +β Among them, M node Indicates the amount of memory occupied by each node, M edge Indicates the amount of memory occupied by each edge, M feat Indicates the amount of memory occupied by each feature; |V i | represents the partition subgraph S i The number of nodes, |E i | represents the partition subgraph S i The number of edges, |d i | represents the partition subgraph S i The feature dimension of ; β represents the reserved memory, which is set by the user.

2. According to claim 1, the graph neural network partitioning method based on heterogeneous GPU load balancing is characterized in that: In step S1, the original graph G and the target number of partitions N are input by the user, and the target number of partitions N is the same as the number of GPUs; the GPU set P is represented as: P = {P1, P2, ...P i ,...,P N }, where P i Represents the i-th GPU.

3. According to claim 2, the graph neural network partitioning method based on heterogeneous GPU load balancing is characterized in that: In step S2, the original graph G is divided into N partition subgraphs, and the set table of the N partition subgraphs is It is shown as S = {S1, S2, ... S i ,...,S N },S i represents the i-th partition subgraph; The partitioned subgraph contains two types of nodes: one is the internal node, which is the node belonging to the current partition; the other is the external node, which only contains the structural information of the external node in the current partition.

4. According to claim 3, the graph neural network partitioning method based on heterogeneous GPU load balancing is characterized in that: The step S7 comprises the following sub-steps: S71. Initialize sequence number i=1; S72. Select partition subgraph S i , calculate the total cost and corresponding cost variance of the current partition subgraph; when the cost variance is greater than the variance threshold ε, proceed to sub-step S73, otherwise proceed to sub-step S74; S73. Sort the external points in the current partition subgraph by the number of edges, remove the external point with the smallest number of edges, and jump back to substep S72; if the current partition subgraph does not contain external points, proceed directly to substep S74; S74. Set i+1 and return to sub-step S72; if i=N, end the process.

5. According to claim 4, the graph neural network partitioning method based on heterogeneous GPU load balancing is characterized in that: The step S8 is specifically as follows: directly allocating the adjusted partition subgraph to the memory of the corresponding GPU; each GPU starts to execute the corresponding graph computing task locally; finally, aggregating the calculation results on each GPU to form the output result of the graph computing task.

Images