An automatic parallel method based on double graph manhattan network and policy gradient optimization

By employing a dual-graph Mamba network and an automatic parallel approach for policy gradient optimization, the problems of insufficient computational graph feature representation and policy fluctuations in reinforcement learning are addressed, enabling more efficient model-parallel policy search and device placement.

CN119558378BActive Publication Date: 2025-11-21HANGZHOU DIANZI UNIV
View PDF 2 Cites 0 Cited by

Patent Information

Application Number
CN202411698245.6
Authority / Receiving Office
CN · China
Patent Type
Patents(China)
Current Assignee / Owner
Filing Date
2024-11-26
Publication Date
2025-11-21
Estimated Expiration
2044-11-26

AI Technical Summary

Technical Problem

Existing equipment placement methods suffer from insufficient computational graph feature representation capabilities and performance fluctuations when solving strategies using reinforcement learning, resulting in high time costs and low efficiency in parallel strategy search.

Method used

An automatic parallel approach based on dual graph Mamba network and policy gradient optimization is adopted. The computation graph features are extracted through dual graph Mamba network, and combined with node selection mechanism and semantic features, a reinforcement learning iterative solution module is generated to optimize the device placement strategy.

Benefits of technology

It significantly reduces the time cost of parallel policy search in the model, improves the efficiency and stability of policy learning, and obtains better parallel policies.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN119558378B_ABST
    Figure CN119558378B_ABST
Patent Text Reader

Abstract

The application discloses an automatic parallel method based on a double graph manba network and a policy gradient optimization, which first acquires a public AI model data set, fuses operators for each calculation graph G in the data set, and acquires a feature matrix X.Secondly, a double graph manba network containing a graph manba branch and a decoupling graph neural network branch is constructed, and the feature matrix X is input into the two branches to obtain feature representations X M and X D Then, X M and X D are averaged, input into an LSTM network for feature fusion, and the feature representation of the double graph manba network is obtained.Finally, the feature representation is input into a reinforcement learning iteration solving module, the optimal parallel strategy is obtained by continuously iterating and optimizing the intermediate reward value, and the expression capacity of the graph embedding is enhanced, and the inaccuracy and stability problems caused by the large difference between the new and old strategies of the traditional policy gradient algorithm in the device placement problem are effectively solved.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] This invention belongs to the field of automatic parallel technology optimization for AI models, specifically involving an automatic parallel method based on dual graph Mamba networks and policy gradient optimization. Background Technology

[0002] In recent years, the size of basic models has increased dramatically, but the growth in computing power and storage capacity has failed to keep pace, leading to the need for multi-device collaboration to train large-scale AI models. This challenge has prompted the research and industry communities to continuously explore new strategies for efficiently training large-scale models.

[0003] Automated parallelism offers an effective solution to the challenges of large-scale model training. It automatically generates parallel strategies, efficiently designs distributed training strategies for models, and accelerates distributed computation across multiple devices, allowing developers to focus more on algorithmic model innovation. Currently, core methods for achieving distributed model training include data parallelism and model parallelism. Data parallelism involves partitioning and distributing training data across multiple devices while replicating the complete AI model on each. Each device independently processes its assigned subset of data and periodically synchronizes model parameters via communication mechanisms to ensure global consistency. Model parallelism, on the other hand, is specifically designed for handling ultra-large-scale models. When model complexity exceeds the capacity of a single device, this method partitions the model and its computational tasks into multiple parts and flexibly deploys them across different devices. This process is figuratively known in academia as "device placement," which allows each device to focus on a sub-segment of the model, completing the overall training task through efficient collaboration.

