A parallel strategy search method for efficient training of large artificial intelligence models

Abstract

The present invention discloses a parallel strategy search method for efficient training of large artificial intelligence models. The method comprises the following steps: inputting a large model; obtaining system execution performance information and the large model's execution performance information; using a cost model to estimate the execution time and storage overhead of the large model on the current system; establishing and solving a mixed integer quadratic programming mathematical model for automatically searching for parallel training strategies for the large model, taking the time spent on each iteration of pipeline parallel training as the optimization target; and outputting the optimal parallel training strategy for the large model. The method supports pipeline parallelism, data parallelism, tensor parallelism, and optimizer parallelism, has a large strategy space, and shortens search time. It can be applied to single-machine multi-GPU parallel training of large models and distributed training of multi-machine clusters, thereby improving the training efficiency of large models.

Description

Technical Field

[0001] The present invention relates to artificial intelligence technology, and in particular to a parallel strategy search method for efficient training of artificial intelligence large models (hereinafter referred to as large models). Background Art

[0002] In recent years, the number of parameters in large AI models has increased rapidly, achieving state-of-the-art results across a wide range of tasks in fields like computer vision and natural language processing. However, computing hardware development has been very slow, necessitating the use of multiple GPUs on a single machine or in clusters for parallel training.

[0003] Currently, the main parallel training modes for large models include pipeline parallelism, data parallelism, tensor parallelism, and optimizer parallelism. Pipeline parallelism primarily partitions the layers of a large model and places several adjacent layers on the same pipeline stage. During computation, pipeline parallelism divides a mini-batch into multiple micro-batches and organizes these micro-batches into pipelines, thereby improving overall training efficiency. Data parallelism, tensor parallelism, and optimizer parallelism perform parallel computations within each layer of a large model, improving intra-layer computation efficiency. Data parallelism requires copying the entire large model and optimizer state across different nodes, resulting in lower communication traffic but higher memory overhead. Tensor parallelism and optimizer parallelism, respectively, partition the tensors and optimizer state of a large model. While this reduces memory overhead, the communication traffic is higher than that of data parallelism. It is important to note that each of these parallel computing strategies has its own advantages and disadvantages for different models and computing environments. Therefore, combining these parallel computing strategies to maximize training efficiency for specific models and given computing environments has become an important research direction in parallel training of large models.

[0004] However, due to the large number of parallel strategies, manual combination and evaluation are expensive. At the same time, for different models and computing environments, selecting the appropriate combination of parallel strategies requires expert-level knowledge, which is difficult to apply in real-world scenarios. Therefore, some researchers have proposed methods for automatically searching for parallel training strategies for large models. However, to date, existing methods either focus on only one or several of the four parallel strategies mentioned above, or consider the four parallel strategies in a layered manner and solve them layer by layer. The strategy space of these methods is limited, and they are prone to obtaining locally optimal parallel training strategies rather than globally optimal parallel training strategies. Because training a large model with a large parameter scale often takes weeks or even months and consumes a lot of energy, how to automatically search for a globally optimal parallel training strategy to maximize the acceleration of large model training has become an important issue. Summary of the Invention

[0005] Purpose of the invention: The purpose of the present invention is to provide a parallel strategy search method for efficient training of large artificial intelligence models, thereby providing a parallel strategy search method for efficient training of large models based on mixed integer quadratic programming, so as to solve the local optimality problem of the parallel strategies searched by the prior art, thereby further improving the training efficiency of large models and reducing energy consumption.

[0006] Technical solution: The present invention provides a parallel strategy search method for efficient training of large artificial intelligence models, comprising the following steps:

[0007] Step S1, input the large model;

[0008] Step S2: Obtaining system execution performance information and large model execution performance information;

[0009] Set the All-Reduce communication volume a ar , statistics of All-Reduce communication time t ar , thereby obtaining the system's All-Reduce communication efficiency e ar ;

