A model training method and system based on reinforcement learning memory scheduling decision

Abstract

The application provides a model training method and system based on reinforcement learning memory scheduling decision, the application is based on a reinforcement learning algorithm, analyzes training information generated in a training process, and updates a decision scheme according to feedback to determine which data to transfer, thereby further optimizing memory space and improving the overall performance of deep learning model training.

Description

Technical Field

[0001] This invention relates to the field of computer science and artificial intelligence, and in particular to a model training method and system for memory scheduling decisions in deep learning GPU memory management based on reinforcement learning. Background Technology

[0002] The innovation of deep learning technology has greatly promoted the development of computer vision, natural language processing, medicine, and other fields. To further improve accuracy, deep learning models are becoming increasingly larger and more complex, placing higher demands on the storage space of training devices. However, the storage space of deep learning accelerators, represented by GPUs, is limited, restricting researchers from developing and exploring larger-scale networks and hindering the iteration of deep learning technology. Currently, common memory optimization schemes involve transferring temporarily unnecessary layer data from GPU memory to system memory during training to alleviate GPU memory pressure, and then transferring the required data back to GPU memory before later access. The challenge of this technique lies in selecting which data to exchange at what stage. However, the transfer bandwidth between GPU and CPU is limited, and coarse-grained exchanges of large amounts of data result in long transfer times, introducing non-negligible communication latency. Existing solutions design transfer strategies using traditional heuristics or expert experience. However, it is difficult to determine the optimal data transfer scheme on large-scale neural networks, leading to significant performance degradation in model training. Summary of the Invention

[0003] The purpose of this invention is to address the shortcomings of existing technologies by proposing a model training method and system based on reinforcement learning memory scheduling decision-making.

[0004] The specific technical solution adopted in this invention is as follows:

[0005] A model training method based on reinforcement learning memory scheduling decision-making includes the following steps:

[0006] Step 1: During the first round of iterative training, initialize the scheduling scheme;

[0007] Step 2: Collect the GPU memory space, transfer bandwidth, and network layer information of the neural network model after the previous round of neural network model iteration training, record the data size and the dependencies between each layer; and calculate the data exchange time and training time based on the collected information.

[0008] Step 3: Based on the information obtained in Step 2, dynamically decide on the scheduling plan for each round, including the following sub-steps:

[0009] (3.1) After normalizing the training time and memory space consumption obtained in step two, a weighted evaluation is performed as the reward for the previous round.

[0010] (3.2) Based on the reward income of the previous round obtained in step (3.1) and the feedback value Q obtained in the previous round, the ε-greedy algorithm is used to select the tensor waiting to be exchanged in the current round and the scheduling scheme of the current round. The neural network model is trained in the current round according to the scheduling scheme.

[0011] (3.3) Update the feedback value Q and store it in the Q table based on the feedback value Q and reward income of the previous round and the scheduling plan of the current round;

[0012] Step 4: Repeat steps 2 and 3 for iterative training; as the number of iterations increases, the scheduling scheme will converge to the optimal scheduling solution; when the expected training effect is achieved, end the training.

[0013] Furthermore, in step (3.1), the reward R for the current round plan Specifically as follows:

[0014] R plan =R time +w(R mem -R time )

[0015] Among them, R time R represents the training time of the previous round. mem w represents the GPU memory consumed in the previous training round, and w represents the reward R. plan In the previous training round, the GPU memory space R consumed... mem The proportion it occupies.

[0016] Furthermore, in step (3.3), updating the feedback value Q specifically involves:

[0017] Q k (s,t)=Q k-1 (s,t)+α[R plan +γQ k-1 (s',t')-Q k-1 (s,t)]

[0018] Where k represents the current round, k-1 represents the previous round, t is the data tensor swapped in the previous round, s is the scheduling scheme used in the previous round, and t' and s' are the data tensor swapped in the current round and the scheduling scheme used, respectively. R plan This represents the reward from the previous round, where α is the learning rate and γ is the discount factor; Q k-1 (s,t) represents the feedback value in the (k-1)th round when the data tensor being exchanged is t and the scheduling scheme used is s. k-1(s',t') represents the feedback value in the (k-1)th round when the data tensor being exchanged is t' and the scheduling scheme is s'. k-1 (s,t), Q k-1 (s',t') is obtained by looking up the Q table.

[0019] Furthermore, in step (3.2), when using the ε-greedy algorithm to select the tensors waiting to be exchanged in the current round and the scheduling scheme for the current round, only the tensors corresponding to the same network layer type are selected as the tensors waiting to be exchanged in each round.

[0020] Furthermore, the model training method includes pre-training and formal training. Pre-training uses a portion of the training data for training and dynamically decides the scheduling scheme for each round. Formal training uses all the training data to continue training based on the pre-training and dynamically decides the scheduling scheme for each round.