[0004] The learning ability of reinforcement learning and the effective representation of computational graph features are key factors driving the efficiency and performance of automatic parallel policy search. Regarding the optimization of reinforcement learning modules, while ColocRL, Hierarchical Device Placement (HDP), and Reinforcement Genetic Algorithm Learning (REGAL) accelerate parallel policy search in AI models using the Reinforce algorithm, their low sample efficiency significantly prolongs the training cycle. Methods such as Placeto, GraphSAGE, and P-GNN employ Markov Decision Processes (MDPs) to optimize device placement, precisely capturing state transition probabilities and immediate rewards to find the optimal path to maximize long-term returns. However, MDPs require the Markov property that the next state depends only on the current state and action, and is independent of historical states. Post, Spolight, GDP, Mars, and Trinity, on the other hand, utilize Proximal Policy Optimization (PPO) algorithms to optimize device placement, enhancing training stability by improving sampling efficiency and limiting the magnitude of new policy updates. In graph embedding techniques, Placeto and Trinity, for example, neglect crucial node location information when constructing their node embedding systems. This can lead to the marginalization of key features, causing confusion in the vector space between structurally similar computational graphs and severely interfering with the accuracy and effectiveness of operator placement strategies. While GraphSAGE attempts to aggregate local information from neighboring nodes through random sampling to capture subtle differences in the computational graph, this approach carries the risk of random loss of important information. As model complexity increases, the number of operators surges, and the sampling scale expands accordingly. This not only slows down the training process but also significantly increases the computational burden and execution time, becoming another bottleneck restricting performance. P-GNN and Aware attempt to integrate location encoding into node representations to compensate for the lack of location information. However, their methods rely on aggregating the representations of nodes farthest from the target node to integrate location information. This strategy is limited by the inherent defects of message passing mechanisms (MPNNs), resulting in limited overall feature representation capabilities of the computational graph, making it difficult to fully reflect its internal structure and complexity, ultimately affecting the efficiency and depth of parallel strategy search. More importantly, all current graph embedding-based device placement methods, including Placeto, GraphSAGE, P-GNN, Trinity, and Aware, neglect the potential impact of semantic information in the computation graph on device placement decisions. This neglect may not only limit further improvements in model parallel performance but also hinder the maximization of overall efficiency. Summary of the Invention

[0005] This invention addresses the shortcomings of existing equipment placement methods in terms of the insufficient expressive power of computational graph features and the performance fluctuations in reinforcement learning strategy solving. It designs and implements an automatic parallel method based on a dual-graph Mamba network and policy gradient optimization. This method effectively reduces the time cost of model parallel policy search and obtains a better parallel strategy. First, computational graph features are extracted using a dual-graph Mamba network. This network comprises two branches: one branch (graph Mamba branch) enhances long-distance dependencies using a node selection mechanism, thereby improving the expressive power of graph embedding; the other branch (decoupled graph neural network branch) combines node attributes, topological space, and semantic features to generate a powerful and complete node representation. The semantic features are constructed from computational graph operator attributes, providing supplementary information to the topological graph. Next, the features extracted from the two branches are fused to generate the final computational graph features. This process improves the expressive power of the computational graph by deeply mining the semantic information within the graph. Finally, the obtained computational graph features are input into the reinforcement learning iterative solving module. In this module, an additional loss layer is introduced for each node, enabling the algorithm to quickly capture subtle changes in the environment and allowing for finer-grained learning and decision adjustment. By adjusting the strategy through feedback, this method significantly improves the learning efficiency, stability, and exploratory ability of the strategy, ultimately obtaining the optimal parallel strategy.

[0006] This invention provides an automatic parallel method based on dual graph Mamba networks and policy gradient optimization, with the following steps:

[0007] Step 1: Obtain a publicly available AI model dataset, perform operator fusion on each computation graph G in the dataset, and obtain the feature matrix of the computation graph. Where V represents the set of nodes in computation graph G, and F represents the features of the nodes in computation graph G, as follows:

[0008] Obtain a publicly available AI model dataset. For each computation graph G in the dataset, perform operator fusion using the ColocRL algorithm. The fused computation graph contains relevant attributes of all operators in that computation graph, including operator type, input, output, and access identifier. Based on these operator attributes, construct the original feature vectors of the operators, i.e., initialize the node encoding, and obtain the complete feature matrix X of the computation graph. Where V represents the set of nodes in the computation graph G, and F represents the features of the nodes in the graph.