[0010] Set the P2P communication volume a p2p , statistics of P2P communication time t p2p , thereby obtaining the system's P2P communication efficiency e p2p ;

[0011] Set the All-Reduce communication volume a ar , respectively count the All-Reduce communication time t of executing All-Reduce separately n , and the All-Reduce communication time t when performing All-Reduce and matrix multiplication calculations simultaneously o , thereby obtaining the system's computation-communication overlap factor e o ;

[0012] Statistical model G (V, E) each layer u (u ∈ V) forward execution time fp u ;

[0013] Statistics of the activation memory of each layer of the large model G(V,E) in addition to the memory occupied by the model and optimizer And the additional video memory m required to load the large model G(V,E) other .

[0014] The execution time overhead described in step S3 includes the sum of the calculation and communication time of all candidate strategies in each layer of the large model, the communication time of all candidate strategy pairs in adjacent layers in the same pipeline stage, and one or more of the P2P communication time across pipeline stages.

[0015] Step S3: using the cost model to estimate the execution time and storage overhead of the large model on the current system;

[0016] The execution time overhead described in step S3 specifically includes the following steps:

[0017] Use the cost model to estimate the sum of the computation and communication time A of all candidate data parallel, tensor parallel, and optimizer parallel strategies for layer u (u∈V) of the large model G(V,E) u ;

[0018] Use the cost model to estimate the edge <u,v>(<u,v> The communication time R in the same pipeline stage between all candidate data parallel, tensor parallel, and optimizer parallel strategy pairs corresponding to layer u and layer v ∈ E) uv ;

[0019] Use the cost model to estimate the edge of the large model G(V,E)<u,v> The cross-pipeline stage P2P communication time R′ between all candidate data parallel, tensor parallel, and optimizer parallel strategy pairs corresponding to layer u and layer v (<u,v>∈E) uv .

[0020] The storage overhead described in step S3 includes the memory overhead of all candidate strategies in each layer of the large model.

[0021] The storage cost described in step S3 specifically includes the following steps: using the cost model to estimate the memory cost M of all candidate data parallel, tensor parallel and optimizer parallel strategies of layer u (u∈V) of the large model G(V,E) u .

[0022] Step S4: Taking the time spent on each iteration of pipeline parallel training as the optimization target, a mixed integer quadratic programming mathematical model for automatically searching for a parallel training strategy for a large model is established and solved;

[0023] Step S4.1: Enumerate all legal numbers of pipeline stages deg and legal numbers of micro-batches c, select the minimum time per iteration TPI, and its corresponding one or more of the pipeline parallelism strategy, data parallelism strategy, tensor parallelism strategy, and optimizer parallelism strategy;

[0024] Step S4.2: Based on a non-pipelined parallel strategy (deg = 1), establish and solve a mixed integer quadratic programming mathematical model with the optimization objective of the time per iteration (TPI) of data parallelism, tensor parallelism, and optimizer parallelism. The search content is one or more of the data parallel strategy, tensor parallel strategy, and optimizer parallel strategy.

[0025] Step S4.3: Based on a given pipeline parallel strategy (deg>1), establish and solve a mixed integer quadratic programming mathematical model with the time TPI spent on each iteration of pipeline parallelism, data parallelism, tensor parallelism and optimizer parallel training as the optimization target. The search content is one or more of the pipeline parallel strategy, data parallel strategy, tensor parallel strategy and optimizer parallel strategy.

[0026] Step S5: output the optimal parallel training strategy for the large model.

[0027] Based on the solution results of the mixed integer quadratic programming mathematical model, one or more combinations of the optimal pipeline parallel strategy, data parallel strategy, tensor parallel strategy and optimizer parallel strategy for the large model under the current system are output.

[0028] A computer storage medium stores a computer program, which, when executed by a processor, implements the above-mentioned parallel strategy search method for efficient training of large artificial intelligence models.

[0029] A computer device includes a memory, a processor, and a computer program stored in the memory and executable on the processor. When the processor executes the computer program, the method for searching for a parallel strategy for efficient training of large artificial intelligence models is implemented.

