A neural network distributed automatic parallel training method based on AC reinforcement learning

Abstract

The application discloses a neural network distributed automatic parallel training method based on AC reinforcement learning. Firstly, the application performs performance data sampling analysis on a neural network model based on a performance analysis method, extracts performance data of model operators, and realizes operator grouping fusion. Secondly, global feature vector representation of a calculation graph is realized, and a state search space of reinforcement learning is constructed. Then, an AC algorithm based on time series difference sampling is adopted to complete iterative optimization of the reinforcement learning model, so as to search for an optimal distributed parallel strategy. Finally, a multi-thread simulation executor based on event driving is constructed to simulate the calculation process of the neural network model. The application reduces the state search space scale of the AC reinforcement learning, improves the search performance and universality of the parallel strategy, reduces the strategy iteration execution time, and reduces the hardware execution cost in the iterative search process.

Description

Technical Field

[0001] This invention belongs to the field of parallel training of large-scale complex neural networks, and specifically relates to a distributed automatic parallel training method for neural networks based on AC reinforcement learning. Technical Background

[0002] With the accelerating pace of research and innovation in deep learning, deep learning models are widely used in scenarios such as computer vision (CV), natural language processing (NLP), and search and recommendation. A common characteristic of deep learning models is the use of complex hierarchical neural networks. For example, the BERT (Bidirectional Transformer Language Model) network is constructed from composable modules such as encoders, decoders, and attention mechanisms; CNN networks are built based on components such as convolutional layers and pooling layers. Neural networks have proven to have significant combinatorial scalability, allowing for the training of large-scale model parameters on larger datasets to improve prediction accuracy. However, due to the limited resources of single devices, they cannot handle large volumes of input data and complex model parameters. Therefore, to effectively train such complex neural networks, it is necessary to partition and schedule large-scale datasets (data parallelism) or partition, schedule, and execute neural network models across multiple devices (model parallelism) to fully utilize the performance of computing devices.

[0003] Currently, there are many frameworks for distributed training, such as TensorFlow, PyTorch, and MindSpore. However, existing methods mainly rely on expert experience to manually search for parallel strategies. This requires developers to possess expertise in AI, distributed computing, and architecture, as well as the ability to make professional choices within these domains. Manually searching for the optimal parallel strategy is inherently difficult. Therefore, to simplify the design and implementation of parallel methods for neural network models and improve the versatility of parallel strategy design, the industry has begun to research automatic parallel training methods for neural networks, enabling automatic search and tuning of distributed parallel strategies. Reinforcement learning, in particular, has achieved outstanding performance in recent years in complex decision-making problems such as games and autonomous driving, reaching or even surpassing human decision-making levels, thus becoming a key research focus for automatic parallel strategy search.

[0004] Google first proposed the Hierarchical method, which extracts features from the neural network model and clusters, and uses reinforcement learning (RL) to guide the parallel policy search. However, this method requires frequent sampling and has a large search space, making the policy search process expensive. Therefore, its performance improvement compared to expert-based model parallel methods is limited. Gao et al. proposed Spotlight, which for the first time modeled the neural network operator scheduling problem as a Markov decision process (MDP). However, these methods are only effective for the current specific network model. When encountering a new network model, a new parallel policy search is needed, lacking the ability to be ported to other similar networks. For different network models, the cost of designing and implementing parallel policies remains high. To address this issue, Addanki et al. proposed Placeto, which introduces graph embedding encoding to give parallel policies portability and avoid repeated training on similar unknown networks. Wang Siyu et al. proposed AutoMap, which performs automatic parallel policy search based on finer-grained XLA-IR graphs, but its DQN sampling efficiency is low and requires a large amount of historical experience storage, resulting in low overall execution efficiency. Baidu Labs combined pipeline technology with reinforcement learning methods to achieve coarse-grained layer-level scheduling to improve training throughput and reduce model training costs. However, it is still based on Monte Carlo full-process sampling reinforcement learning, which has low sampling efficiency. When the complexity of the neural network model increases exponentially, the efficiency of a single round of sampling also decreases exponentially, leading to a decrease in model convergence rate, getting stuck in local optima, and difficulty in coping with the ever-expanding neural network model. Summary of the Invention

