A data allocation optimization method and device based on calculation graph and recalculation

Abstract

The present invention provides a data allocation optimization method and device based on calculation graphs and recalculations, and relates to the field of optimization technology. The method comprises: obtaining storage overhead and calculation overhead; determining the optimal recalculation strategy corresponding to each calculation unit according to the storage overhead and calculation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of calculation units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit, and the first goal of saving the most calculation time for all calculation units in each subgraph; determining the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline node corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline node corresponding to each subgraph, with the shortest total training time as the second goal. The device executes the above method. The method and device provided by the embodiment of the present invention can improve the overall calculation efficiency of the model.

Description

Technical Field

[0001] The present invention relates to the field of optimization technology, and in particular to a data allocation optimization method and device based on calculation graph and recalculation. Background Art

[0002] Currently in the field of deep learning, the scale of models continues to grow, far exceeding the storage capacity of a single card, and parallel training is needed. Existing parallel schemes include data parallelism, tensor parallelism, and pipeline parallelism. In order to efficiently train large models, a hybrid parallel scheme combining the above three parallel schemes is often used. Among them, pipeline parallelism is widely used in large model training because it can split model parameters and introduces less communication. However, the pipeline parallel strategy requires a large amount of video memory, so it is often used together with the recalculation mechanism. The recalculation mechanism reduces storage overhead by discarding the intermediate results of the forward process, but the discarded intermediate results need to be recalculated in the reverse process, which introduces additional computing overhead.

[0003] In terms of pipeline parallelism, existing work mainly includes two parts: scheduling strategy and computational graph segmentation. In terms of scheduling strategy, the scheduling strategy proposed by GPipe can reduce the proportion of bubbles during training. PipeDream and DAPPLE further proposed the 1F1B strategy to reduce the overhead of video memory during training. Interleaved 1F1B divides the computational graph into more stages to reduce the number of bubbles, but will introduce greater video memory overhead. In the computational graph segmentation part, the main optimization goal of the current PipeDream and Alpa works is to reduce the longest computing time in all stages. However, all of the above works have the problems of uneven video memory utilization and large video memory requirements, and the optimization goals of the computational graph segmentation part are too simple, leaving a lot of room for optimization.

[0004] In terms of recalculation mechanism, existing frameworks such as Megatron-LM support recalculation in layers, and the latest selective recalculation technology only recalculates several operators with large intermediate results in the decoding layer. These works use the same recalculation strategy for different stages of pipeline parallelism and different layers of the model, and cannot dynamically adjust their respective recalculation strategies according to the video memory capacity of the accelerator card itself, thus failing to fully utilize the video memory to reduce recalculation overhead.

[0005] Existing technology 1 includes pipeline parallel technology: Currently, 1F1B is the mainstream pipeline parallel scheduling algorithm, supported by multiple frameworks such as Megatron-LM and Alpa. Among them, frameworks such as Megatron-LM manually configure the segmentation strategy of the calculation graph, while frameworks such as Alpa automatically find the best segmentation strategy with the goal of optimizing the longest calculation time of all pipeline stages.

[0006] For frameworks that require manual configuration, such as Megatron-LM, the splitting strategy is limited by the framework restrictions that require the computational graph to be evenly split, and therefore it is often impossible to find the optimal strategy. For automatic parallel frameworks such as Alpa, their optimization goals can only solve the situation where the number of samples far exceeds the number of pipeline parallel stages, and the effect is not good in other cases. In addition, automatic parallel work such as Alpa does not support the recalculation mechanism, and cannot uniformly solve the problem of insufficient and unbalanced video memory utilization.

[0007] Existing technology 2 includes recalculation mechanism: the recalculation mechanism is widely used in large model training. For large models with Transformer structure, there are currently two mainstream recalculation mechanisms. One is recalculation by layer, which only retains the input tensor of each Transformer layer. The other is selective recalculation, which only recalculates the Softmax operator and some adjacent operators in the Transformer layer, and does not recalculate other operators.

[0008] The layer-based recalculation mechanism cannot fully utilize the storage resources of the graphics card and introduces a large amount of recalculation overhead. Although the selective recalculation mechanism reduces the recalculation overhead, it still needs to save a large number of intermediate results, which occupies a large amount of video memory and is not applicable in many scenarios. In addition, the above strategies are all static, that is, different layers and different pipeline stages use the same recalculation strategy, which cannot solve the problem of video memory imbalance. Summary of the invention

[0009] In response to the problems in the prior art, an embodiment of the present invention provides a data allocation optimization method and device based on a calculation graph and recalculation, which can at least partially solve the problems existing in the prior art.