[0030] Beneficial effects: Compared with the prior art, the present invention has the following advantages:

[0031] 1. The present invention supports four types of parallel training strategies: pipeline parallelism, data parallelism, tensor parallelism, and optimizer parallelism, and its strategy space is large;

[0032] 2. The present invention uses mixed integer quadratic programming, which shortens the search time.

[0033] 3. The present invention can be applied to single-machine multi-card parallel training of large models and distributed parallel training of multi-machine clusters, which helps to improve the training efficiency of large models, thereby reducing training time and energy consumption. BRIEF DESCRIPTION OF THE DRAWINGS

[0034] Figure 1 A flowchart of the steps of the method of the present invention;

[0035] Figure 2 A schematic diagram of a system environment according to an embodiment of the present invention;

[0036] Figure 3 This is a flowchart of step S4.1 of an embodiment of the present invention. DETAILED DESCRIPTION

[0037] The technical solution of the present invention will be further described below with reference to the accompanying drawings.

[0038] The parallel strategy search method for efficient training of large artificial intelligence models provided by the present invention supports pipeline parallelism, data parallelism, tensor parallelism, and optimizer parallel training of various large models. It can be applied to single-machine multi-card parallel training of large models and distributed parallel training of multi-machine clusters.

[0039] This example uses two servers, Server0 and Server1, each with four GPUs numbered GPU0, GPU1, GPU2, and GPU3. The two servers are isomorphic, including but not limited to identical hardware, software, network, and protocol environments. In this example, the global GPU ID is defined as globalID = serverID * gpuPerServer + localID, where serverID is the server ID, gpuPerServer is the number of GPUs per server, and localID is the local ID of the current GPU. For example, the global ID of GPU1 on Server1 is 1 * 4 + 1 = 5.

[0040] To help understand the relationship between the two servers in this embodiment, Figure 2 The topology of the two servers of this embodiment is shown. GPU0 and GPU1 of the two servers are connected to CPU0 via the PCIe bus, and GPU2 and GPU3 are connected to CPU1 via the PCIe bus. CPU0 and CPU1 of the two servers are connected via the QPI bus. The two servers are connected via Ethernet. This embodiment will apply a parallel strategy search method for efficient training of large artificial intelligence models to search for a parallel strategy that can be used to distribute the training of the BERT-Huge model on Server0 and Server1 using the Adam optimizer. In this embodiment, the specific workflow of the method of the present invention is as follows:

[0041] The workflow of the parallel strategy search method for efficient training of large artificial intelligence models is as follows: Figure 1 shown.

[0042] Step S1: Input the large model.

[0043] In this embodiment, the large model is the BERT-Huge model.

[0044] Step S2: Obtain system execution performance information and large model execution performance information.

[0045] In this embodiment, the acquisition of system execution performance information and large model execution performance information includes:

[0046] Set the All-Reduce communication volume aar , statistics of All-Reduce communication time t ar , thereby obtaining the system's All-Reduce communication efficiency

[0047] Set the P2P communication volume a p2p , statistics of P2P communication time t p2p , thereby obtaining the P2P communication efficiency of the system

[0048] Set the All-Reduce communication volume a ar , respectively count the All-Reduce communication time t of executing All-Reduce separately n , and the All-Reduce communication time t when performing All-Reduce and matrix multiplication calculations simultaneously o , thereby obtaining the system's computation-communication overlap factor

[0049] On Server0, using GPU0, setting the batch size to 1, calculate the forward execution time fp of each layer u (u∈V) of the BERT-Huge model G(V,E) u ;

[0050] On Server0, using GPU0, with a batch size of 1, we count the memory used by the BERT-Huge model and the Adam optimizer, and the memory used by each layer of the BERT-Huge model G(V,E). And the additional video memory m required to load the BERT-Huge model G(V,E) other .

[0051] Step S3: Use the cost model to estimate the execution time overhead and storage overhead of the large model on the current system.