[0005] To address the aforementioned problems, this invention designs and implements a distributed automatic parallel training method for neural networks based on AC (Actor and Critic) reinforcement learning.

[0006] This invention first performs performance data sampling and analysis on a neural network model using profiling methods, extracting multi-dimensional performance data such as the computation of model operators and the communication costs between different devices. It then constructs constrained grouping rules for the parallel computation benefits and communication costs of operators, achieving operator grouping and fusion to reduce the size of the computation graph. Second, based on the original features of the fused computation graph, it implements a global feature vector representation of the computation graph for message passing, constructing a state search space for reinforcement learning. Finally, based on the constructed state search space, it uses the AC algorithm based on temporal difference sampling to iteratively optimize the reinforcement learning model, searching for the optimal distributed parallel strategy and improving the stability, convergence rate, and policy search performance of the reinforcement learning algorithm. Furthermore, it constructs an event-driven multi-threaded simulator to simulate the computation process of the neural network model, reducing hardware execution costs and shortening the time for outputting policy rewards through multi-threading, providing an efficient execution and optimization environment for the entire method.

[0007] The specific steps of this invention are as follows:

[0008] Step 1: Perform performance data sampling and analysis on the neural network model using the performance analysis profiling method, extract multi-dimensional performance data such as computational cost, memory cost, and tensor transfer size of the model operators, and fit the communication cost function in the real execution environment to realize the grouping and fusion of computation graph operators based on computation-communication cost constraints, thereby reducing the size of the computation graph.

[0009] Step 2: Implement message-passing-based computation graph feature encoding based on the original features of the fused computation graph. First, extract the original features of the fused computation graph after operator grouping, mainly including computational cost, in-degree, and operator output tensor size. Second, construct the original feature vector of the computation graph using these original feature information. Finally, based on the original feature vector, obtain the nearest neighbor information of the operators using the message-passing mechanism. By capturing more local information, realize the global feature representation of the computation graph and construct the state search space for reinforcement learning, i.e., computation graph feature encoding.

[0010] Step 3: For the constructed state search space, the AC algorithm based on temporal differential sampling is used to iteratively optimize the reinforcement learning model and output the optimal scheduling strategy.

[0011] Reinforcement learning utilizes the interaction between the agent and the critic, i.e., the constructed state search space, to sample quintuples (...). Using P, r, V) data, an automatic optimal parallel strategy search is achieved, where This indicates the state of reinforcement learning in the current environment. Let P represent the actions that the Agent can take in the environment, let r represent the reward generated by the current action, and let V represent the value of the current state. The Agent samples the above data set multiple times and iteratively updates its parameters based on the reward and value to achieve the optimal automatic parallel strategy output.

[0012] First, construct the reinforcement learning state search space. This state space is the computation graph feature encoding constructed in step 2. During the reinforcement learning iteration process, the agent changes the current state of the computation graph by interacting with the environment. This state is the original feature vector in step 2. Encoding the graph features of these different original feature vectors constitutes the overall state search space of reinforcement learning.

[0013] Secondly, a Markov decision process is implemented based on TD(k) (Temporal Difference k, k-step temporal difference learning). During single-step iteration, the Agent is divided into an Actor Agent and a Critic Agent. The Actor is a feedforward neural network with a SoftMax layer, responsible for iterative policy optimization and outputting actions based on the current action probability distribution P. Critic is an MLP (Multi-Layer Perceptron) network that outputs a value estimate V for two adjacent states before and after an action. t V t+1 The value estimate of this state represents the magnitude of the current state's value; a higher value indicates a better current action. The Actor Agent implements a temporal difference parameter iterative update algorithm by sampling the values ​​of multiple adjacent states and the rewards for state transitions. Its parameter iteration formula is:

[0014]

[0015] Among them, s t π(a) represents the state of the environment at time t. t |s t ;θ) represents the policy network of the current Actor Agent, a t θ is the action taken by the Actor Agent based on the current policy at time t, where θ is the parameter of the policy network; It is an entropy regularization term used to enhance the Actor Agent's exploratory capabilities; A(s t a t ω) is the advantage function, indicating the dominance in the current state s. t Next, action a tThe advantage function is the size of the advantage; ω is a parameter of the Critic Agent's value network. The advantage function is constructed based on the estimated state value and the reward for state transitions from the Critic Agent's value network, where the reward value is defined by the following formula:

[0016] Among them, R(a) t ) represents the current action a t The reward received, r t M represents the execution time of the current strategy, and M represents the memory limit of the hardware device. cost This represents the memory consumption of the current policy, and α represents the penalty factor. If the memory consumption of the current policy exceeds the device's memory limit, a memory penalty term will be added.

[0017] Critic itself uses bootstrapping to iterate network parameters by using the mean square error of the values ​​of two adjacent states as the loss function.

[0018] Finally, the parallel strategy is input into the simulation execution engine for simulation execution, and the output execution reward value is used for iterative optimization of Actor and Critic to achieve automatic search for the optimal parallel strategy.

[0019] Step 4: Construct an event-driven, multi-threaded simulator to simulate the execution process of a neural network model in a real environment, reduce hardware execution costs, reduce the time for outputting policy rewards, and provide an efficient execution and optimization environment for the entire method.

[0020] First, to address the computational and communication costs of the execution entity in the neural network model, a device execution queue and a device communication queue are constructed to simulate the computation and communication processes of operators in the neural network model. Second, an event queue is constructed using various event mechanisms (such as computation events, communication events, topology refresh events, etc.) to enable event-driven interaction between the device execution queue and the device communication queue. Finally, the states of the three queues—the device execution queue, the device communication queue, and the event queue—are traversed. When all queues are empty, the neural network model simulation is complete, and training is finished.

[0021] The beneficial effects of this invention are as follows:

[0022] First, performance data sampling of the neural network model is performed using profiling methods. Based on the sampled data of the model operators, an operator constraint grouping and fusion method for computation-communication cost is implemented to reduce the state search space size of AC reinforcement learning and improve the search performance of parallel policies. Second, feature encoding of the computation graph is performed using a message-passing-based approach to achieve comprehensive feature representation of the computation graph, thereby improving the versatility of parallel policies. Finally, an AC algorithm based on temporal difference sampling is implemented to iteratively optimize the reinforcement learning model for efficient searching of the optimal distributed parallel policy. Furthermore, a simulated executor based on an event-driven method is constructed to simulate the real execution environment. Simultaneously, the accuracy of the simulated executor is improved by fitting the communication cost function, making the simulated execution environment closer to the real environment. In this simulated environment, the policy iteration execution time can be reduced, and the hardware execution cost during the iterative search process can be lowered, providing an efficient execution and optimization environment for the entire method. Attached Figure Description

[0023] Figure 1 This is a schematic diagram of the overall process of searching for the optimal parallel strategy based on AC reinforcement learning.

[0024] Figure 2 This is a schematic diagram of computational graph feature encoding based on nearest neighbor message passing;

[0025] Figure 3 This is a schematic diagram of the iterative optimization process of the AC reinforcement learning agent;

[0026] Figure 4 This is a schematic diagram of the execution process of a high-precision simulated actuator. Detailed Implementation

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

[0028] A distributed automatic parallel training method for neural network models based on AC reinforcement learning, such as... Figure 1 As shown, the specific steps include the following:

[0029] Step 1: Construct the computation graph of the neural network model based on mainstream deep learning frameworks (Tensorflow, PyTorch, etc.). The computation graph is a directed acyclic graph defined according to the data flow of the neural network, containing computation nodes and resource dependencies. Therefore, the computation graph is an implementation form of the neural network model. Performance data sampling and analysis of the computation graph constructed by the neural network is performed using profiling methods to extract multi-dimensional performance data such as computational cost, memory cost, and tensor transfer size of the computation graph operators. A communication cost function under a real execution environment is then fitted to achieve computation graph operator grouping and fusion based on computation-communication cost constraints, reducing the size of the computation graph.

[0030] First, the neural network model is trained in multiple rounds to sample the computational cost, memory cost, tensor transfer size, and tensor communication cost of the neural network model operators. Based on this, the tensor size and tensor communication cost during the tensor transfer process are used to fit the communication cost function. Based on the sampled data, the constrained operator grouping based on computation-communication cost is implemented, and the constrained grouping rule is constructed, that is, when the parallel computation time of the operator is shorter than the communication time, constrained grouping is performed.