[0009] Step 2: Construct a dual graph mamba network containing a graph mamba branch and a decoupled graph neural network branch, and combine the feature matrix... The inputs are fed into the two branches respectively to obtain the feature representation X. M With X D ;

[0010] Step 2.1: The Graph Mamba branch generates the node encoding matrix M based on the feature matrix X. 0The edge encoding matrix Q is used to obtain the intermediate representation N through the message passing network MPNN. t Node encoding matrix M 0 The input is fed into the Graph Mamba module to extract features and obtain the intermediate feature matrix M. t M t and N t The characteristic matrix M is obtained by averaging the elements one by one. T M T The input is fed into a multilayer perceptron to obtain M. T′ M T′ With N t The final feature representation X of the graph mamba branch is obtained through a message passing network. M .

[0011] Sub-step 2.1.1: Obtain the intermediate feature matrix by adding the adjacency matrix D of X and the identity matrix I. By analyzing the matrix Laplace normalization is performed to obtain the edge encoding matrix. The edge encoding matrix Q is input into the message passing network to obtain the intermediate representation.

[0012] Sub-step 2.1.2: Feature matrix X and intermediate feature matrix Multiplying them yields the node encoding matrix M. 0 Node encoding matrix M 0 Inputting into the TuMamba module yields the intermediate node code M t Node encoding M t The next state M is obtained through the Selective State Management (SSM) mechanism of the Graph Mamba module. t+1 ;

[0013] The implementation process of the Selective State Management (SSM) mechanism first introduces three learnable weight matrices W. B W c and W Δ For M respectively t A linear transformation is performed to obtain the projection parameters B and C, and the discrete step size Δ after Softplus activation. The evolution parameters are then calculated using the evolution parameters A obtained at the previous time step, the projection parameters B at the current time step, and the discrete step size Δ. and in It is related to Δ*A through an exponential function. It is then calculated based on the difference between the inverse matrix of Δ*A and the exponential function; subsequently, it is combined with M... t-1 and M t And projection parameter C and evolution parameter and Update the current state to get M t+1 .

[0014] Sub-step 2.1.3: Process the feature matrix X using a message passing network to obtain the intermediate representation N. t+1 ; N t+1 and M t+1 By accumulating the results, we obtain the fusion feature matrix M at time t. t+1′ and M t+1′ The input is fed into a multilayer perceptron to obtain the feature matrix M at the next time step. t+2 After T identical iterations, the final feature matrix M is obtained. T .

[0015] Sub-step 2.1.4: Analyze the feature matrix M after T iterations... T With the intermediate representation N T Simultaneously, this is used as input again and processed through a message passing network to obtain the fused feature matrix M at time T. T′ Ultimately, M T′ The computational graph characteristic representation of the graph mamba branch is denoted as...

[0016] Step 2.2: Use the feature matrix X as the node attribute feature A 0 and A 0 Semantic features C are obtained by passing the edge encoding matrix Q through a GCN network. 0 Meanwhile, A 0 The semantic intermediate representation B is obtained after passing through a multi-layered graph mamba module. 0 A 0 C 0 B 0 AT, CT, and BT are obtained through weight sharing and structural consistency constraints. Finally, the weights of these three components are used to obtain the final feature representation of the decoupled graph neural network branch.

[0017] Sub-step 2.2.1: Use the feature matrix X as the node attribute features. Initial node attribute feature A 0 The edge encoding matrix Q is input into the graph convolutional network GCN to obtain the initial structural features. And use the Graph Mamba module to analyze node attribute feature A 0 Processing was performed; initial semantic features were obtained.