[0052] In this embodiment, the execution time overhead of the large model on the current system includes:

[0053] Use the cost model to estimate the sum of the computation and communication time A of all candidate data parallel, tensor parallel, and optimizer parallel strategies for layer u (u∈V) of the BERT-Huge model G(V,E) with a batch size of bs. u In this embodiment, a candidate strategy S for layer u of BERT-Huge is ui For example, let the data parallelism / optimizer parallelism be dp ui , the tensor parallelism is tp ui , the data parallel / optimizer parallel communication volume is a dp, the forward and reverse communication volume of tensor parallelism is a tp Then apply strategy S ui When , the sum of the forward computation and communication time of layer u is The reverse computation and communication time includes the time when the reverse computation and data parallel / optimizer parallel communication overlap, the time when the reverse computation and data parallel / optimizer parallel communication do not overlap, and the reverse tensor parallel communication time, which is Therefore, a candidate strategy S for layer u of BERT-Huge in this embodiment is ui For example, when the batch size is bs, the sum of the computation and communication time of layer u is A ui =FP ui +BP ui ;

[0054] Use the cost model to estimate the edge of the BERT-Huge model G(V,E)<u,v> (<u,v> ∈E) corresponds to the layer u and layer v The communication time R between all candidate data parallel, tensor parallel, and optimizer parallel strategy pairs in the same pipeline stage with batch size bs uv In this embodiment, let the candidate strategy of layer u be S ui , the candidate strategy of layer v is S vj ,<u,v> ∈E, the communication volume in the same pipeline stage between layer u and layer v is a uvij , then the communication time between layers u and v in the same pipeline stage is

[0055] Use the cost model to estimate the edge of the BERT-Huge model G(V,E)<u,v> (<u,v> The P2P communication time R′ across pipeline stages between all candidate data parallel, tensor parallel, and optimizer parallel strategy pairs corresponding to layer u and layer v ∈ E) with batch size bs uv In this embodiment, let the candidate strategy of layer u be S ui , the candidate strategy of layer v is S vj ,<u,v> ∈E, the amount of P2P communication between layer u and layer v across pipeline stages is a′ uvij , then the P2P communication time between layer u and layer v across pipeline stages is

[0056] In this embodiment, the storage overhead of the large model on the current system includes:

[0057] Use the cost model to estimate the memory overhead M of all candidate data parallel, tensor parallel, and optimizer parallel strategies for layer u (u∈V) of the BERT-Huge model G(V,E) with a batch size of bs. u In this embodiment, the parameter of layer u is set to C u , data parallelism / optimizer parallelism is dp ui , the tensor parallelism is tp ui , trained with the Adam optimizer, then on a GPU, the memory overhead of layer u is

[0058] Step S4: Taking the time spent on each iteration of pipeline parallel training as the optimization target, a mixed integer quadratic programming mathematical model for automatically searching for a parallel training strategy for a large model is established and solved:

[0059] In this embodiment, the time spent on each iteration of pipeline parallel training is used as the optimization target, and a mixed integer quadratic programming mathematical model for automatically searching for a parallel training strategy for a large model is established and solved, including:

[0060] Step S4.1: Enumerate all legal numbers of pipeline stages deg and legal numbers of micro-batches c, select the minimum time per iteration TPI, and its corresponding one or more of the pipeline parallelism strategy, data parallelism strategy, tensor parallelism strategy, and optimizer parallelism strategy;

[0061] Step S4.2: Based on a non-pipelined parallel strategy (deg = 1), establish and solve a mixed integer quadratic programming mathematical model with the optimization objective of the time per iteration (TPI) of data parallelism, tensor parallelism, and optimizer parallelism. The search content is one or more of the data parallel strategy, tensor parallel strategy, and optimizer parallel strategy.