[0031] Secondly, during the operator fusion process based on the above rules, it is necessary to ensure the directed acyclicity of the computation graph; otherwise, circular dependencies will occur, leading to model execution failure. However, since the computation graph is too large, using conventional methods to determine whether a cycle has formed is too costly. Therefore, strict ablation conditions are set: operator ablation requires either the out-degree of the predecessor operator to be 1 or the in-degree of the successor operator to be 1. This ensures smooth model execution while fusing nodes based on the aforementioned constraint rules as much as possible, thereby reducing the size of the search space. Simultaneously, an ablation threshold T is set to control the ablation range of the computation graph. Therefore, the final grouping rules are as shown in Formula 1:

[0032]

[0033] Among them, o i and o j "Don't" represents different operator nodes in the computation graph, succ(o i ) represents o i The successor node, outdegree(fore(o i )) represents o i The out-degree of the predecessor node; indegree(succ(o i )) represents o i The subsequent in-degree; o i and o j Inter-connection communication cost; C avg (o j ) represents o j Average computational cost.

[0034] Finally, extract the computation graph of the neural network model, traverse each node of the computation graph according to the topological order, and find the current driver node. i and successor node o j Nodes that meet the constraint rules are grouped together. Ablation ends when the number of nodes in the computation graph is below the set threshold T or when there are no nodes that meet the conditions. Node ablation is shown in Algorithm 1:

[0035]

[0036] Step 2: Based on the original features of the fused computation graph, implement message-passing-based computation graph feature encoding to represent the computation graph and complete the construction of the state search space for reinforcement learning.

[0037] First, the original features of the computation graph after operator grouping and fusion are extracted, mainly including features such as computational cost, in-degree and out-degree, operator output tensor size, hardware device label, and access label.

[0038] Secondly, the original feature vectors of the computational graph are constructed using this raw information.

[0039] Finally, based on the original feature vectors, a message-passing-based computational graph feature encoding method is implemented. This process consists of two stages: message passing and neighbor information aggregation, as detailed below. Figure 2 As shown.

[0040] The first phase, the message passing phase, obtains the current computing node o through two processes: forward computation and backward computation. i Original feature vector of the predecessor node Original feature vector of successor node and the original feature vectors of parallel nodes The global environment characteristics of the current node are obtained through message passing, thereby enhancing the node's feature representation. The message passing process is shown in the formula:

[0041]

[0042] Among them, P μ With S θ These are multilayer perceptron networks for forward and backward information transmission, respectively, where θ and μ are trainable parameters. This represents the original feature vector of the current node i after aggregating information from its neighboring nodes. It is a feature vector generated after information transmission between nodes. The network parameters can be reused in different neural networks to improve the generalization ability of the reinforcement learning model's policy search.

[0043] The second stage, the neighbor information aggregation stage, involves combining the feature vectors obtained from the neighbor information acquisition for each operator in the first stage with their parent node groups. child node group and parallel groups Aggregation is performed to complete the feature encoding of the computational graph. Specifically, a comprehensive combinatorial embedding process is implemented for each node. During the aggregation process, the parent node group, child node group, and parallel group are searched and determined. The parent node group represents the set of parent nodes that the current node can reach, the child node group represents the set of child nodes that the current node can reach, and the parallel group represents the set of nodes that cannot be reached from the current node or reach the current node. The aggregation process formula is as follows:

[0044]

[0045] in f represents the feature encoding of the current three node groups. μ and g θ For a multilayer perceptron network with trainable parameters, the three embedding encoding results are finally combined with those from the first stage using a vector concatenation operation. The connections are made to form a computational graph feature encoding, which serves as the state search space for reinforcement learning. The algorithm for the neighbor information aggregation stage is shown in the table below:

[0046]

[0047]

[0048] Step 3: For the constructed state search space, the AC algorithm based on temporal difference sampling is used to iteratively optimize the reinforcement learning model, and the optimal scheduling policy is output, such as... Figure 3 As shown.