[0010] On the one hand, the present invention proposes a data allocation optimization method based on a computation graph and recomputation, comprising:

[0011] Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer;

[0012] According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit;

[0013] According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

[0014] The method of determining the optimal recalculation strategy corresponding to each computing unit according to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and saving the most computation time for all computing units in each subgraph as the first goal, includes:

[0015] A first dynamic programming algorithm based on the knapsack problem is used, with the first goal as the dynamic programming goal, and the storage overhead, the computing overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node, the number of computing units corresponding to each subgraph and the upper limit of the video memory of the intermediate result as algorithm inputs, to obtain the optimal recalculation strategy corresponding to each computing unit.

[0016] Wherein, determining the upper bound of the video memory of the intermediate result includes:

[0017] Obtaining first video memory overhead information related to the parameter quantity in the subgraph corresponding to the current pipeline node, and second video memory overhead information related to the layer-by-layer recalculation strategy;

[0018] The subtraction result of the total video memory capacity and the first video memory overhead information, and the subtraction result of the total video memory capacity and the second video memory overhead information is used as the video memory upper bound of the intermediate result.

[0019] The parameter quantity includes model parameters, model gradients, and optimizer states; accordingly, obtaining first video memory overhead information related to the parameter quantity in the subgraph of the current pipeline node includes:

[0020] Calculate the model parameter memory overhead corresponding to the model parameter according to the model parameter quantity and the size of tensor parallelism;

[0021] Calculate the model gradient memory overhead corresponding to the model gradient according to the model gradient parameter quantity and the size of tensor parallelism;

[0022] Calculate the optimizer state memory overhead corresponding to the optimizer state according to the optimizer parallel size, optimizer type parameter, optimizer state parameter quantity and tensor parallel size;

[0023] The sum of the model parameter memory overhead, the model gradient memory overhead and the optimizer state memory overhead is used as the first memory overhead information.

[0024] The step of obtaining the second video memory overhead information related to the layer-by-layer recalculation strategy includes:

[0025] The cache size required by the layer recalculation strategy is used as the second video memory overhead information.

[0026] The method of determining the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal, includes:

[0027] A second dynamic programming algorithm based on the 1F1B scheduling strategy is used, with the starting phase time, cooling phase time, the sum of the longest forward time and reverse time from the current pipeline node to the last pipeline node, the forward time and reverse time of the current pipeline node, and the total calculation time starting from the current pipeline node and the second target as the dynamic programming target, and the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph in the stable stage as the algorithm input to obtain the number of decoding layers contained in the subgraph of each pipeline node.

[0028] The step of determining the sum of the longest forward time and the longest reverse time from the current flow node to the last flow node includes:

[0029] From back to front, use the following formula to obtain the sum of the longest forward time and reverse time M from the current flow node to the last flow node. i :

[0030] M i =max(M i+1 ,F i +B i )

[0031] Among them, B i is the forward time of the current flow node, F i It is the backward time of the current flow node.

[0032] On the one hand, the present invention provides a data allocation optimization device based on a computation graph and recomputation, comprising:

[0033] An acquisition unit, used to acquire storage overhead and computation overhead of computation units corresponding to each subgraph in a computation graph, wherein the computation graph includes a decoding layer;

[0034] a determination unit, configured to determine the optimal recalculation strategy corresponding to each computing unit according to the storage overhead and the computation overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and the first goal of saving the most computation time for all computing units in each subgraph;

[0035] The optimization unit is used to determine the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal.

[0036] In another aspect, an embodiment of the present invention provides a computer device, including a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein the processor implements the following method when executing the computer program:

[0037] Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer;

[0038] According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit;

[0039] According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

[0040] An embodiment of the present invention provides a computer-readable storage medium, including:

[0041] The computer readable storage medium stores a computer program, and when the computer program is executed by a processor, the following method is implemented:

[0042] Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer;

[0043] According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit;

[0044] According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

[0045] The data allocation optimization method and device based on calculation graph and recalculation provided by the embodiment of the present invention obtain the storage overhead and calculation overhead of the calculation units corresponding to each subgraph in the calculation graph; wherein the calculation graph includes a decoding layer; according to the storage overhead and the calculation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of calculation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and all the calculation units in each subgraph save the most calculation time as the first goal, determine the optimal recalculation strategy corresponding to each calculation unit; according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline node corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline node corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal, which can optimize data allocation and thus improve the overall calculation efficiency of the model. BRIEF DESCRIPTION OF THE DRAWINGS