[0062] Step S4.3: Based on a given pipeline parallel strategy (deg>1), establish and solve a mixed integer quadratic programming mathematical model with the time TPI spent on each iteration of pipeline parallelism, data parallelism, tensor parallelism and optimizer parallel training as the optimization target. The search content is one or more of the pipeline parallel strategy, data parallel strategy, tensor parallel strategy and optimizer parallel strategy.

[0063] For ease of understanding, Figure 3 The execution flow of step S4.1 in this embodiment is shown. In this embodiment, the pipeline is a synchronous pipeline in GPipe mode.

[0064] Step S4.1.1. Input the sum of the computation and communication time A and the memory overhead M of each layer of the large model, the communication time R of the unified pipeline stage of adjacent layers and the P2P communication time R′ of different pipeline stages, as well as the number of GPUs n and the number of small batches B.

[0065] Step S4.1.2, let the number of pipeline stages deg = 1, the minimum execution time per iteration minTPI = ∞, and its corresponding minimum number of pipeline stages minDeg = -1, the number of micro-batches minC = -1, and the placement strategy of each layer of the model and parallel strategies

[0066] Step S4.1.3: Determine whether the number of pipeline stages deg is not greater than the number of GPUs n. If the judgment in S4.1.3 is true, enter the outer loop (entry is step S4.1.4).

[0067] Step S4.1.4, determine whether the pipeline degree deg is 1. If so, proceed to step S4.2, establish and solve a mixed integer quadratic programming mathematical model with the time TPI spent on each iteration of data parallelism, tensor parallelism, and optimizer parallelism as the optimization target. After the solution is completed, proceed to step S4.1.5, determine whether the optimization target TPI is less than minTPI. If the judgment in step S4.1.5 is established, the program will execute step S4.1.6, set minTPI = TPI, minDeg = deg, minC = c, minP = P, minS = S. Finally, the program will execute step S4.1.7, set deg = deg*2, and then judge the condition of the outer loop in step S4.1.3.

[0068] If the result of step S4.1.4 is negative, execute step S4.1.8, set the micro-batch size c to 2, and then execute step S4.1.9 to determine whether the micro-batch size c does not exceed the mini-batch size B. If this is not true, the program executes step S4.1.7, deg = deg * 2; if so, the program enters the inner loop. Within the inner loop, the program first executes step S4.1.10, determining whether B is not divisible by c. If so, execute step S4.1.14, c = c + 1, and continue to determine the inner loop condition in step S4.1.9. If step S4.1.10 is not true, execute step S4.1.11, set the micro-batch size b = B / c, and then execute step S4.3 to establish and solve a mixed integer quadratic programming mathematical model with the optimization objective of the time per iteration (TPI) of pipeline parallelism, data parallelism, tensor parallelism, and optimizer parallelism. After obtaining the solution result, execute step S4.1.12 to determine whether the optimization target TPI is less than minTPI. If the judgment is true, then execute step S4.1.13 and set minTPI = TPI, minDeg = deg, minC = c, minP = P, minS = S. Then the program will execute step S4.1.14, c = c+1, and continue to judge the inner loop condition in step S4.1.9.

[0069] If the judgment in step S4.1.3 is not established, go to step S4.1.15, output the minimum execution time per iteration minTPI and its corresponding minimum number of pipeline stages minDeg, number of micro-batches minC, placement strategy minP for each layer of the model and parallel strategy minS for each layer of the model, and end.

[0070] In step S4.2 provided in this embodiment, to help understand the mixed integer quadratic programming mathematical model with the time TPI spent per iteration of data parallelism, tensor parallelism, and optimizer parallelism as the optimization target, this embodiment will exemplify a mixed integer quadratic programming mathematical model:

[0071] Assuming that the objective function minimizes the time spent in each iteration TPI can be expressed as formula (1):

[0072] min TPI=p1 (1)