[0049] First, the Actor encodes the computational graph features of the neural network model as state S. t The inputs are fed into their respective networks; the Actor is based on the current system state features S. t The scheduling action 'a' of the current operator is output after passing through the feedforward neural network (FNN) and the softmax layer. t That is, which device to plan for execution, and at the same time, based on the current a t Regarding the state S of the system t Update to reach S t+1 And save the changed state.

[0050] Secondly, Critic utilizes an MLP (Multi-Layer Perceptron) to process the initial input state S. t And the new state S obtained in the previous step t+1 Perform a value assessment and output the state value V(s) t ;ω) and V(s t+1 ;ω).

[0051] Then, using the obtained reward value, V(s) t+1 ;ω) and V(s t ;ω) calculates TD-error (Temporal Difference-error), which is an advantage function built based on the AC algorithm, used to evaluate the magnitude of the advantage of the current agent's action, where V(st+1 ;ω) and V(s t ω) is estimated by the Critic value network, while the reward value is defined by the following formula:

[0052]

[0053] Among them, R(a) t ) represents the current action a t The reward received, r t M represents the execution time of the current strategy, and M represents the memory limit of the hardware device. d This represents the memory consumption of the current strategy, and c represents the penalty factor. If the memory consumption of the current strategy exceeds the device's memory limit, a memory penalty item will be added to guide the Agent to explore in a more balanced direction in terms of memory usage.

[0054] Finally, using the above methods, we complete the parameter updates for both the Actor and the Critic. The parameter update method for the Actor is as follows:

[0055]

[0056] in Defined by the following formula:

[0057]

[0058] π(a t |s t ;θ) represents the policy network of the current Actor Agent, A(s t a t ;ω) represents the TD-error, which is calculated as follows:

[0059]

[0060] Where k represents the interval of the current sampling data, i.e. how many actions the Agent performs before iterating the policy; γ is the decay coefficient, a hyperparameter used to control the importance of reward and state value; R represents the reward obtained by the current action; V represents the value of the current state of the Agent, where the state value is predicted by the Critic network; and ω is its network parameter.

[0061] As the model training progresses, Actors will increasingly choose actions with higher value, thus abandoning exploration of the environment. Therefore, an entropy regularization term is used to guide Actors to explore actions with greater uncertainty and higher information content, preventing the entire model from getting trapped in local optima. Based on the definition of Shannon entropy (i.e., information entropy), an entropy regularization term is introduced.

[0062]

[0063] in, The set of all policies is represented by π(u), and π(u) represents the currently searched policy. Shannon entropy can be used to calculate the information gain of the current system. Therefore, the essence of this regularization term is to allow the agent to explore in directions that acquire more information during the iteration process, thereby preventing the model from getting trapped in local optima. The Critic network parameter iteration uses the bootstrapping method, using the mean squared error of the two state values ​​as the loss function for iterative parameter optimization.

[0064]

[0065] Among them, V s and V s+k These represent the value estimates made by the Critic network for state s and state s+k after k steps, respectively, when the parameter is ω. This represents the partial derivative of the mean squared error with respect to the parameter ω. Finally, ω is updated using the gradient descent method.

[0066] The specific iterative process is shown in Algorithm 3:

[0067]

[0068]

[0069] Currently, model developers build neural network models based on general frameworks (such as Tensorflow, PyTorch, etc.). However, improving model accuracy requires repeated debugging and training, which is very time-consuming. The method of this invention can automatically and efficiently find the optimal parallel training strategy for various neural network models, reducing the training time for neural network model developers. Taking the classic neural network model InceptionV3 in the field of image classification as an example, this model is one of the most widely used image classification models. Most image classification model developers extend and develop based on this model. Therefore, by improving the training efficiency of this model, the overall development efficiency can be improved. After operator fusion and grouping, the computation graph of this neural network model has approximately 1400 operators. Based on existing methods, all operators need to be scheduled before a policy iteration can be performed, ultimately requiring 30 iterations to converge to the optimal policy, with an average single iteration time of 95 seconds. However, the reinforcement learning model based on AC proposed in this invention can achieve single-step policy iteration, that is, in each iteration, each operator scheduling allows for a policy iteration, thus improving the reinforcement efficiency. The learning model has high policy search efficiency. Under the same environment and model, the method of this invention converges to the same optimal policy as the former in the 14th round. Although the average iteration time of this method is 104s, which is longer than the former, the overall search time for the same optimal policy is reduced by 48.9%. Finally, the distributed parallel strategy is applied to the training and debugging of the InceptionV3 model, which can reduce the training time per round from 2.65s to 2s. When model developers need to group the model for multiple rounds of training to improve the model accuracy, for example, dividing it into 10 groups and performing 1000 rounds of iteration per group, the training and debugging time can be reduced by 1.8 hours. This can greatly improve the model development efficiency of various neural network model developers and train the same high-precision model in a shorter time.