[0018] Sub-step 2.2.2: Obtain A 0 C 0 B 0 The node attribute feature matrix A at the final moment is obtained by combining weight sharing and structural consistency constraints. T Topological feature matrix C Tand semantic information feature matrix B T The final feature representation is obtained. To achieve this goal, a learnable shared reconstruction factor matrix W is first introduced. S Calculate the reconstructed feature matrix A t B t and C t The relationship between the three is optimized by minimizing the error. During the optimization process, a nonlinearly transformed sigmoid function and a series of hyperparameters are used to balance the contributions of different eigenma matrices. Finally, after multiple iterations, A... t B t and C t After optimization, the feature matrix A is obtained. T B T and C T , will A T B T and C T By connecting these components and constraining them with weights, the final feature representation of the decoupled graph neural network branches can be obtained.

[0019] Step 3: Extract the graph feature representation X from the graph mamba branches. M The final feature representation X of the decoupled graph neural network branch D The summation and averaging are used as input to an LSTM network for feature fusion, resulting in the final feature representation of the dual graph mamba network.

[0020] Step 4: Represent the features obtained in Step 3 The input is fed into the reinforcement learning iterative solution module, which utilizes the characteristics of reinforcement learning (proximal policy gradient optimization) to obtain the optimal parallel policy by continuously iterating and optimizing the intermediate reward value.

[0021] The beneficial effects of this invention are:

[0022] (1) The use of dual graph mamba graph neural networks effectively improves the expressive power of computation graphs. Specifically, the decoupled graph neural network branch constructs a semantic graph through the feature space of computation graph operator attributes, which provides supplementary information to the topological graph of the computation graph and significantly improves the expressive power of nodes. The graph mamba branch enhances long-distance dependencies in graph encoding through the node selection mechanism, further enhancing the expressive power of graph embedding.

[0023] (2) By using the policy gradient reinforcement learning method, the inaccuracy and stability problems caused by the large difference between the old and new policies in the traditional policy gradient algorithm in the device placement problem are effectively solved. Attached Figure Description

[0024] Figure 1 This is the overall architecture diagram;

[0025] Figure 2 It is based on the dual graph Mamba coding module graph. Detailed Implementation

[0026] The present invention will be further described below with reference to the accompanying drawings and specific implementation steps:

[0027] Figure 1 The diagram shows the overall flow of an automated parallel method based on dual graph Mamba networks and policy gradient optimization, such as transforming an AI model into a graph. Where V represents the set of computation graph operators (nodes), and operator v i ∈V represents a single operator (e.g., matrix multiplication, convolution, etc.); E represents the set of directed edges between nodes, e i,j ∈E represents the operator v i and v j The data communication dependencies between them. Given D = {d1, d2, d3, ..., d...} m} represents the devices that can be used in the cluster, where d i Let D represent a computing device (e.g., CPU or GPU). Given computing resources D and a computing graph... Find mappings i (i.e., parallel strategy, where s) i ∈S policy set) such that every operator v i Corresponding to a device d i The ultimate goal is to find a device placement strategy. i When the computation graph The middle operator follows strategy s i Minimize the computation graph during placement. Execution time Neural network computation graph and random device placement strategies i As input, then iterate through all nodes v in the computation graph. i ∈V, firstly, it uses a dual graph mamba network to obtain graph embedding representations. Secondly, for each selected node v i It applies a policy based on reinforcement learning (proximal policy optimization). i v i Assigned to device d i Then, it depends on the execution time. Update reinforcement learning strategies i And continuously optimize the placement strategy for subsequent nodes. Figure 2 The diagram illustrates a process for extracting computational graph feature representations based on a dual-graph Mamba network, comprising the following steps:

[0028] Step 1: Obtain a publicly available AI model dataset. Perform operator fusion on each computation graph G in the dataset. The fused computation graph contains the relevant attributes of all operators in the computation graph (including operator type, input, output, access identifier, etc.). Construct the original feature vector of the operator based on the operator-related attributes, i.e., node encoding initialization, to obtain the complete feature matrix X of the computation graph. Where V represents the set of nodes in the computation graph G, and F represents the features of the nodes in the graph.