[0073] In formula (1), TPI is the time spent on each iteration, so this formula is equivalent to maximizing the number of iterations per unit time, that is, equivalent to maximizing throughput. iis the total computation and communication time of the i-th pipeline parallel stage. Because step S4.2 does not require pipeline parallelism, one or more combinations of data parallelism, tensor parallelism, and optimizer parallelism can be considered as parallel computation performed in a single pipeline stage. Without loss of generality, we assume that data parallelism, tensor parallelism, and optimizer parallelism are performed in the first stage of the pipeline, and denote the total computation and communication time of the first stage of the pipeline as p1.

[0074] For a better understanding, the following introduces the constraints in combination with the formula:

[0075] Constraint 1: The computation time of the first stage of the pipeline includes the sum of the computation and communication time of all layers u of BERT-Huge, as well as the sum of the computation and communication time of all edges of BERT-Huge.<u,v> (<u,v> ∈E) is the sum of the communication costs of the same pipeline stage between layer u and layer v. Constraint 1 can be expressed as formula (2):

[0076]

[0077] Among them S uk Indicates whether layer u selects the k-th candidate data parallel, tensor parallel, and optimizer parallel strategy. It is a 0-1 variable.

[0078] Constraint 2: The storage cost of the first stage of the pipeline includes the sum of the memory costs of all candidate data parallel, tensor parallel, and optimizer parallel strategies for all layers u (u∈V) of BERT-Huge. Constraint 2 can be expressed as formula (3):

[0079]

[0080] In this embodiment, m represents the memory capacity of GPU0 minus the additional memory m required to load the BERT-Huge model G(V,E) other In this embodiment, Server0 and Server1 are isomorphic, so m is a constant for different GPUs.

[0081] Constraint 3: All layers u(u∈V) of BERT-Huge can and only can select one candidate data parallel, tensor parallel and optimizer parallel strategy. Constraint 3 can be expressed as formula (4)(5):

[0082]

[0083]

[0084] where g urepresents the set of candidate data parallel, tensor parallel, and optimizer parallel strategies for BERT-Huge layer u, |g u | indicates g u The momentum of the collection.

[0085] In step S4.3 provided in this embodiment, to help understand the mixed integer quadratic programming mathematical model with the time TPI spent in each iteration of pipeline parallelism, data parallelism, tensor parallelism, and optimizer parallel training as the optimization target, this embodiment will exemplify a mixed integer quadratic programming mathematical model:

[0086] Assuming that the objective function minimizes the time spent in each iteration TPI can be expressed as formula (6):

[0087]

[0088] In formula (6), p i is the total computation and communication time of the ith stage of the pipeline parallelism, ps is the total number of stages of the pipeline parallelism, and ps = deg. j is the total communication time of the jth communication phase across pipeline stages, os is the total number of communication phases across pipeline stages in pipeline parallelism, satisfying os = ps-1. c is the number of micro-batches in pipeline calculation. represents the sum of computation and communication time of all stages in the pipeline, represents the sum of all P2P communication times across pipeline stages in the pipeline, max{p1,…,p ps}*(c-1) models the impact of the slowest stage in all stages of the pipeline and the number of micro-batches on the overall pipeline computation time.

[0089] For a better understanding, the following introduces the various constraints in combination with the formula:

[0090] Constraint 1: The computation and communication time of the i-th stage of the pipeline, including: (1) the sum of the computation and communication time of the BERT-Huge layers placed in the i-th stage; (2) the sum of the edges between layers u and v of BERT-Huge placed in the i-th stage.<u,v> (<u,v> ∈E) corresponds to the sum of the communication costs of the same pipeline stage. Constraint 1 can be expressed as formula (7):

[0091]

[0092] Among them, P ui Indicates whether layer u is placed in the i-th stage of the pipeline, which is a 0-1 variable. uk Indicates whether layer u selects the k-th candidate data parallel, tensor parallel, and optimizer parallel strategy. It is a 0-1 variable.

[0093] Constraint 2: The communication time of the jth cross-pipeline stage in the pipeline parallelism. Constraint 2 can be expressed as formula (8):

[0094]