[0046] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the following briefly introduces the drawings required for use in the embodiments or the prior art descriptions. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work. In the drawings:

[0047] Figure 1 It is a flowchart of a data allocation optimization method based on a calculation graph and recalculation provided by an embodiment of the present invention.

[0048] Figure 2 It is a schematic diagram illustrating a 1F1B scheduling strategy provided by an embodiment of the present invention.

[0049] Figure 3 It is a schematic diagram for explaining the implementation based on the second dynamic programming algorithm provided by an embodiment of the present invention.

[0050] Figure 4It is a structural diagram of a data allocation optimization device based on calculation graph and recalculation provided by an embodiment of the present invention.

[0051] Figure 5 A schematic diagram of the physical structure of a computer device provided in an embodiment of the present invention. DETAILED DESCRIPTION

[0052] In order to make the purpose, technical scheme and advantages of the embodiments of the present invention more clear, the embodiments of the present invention are further described in detail below in conjunction with the accompanying drawings. Here, the illustrative embodiments of the present invention and their descriptions are used to explain the present invention, but are not intended to limit the present invention. It should be noted that, in the absence of conflict, the embodiments in this application and the features in the embodiments can be combined with each other arbitrarily.

[0053] The professional terms of the embodiments of the present invention are described as follows:

[0054] Pipeline Parallelism: The computation graph is divided into multiple stages, each of which is processed by a different device.

[0055] Recomputation Strategy: It is used to describe which operators or computing units in the computation graph need to be recomputed.

[0056] Graph Partitioning Strategy: It is used to describe the strategy of dividing the computational graph into subgraphs of different lengths.

[0057] 1F1B scheduling strategy (one-forward-one-backward scheduling mechanism): a pipeline scheduling strategy that performs forward and reverse interleaving execution in the stable phase.

[0058] Pipeline stage: Pipeline parallelism divides the computational graph into multiple serially connected subgraphs, each of which is a pipeline stage.

[0059] Micro-batch: In pipeline parallelism, each batch of samples is evenly divided into multiple parts, each of which is a micro-sample. After all micro-samples in each batch are calculated, the parameters are updated uniformly.

[0060] Figure 1 is a flow chart of a data allocation optimization method based on a calculation graph and recalculation provided by an embodiment of the present invention, such as Figure 1 As shown, the data allocation optimization method based on calculation graph and recalculation provided by the embodiment of the present invention includes:

[0061] Step S1: Obtain storage overhead and computational overhead of a computational unit corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer.

[0062] Step S2: According to the storage overhead and the computing overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and all the computing units in each subgraph saving the most computing time as the first goal, determine the optimal recalculation strategy corresponding to each computing unit.

[0063] Step S3: According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each sub-graph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each sub-graph, the number of decoding layers contained in the sub-graph corresponding to each pipeline node is determined with the second goal of minimizing the total training time.

[0064] In the above step S1, the device obtains the storage overhead and computational overhead of the computing units corresponding to each subgraph in the computational graph; wherein the computational graph includes a decoding layer. The device may be a computer device that executes the method, such as a server. The acquisition, storage, use, and processing of data in the technical solution of this application are in compliance with relevant regulations. The method of the embodiment of the present invention can be applied to the field of deep learning models, specifically including the field of pre-trained language models.

[0065] The pre-trained language model mainly includes multiple isomorphic decoding layers, each of which can be divided into an attention layer and a feedforward layer. According to the structural characteristics of these two layers, they can be further divided into more fine-grained computing units, where operators such as addition and transposition do not need to save intermediate results, so these operators can be omitted in the diagram, but each computing unit may actually include multiple such operators. All operators in each computing unit are either recalculated or not recalculated.

[0066] First, the storage cost and computation cost of each computing unit are constructed through performance analysis. For computing unit U, use T f (U), T b (U) represents the forward calculation time of the computing unit and the reverse calculation time of the computing unit respectively, and Mem(U) represents the storage overhead.

[0067] For a subgraph of a computational graph, it can be regarded as a collection of a series of computational units, which can be written as The recalculation strategy for this subgraph can be expressed as Indicates which computational units in the subgraph can be recomputed.

[0068] Then we can get the subgraph in the recalculation strategy as Calculation time

[0069]

[0070] In the above step S2, the device determines the optimal recalculation strategy corresponding to each computing unit according to the storage cost and the computational cost, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and the first goal of saving the most computational time for all computing units in each subgraph. Determining the upper limit of the video memory of the intermediate result includes:

[0071] Obtain first memory overhead information related to the parameter quantity in the subgraph corresponding to the current pipeline node, and second memory overhead information related to the layer-by-layer recalculation strategy; the parameter quantity includes model parameters, model gradients and optimizer states, and the model parameter quantity corresponding to the model parameters is denoted by N; one parameter corresponds to one model gradient, and one parameter corresponds to one optimizer state.

[0072] Obtaining first video memory overhead information related to the parameter quantity in the subgraph of the current pipeline section, including:

[0073] According to the model parameter quantity and the size of tensor parallelism, the model parameter memory overhead corresponding to the model parameter is calculated; the model parameter memory overhead is calculated according to the following expression:

[0074] t represents the size of tensor parallelism.

[0075] According to the model gradient parameter quantity and the size of tensor parallelism, the model gradient memory overhead corresponding to the model gradient is calculated; the model gradient memory overhead is calculated according to the following expression:

[0076] t represents the size of tensor parallelism.

[0077] According to the size of data parallelism, optimizer type parameter, optimizer state parameter quantity and tensor parallelism, calculate the optimizer state memory overhead corresponding to the optimizer state; calculate the optimizer state memory overhead according to the following expression:

[0078] t represents the size of tensor parallelism, d represents the size of optimizer parallelism. The optimizer type parameter k is related to the type of optimizer. For the current mainstream Adam optimizer, k is 8.

[0079] The sum of the model parameter memory overhead, the model gradient memory overhead and the optimizer state memory overhead is used as the first memory overhead information.

[0080] Get the secondary memory overhead information related to the layer recalculation strategy, including:

[0081] The cache size required by the layer recalculation strategy is used as the second video memory overhead information. It is mainly to save the intermediate results obtained by recalculation during the reverse calculation process. The specific size of the cache is directly related to the recalculation strategy. For the traditional layer recalculation strategy, the cache size is the sum of the storage of all intermediate results of a single decoding layer. For simplicity, let the operator that is not recalculated be a superset of the layer recalculation, so the required cache must not be larger than the cache required for layer recalculation, and the cache size of this part can be set to the cache size required by the layer recalculation strategy.

[0082] It should be noted that when recalculating by layer, only the input tensor of each layer needs to be retained. When recalculating in reverse, other tensors in each layer need to be calculated and saved in the cache during the reverse process, and then discarded after use. In the recalculation strategy, the tensors that are not recalculated must include the input tensors of each layer, so the sum of other tensors that need to be retained must be less than the recalculation by layer.

[0083] The subtraction result of the total video memory capacity and the first video memory overhead information, and the subtraction result of the second video memory overhead information is used as the video memory upper limit of the intermediate result. That is:

[0084] The upper limit M of the intermediate result's video memory = total video memory capacity - first video memory overhead information - second video memory overhead information.

[0085] The memory occupied by the intermediate results is related to the recalculation strategy. It is related to the current sequence number s of the pipeline node where the sub-graph is located. This part of the video memory can be expressed as Where p represents the size of the pipeline parallelism.

[0086] The optimal recalculation strategy corresponding to each computing unit is determined according to the storage overhead and the computation overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and the first goal of saving the most computation time for all computing units in each subgraph, including:

[0087] A first dynamic programming algorithm based on the knapsack problem is used, with the first goal as the dynamic programming goal, and the storage overhead, the computing overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node, the number of computing units corresponding to each subgraph and the upper limit of the video memory of the intermediate result as algorithm inputs, to obtain the optimal recalculation strategy corresponding to each computing unit.

[0088] The Knapsack problem is an NP-complete problem of combinatorial optimization. The problem can be described as: given a set of items, each with its own weight and price, how to choose items within a limited total weight to maximize the total price of the items.

[0089] Each computing unit in the embodiment of the present invention corresponds to each item, the video memory overhead occupied by the computing unit corresponds to the weight of the item, the computing overhead of the computing unit corresponds to the price of the item, M corresponds to the specified total weight, and the maximum computing time saved by all computing units corresponds to the highest total price of the item.

[0090] Hypothetical calculation Figure 1 There are NX computing units in total, and the upper limit of the video memory for the intermediate results is M. i (m) represents the maximum computing time that can be saved for the first i computing units when the video memory limit is m. Then the following state transition equation can be listed:

[0091] S i (m) = max(S i-1 (m),S i-1 (m-(ps)×Mem(U i ))+T f (U i ))

[0092] The parameters can refer to the above description to find the optimal recalculation strategy S NX (M).

[0093] In the above step S3, the device determines the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and the reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal. The method of determining the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and the reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal, includes:

[0094] A second dynamic programming algorithm based on the 1F1B scheduling strategy is used, with the starting phase time, cooling phase time, the sum of the longest forward time and reverse time from the current pipeline node to the last pipeline node, the forward time and reverse time of the current pipeline node, and the total calculation time starting from the current pipeline node and the second target as the dynamic programming target, and the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph in the stable stage as the algorithm input to obtain the number of decoding layers contained in the subgraph of each pipeline node.

[0095] like Figure 2 As shown, the 1F1B scheduling strategy includes three stages, namely the starting stage, the stable stage and the cooling stage.

[0096] The initial stage starts from the forward direction of sample No. 0 and ends before the reverse direction of sample No. 0.

[0097] The stable phase starts from the reverse direction of sample No. 0 until the forward direction of the last sample is completed.

[0098] The remaining part is the cooling phase. i , S i and E i They represent the starting phase time, stable phase time and cooling phase time of the i-th flow node respectively.

[0099] For W i , we can get the following recursive relationship from back to front:

[0100] W i-1 =Max(W i +B i ,(pi)F i-1 )+F i-1

[0101] Among them, p represents the size of the pipeline parallelism, B i represents the forward time of the ith flow node, F i Represents the reverse time of the i-th flow node. In this way, the start time of all flow nodes can be obtained from the back to the front. The time of the cooling stage is similar to that of the starting stage, and can be expressed by the following recursive relationship:

[0102] E i-1 =Max(E i +F i ,(pi)B i-1 )+B i-1

[0103] For S i, let M i It represents the sum of the longest forward time and reverse time from the i-th flow node to the last flow node. Similarly, the stable phase time of each flow node can be obtained from the back to the front using the following formula:

[0104] M i =max(M i+1 ,F i +B i )

[0105] S i =(n-p+i)×M i

[0106] Where n represents the number of samples. The samples in the embodiment of the present invention are micro-samples, i.e., micro-batches. i The following instructions are given:

[0107] For M p-1 , which is the sum of the longest forward time and the longest reverse time in the last flow node.

[0108] For M p-2 , which is the above M p-1 The larger one is the sum of the longest forward time and the longest reverse time in the last flow node and the sum of the longest forward time and the longest reverse time in the second to last flow node.

[0109] For M p-3 , which is the above M p-2 The larger one of the longest forward time and the longest reverse time in the third-to-last flow node, and so on.

[0110] The starting time is recorded as W, the cooling time is recorded as E, the sum of the longest forward time and reverse time from the current flow node to the last flow node is recorded as M, the forward time of the current flow node is recorded as F, the reverse time of the current flow node is recorded as B, and the total calculation time starting from the current flow node is recorded as T, which can be obtained by Figure 3 The program code of the second dynamic programming algorithm shown is used to output the number of decoding layers contained in the subgraph of each pipeline node.

[0111] Description of the beneficial technical effect of the second dynamic programming algorithm: The algorithm automatically adjusts the segmentation strategy of the calculation graph according to the number of micro-batches to achieve the optimal total training time. As the number of micro-batches increases, the algorithm will gradually tend to the following segmentation method: the decoding layers of the first few pipeline sections are fewer, and the decoding layers of the last few pipeline sections are more. Since the recalculation of the first few pipeline sections is less and the recalculation of the last few pipeline sections is more, the calculation time of different pipeline sections is more balanced, thereby reducing the overall time.

[0112] First, according to the designed adaptive recalculation algorithm, the computational graph is divided into finer-grained units, supporting different recalculation strategies for different computational units in different pipeline nodes, which can fully utilize the device's video memory to reduce the amount of recalculation. In addition, the adaptive computational graph splitting algorithm designed according to the pipeline parallel 1F1B scheduling strategy is used to balance the computational load of different pipeline nodes to achieve the optimal overall computing efficiency.

[0113] The data allocation optimization method based on computational graph and recomputation provided by the embodiment of the present invention obtains the storage overhead and computational overhead of the computational units corresponding to each subgraph in the computational graph; wherein the computational graph includes a decoding layer; according to the storage overhead and the computational overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computational units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and all computational units in each subgraph saving the most computational time as the first goal, the optimal recomputation strategy corresponding to each computational unit is determined; according to the total number of pipeline nodes, the number of micro-samples, the optimal recomputation strategy in the pipeline nodes corresponding to each subgraph, and the forward computation time and reverse computation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined, which can optimize data allocation and thus improve the overall computational efficiency of the model.

[0114] Further, the optimal recalculation strategy corresponding to each computing unit is determined based on the storage overhead and the computation overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and all computing units in each subgraph saving the most computing time as the first goal, including:

[0115] By using the first dynamic programming algorithm based on the knapsack problem, taking the first target as the dynamic programming target, taking the storage cost, the computational cost, the size of the pipeline parallelism, the current sequence number of the pipeline node, the number of computational units corresponding to each subgraph, and the upper bound of the video memory of the intermediate result as the algorithm input, the optimal recomputation strategy corresponding to each computational unit is obtained. The above-mentioned embodiment can be referred to for explanation and will not be repeated here.

[0116] Further, determining the upper bound of the video memory of the intermediate result includes:

[0117] Obtain first video memory overhead information related to the parameter quantity in the subgraph corresponding to the current pipeline node, and second video memory overhead information related to the layer-by-layer recalculation strategy; refer to the above embodiment for description, which will not be repeated here.

[0118] The subtraction result of the total video memory capacity and the first video memory overhead information, and the subtraction result of the total video memory capacity and the second video memory overhead information is used as the video memory upper limit of the intermediate result. The above embodiment can be referred to for explanation, and no further description is given.

[0119] Further, the parameter quantity includes model parameters, model gradients and optimizer states; accordingly, obtaining first video memory overhead information related to the parameter quantity in the subgraph of the current pipeline node includes:

[0120] According to the model parameter quantity and the size of tensor parallelism, the model parameter memory overhead corresponding to the model parameter is calculated; the above-mentioned embodiment can be referred to for description and will not be repeated here.

[0121] According to the model gradient parameter quantity and the size of tensor parallelism, the model gradient memory overhead corresponding to the model gradient is calculated; the description may be made with reference to the above embodiment and will not be repeated here.

[0122] According to the size of optimizer parallelism, optimizer type parameters, optimizer state parameter quantity and tensor parallelism, the optimizer state memory overhead corresponding to the optimizer state is calculated; the above-mentioned embodiment can be referred to for description and will not be repeated here.

[0123] The sum of the model parameter memory overhead, the model gradient memory overhead and the optimizer state memory overhead is used as the first memory overhead information. The above embodiment can be referred to for explanation and will not be described in detail.

[0124] Furthermore, obtaining the second video memory overhead information related to the layer-by-layer recalculation strategy includes:

[0125] The cache size required by the layer recalculation strategy is used as the second video memory overhead information. The above embodiment can be referred to for description, which will not be repeated here.

[0126] Furthermore, the method determines the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal, including:

[0127] The second dynamic programming algorithm based on the 1F1B scheduling strategy is used, with the starting phase time, cooling phase time, the sum of the longest forward time and reverse time from the current flow node to the last flow node, the forward time and reverse time of the current flow node, and the total calculation time starting from the current flow node and the second target as the dynamic programming target, and the total number of flow nodes, the number of micro-samples, the optimal recalculation strategy in the flow nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the flow nodes corresponding to each subgraph in the stable stage as the algorithm input, to obtain the number of decoding layers contained in the subgraph of each flow node. The above-mentioned embodiment can be referred to for description and will not be repeated.

[0128] Further, determining the sum of the longest forward time and the longest reverse time from the current flow node to the last flow node includes:

[0129] From back to front, use the following formula to obtain the sum of the longest forward time and reverse time M from the current flow node to the last flow node. i :

[0130] M i =max(M i+1 ,F i +B i )

[0131] Among them, B i is the forward time of the current flow node, F i The backward time of the current flow node can be referred to the above embodiment, which will not be described in detail.

[0132] Figure 4 is a schematic diagram of the structure of a data allocation optimization device based on a calculation graph and recalculation provided by an embodiment of the present invention, such as Figure 4 As shown, the data allocation optimization device based on calculation graph and recalculation provided by the embodiment of the present invention includes an acquisition unit 401, a determination unit 402 and an optimization unit 403, wherein:

[0133] The acquisition unit 401 is used to obtain the storage overhead and computational overhead of the computational units corresponding to each subgraph in the computational graph; wherein the computational graph includes a decoding layer; the determination unit 402 is used to determine the optimal recalculation strategy corresponding to each computational unit based on the storage overhead and the computational overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computational units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and all computational units in each subgraph saving the most computational time as the first goal; the optimization unit 403 is used to determine the number of decoding layers contained in the subgraph corresponding to each pipeline node based on the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward computation time and reverse computation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal.