[0029] Step 2: Construct a dual graph mamba network containing a graph mamba branch and a decoupled graph neural network branch, and combine the feature matrix... The inputs are fed into the two branches respectively to obtain the feature representation X. M With X D .

[0030] Step 2.1: The Graph Mamba branch generates the node encoding matrix M based on the feature matrix X. 0 The edge encoding matrix Q is used to obtain the intermediate representation N through the message passing network MPNN. t Node encoding matrix M 0 The input is fed into the Graph Mamba module to extract features and obtain the intermediate feature matrix M. t M t and N t The characteristic matrix M is obtained by averaging the elements one by one. T M T The input is fed into a multilayer perceptron to obtain M. T′ M T′ With N t The final feature representation X of the graph mamba branch is obtained through a message passing network. M .

[0031] Sub-step 2.1.1: Obtain the intermediate feature matrix by adding the adjacency matrix D of X and the identity matrix I. By analyzing the matrix Laplace normalization is performed to obtain the edge encoding matrix. The edge encoding matrix Q is input into the message passing network to obtain the intermediate representation.

[0032] Sub-step 2.1.2: Feature matrix X and intermediate feature matrix Multiplying them yields the node encoding matrix M. 0 Node encoding matrix M 0 Inputting into the TuMamba module yields the intermediate node code M t Simultaneously, node encoding M t M is obtained by extracting features using the Graph Mamba module. t+1 First, we introduce a learnable weight matrix W. B W c and WΔ The projection parameters B and C and the deviation step length Δ are obtained. Then, the evolution parameters are used... and projection parameters The continuous state is transformed into a discrete state, and finally, the Graph Mamba module is used for selective state management (SSM) to update the current state and obtain M. t+1 .

[0033] B = W B M t C = W c M t ,Δ=Softplus(W Δ M t )

[0034]

[0035] Sub-step 2.1.3: Process the feature matrix X using a message passing network to obtain the intermediate representation N. t+1 ; N t+1 and M t+1 By accumulating the results, we obtain the fusion feature matrix M at time t. t+1′ and M t+1′ The input is fed into a multilayer perceptron to obtain the feature matrix M at the next time step. t+2 .

[0036] Sub-step 2.1.4: Convert the intermediate representation N T The feature matrix M is obtained after passing through a multilayer perceptron. T As input, after passing through a message passing network, the final computational graph feature representation of the graph mamba branch is obtained.

[0037] X M =Combine(N) T M T )

[0038] Step 2.2: Extract features for the decoupled graph neural network branches, and use the feature matrix X as the node attribute feature A. 0 and A 0 Semantic features C are obtained by passing the edge encoding matrix Q through a GCN network. 0 Meanwhile, A 0 The semantic intermediate representation B is obtained after passing through a multi-layered graph mamba module. 0 A 0 C 0 B 0 A is obtained through weight sharing and structural consistency constraints. T C T and B T Finally, the three factors are weighted to obtain the final feature representation of the decoupled graph neural network branch.

[0039] Sub-step 2.2.1: Extract node attributes and topological features. First, extract node attribute features from the complete information (i.e., feature matrix) of the computation graph X. Simultaneously, the topological structure of the computation graph is analyzed, and its features are extracted. Based on the extracted node attribute features A 0 To obtain the initial semantic features The input is then fed into a Graph Convolutional Network (GCN), where it undergoes multiple convolutional operations and possible nonlinear activation functions to obtain an intermediate representation of the topological feature matrix.

[0040] Sub-step 2.2.2: Transfer node attribute feature A 0 The input is fed into the Graph Mamba neural network module for further processing of node attribute features, resulting in a semantic information feature matrix. Subsequently, an intermediate representation of the semantic information feature matrix is ​​obtained through a multi-layer graph mamba neural network.