[0095] Constraint 3: The storage cost of the i-th stage of the pipeline includes the sum of the memory costs of all candidate data parallel, tensor parallel, and optimizer parallel strategies for the BERT-Huge layer u (u∈V) placed in the i-th stage. Constraint 3 can be expressed as Equation (9):

[0096]

[0097] In this embodiment, m represents the memory capacity of GPU0 minus the additional memory m required to load the BERT-Huge model G(V,E) other In this embodiment, Server0 and Server1 are isomorphic, so m is a constant for different GPUs.

[0098] Constraint 4: Each pipeline stage must contain a continuous portion of the BERT-Huge model. Constraint 4 can be expressed as formula (10)(11)(12):

[0099]

[0100]

[0101]

[0102] where Z vi is a 0-1 variable representing whether there is a layer w in the successor of layer v, so that layer w and a predecessor layer u of layer v are placed on the i-th pipeline stage. If layer w exists, then Z vi =1, otherwise it is equal to 0. In this embodiment, formulas (10)(11)(12) jointly constrain the different layers of the BERT-Huge model not to be placed in the pipeline stages out of order, but to be placed sequentially.

[0103] Constraint 5: All layers u(u∈V) of BERT-Huge can and can only be placed on one pipeline stage, and all pipeline stages have at least one layer. Constraint 5 can be expressed as formula (13)(14)(15):

[0104]

[0105]

[0106]

[0107] Constraint 6: All layers u(u∈V) of BERT-Huge can and only can select one candidate data parallel, tensor parallel, and optimizer parallel strategy. Constraint 6 can be expressed as formula (4)(5).

[0108] Finally, based on the solution of the above mixed integer quadratic programming model, the optimal parallel training strategy for the large model is output (step S5):

[0109] In this embodiment, the parallel training strategy for the optimal large model output specifically includes:

[0110] Based on the output of step S4, step S4.1.15 outputs one or more combinations of the optimal pipeline parallel strategy, data parallel strategy, tensor parallel strategy, and optimizer parallel strategy for the BERT-Huge model under the current system.

[0111] More specifically, minTPI represents the minimum execution time per iteration, minDeg represents the minimum number of pipeline stages corresponding to the current minTPI, minC represents the number of micro-batches corresponding to the current minTPI, minP represents the placement strategy for each layer of the model corresponding to the current minTPI, and minS represents the parallel strategy for each layer of the model corresponding to the current minTPI. Therefore, the final pipeline parallel strategy is: set the number of pipeline stages to minDeg, and split the mini-batch into minC micro-batches during pipeline calculation. Each layer u of the model should be placed on the pipeline stage corresponding to each layer in minP. At the same time, the data parallelism, tensor parallelism, and optimizer parallelism strategies for each layer u of the model are the parallel strategies adopted for each layer in minS. In addition, for the mapping relationship between parallel computing strategies and systems, the GPU globally numbered l in the system will be responsible for calculating the first layer of the pipeline. For example, if minDeg = 4, then the global number of GPU3 of Server0 is 0*4+3=3, which will be responsible for calculating the first stage of the pipeline. The global number of GPU0 of Server1 is 1*4+0=4, which will be responsible for calculating the first stage of the pipeline. stage.

[0112] The method of the present invention was experimented on multiple models and systems. During the experiment, the search time of the parallel strategy search method for efficient training of large artificial intelligence models was counted on the server numbered 0. After the parallel training strategy was searched by the method, the throughput of the parallel strategy under the current system was counted on the server numbered 0. The experimental results on the four models of BERT-Huge, T5-Large, ViT-Huge, and Swin-Huge showed that compared with the best existing methods, the method of the present invention can improve the strategy search speed by about 16 times at the fastest and the training throughput by about 1.7 times at the maximum.

Claims