[0070] Step 4: Construct an event-driven simulation executor to simulate the execution process of a neural network model in a real-world environment. This reduces hardware execution costs, shortens the time spent outputting policy rewards, and provides an efficient execution and optimization environment for the entire method. Figure 4 As shown.

[0071] First, based on the DAG (Directed Acyclic Graph) computation graph of the neural network model, device execution events are initialized, and initial executable operators are placed into the device execution queue, completing the initialization of the simulator's overall state. Second, by continuously retrieving corresponding events from the event queue and selecting the appropriate processor through a selector, the execution process of the entire computation graph is simulated. Each processor contains a scheduler and an event generator. The specific processing procedure for each event is as follows:

[0072] 1. Device execution processor:

[0073] a. Device execution scheduler: Retrieves the head task from the current device execution queue, completes operator execution, and triggers the event generator.

[0074] b. Event Generator

[0075] i. Device execution event: The current scheduling operator has completed execution. The device that has completed the current operator is rescheduled, that is, the current device is idle and the next device execution event is triggered.

[0076] ii. Device Communication Event: When the current scheduling operator has completed execution, and its successor operator's execution device is different from its execution device, a device communication event is triggered, and the current operator's execution device is changed. i Add to the device communication queue.

[0077] iii. Topology update event: The current scheduling operator has completed execution, and a topology update event is triggered when its successor operator has the same execution device.

[0078] 2. Communication Execution Processor

[0079] a. Communication Execution Scheduler: Retrieves the head task from the current device communication queue, calculates the communication cost between nodes using a fitting function, completes simulated communication between devices, and triggers the event generator.

[0080] b. Event Generator

[0081] i. Device Communication Event: If the end time of this communication is less than the global execution time, the current device communication can be executed normally, and the communication state of the current device will be rescheduled, meaning the current device communication is in an idle state and can trigger the next device communication event; otherwise, the current operator will be... i The communication task is re-added to the device execution queue to await the next scheduling.

[0082] ii. Topology refresh event: If the current device completes communication, it means that the current operator information has been transmitted to the device where its successor node is located, thus triggering a topology refresh event.

[0083] 3. Topology processor

[0084] a. Topology Scheduler: In the DAG graph of the model, operators that have completed computation and communication are marked as visited, the in-degree of their child nodes is decremented by 1, and the event generator is triggered.

[0085] b. The event occurs:

[0086] i. Device Execution Event: Operators to be executed are simultaneously inserted into the execution queue, and an event is generated. When the in-degree of a child node is 0, the current operator enters a schedulable state, thus generating a device execution event and setting the current operator's in-degree to 0. i The child nodes are placed in the device execution queue to await scheduling.

[0087] Finally, after the three processors have completed all event processing, that is, after scheduling all operators in the model graph and all devices are in an idle state, no new events will be generated. At this point, the execution is complete, and the final execution time is output as the simulation execution time. The specific execution process is shown in Algorithm 4:

[0088]

[0089]

Claims

1. A distributed automatic parallel training method for neural networks based on AC reinforcement learning, characterized in that, Includes the following steps: Step 1: Perform performance data sampling analysis on the neural network model using the performance analysis profiling method, extract multi-dimensional performance data of the model operators, and fit the communication cost function in the real execution environment to realize the grouping and fusion of computation graph operators based on computation-communication cost constraints; Step 2: Based on the original features of the fused computation graph, implement message-passing-based computation graph feature encoding, i.e., state search space; Step 3: For the state search space, the AC algorithm based on temporal differential sampling is used to iteratively optimize the reinforcement learning model and output the optimal scheduling strategy. Step 4: Construct an event-driven, multi-threaded simulator to simulate the execution process of a neural network model in a real environment, providing an efficient execution and optimization environment; The multi-dimensional performance data mentioned in step 1 includes: computational cost, memory cost, and tensor transfer size.