[0021] A model training system based on reinforcement learning memory scheduling decision-making, used to implement the aforementioned model training method based on reinforcement learning memory scheduling decision-making, includes:

[0022] The acquisition module is used to collect the GPU's memory space, transfer bandwidth, and network layer information of the neural network model after the previous round of neural network model iteration training, record the data size and the dependencies between each layer; and calculate the corresponding time of data exchange and training time based on the acquired information.

[0023] The decision-making module is used to dynamically decide the scheduling scheme for each round based on the information obtained by the acquisition module;

[0024] The training module iteratively trains based on the data from the acquisition and decision-making modules; as the number of iterations increases, the scheduling scheme will converge to the optimal scheduling solution; when the expected training effect is achieved, the training ends.

[0025] The beneficial effects of this invention are: Based on reinforcement learning algorithms, this invention analyzes the training information generated during the training process and updates the decision-making scheme according to feedback, determining which data to transfer. This further optimizes memory space and improves the overall performance of deep learning model training. Attached Figure Description

[0026] Figure 1 It is a system architecture diagram.

[0027] Figure 2 It is a flowchart of the decision-making process.

[0028] Figure 3 This is a statistical chart of the maximum batch growth rate supported during training on CIFAR10.

[0029] Figure 4 This is a graph showing the results of batch learning optimization of a subset of datasets. Detailed Implementation

[0030] The present invention will be further described below with reference to the accompanying drawings and specific embodiments.

[0031] This invention provides a model training method based on reinforcement learning memory scheduling decision-making, such as... Figure 2 As shown, it includes the following steps:

[0032] Step 1: After the user implements the model code through the interface, the deep learning framework initializes the scheduling scheme 's' and performs training during the first round of iteration training. The scheduling scheme describes whether each tensor in the training data is swapped in each iteration, including the selected tensors to be swapped and those not to be swapped. Generally, the scheduling scheme can be represented using a two-dimensional array, that is, the state of the data is represented using a two-dimensional array to facilitate policy expansion. In the scheduling scheme, each two-dimensional array has a dimension of 2, where a scheduling policy of 0 indicates that the data is kept in GPU memory, and 1 indicates that data is swapped. A series of arrays constitutes a state, and the change in decision-making is reflected in the transition between states. The ultimate goal of the system is to solve for a set of arrays to achieve the desired training effect. The initialization of the scheduling scheme 's' can be set by the user.

[0033] Step Two: Starting from the second round of training, each training round first collects the GPU memory space, transfer bandwidth, and network layer information of the neural network model after the previous round of neural network model iteration training, recording the data size and the dependencies between layers; and calculates the data exchange time and training time based on the scheduling scheme of the previous round and the bandwidth; among them, the dependencies between layers are used to ensure the correctness of data exchange during training. For example, ensuring that when the current layer starts reverse computation, the data of the dependent layers are in GPU memory.

[0034] Step 3: Make dynamic decisions based on the information obtained in Step 1. This step is the core of the invention and can be divided into the following sub-steps:

[0035] (3.1) After normalizing the training time and memory consumption obtained in step two, a weighted evaluation is performed as the reward for the previous round. The reward for the previous round is the action reward of the scheduling scheme decided in the previous round. For example, it can be calculated using the following method:

[0036] R plan =R time +w(R mem -R time )

[0037] Among them, R timeR represents the training time of the previous round (in milliseconds). mem w represents the GPU memory consumed in the previous training round, and w represents the reward R. plan In the previous training round, the GPU memory space R consumed... mem The proportion (weight) it occupies takes a value between 0 and 1. The learning objective of dynamic decision-making is to minimize this reward objective value.

[0038] (3.2) Based on the reward income of the previous round obtained in step (3.1) and the feedback value Q obtained in the previous round, the ε-greedy algorithm is used to select the tensor waiting to be exchanged in the current round and the scheduling scheme of the current round and record them. The neural network model is trained in the current round according to the scheduling scheme.

[0039] (3.3) Based on the feedback value Q, reward income and scheduling scheme of the previous round, evaluate the strategy selection of the previous round, update the feedback value Q and store it in the Q table to provide a basis for scheduling decision.

[0040] The specific steps to update the feedback value Q are as follows:

[0041] Q k (s,t)=Q k-1 (s,t)+α[R plan +γQ k-1 (s',t')-Q k-1 (s,t)]