1. A parallel strategy search method for efficient training of large artificial intelligence models, characterized by: The following steps are involved: Step S1: inputting the artificial intelligence model; Step S2: Obtaining system execution performance information and execution performance information of the artificial intelligence large model; Step S3: using the cost model to estimate the execution time and storage overhead of the large artificial intelligence model on the current system; Step S4: Taking the time spent on each iteration of pipeline parallel training as the optimization target, a mixed integer quadratic programming mathematical model for automatically searching for a parallel training strategy for a large model is established and solved; Step S5: outputting the optimal parallel training strategy for the large artificial intelligence model; The step S2 is specifically as follows: Setting the All-Reduce Traffic Volume , statistics of All-Reduce communication time , thereby obtaining the system's All-Reduce communication efficiency ; Set P2P communication volume , statistics of P2P communication time , thereby obtaining the P2P communication efficiency of the system ; Setting the All-Reduce Traffic Volume , respectively count the All-Reduce communication time of executing All-Reduce separately , and the All-Reduce communication time when performing All-Reduce and matrix multiplication calculations simultaneously , thereby obtaining the system's computation-communication overlap factor ; Statistical large models Each layer Forward execution time ; Statistics In addition to the memory occupied by the model and optimizer, large models Active memory for each layer and loading large models Additional video memory required ; The execution time overhead described in step S3 specifically includes the following steps: Estimating Large Models Using Cost Models Layer The sum of the computation and communication time of all candidate data parallel, tensor parallel, and optimizer parallel strategies ; Estimating Large Models Using Cost Models edge The corresponding layer and layer Communication time within the same pipeline stage between all candidate data parallelism, tensor parallelism, and optimizer parallelism strategy pairs ; Estimating Large Models Using Cost Models edge The corresponding layer and layer P2P communication time across pipeline stages between all candidate data parallelism, tensor parallelism, and optimizer parallelism strategy pairs ; The step S4 is specifically as follows: Step S4.1: Enumerate all legal pipeline stage numbers and legal micro-batch quantities , select the minimum time spent on each iteration , and one or more of its corresponding pipeline parallel strategy, data parallel strategy, tensor parallel strategy and optimizer parallel strategy; Step S4.2: Based on a non-pipelined parallel strategy, establish and solve the time taken for each iteration in data parallelism, tensor parallelism, and optimizer parallelism A mixed integer quadratic programming mathematical model for optimizing the objective, wherein the search content is one or more of a data parallel strategy, a tensor parallel strategy, and an optimizer parallel strategy; Step S4.3: Based on a given pipeline parallel strategy, establish and solve the time taken for each iteration of pipeline parallelism, data parallelism, tensor parallelism and optimizer parallelism training A mixed integer quadratic programming mathematical model is used to optimize the target, and the search content is one or more of pipeline parallel strategy, data parallel strategy, tensor parallel strategy and optimizer parallel strategy.

2. A parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The execution time overhead described in step S3 includes the sum of the calculation and communication time of all candidate strategies in each layer of the large model, the communication time of all candidate strategy pairs in adjacent layers in the same pipeline stage, and one or more of the P2P communication time across pipeline stages.

3. A parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The storage overhead described in step S3 includes the memory overhead of all candidate strategies in each layer of the large model.

4. A parallel strategy search method for efficient training of large artificial intelligence models according to claim 1, characterized in that: The storage cost described in step S3 specifically includes the following steps: using the cost model to estimate the large model Layer The memory overhead of all candidate data parallel, tensor parallel, and optimizer parallel strategies .

5. The parallel strategy search method for efficient training of large artificial intelligence models according to claim 1 is characterized in that: The step S5 is specifically as follows: Based on the solution results of the mixed integer quadratic programming mathematical model, one or more combinations of the optimal pipeline parallel strategy, data parallel strategy, tensor parallel strategy and optimizer parallel strategy for the large model under the current system are output.

6. A computer storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, it implements a parallel strategy search method for efficient training of large artificial intelligence models as described in any one of claims 1 to 5.

7. A computer device comprising a memory, a processor, and a computer program stored in the memory and executable on the processor, wherein: When the processor executes the computer program, it implements a parallel strategy search method for efficient training of large artificial intelligence models as described in any one of claims 1-5.

Images