2. The distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 1, characterized in that, In step 1, the grouping rules for group fusion are as follows: in, and Each represents a different operator node in the computation graph. express The successor node, express The out-degree of the predecessor node; express The subsequent in-degree; express and The cost of fitting communication between them express Average computational cost.

3. The distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 2, characterized in that, Step 2 is as follows: 2.1 Extracting original features from the computational graph after operator grouping and fusion; 2.2 Constructing the original feature vector of the computational graph using the original features; 2.3 Based on the original feature vector, the nearest neighbor information of the operator is obtained by using the message passing mechanism. By capturing local information, the global feature representation of the computation graph is realized, and the state search space of reinforcement learning is constructed, namely the feature encoding of the computation graph.

4. The distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 3, characterized in that, The original features described in 2.1 include computational cost, in-degree size, and operator output tensor size.

5. The distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 4, characterized in that, Step 3 is as follows: 3.1 Constructing the State Search Space for Reinforcement Learning The state space is the computation graph feature encoding constructed in step 2. During the reinforcement learning iteration process, the agent changes the current state of the computation graph by interacting with the environment. This state is the original feature vector in step 2. Encoding the graph features of these different original feature vectors constitutes the state search space of reinforcement learning. 3.2 Implement Markov decision process, based on k-step temporal difference learning, and divide the Agent into actor agent and critic agent during single-step iteration; The Actor Agent implements a temporal difference parameter iterative update algorithm by sampling the values ​​of multiple adjacent states and the rewards for state transitions. The parameter iteration formula is as follows: ， in, express The state of the environment at any time, The policy network representing the current Actor Agent. Is At any given moment, the action taken by the Actor Agent based on the current policy; These are the parameters of the policy network; It is an entropy regularization term, used to improve the exploration capabilities of the Actor Agent; This is the dominance function, indicating the current state. Next, action The extent of the advantage; These are parameters of the Critic Agent's value network; Critic itself uses the mean square error of the values ​​of two adjacent states as a loss function to iterate the network parameters through bootstrapping. 3.3 Input the parallel strategy into the simulation execution engine to simulate execution, and output the execution reward value for iterative optimization of Actor and Critic to achieve automatic search for the optimal parallel strategy.

6. The distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 5, characterized in that, In section 3.2, the Actor is a feedforward neural network with a SoftMax layer, responsible for iterative policy optimization based on the current action probability distribution. Output Action ; The Critic is a multilayer perceptron (MLP) network that outputs value estimates for two adjacent states before and after an action. The value estimate of this state represents the value of the current state; the higher the value, the better the current action.

7. The distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 6, characterized in that, The advantage function described in section 3.2 is constructed based on the state value estimated from the Critic Agent's value network and the reward value for state transitions, where the reward value is defined by the following formula: in, Represents the current action The rewards received This represents the execution time of the current strategy. This represents the maximum memory capacity of the hardware device. This represents the memory consumption of the current strategy. This represents the penalty factor. If the current policy consumes more memory than the device's memory limit, a memory penalty will be added.

8. A distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 7, characterized in that, Step 4 is as follows: 4.1 To address the computational and communication costs of the execution entity in the neural network model, a device execution queue and a device communication queue are constructed to simulate the computation and communication processes of operators in the neural network model. 4.2 Construct event queues using multiple event mechanisms to enable interaction between event-driven device execution queues and device communication queues; 4.

3. Traverse the states of the three queues: device execution queue, device communication queue, and event queue. When all queue states are empty, the neural network model simulation is complete, and training is finished.

9. A distributed automatic parallel training method for neural networks based on AC reinforcement learning according to claim 8, characterized in that, The event mechanisms described in 4.2 include: computation events, communication events, and topology refresh events.

Images