[0042] Where k represents the current round, k-1 represents the previous round, t is the data tensor swapped in the previous round, s is the scheduling scheme used in the previous round, and t' and s' are the data tensor swapped in the current round and the scheduling scheme used, respectively. R plan This represents the reward from the previous round. α and γ are common parameters in reinforcement learning. α is the learning rate, representing the step size of policy learning, and γ is the discount factor, representing the decay coefficient of the current reward's impact on subsequent time steps. Both parameters take values ​​between 0 and 1, with common configurations being α = 0.5 and γ = 0.9. k-1 (s,t) represents the feedback value in the (k-1)th round when the data tensor being exchanged is t and the scheduling scheme used is s. k-1 (s',t') represents the feedback value in the (k-1)th round when the data tensor being exchanged is t' and the scheduling scheme is s'. k-1 (s,t), Q k-1 (s',t') are obtained by looking up the Q-table. The Q-table records all the exchanged data tensors and the feedback values ​​under the adopted scheduling scheme, and is updated with each iterative training. At the beginning of iterative training, all values ​​in the Q-table are initialized.

[0043] Step 4: Repeat steps 2 and 3. The training time and space consumption will change in each iteration. As the number of iterations increases, the scheduling scheme will converge to the optimal solution. When the number of iterations is reached, or the training effect reaches the expected standard, the expected scheduling decision is obtained. When the expected training effect is achieved, the system ends training and saves the learning results of the current decision module.

[0044] Furthermore, during the decision-making process, based on the calculated historical Q-values, the ε-greedy algorithm is invoked to derive the next decision plan. For adjustments to the decision, the system modifies only one type of tensor of the neural network layer at a time, treating it as a single action. Throughout the process, the model makes decisions regarding new scheduling schemes based on prior experience and real-time training feedback. Once a decision is made, the system sends the specific decision to the training module for the next iteration.

[0045] Furthermore, dynamic decision-making can achieve better scheduling schemes through finite-state (decision) learning. However, after selecting a strategy, a long training iteration is required to obtain feedback, resulting in slow convergence speed in the search solution. Therefore, this invention, combining the training characteristics of neural networks, also provides an optimized scheme. During the iterative process of multi-round neural network training, the data access pattern is the same. For the same neural network, the access pattern of training data is also more similar. Therefore, training can be divided into two processes: pre-training and real training. Before training begins, the original training data is split. During pre-training, training decisions are made using sub-batch data of the original training data, and dynamic adjustments continue after formal training, thereby accelerating convergence. Experimental results show that, while ensuring good memory optimization, this scheme significantly improves the performance of network training.

[0046] Corresponding to an embodiment of a model training method based on reinforcement learning memory scheduling decision, the present invention also provides an embodiment of a model training system based on reinforcement learning memory scheduling decision.

[0047] The system architecture designed in this invention includes a training module, a data acquisition module, and a decision-making module, such as... Figure 1As shown, the training module completes the entire training task, including data input, model forward propagation, back propagation, gradient update, and other processes. Furthermore, this invention adds a data exchange technology to the traditional training framework. This technology uses a transfer function (CudaMemcpy()) to support data exchange between the GPU and CPU at specific stages according to the transfer strategy. The acquisition module is responsible for collecting training information, including device information such as GPU memory space and transfer bandwidth, network information such as the type and parameters of each layer of the model, and model training time. This information is promptly fed back to the decision module after each iteration. The decision module analyzes the training information and provides a memory scheduling scheme for the next iteration.

[0048] For the system embodiments, since they basically correspond to the method embodiments, the relevant parts can be referred to in the description of the method embodiments. The device embodiments described above are merely illustrative. The units described as separate components may or may not be physically separate, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed across multiple network units. Some or all of the modules can be selected to achieve the purpose of the present invention according to actual needs. Those skilled in the art can understand and implement this without creative effort.

[0049] The present invention also provides a specific embodiment to further illustrate the effects of the present invention:

[0050] The specific experiment is as follows:

[0051] Experimental setup:

[0052] (1) Operating system: Ubuntu 18.04.3LTS;

[0053] (2) CPU: 2.10GHz Intel(R)Xeon(R)Gold 521R, equipped with 128GB DRAM;

[0054] (3) GPU: RTX 2080Ti 11GB VRAM;

[0055] Model configuration:

[0056] (1) Models: AlexNet, Vgg16, ResNet50, InceptionV3;

[0057] (2) Dataset: Cifar10, containing 60,000 color images, each 32*32 in size, divided into 10 classes, with 6,000 images per class, of which 50,000 images are used for training and 10,000 images are used for testing;

[0058] (3) The formulas for reward assessment and strategy assessment are based on experience, with α set to 0.5 and γ set to 0.9;

[0059] Introduction of the comparison objects:

[0060] 1. PyTorch: The industry-standard deep learning training framework that does not involve GPU memory optimization techniques;

[0061] 2. vDNN and Capuchin: Currently the best deep learning memory optimization frameworks;

[0062] Final test results:

[0063] Space testing: such as Figure 3 As shown, under the experimental conditions described above, the maximum batch size supported by training AlexNet using the PyTorch framework exceeds 3800, while our system supports a maximum batch size exceeding 5500. For the largest neural network, InceptionV3, our system supports a maximum batch size exceeding 360, more than twice the size supported by PyTorch. Compared to the PyTorch framework, our system improves the maximum batch size supported for training on the CIFAR10 dataset by 1.4× to 2.0×. Compared to the vDNN approach, our system supports a maximum batch size improvement of 20% to 70%. Compared to the world's best Capuchin approach, our system achieves a 3% to 5% improvement in space optimization without significant performance loss.

[0064] Figure 4 The study demonstrates the learning optimization effect of using a subset of the dataset for pre-training. Experiments show that selecting 80% of the CIFAR dataset can improve learning efficiency by 30% and achieves an optimal training performance of 85.6%. Therefore, selecting an appropriately sized subset of the dataset can significantly improve training efficiency.

[0065] Obviously, the above embodiments are merely illustrative examples for clear explanation and are not intended to limit the implementation. Those skilled in the art will recognize that other variations or modifications can be made based on the above description. It is neither necessary nor possible to exhaustively list all possible implementations here. However, obvious variations or modifications derived therefrom are still within the scope of protection of this invention.

Claims

1. A model training method based on reinforcement learning memory scheduling decision-making, characterized in that, Includes the following steps: Step 1: During the first round of iterative training, initialize the scheduling scheme; Step 2: Collect the GPU memory space, transfer bandwidth, and network layer information of the neural network model after the previous round of neural network model iteration training, and record the size of the data and the dependencies between each layer; And calculate the data exchange time and training time based on the collected information; Step 3: Based on the information obtained in Step 2, dynamically decide on the scheduling plan for each round, including the following sub-steps: (3.1) After normalizing the training time and memory space consumption obtained in step two, a weighted evaluation is performed as the reward for the previous round; (3.2) Based on the reward income of the previous round obtained in step (3.1) and the feedback value Q obtained in the previous round, the ε-greedy algorithm is used to select the tensor waiting to be exchanged in the current round and the scheduling scheme of the current round. The neural network model is trained in the current round according to the scheduling scheme. (3.3) Update the feedback value Q and store it in the Q table based on the feedback value Q and reward income of the previous round and the scheduling plan of the current round; Step 4: Repeat steps 2 and 3 for iterative training; as the number of iterations increases, the scheduling scheme will converge to the optimal scheduling solution; when the expected training effect is achieved, end the training.

2. The model training method according to claim 1, characterized in that, In step (3.1), the reward R from the previous round plan Specifically as follows: R plan =R time +w(R mem -R time ) Among them, R time R represents the training time of the previous round. mem w represents the GPU memory consumed in the previous training round, and w represents the reward R. plan In the previous training round, the GPU memory space R consumed... mem The proportion it occupies.

3. The model training method according to claim 1, characterized in that, In step (3.3), updating the feedback value Q specifically involves: Q k (s,t)=Q k-1 (s,t)+α[R plan +γ Q k-1 (s', t')-Q k-1 (s,t)] Where k represents the current round, k-1 represents the previous round, t is the data tensor swapped in the previous round, s is the scheduling scheme used in the previous round, and t' and s' are the data tensor swapped in the current round and the scheduling scheme used, respectively. R plan Q represents the reward from the previous round, α is the learning rate, and γ is the discount factor; k-1 (s,t) represents the feedback value Q in the (k-1)th round when the data tensor being exchanged is t and the scheduling scheme used is s. k-1 (s', t') represents the feedback value in round k-1 when the data tensor being exchanged is t' and the scheduling scheme used is s'. k-1 (s,t), Q k-1 (s', t') are obtained by looking up the Q table.

4. The model training method according to claim 1, characterized in that, In step (3.2), the ε-greedy algorithm is used to select the tensors waiting to be exchanged in the current round and the scheduling scheme for the current round. In each round, only the tensors corresponding to the same network layer type are selected as the tensors waiting to be exchanged.

5. The model training method according to claim 1, characterized in that, The model training method includes pre-training and formal training. Pre-training uses a portion of the training data for training and dynamically decides the scheduling scheme for each round. Formal training uses all the training data to continue training based on the pre-training and dynamically decides the scheduling scheme for each round.

6. A model training system based on reinforcement learning memory scheduling decision-making, characterized in that, The method for training a model based on reinforcement learning memory scheduling decision as described in any one of claims 1-5 includes: The acquisition module is used to collect the GPU's memory space, transfer bandwidth, and network layer information of the neural network model after the previous round of neural network model iteration training, record the data size and the dependencies between each layer; and calculate the corresponding time of data exchange and training time based on the acquired information. The decision-making module is used to dynamically decide the scheduling scheme for each round based on the information obtained by the acquisition module; The training module iteratively trains based on the data from the acquisition and decision-making modules; as the number of iterations increases, the scheduling scheme will converge to the optimal scheduling solution; when the expected training effect is achieved, the training ends.

Images