[0041] Sub-step 2.2.3: Obtain the node attribute feature matrix A at time t. t The intermediate representation C of the topological feature matrix t intermediate representation B of semantic information feature matrix t During the process, the three components, through weight sharing and structural consistency constraints, obtain the final value of A at time T. T B T and C T The final feature representation of the decoupled graph neural network branch is obtained by constraining it with ε weights.

[0042] Sub-step 2.2.3.1: Introduce the learnable shared reconstruction factor matrix W S This is used to reconstruct the potential correlations between adjacency matrices. Simultaneously, to ensure the symmetry of the reconstructed matrix, constraints are imposed on W. S =WW T .in This is achieved by using the node attribute feature matrix A and the learnable shared reconstruction factor matrix W. S and its transpose A T The product of these is used to capture potential inter-node correlations in A. and Similarly, they represent the semantic information feature matrix B and the topological feature matrix C after passing through W. S The reconstructed adjacency representation. Furthermore... and These items are based on A, B, and C, with the addition of W. Sand its transpose-related additional items (such as and These additional terms introduce extra information to capture the complex relationships between the feature matrices. The specific calculation and update formulas are as follows:

[0043]

[0044] Where A is the node attribute feature matrix, B is the semantic information feature matrix, C is the topological feature matrix, and W, F, K and H are bias matrices.

[0045] Sub-step 2.2.3.2: A t B t and C t Error between and By reconstructing factor W S Minimize. Where ε is the node attribute feature A. t Semantic feature B t With topology C t The balance coefficient is used to balance the weight relationship among the three, and its value ranges from 0 to 1.

[0046]

[0047] Sub-step 2.2.3.3: σ is the Sigmoid function of the nonlinear transformation. L1, L2, and L3 are used to balance the hyperparameters contributing to the three graph encodings (A, B, C) in the decoupled graph neural network. After T iterations, the final A... T B T C T By concatenating them, we obtain the final feature representation.

[0048]

[0049] X D =Concat(A T B T C T )

[0050] Step 3: Extract the graph feature representation X from the graph mamba branches. M The final feature representation X of the decoupled graph neural network branch D As input, it is fed into an LSTM network for feature fusion, resulting in the final feature representation of the dual graph mamba network.

[0051] Step 4: Represent the features obtained in Step 5 The input is fed into the policy gradient (PPO) iterative solution module, which utilizes the characteristics of the policy gradient method to obtain the optimal parallel policy by continuously iterating and optimizing the reward value.

[0052] Example:

[0053] To verify the effectiveness of an automated parallel method based on dual graph Mamba networks and policy gradient optimization, three carefully constructed datasets—ptb, cifar10, and nmt—were selected for testing. The ptb dataset focuses on recurrent neural network architectures, cifar10 is based on convolutional neural networks, and the nmt dataset was generated by adjusting the parameters of a neural machine translation model. Each dataset contains 17 randomly selected computation graphs, with average operator counts of 190, 300, and 500 for the computation graphs in nmt, cifar10, and ptb, respectively. Performance was evaluated using two metrics: single-round training time and policy search time. Shorter single-round training time indicates more efficient parallel policies, while shorter policy search time indicates faster policy search speed. Experimental data reveals that compared to existing device placement methods such as Placeto, GraphSAGE, P-GNN, and CP-GNNAK, the proposed and implemented automated parallel method based on dual graph Mamba networks and policy gradient optimization demonstrates significant advantages in both policy search speed and parallel policy quality.

[0054] Table 1 visually illustrates the execution time of each method under different datasets and the number of GPUs. The parallel strategy implemented in this invention is significantly superior to other methods. Compared with the Placeto method, the parallel strategy obtained by applying this invention reduces the average training time of the model by 16.58%. Table 2 details the computation time of each method under different datasets and the number of GPUs, further demonstrating the superior performance of this invention in strategy search speed. Compared with the Placeto method, the average search speed of obtaining the optimal parallel strategy by this invention is improved by 31.53 times.

[0055] Table 1. Comparison of single-round training time for different methods on different datasets with different numbers of cards (unit: sec)

[0056]

[0057] Table 2 Comparison of strategy search time for different methods on different datasets with different numbers of cards (unit: sec)

[0058]

Claims

1. An automatic parallel method based on dual graph Mamba networks and policy gradient optimization, characterized in that, Includes the following steps: Step 1: Obtain a publicly available AI model dataset, perform operator fusion on each computation graph G in the dataset, and obtain the feature matrix of the computation graph. Where V represents the set of nodes in computation graph G, and F represents the features of the nodes in computation graph G; Transform an AI model into a graph Where V represents the set of computation graph operators, i.e., the set of nodes, and operator v i ∈V represents a single operator; E represents the set of directed edges between nodes, where element e i,j ∈E represents the operator v i and v j Data communication dependencies between them; given D = {d1, d2, d3, ..., d...} m } represents the devices in the cluster, where d i ∈D represents a computing device; given D and a computing graph Find mappings i Make each operator v i Corresponding to a device d i The ultimate goal is to find a device placement strategy. i When the computation graph The middle operator follows strategy s i Minimize the computation graph during placement. Execution time Step 2: Construct a dual graph mamba network containing a graph mamba branch and a decoupled graph neural network branch, and combine the feature matrix... The inputs are fed into the two branches respectively to obtain the feature representation X. M With X D The specific implementation process is as follows: Step 2.1: The Graph Mamba branch generates the node encoding matrix M based on the feature matrix X. 0 The edge encoding matrix Q is used to obtain the intermediate representation N through the message passing network MPNN. t Node encoding matrix M 0 The input is fed into the Graph Mamba module to extract features and obtain the intermediate feature matrix M. t M t and N t The characteristic matrix M is obtained by averaging the elements one by one. T ; M T The input is fed into a multilayer perceptron to obtain M. T′ ; M T′ With N t The final feature representation X of the graph mamba branch is obtained through a message passing network. M ; Step 2.2: Use the feature matrix X as the node attribute feature A 0 and A 0 Semantic features C are obtained by passing the edge encoding matrix Q through a GCN network. 0 Meanwhile, A 0 The semantic intermediate representation B is obtained after passing through a multi-layered graph mamba module. 0 A 0 C 0 B 0 A is obtained through weight sharing and structural consistency constraints. T C T and B T Finally, the three factors are weighted to obtain the final feature representation of the decoupled graph neural network branch. Step 3: Represent the computational graph features X extracted from the graph mamba branch. M Feature representation X of decoupled graph neural network branches D The mean value is calculated and fed into an LSTM network for feature fusion, resulting in the final feature representation of the dual graph mamba network. Step 4: Represent the features The input is fed into the reinforcement learning iterative solution module, which uses the near-end policy gradient optimization method of reinforcement learning to obtain the optimal parallel policy by continuously iterating and optimizing the intermediate reward value.