[0134] Specifically, the acquisition unit 401 in the device is used to obtain the storage overhead and the computational overhead of the computing units corresponding to each subgraph in the computational graph; wherein the computational graph includes a decoding layer; the determination unit 402 is used to determine the optimal recalculation strategy corresponding to each computing unit based on the storage overhead and the computational overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and all computing units in each subgraph saving the most computing time as the first goal; the optimization unit 403 is used to determine the number of decoding layers contained in the subgraph corresponding to each pipeline node based on the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal.

[0135] The data allocation optimization device based on calculation graph and recalculation provided by the embodiment of the present invention obtains the storage overhead and calculation overhead of the calculation units corresponding to each subgraph in the calculation graph; wherein the calculation graph includes a decoding layer; according to the storage overhead and the calculation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of calculation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and all the calculation units in each subgraph save the most calculation time as the first goal, the optimal recalculation strategy corresponding to each calculation unit is determined; according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline node corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline node corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal, which can optimize data allocation and thus improve the overall calculation efficiency of the model.

[0136] An embodiment of the present invention provides an embodiment of a data allocation optimization device based on a computation graph and recalculation, which can be specifically used to execute the processing flow of the above-mentioned method embodiments. Its functions are not repeated here, and reference can be made to the detailed description of the above-mentioned method embodiments.

[0137] Figure 5 A schematic diagram of the physical structure of a computer device provided in an embodiment of the present invention is shown in FIG. Figure 5 As shown, the computer device includes: a memory 501, a processor 502, and a computer program stored in the memory 501 and executable on the processor 502. When the processor 502 executes the computer program, the following method is implemented:

[0138] Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer;

[0139] According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit;

[0140] According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

[0141] This embodiment discloses a computer program product, the computer program product including a computer program, and when the computer program is executed by a processor, the following method is implemented:

[0142] Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer;

[0143] According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit;

[0144] According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

[0145] This embodiment provides a computer-readable storage medium, wherein the computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the following method is implemented:

[0146] Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer;

[0147] According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit;

[0148] According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

[0149] Compared with the technical solutions in the prior art, the data allocation optimization method based on calculation graph and recalculation provided by the embodiment of the present invention obtains the storage overhead and the calculation overhead of the calculation units corresponding to each subgraph in the calculation graph; wherein the calculation graph includes a decoding layer; according to the storage overhead and the calculation overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of calculation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and all the calculation units in each subgraph save the most calculation time as the first goal, the optimal recalculation strategy corresponding to each calculation unit is determined; according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline node corresponding to each subgraph, and the forward calculation time and the reverse calculation time of the pipeline node corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal, which can optimize data allocation and thus improve the overall calculation efficiency of the model.

[0150] Those skilled in the art will appreciate that embodiments of the present invention may be provided as methods, systems, or computer program products. Therefore, the present invention may take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Moreover, the present invention may take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program code.

[0151] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowchart and / or block diagram, as well as the combination of processes and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the flowchart and / or block diagram. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0152] These computer program instructions may also be stored in a computer-readable memory capable of directing a computer or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer-readable memory produce an article of manufacture comprising an instruction device, which implements the process Figure 1 A process or multiple processes and / or boxes Figure 1 A function specified in one or more boxes.

[0153] These computer program instructions can also be loaded onto a computer or other programmable data processing device so that a series of operating steps are executed on the computer or other programmable device to produce a computer-implemented process, thereby providing instructions for implementing the process. Figure 1 A process or multiple processes and / or boxes Figure 1 A step that specifies a function in one or more boxes.

[0154] In the description of this specification, the description with reference to the terms "one embodiment", "a specific embodiment", "some embodiments", "for example", "example", "specific example", or "some examples" means that the specific features, structures, materials or characteristics described in conjunction with the embodiment or example are included in at least one embodiment or example of the present invention. In this specification, the schematic representation of the above terms does not necessarily refer to the same embodiment or example. Moreover, the specific features, structures, materials or characteristics described may be combined in any one or more embodiments or examples in a suitable manner.

[0155] The specific embodiments described above further illustrate the objectives, technical solutions and beneficial effects of the present invention in detail. It should be understood that the above description is only a specific embodiment of the present invention and is not intended to limit the scope of protection of the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.

Claims

1. A data allocation optimization method based on computation graph and recomputation, characterized in that: include: Obtaining storage overhead and computational overhead of computational units corresponding to each subgraph in a computational graph, wherein the computational graph includes a decoding layer; According to the storage overhead and the computation overhead, the size of pipeline parallelism, the current sequence number of the pipeline node and the number of computation units corresponding to each subgraph, the upper limit of the video memory of the intermediate result is taken as the upper limit, and the first goal is to save the most computation time for all computation units in each subgraph, determine the optimal recomputation strategy corresponding to each computation unit; According to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, the number of decoding layers contained in the subgraph corresponding to each pipeline node is determined with the shortest total training time as the second goal.