2. The automatic parallel method based on dual graph Mamba network and policy gradient optimization according to claim 1, characterized in that, The specific implementation process of step 1 is as follows: Obtain a publicly available AI model dataset. For each computation graph G in the dataset, perform operator fusion using the ColocRL algorithm. The fused computation graph contains relevant attributes of all operators in that computation graph, including operator type, input, output, and access identifier. Based on these operator attributes, construct the original feature vectors of the operators, i.e., initialize the node encoding, and obtain the complete feature matrix X of the computation graph. Where V represents the set of nodes in the computation graph G, and F represents the features of the nodes in the graph.

3. The automatic parallel method based on dual graph Mamba network and policy gradient optimization according to claim 2, characterized in that, The specific implementation process of step 2.1 is as follows: Step 2.1.1: Obtain the intermediate feature matrix by adding the adjacency matrix D of X and the identity matrix I. By analyzing the matrix Laplace normalization is performed to obtain the edge encoding matrix. The edge encoding matrix Q is input into the message passing network to obtain the intermediate representation. Step 2.1.2: Feature matrix X and intermediate feature matrix Multiplying them yields the node encoding matrix M. 0 Node encoding matrix M 0 Inputting into the TuMamba module yields the intermediate node code M t Node encoding M t The next state M is obtained through the Selective State Management (SSM) mechanism of the Graph Mamba module. t+1 ; Step 2.1.3: Process the feature matrix X using a message passing network to obtain the intermediate representation N. t+1 ; N t+1 and M t+1 By accumulating the results, we obtain the fused feature matrix M at time t. t+1′ and M t+1′ The input is fed into a multilayer perceptron to obtain the feature matrix M at the next time step. t+2 ; Step 2.1.4: The feature matrix M after T iterations... T With the intermediate representation N T Simultaneously, this is used as input again and processed through a message passing network to obtain the fused feature matrix M at time T. T′ Ultimately, M T′ The computational graph characteristic representation of the graph mamba branch is denoted as...

4. The automatic parallel method based on dual graph Mamba network and policy gradient optimization according to claim 3, characterized in that, The selective state management (SSM) mechanism is implemented as follows: First, three learnable weight matrices W are introduced. B W c and W Δ For M respectively t A linear transformation is performed to obtain the projection parameters B and C, and the discrete step size Δ after Softplus activation. The evolution parameters are then calculated using the evolution parameters A obtained at the previous time step, the projection parameters B at the current time step, and the discrete step size Δ. and in It is related to Δ*A through an exponential function. Then it is calculated based on the difference between the inverse matrix of Δ*A and the exponential function; subsequently, it is combined with M t-1 and M t And projection parameter C and evolution parameter and Update the current state to get M t+1 The specific calculations are as follows: B=W B M t ,C=W c M t ,Δ=Softplus(W Δ M t ) 5. The automatic parallel method based on dual graph Mamba network and policy gradient optimization according to claim 4, characterized in that, The specific implementation process of step 2.2 is as follows: Step 2.2.1: Use the feature matrix X as the node attribute features. Initial node attribute feature A 0 The edge encoding matrix Q is input into the graph convolutional network GCN to obtain the initial structural features. And use the Graph Mamba module to analyze node attribute feature A 0 Processing was performed; initial semantic features were obtained. Step 2.2.2: Obtain A 0 C 0 B 0 The node attribute feature matrix A at the final time is obtained by combining weight sharing and structural consistency constraints. T Topological feature matrix C T and semantic information feature matrix B T The final feature representation is obtained.

6. The automatic parallel method based on dual graph Mamba network and policy gradient optimization according to claim 5, characterized in that, The specific implementation process of step 2.2.2 is as follows: Step 2.2.3.1: Introduce the learnable shared reconstruction factor matrix W S Reconstruct the potential correlations between adjacency matrices and constrain W. S =WW T The specific calculation and update formulas are as follows: Where A is the node attribute feature matrix, B is the semantic information feature matrix, C is the topological feature matrix, and W, F, K and H are bias matrices; Step 2.2.3.2: A t B t and C t Error between and By reconstructing factor W S Minimize; where ε is the node attribute feature A t Semantic feature B t With topology C t The balance coefficient is used to balance the weight relationship among the three factors. Its value ranges from 0 to 1, and the specific formula is as follows: Step 2.2.3.3: After T iterations, the final A... T B T C T By concatenating the sequences, the final feature representation is obtained. The specific formula is as follows: X D =Concat(A T ,B T ,C T ) σ is the sigmoid function of the nonlinear transformation; L1, L2, and L3 are hyperparameters.

Citation Information

Patent Citations

  • Neural network distributed automatic parallel training method based on AC reinforcement learning

    CN116306897A

  • Multi-target chip layout optimization method based on hierarchical reinforcement learning

    CN117787186A