2. The data allocation optimization method based on calculation graph and recalculation according to claim 1 is characterized in that: The optimal recalculation strategy corresponding to each computing unit is determined according to the storage overhead and the computation overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and the first goal of saving the most computation time for all computing units in each subgraph, including: A first dynamic programming algorithm based on the knapsack problem is used, with the first goal as the dynamic programming goal, and the storage overhead, the computing overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node, the number of computing units corresponding to each subgraph and the upper limit of the video memory of the intermediate result as algorithm inputs, to obtain the optimal recalculation strategy corresponding to each computing unit.

3. The data allocation optimization method based on calculation graph and recalculation according to claim 2 is characterized in that: Determining the upper bound of the video memory of the intermediate result includes: Obtaining first video memory overhead information related to the parameter quantity in the subgraph corresponding to the current pipeline node, and second video memory overhead information related to the layer-by-layer recalculation strategy; The subtraction result of the total video memory capacity and the first video memory overhead information, and the subtraction result of the total video memory capacity and the second video memory overhead information is used as the video memory upper bound of the intermediate result.

4. The data allocation optimization method based on calculation graph and recalculation according to claim 3 is characterized in that: The parameter quantities include model parameters, model gradients, and optimizer states; accordingly, obtaining first video memory overhead information related to the parameter quantities in the subgraph of the current pipeline node includes: Calculate the model parameter memory overhead corresponding to the model parameter according to the model parameter quantity and the size of tensor parallelism; Calculate the model gradient memory overhead corresponding to the model gradient according to the model gradient parameter quantity and the size of tensor parallelism; Calculate the optimizer state memory overhead corresponding to the optimizer state according to the optimizer parallel size, optimizer type parameter, optimizer state parameter quantity and tensor parallel size; The sum of the model parameter memory overhead, the model gradient memory overhead and the optimizer state memory overhead is used as the first memory overhead information.

5. The data allocation optimization method based on calculation graph and recalculation according to claim 3 is characterized in that: Get the secondary memory overhead information related to the layer recalculation strategy, including: The cache size required by the layer recalculation strategy is used as the second video memory overhead information.

6. The data allocation optimization method based on calculation graph and recalculation according to any one of claims 1 to 5, characterized in that: The method determines the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal, including: A second dynamic programming algorithm based on the 1F1B scheduling strategy is used, with the starting phase time, cooling phase time, the sum of the longest forward time and reverse time from the current pipeline node to the last pipeline node, the forward time and reverse time of the current pipeline node, and the total calculation time starting from the current pipeline node and the second target as the dynamic programming target, and the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph in the stable stage as the algorithm input to obtain the number of decoding layers contained in the subgraph of each pipeline node.

7. The data allocation optimization method based on calculation graph and recalculation according to claim 6 is characterized in that: Determining the sum of the longest forward time and the longest reverse time from the current flow node to the last flow node, including: From back to front, use the following formula to obtain the sum of the longest forward time and reverse time M from the current flow node to the last flow node. i : M i =max(M i+1 ,F i +B i ) Among them, B i is the forward time of the current flow node, F i It is the backward time of the current flow node.

8. A data allocation optimization device based on calculation graph and recalculation, characterized in that: include: An acquisition unit, used to acquire storage overhead and computation overhead of computation units corresponding to each subgraph in a computation graph, wherein the computation graph includes a decoding layer; a determination unit, configured to determine the optimal recalculation strategy corresponding to each computing unit according to the storage overhead and the computation overhead, the size of the pipeline parallelism, the current sequence number of the pipeline node and the number of computing units corresponding to each subgraph, with the upper limit of the video memory of the intermediate result as the upper limit and the first goal of saving the most computation time for all computing units in each subgraph; The optimization unit is used to determine the number of decoding layers contained in the subgraph corresponding to each pipeline node according to the total number of pipeline nodes, the number of micro-samples, the optimal recalculation strategy in the pipeline nodes corresponding to each subgraph, and the forward calculation time and reverse calculation time of the pipeline nodes corresponding to each subgraph, with the shortest total training time as the second goal.

9. A computer device comprising a memory, a processor and a computer program stored in the memory and executable on the processor, characterized in that: When the processor executes the computer program, the method according to any one of claims 1 to 7 is implemented.

10. A computer-readable storage medium, characterized in that: The computer-readable storage medium stores a computer program, and when the computer program is executed by a processor, the method according to any one of claims 1 to 7 is implemented.

Images