Recomputation-Aware Model Partitioning Method and System for Pipeline Parallel Training

Abstract

The present invention discloses a recomputation-aware model splitting method and system for pipeline parallel training, belonging to the field of artificial intelligence in computer science. The method of the present invention specifically includes: using the Dijkstra algorithm and two-layer binary search to calculate the minimum value of the sum of the forward propagation and backward propagation times of the pipeline microbatch under recomputation-aware model splitting; and using the intermediate results obtained by the Dijkstra algorithm during the search process to calculate and output a recomputation-aware approximately optimal model splitting strategy. The present invention utilizes the model characteristics of the Transformer large model and the characteristics of the pipeline parallel mode, designs a model splitting method and system, non-uniformly splits the model into each stage of the pipeline, and ensures the computational load balance of each pipeline stage when recomputation is enabled, thereby improving the overall performance of large model pipeline parallel training and the utilization efficiency of cluster devices.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence in computer science, and particularly to a recomputation-aware model splitting method and system for pipeline parallel training. Background Art

[0002] With the development of artificial intelligence technology, the scale and complexity of models have been continuously increasing, and more powerful performance and more accurate prediction results have been obtained. These large models, such as Transformer-based models, usually contain hundreds of millions or even billions of parameters. For example, one of the latest open-source large models, Llama3.1-405B, contains more than 400 billion parameters. However, the increase in these parameters also poses a huge higher demand for the capacity of video memory, making it very difficult to train these models on a single GPU or CPU. To overcome the limitation of video memory capacity, researchers and engineers have developed model splitting technologies. The core idea of this technology is to split a huge model into multiple smaller sub-models, and each sub-model can be independently computed on different hardware devices, and they form a pipeline with each other to efficiently utilize the potential of the devices. However, model splitting needs to satisfy the computational load balance on each device, otherwise the synchronization operation between devices will reduce the overall training efficiency.

[0003] In a distributed training environment with pipeline parallelism, due to the application of recomputation technology, the computational loads on different devices may become unbalanced.

[0004] In the case of uniform slicing, the computational loads on each device under pipeline parallelism are equal, and there will be no device idle time (bubbles). However, due to the different pipeline stages of the devices, their video memory loads are also different. To prevent video memory overflow, the latest training frameworks (such as Megatron, Deepspeed, etc.) allow users to use recomputation technology, and some new technologies (such as Lynx, AdaPipe, etc.) also propose to use corresponding recomputation strategies according to different pipeline stages. However, directly using different recomputation strategies under uniform slicing of the pipeline will result in different recomputation time overheads in different stages, further leading to a large amount of idle time for the devices in the pipeline.

[0005] To solve this problem, a recomputation-aware model splitting technology is needed, which can adaptively adjust the model splitting strategy by combining information such as the storage requirements and recomputation time on each device to optimize the overall training efficiency. Summary of the Invention

[0006] To solve the problem that the inefficient splitting in the training of existing large Transformer models leads to low utilization of the hardware performance for model training, the purpose of the present invention is to provide a recomputation-aware model splitting method and system for pipelined parallel training, which can balance the time overhead of each stage in pipelined parallelism, reduce the idle time of the device when using pipelined parallel training, and improve the overall device utilization rate.

[0007] The purpose of the present invention is achieved through the following technical solutions: A recomputation-aware model splitting system for pipelined parallel training, the system includes:

[0008] An information collection module: used to configure the structure of the model, parallel training strategies, and hyperparameters (GlobalBatchSize, Micro Batch Size), deploy the defined model on a real GPU, pre-run several training iterations, and collect data; the data includes the number of model operators, computation time, output size, and operator dependencies; the parallel training strategies include settings of Tensor parallelism and pipelined parallelism.

[0009] A model splitting module: used to provide a model splitting mechanism with recomputation strategy awareness to achieve uniform distribution of workload among multiple devices; the model splitting mechanism with recomputation strategy awareness is implemented through the dijkstra algorithm and two-layer binary search. Among them, the first-layer binary search performs a binary search on an approximate solution of the theoretical minimum value of the maximum sum of the forward propagation and backward propagation times (including recomputation time) for training one microbatch in each stage of the pipeline; the second-layer binary search is the calculation method of "calculating the maximum sub-model scale that each stage can support with T mid as a constraint condition". According to the structure of the last computing unit of the sub-model, a binary search is respectively performed on the maximum sub-model scale that each stage can support under the condition of "with T mid as a constraint condition", where T mid is a temporary variable representing the intermediate value of the interval in the first-layer binary search process; the dijkstra algorithm is used to determine whether the intermediate results in the process of the first-layer binary search meet the requirements and for the calculation of the final result.

[0010] A model training module: configures the model splitting strategy and the recomputation strategy of each device into the training strategy according to the model splitting mechanism.

[0011] Further, the model splitting module includes a sub-module, which is a memory module containing two tables. One table is used to save the intermediate results of the sub-model scale search ending with a FeedForward unit, and the other table is used to save the intermediate results of the sub-model scale search ending with an Attention unit. The intermediate results include the number of consecutive computing units included in the sub-model, the recomputation strategy, and the estimated computing time of the sub-model.

[0012] Further, the recomputation strategy-aware model splitting mechanism is specifically as follows:

[0013] S1. Initialize the search interval (T low , T high ) for the sum of the maximum forward propagation time and the backward propagation time (including recomputation time) of training a single microbatch in the pipelining stage; initially set the search interval to (T min , T max ), where T min is the time taken for a single microbatch when the current complete model is evenly sliced by layer into each stage and no recomputation strategy is used, and T max is the time taken for a single microbatch when the current complete model uses the full recomputation strategy under the current training and device configuration. The time taken for a single microbatch represents the sum of the forward propagation and backward propagation times of training a single microbatch in each stage of the pipeline. T low represents the lower bound of the search interval in the first-layer binary search process, and T high represents the upper bound of the search interval in the first-layer binary search process;

[0014] S2. Use the middle value T low , T high of the interval (T mid ) as the constraint for the maximum single microbatch time in the pipelining stage. For each stage i, search for the maximum number of consecutive computing units that can be trained under the corresponding constraint T mid ;

[0015] S3. To verify whether the overall pipelining slicing is feasible under the constraint of T mid , for all l i，s (i ∈ {0...p}, s ∈ {0, 1}) of all stages obtained in S2, establish an edge from l i-1，j to l i，k between the (i - 1)-th stage and the i-th stage according to the following edge weight ruleThe edges j, k ∈ {0, 1}, where i represents the stage number of pipeline parallelism, s represents the type of computing unit at the end of this stage (0 represents ending with a FeedForward unit, 1 represents ending with an Attention unit), and l i，s represents the maximum number of consecutive computing units that can be trained under the constraint T of the maximum time-consuming of the largest single microbatch at the i-th stage obtained by S2, where the computing unit at the end is represented by s. j and k represent the types of computing units at the corresponding stages (0 represents ending with a FeedForward unit, 1 represents ending with an Attention unit). Let l mid represent the node number (i - 1, j), and let l i-1，j represent the node number (i, k), and finally connect them into a directed acyclic graph: i，k

[0016] When j = 0 and k = 0, the edge weight is less than or equal to the largest even number less than or equal to L i，0 ;

[0017] When j = 0 and k = 1, the edge weight is less than or equal to the largest odd number less than or equal to l i，0 ;

[0018] When j = 1 and k = 0, the edge weight is less than or equal to the largest odd number less than or equal to l i，1 ;

[0019] When j = 1 and k = 1, the edge weight is less than or equal to the largest even number less than or equal to l i，1 ;

[0020] Execute the Dijkstra single-source path-finding algorithm on the directed acyclic graph to find the longest path (the length is the sum of the weights of the edges on the path) from the node numbered ((0, 0) to the node numbered (p, 0). If the length of the longest path is greater than or equal to 2L, it is considered that the pipeline splitting is feasible under the constraint of T mid , otherwise it is not feasible. Here, L is the number of layers of the Transformer large model to be trained. If the splitting is feasible, update the search interval of T (T low , T high ) to (T low , T high ). If the splitting is not feasible, update the search interval of T (T low , T high ) to (T mid , T high ). Here, T is the target to be found: the theoretical minimum value of the maximum time-consuming of a single microbatch in the pipeline stage;

[0021] Judge the termination condition T high - T low ​≤δ, where δ is a hyperparameter set by the user; if the termination condition is met, return T high as If not, repeat S2 - S3 with the updated search interval and obtain an approximate optimal solution of T through multiple iterations

[0022] S4. Select any path {(0, 0), (1, s1),...(i, s ),...(p, 0)} (s i ) in the directed acyclic graph in the last round of iteration where the pipelining splitting is feasible under the condition that the longest time consumption is less than or equal to i the limit and with a length greater than or equal to 2L as the basis for splitting and split the model

[0023] Furthermore, S1 includes:

[0024] Search for the minimum value T of the time consumption t(l q ) of the stage with the longest time consumption under the training video memory limit, and calculate the number of Transformer computing units obtained by splitting each stage under this minimum value limit in reverse; where the objective and constraints are represented by the following formula, and at this time T participates in the optimization as a parameter:

[0025]

[0026] Find an approximate solution of the minimum value T of t(l q ) through binary search of t(l q )

[0027] Furthermore, S2 includes:

[0028] S2.1. Initialize the maximum number of consecutive computing units l mid that can be trained in this stage when the sum of the maximum forward propagation time and the backward propagation time of the pipelining stage i for training a single microbatch is less than T i，0 and l i，1 and the initial search intervals [l i，0 low , l i，0 up and [l i，1 low , l i，1 up , where l i，0 and l' i，1 respectively represent the maximum number of consecutive computing units with the FeedForward unit and the Attention unit as the last computing unit structure of the sub - model in this stage;

[0029] Use l' iDenote the sub-model scale at stage i under the final splitting scheme (the number of consecutive computing units of the sub-model). When the VRAM occupancy during sub-model training is less than or equal to the training VRAM limit and the time consumption is less than or equal to T mid If there exists stage i, and the sum of the sub-model scales l′ i satisfies the following equation, then the pipeline splitting scheme is feasible; otherwise, it is not feasible. Here, L is the number of layers of the large Transformer model to be trained:

[0030]

[0031] By finding the values of l i，0 and l i，1 to determine the range of l′ i under the final splitting scheme;

[0032] Use the binary search method to find the maximum trainable sub-model scale that satisfies the constraints t(l i，0 ) ≤ T mid , t(l i，0 ) ≤ T mid , which is represented by the following formula:

[0033]

[0034] where, and are the specific recomputation strategies obtained by calling the recomputation strategy search method and inputting the model configuration, training configuration, device information, the sub-model scale to be searched, and the information of the current stage number i. and are the time consumption of a single microbatch under the current recomputation strategy configuration;

[0035] By determining an interval [l i，0low , l i，0high and [l i，1low , l i，1high , ensure that the target l i，0 and l i，1 fall within this interval, and then use binary search to find the maximum value that satisfies the constraints as l i，0 and l i，1 ; [[ID=ó2]]

[0036] S2.2. Use the floor value of the middle value of the interval [l i，0 low , l i，0 high initialized in S2.1 as the number of consecutive computing units l i，0 mid, and according to the number of such units, connect the computing units into sub-model M'0 respectively in the structure of interleaved FeedForward units and Attention units, input the corresponding model configuration, model scale, current device, training configuration, and current stage number of sub-model M'0 into the recomputation strategy search unit, and call the recomputation strategy search unit and its corresponding Cost Model to obtain the corresponding recomputation strategy and the time consumption of a single microbatch of sub-model M'0 calculated according to the Cost Model under the corresponding recomputation strategy

[0037] S2.3. Use the middle value of the initialized interval [l i，1low , l i，1 high ) rounded down as the number of consecutive computing units included in the sub-model of the pipeline stage l i，1 mid , and according to the number of such units, connect the computing units into sub-model M'1 respectively in the structure of interleaved FeedForward units and Attention units, and input the corresponding model configuration, model scale, current device, training configuration, and current stage number of sub-model M'1 into the recomputation strategy search unit, and call the recomputation strategy search unit and its corresponding Cost Model to obtain the corresponding recomputation strategy and the time consumption of a single microbatch of sub-model M'1 calculated according to the Cost Model under the corresponding recomputation strategy

[0038] Furthermore, the determination of an interval [l i,0low , l i,0high ) and [l i，1 low , l i，1 high ), ensuring that the target l i，0 and l i，1 fall within this interval, and then using binary search to find the maximum value that satisfies the constraint as l i，0 and l i,1 includes:

[0039] Initialize the search intervals [l i,0 , l i,1 ) and [l i,0low , l i,0high ) for l i,1low , l i,1high , set l i,0low and l i,1low to 0, and set l i,0high and L i，1high to twice the number of model layers L;

[0040] When repeating the search process, to reduce the search space, through the memory module, during the search for the maximum trainable sub-model scale l i，0 and l i，1 each time the recomputation strategy search unit is called during the process, for the input temporary variable l′ i，0 or l′ i，1 the time consumed by a single microbatch obtained, t(l′ i，0 ) or t(l′ i，1 ) and the corresponding recomputation strategy, after that, l′ i，0 and l′ i，1 as well as the time consumed by a single microbatch corresponding to both and the recomputation strategy are respectively saved in two tables (one table saves the information corresponding to l′ i，0 , and the other saves the information corresponding to l′ i，1 ), each row of the table contains the sub-model scale, the sum of the corresponding forward propagation time and backward propagation time (including recomputation time), and the corresponding recomputation strategy;

[0041] In the initialization process of the second and subsequent second-layer searches, in the table corresponding to l′ i，0 or l′ i，1 in the memory module, respectively find the item with the largest number of computing units where t(l′ i，0 ) or t(l′ i，1 ) is less than T mid as l i，0low or l i，1 low , and the item with the smallest number of computing units where t(l′ i，0 ) or t(l′ i，1 ) is greater than T mid as l i，0high or l i，1high , if the corresponding item cannot be found in the table, it degrades to the initialization strategy during the first search.

[0042] Furthermore, the S2.2 includes:

[0043] Considering the overhead of the recomputation strategy search, during the computation process, before calling the recomputation strategy search unit, by querying the table storing the information of the sub-model ending with the FeedForward unit in the memory module, if the required sub-model scale and the corresponding strategy are already in the table, directly use the data in the table, otherwise, after the search, insert the corresponding information into the table storing the information of the sub-model ending with the FeedForward unit in the memory module;

[0044] If then update the interval [l i，0low , l i，0high ) to [l i，0mid , li，0high ),otherwise the updated interval is [l i，0low , l i，0mid );

[0045] Judge whether [l i，0low , l i，0high ) reaches the termination condition l i，0low + 1 = l i，0high ; if the termination condition is reached, return l i，0low as l i，0 ; if the termination condition is not reached, use the updated new interval to execute S2.2 again.

[0046] Further, the S2.3 includes:

[0047] Considering the overhead of recomputation strategy search, during the calculation process, before calling the recomputation strategy search unit, query the table storing the information of the sub-model ending with the Attention unit in the memory module. If the scale of the required sub-model and the corresponding strategy are already in the table, directly use the data in the table; otherwise, insert the corresponding information into the table storing the information of the sub-model ending with the Attention unit in the memory module after searching;

[0048] If then update the interval [l i，1 low , l i，1 high ) to [l i，1 mid , l i，1 high ), otherwise the updated interval is [l i，1 low , l i，1 mid );

[0049] Judge whether [l i，1 low , l i，1 high ) reaches the termination condition l i，1 low + 1 = l i，1 high ; if the termination condition is reached, return l i，1 low as l i，1 ; if the termination condition is not reached, use the updated new interval to execute S2.3 again.

[0050] Further, in the S4, the splitting of the model includes:

[0051] Satisfy the following constraint conditions:

[0052] The number of units in the i-th stage is less than the length of the edge from (i - 1, s i-1 ) to (i, s i ), and if k i = 0, the last unit structure of the sub-model is the FeedForward unit, if k iIf = 1, the last unit structure of the sub - model is an Attention unit;

[0053] The total number of units obtained by splitting in all stages is 2L;

[0054] The sub - model structure is an interleaved series of Attention units and FeedForward units.

[0055] The present invention also provides a recomputation - aware model splitting method for pipeline parallel training. The method includes:

[0056] Collect data for the defined model. The data includes the number of model operators, computation time, output size, and operator dependencies. The collected data is used for formulating subsequent recomputation strategies and model splitting strategies;

[0057] According to the model splitting mechanism aware of the recomputation strategy, achieve uniform distribution of workload among multiple devices. The model splitting mechanism aware of the recomputation strategy is implemented through the dijkstra algorithm and two - layer binary search. Among them, the first - layer binary search performs a binary search on the minimum value of the sum of the forward propagation and backward propagation times for training one microbatch in each stage of the pipeline; The second - layer binary search is the calculation method of "calculating the maximum sub - model scale that each stage can support with T as the limit condition" in the first - layer binary search. According to the structure of the last computing unit of the sub - model, respectively perform a binary search on the maximum sub - model scale that each stage can support under the condition of "with T as the limit condition"; The dijkstra algorithm is used to determine whether the intermediate results of the first - layer binary search meet the requirements and for calculating the final result; mid as the limit condition, calculate the maximum sub - model scale that each stage can support", according to the structure of the last computing unit of the sub - model, respectively perform a binary search on the maximum sub - model scale that each stage can support under the condition of "with T mid as the limit condition"; The dijkstra algorithm is used to determine whether the intermediate results of the first - layer binary search meet the requirements and for calculating the final result;

[0058] Configure the model splitting strategy and the recomputation strategy of each device into the training strategy.

[0059] The beneficial effects of the present invention are as follows: By establishing a mathematical model using mathematical formulas and graph structures for the pipeline training process of large models, the system can reasonably configure the number of computing units in each stage in a recomputation - aware manner, ultimately achieving load balancing in each stage of the training pipeline, increasing the system throughput by 24% - 38% compared with average splitting, and increasing the system throughput by 3% - 6% compared with the splitting strategy of Lynx, one of the current most advanced technologies. BRIEF DESCRIPTION OF THE DRAWINGS

[0060] Figure 1 is the system architecture diagram of the present invention;

[0061] Figure 2 is the structure diagram of the Attention unit;

[0062] Figure 3 It is the structural diagram of the FeedForward unit;

[0063] Figure 4 It is the flowchart for searching the minimum value of the time-consuming of the longest single microbatch time-consuming stage;

[0064] Figure 5 For T mid It is the flowchart for searching the maximum value of the model scale corresponding to the pipeline stage i under the T

[0065] Figure 6 It is the test result diagram of the present invention. Specific implementation manners

[0066] Here, the exemplary embodiments will be described in detail, and the examples are shown in the drawings. When the following description refers to the drawings, unless otherwise indicated, the same numbers in different drawings represent the same or similar elements. The implementation manners described in the following exemplary embodiments do not represent all implementation manners consistent with the present invention. On the contrary, they are merely examples of devices and methods consistent with some aspects of the present invention as detailed in the appended claims.

[0067] The Transformer model is structurally composed of multiple Transformer Layer structures, and each Transformer structure is in turn composed of two units, Attention and FeedForward. Considering that the pipeline parallelism latency is affected by the communication time between pipeline stages, in the past pipeline splitting methods, the Transformer model often splits by the Transformer Layer structure or the Attention / FeedForward computing unit as the granularity. The present invention considers the final splitting effect and adopts a fine-grained splitting method, that is, the latter splits the pipeline in a fine-grained manner with the Attention / FeedForward computing unit as the granularity.

[0068] The scenario targeted by the present invention is training using distributed GPUs, and at the same time using the pipeline parallelism and recomputation methods during training. The overall structure is designed as Figure 1 shown, specifically including the following three modules:

[0069] Information acquisition module: used to configure the structure of the model, parallel training strategies (including the settings of tensor parallelism and pipeline parallelism), and hyperparameters (Global Batch Size, Micro Batch Size), deploy the defined model on real GPUs, pre-run several training iterations and collect data; the data includes the number of model operators, computing time, output size, and operator dependencies;

[0070] Model splitting module: It is used to provide a model splitting mechanism with recomputation strategy awareness to achieve uniform distribution of workload among multiple devices; the model splitting mechanism with recomputation strategy awareness is implemented through the Dijkstra algorithm and two-layer binary search. Among them, the first-layer binary search performs a binary search on the approximate solution of the theoretical minimum value of the maximum sum of the forward propagation and backward propagation times (including recomputation time) for training one microbatch in each stage of the pipeline; the second-layer binary search is the calculation method of "calculating the maximum sub-model scale that each stage can support with T mid as the constraint condition". According to the structure of the last computing unit of the sub-model, it respectively performs a binary search on the maximum sub-model scale that each stage can support under the condition of "with T mid as the constraint condition", where T mid is a temporary variable representing the intermediate value of the interval in the first-layer binary search process; the Dijkstra algorithm is used to determine whether the intermediate results in the process of the first-layer binary search meet the requirements and is used for the calculation of the final result;

[0071] Model training module: Configure the model splitting strategy and the recomputation strategy of each device into the training strategy, and the system framework deploys the model to the GPU device for training.

[0072] The model splitting mechanism with recomputation strategy awareness is implemented through the Dijkstra algorithm and two-layer binary search: The first-layer binary search performs a binary search on the minimum value of the sum of the forward propagation and backward propagation times (including recomputation time) for training one microbatch in each stage of the pipeline. The process of the first layer is as Figure 4 shown; the second-layer search is the specific calculation method of "calculating the maximum sub-model scale that each stage can support with T mid as the constraint condition" in the first layer. This layer of search respectively performs a binary search on the maximum sub-model scale that each stage can support under the condition of "with T mid as the constraint condition" according to the structure of the last computing unit of the sub-model. The process of the second layer is as Figure 5 shown; the Dijkstra algorithm is mainly used in the model splitting mechanism with recomputation strategy awareness to determine whether the intermediate results in the process of the first-layer binary search meet the requirements and is used for the calculation of the final result. The details of the model splitting mechanism with recomputation strategy awareness are described below: Step 1: Initialize the search interval (T low , T high) Specifically, in pipeline parallelism, the optimal splitting scheme of the model requires that the sum of the forward propagation and backward propagation times (including recomputation time) of training one microbatch in each stage of the pipeline (hereinafter simply referred to as the single microbatch time consumption) be equal, that is: t(l1) = t(l2) =... = t(l p )

[0073] where t(l i ) is the single microbatch time consumption when the submodel scale is l i in stage i, and p is the total number of pipeline parallel stages.

[0074] If a splitting scheme does not satisfy this equation, then in the process of pipeline parallel training, the stage with a shorter single microbatch time consumption needs to wait for the stage with a longer time consumption to complete the computing task. A simple approach is to first split the model through modeling or heuristic methods, and then estimate the time consumption of each stage based on actual training data or CostModel, and then adjust the splitting scheme manually or algorithmically until the pipeline parallelism reaches the expected efficiency. This patent observes that the overall training time is positively correlated with the single microbatch time consumption t(l q ) of the stage q with the longest single microbatch time consumption. Therefore, by minimizing the time consumption t(l q ) of the stage with the longest time consumption, the overall efficiency of the pipeline can be optimized. The difficulty here is that under different splitting schemes, the stage q with the longest time consumption is not a fixed stage, and the stage with the longest time consumption is also different under different pipeline splitting strategies. Therefore, it is necessary to optimize the overall splitting scheme from a global perspective to obtain the submodel scale of each stage. For the global situation, the single microbatch time consumption of each stage must satisfy the following inequality group:

[0075] t(l1) ≤ t(l q )

[0076] t(l2) ≤ t(l q )

[0077] …

[0078] t(l p ) ≤ t(l q );

[0079] Therefore, by searching for the time consumption t(l q) and reversely calculate the number of Transformer computational units obtained by splitting each stage under this minimum constraint, that is, the size of the sub-model at each stage. At this point, the optimization method has changed from first determining the size of the sub-model at each stage and then minimizing T to first finding the minimum value of T and then determining the size of the sub-model at each stage. The specific goal and constraints are expressed in the following formula, where T is also optimized as a parameter:

[0080]

[0081] By directly q )Binary search can find t(l q ) is the approximate solution of the minimum value T

[0082] The first step is to determine a search interval (T low ,T high ), the present invention initially sets this interval to (T min ,T max ), where T min The time taken for a single microbatch when the current complete model is evenly divided into each stage according to the number of layers and no recalculation strategy is used, T mid The time it takes to run a single microbatch for the current complete model using the full recomputation strategy under the current training and device configuration.

[0083] Step 2:

[0084] Use interval (T low , T high )Intermediate value T mid As the constraint on the maximum single microbatch time of the pipeline stage, for each stage i, search for the corresponding constraint T mid The maximum number of consecutive computational units that can be trained under this condition is , and the search process includes the following sub-steps:

[0085] Step 2.1:

[0086] Initialize pipeline stage i in training a single microbatch, the sum of the maximum forward propagation time and the maximum backpropagation time (including recalculation time) is less than T mid The maximum number of continuous computing units that can be trained in this stage is l i，0 and l i，1 The initial search interval [l i，0low , l i，0up ) and [l i，1 low , l i，1up ), where l i，0 and l i，1 Respectively represent the FeedForward unit (structure asFigure 3 as shown) and with an Attention unit (structure as Figure 2 shown) as the maximum continuous computing unit number under the structure of the last computing unit of the sub-model at this stage.

[0087] Use l′ i to represent the scale of the sub-model in stage i under the final splitting scheme (the number of continuous computing units of the sub-model scale). When the video memory occupancy during sub-model training is less than or equal to the training video memory limit and the time consumption is less than or equal to T mid , if there exists a sum of the sub-model scales l′ i of stage i that satisfies the following equation, then the pipeline splitting scheme is feasible; otherwise, it is not feasible. Here, L is the number of layers of the large Transformer model to be trained:

[0088]

[0089] Therefore, by finding the values of l i，0 and l i，1 , the range of l′ i under the final splitting scheme can be determined.

[0090] Considering the problem that the recomputation strategy is affected by the training features of the current pipeline stage i and the scale of the sub-model of the current pipeline stage, resulting in a difficult-to-directly-predict specific mapping relationship between the model scales l i，0 and l i，1 and the corresponding single microbatch time consumption t(l i，0 ) and t(l i，1 ), so that l mid cannot be directly calculated through T i，0 and l i，1 .

[0091] Although it is affected as above and the values of l i，0 and l i，1 cannot be directly predicted, but since there is a positive correlation between l i，0 and l i，1 and the corresponding single microbatch time consumption t(l i，0 ) and t(l i，1 ), the binary search method can be used to find the corresponding maximum trainable sub-model scale that satisfies the constraints t(l i，0 ) ≤ T mid , t(l i，0 ) ≤ T mid . The specific problem and objective are expressed by the following formula, where and The specific recomputation strategy obtained by searching with the input of model configuration, training configuration, device information, the scale of the sub-model to be searched, and the information of the current stage number i for the call to the recomputation strategy search method. and is the time consumption of a single microbatch under the current recomputation strategy configuration:

[0092]

[0093] Specifically, by determining an interval [l i，0 low , l i，0 hight ) and [l i，1 low , l i，1high ), ensuring that the target l i，0 and l i，1 fall within this interval, and then using binary search to find the maximum value that satisfies the constraint as l i，0 and l i，1 .

[0094] In the first step, initialize the search intervals of l i，0 and l i，1 as [l i，0low , l i，0high ) and [l i，1low , l i，1high ). In the first execution of this step in this patent, L i，0low and L i，1low are set to 0, and l i，0high t and l i，1hight [[ID=4,9]]are set to twice the number of model layers L.

[0095] During the repeated search process, to reduce the search space, a memory module is added. During the process of searching for the maximum trainable sub-model scale l i，0 and l' i，1 , each time the recomputation strategy search unit is called, after obtaining the time consumption t(l'[[ID=5,7]] i，0 ) or t(l' i，1 ) of a single microbatch and the corresponding recomputation strategy for the input temporary variable l' i，0 or l' i，1 , l' i，0 and l' i，1 , as well as the time consumption of a single microbatch corresponding to both and the recomputation strategy, are respectively saved in two tables (one table saves the corresponding information of l' i，0 , and the other saves the corresponding information of l' i，1 ). Each row of the table contains the sub-model scale, the sum of the corresponding forward propagation time and backward propagation time (including recomputation time), and the corresponding recomputation strategy.

[0096] In the initialization process of the second and subsequent second-layer searches, it becomes to find in the memory module l′ i，0 or l′ i，1 in the corresponding table to find the number of computing units of the item with the largest number of computing units where t(l′ i，0 ) or t(l′ i，1 ) is less than T mid as l i，0low or l i，1low , and the number of computing units of the item with the smallest number of computing units where t(l′ i，0 ) or t(l′ i，1 ) is greater than T mid as l i，0high or l i，1high . If the corresponding item cannot be found in the table, it degrades to the initialization strategy during the first search.

[0097] Step 2.2:

[0098] Use the floor value of the middle value of the interval [l i，0 low , l i，0high initialized in 2.1 as the number of consecutive computing units l i,0mid included in the sub-model of the pipeline stage. And connect the computing units into the sub-model M′0 according to the structure of interleaved FeedForward units and Attention units respectively according to this number of units, and input the information such as the corresponding model configuration, model scale, current device, training configuration, and current stage number of the sub-model M′0 into the recomputation strategy search unit, and call the recomputation strategy search unit and its corresponding Cost Model to obtain the corresponding recomputation strategy and the time consumption of a single microbatch of the sub-model M′0 calculated according to the Cost Model under the corresponding recomputation strategy

[0099] Considering the overhead of the recomputation strategy search, during the calculation process, before calling the recomputation strategy search unit, query the table storing the information of the sub-model ending with FeedForward units in the memory module. If the required sub-model scale and the corresponding strategy are already in the table, directly use the data in the table. Otherwise, insert the corresponding information into the table storing the information of the sub-model ending with FeedForward units in the memory module after searching.

[0100] If then update the interval [l i，0 low , l i，0 high to [l i，0 mid , l i，0 high , otherwise update the interval to [l i，0low , l i，0mid .

[0101] Judge if i，0 low , i，0 high ) reaches the termination condition i，0 low + 1 = i，0 high . If the termination condition is reached, return i，0 low as i,0 ; if the termination condition is not reached, use the updated new interval to execute step 2.2 again.

[0102] Step 2.3:

[0103] Use the interval [ i,1low , i,1high ) initialized in 2.1, and take the floor of the middle value as the number of consecutive computing units included in the submodel of the pipeline stage i，1 mid , and connect the computing units into submodel M'1 according to the structure of interleaved FeedForward units and Attention units respectively according to the number of units, and input the information such as the corresponding model configuration, model scale, current device, training configuration, and current stage number of submodel M'1 into the recomputation policy search unit, and call the recomputation policy search unit and its corresponding Cost Model to obtain the corresponding recomputation policy and the time consumption of a single microbatch of submodel M'1 calculated according to the Cost Model under the corresponding recomputation policy

[0104] Considering the overhead of recomputation policy search, during the calculation process, before calling the recomputation policy search unit, query the table storing the information of the submodel ending with Attention unit in the memory module. If the required submodel scale and corresponding policy are already in the table, directly use the data in the table. Otherwise, insert the corresponding information into the table storing the information of the submodel ending with Attention unit in the memory module after searching.

[0105] If then update the interval [ i，1 low , i，1 high ) to [ i，1 mid , i，1 high ), otherwise update the interval to [ i，1 low , i，1 mid ).

[0106] Judge if i，1 low , i，1 high ) reaches the termination condition i，1 low + 1 = i，1 high . If the termination condition is reached, return i，1 low as i，1; If the termination condition is not met, then use the updated new interval to execute step 2.3 again.

[0107] Step Three:

[0108] After step two, each stage will obtain two solutions of the maximum trainable sub-model scale generated with the FeedForward unit and the Attention unit as the last layer structures: namely, l ending with the FeedForward unit i，0 and l ending with the Attention unit i，1 , specifically, for the last stage p, l i，1 is not calculated. This patent also adds a blank stage 0, and l i，1 is not calculated in the blank stage 0, i，0 and l is fixed to 0.

[0109] To verify whether the overall pipeline splitting is feasible under the T mid constraint, for all l of all stages obtained in the previous step i，s (i ∈ {0...p}, s ∈ {0, 1}), establish an edge from l i-1，j (numbered (i - 1, j)) to l i，k (numbered (i, k)) between the (i - 1)-th stage and the i-th stage according to the following edge weight rules, where i represents the stage number of pipeline parallelism, s represents the type of computing unit at the end of this stage (0 represents ending with the FeedForward unit, 1 represents ending with the Attention unit), l i，s represents the maximum number of consecutive computing units that can be trained under the constraint of the maximum single microbatch time consumption T mid of the i-th stage obtained in S2 ending with the computing unit represented by s, j, k represent the type of computing unit at the end of the corresponding stage (0 represents ending with the FeedForward unit, 1 represents ending with the Attention unit), the node number represented by l i-1，j is (i - 1, j), the node number represented by l i，k is (i, k), and connect them into a directed acyclic graph:

[0110] When j = 0 and k = 0, the edge weight is the largest even number less than or equal to l i，0 ;

[0111] When j = 0 and k = 1, the edge weight is the largest odd number less than or equal to l i，0 ;

[0112] When j = 1 and k = 0, the edge weight is the largest odd number less than or equal to l i，1 ;

[0113] j=1, k=1, edge weight is less than or equal to l i，1 The largest even number of ;

[0114] Execute the Dijkstra single-source pathfinding algorithm on the directed acyclic graph to find the longest path from the node numbered (0, 0) to the node numbered (p, 0) (the length is the sum of the weights of the edges on the path). If the path length is greater than or equal to 2L, it is considered that the node is in T mid Under the constraint, pipeline splitting is feasible, otherwise it is not feasible, where L is the number of layers of the large Transformer model to be trained. If splitting is feasible, update the search interval of T (T low ,T high ) is (T low ,T mid ), if splitting is not feasible, update the search interval of T (T low ,T high ) is (T mid ,T high ), where T is the desired goal: the theoretical minimum of the longest single microbatch duration in the pipeline stage.

[0115] Determine the termination condition T high -T low ≤δ, where δ is the hyperparameter set by the user. If the termination condition is met, return T high As Otherwise, use the updated search interval to execute steps 2 and 3 again.

[0116] Step 4:

[0117] After iteratively running the above steps, the approximate optimal solution of T can be found And select the one that satisfies the longest time less than or equal to In the last iteration of the pipeline splitting under the constraint, any path {(0,0),(1,s1),...(i,s1) in the directed acyclic graph with a length greater than or equal to 2L is i ),...(p,0)}(s i ∈{0, 1}) as the basis for splitting, and the model is split using this path when the following constraints are met: 1. The number of units in the i-th stage is less than (i-1, s i-1 ) to (i, s i ) and if k i = 0, the last unit structure of the sub-model is the FeedForward unit. If k iIf it is equal to 1, the last unit structure of the sub-model is an Attention unit; 2. The total number of units obtained by slicing in all stages is 2L; 3. The sub-model structure is an interleaved series of Attention units and FeedForward units.

[0118] The present invention also provides a specific embodiment to further illustrate the effect of the present invention. Among them, Megatron-LM is a distributed training framework developed by NVIDIA and most widely used in large model training at present, supporting pipeline parallelism and average splitting strategies; Lynx is one of the most advanced distributed training frameworks at present, supporting pipeline parallelism and an adaptive model splitting method;

[0119] The specific experiments are as follows:

[0120] Experimental configuration:

[0121] (1) Operating system: Ubuntu 22.04 LTS;

[0122] (2) CPU: 2.25GHz AMD(R) EPYC(R) 7742, equipped with 755GB DRAM;

[0123] (3) GPU: 8 * NVIDIA A800 40GB.

[0124] Model configuration:

[0125] (1) Model: Llama2-13B;

[0126] (2) Dataset: Wikitext2-v1, containing 36718 text data.

[0127] The final results are as Figure 6 shown. Compared with the average splitting method provided by the Megatron-LM framework, the method of this patent improves the system throughput by 24% - 38%. Compared with the Lynx framework, the method of this patent can also improve the system throughput by 24% - 38%.

[0128] The above are only the preferred embodiments of the present invention and are not intended to limit the present invention. Any modifications, equivalent replacements, improvements, etc. made within the spirit and principle of the present invention shall be included within the scope of protection of the present invention.

Claims

1. A recomputation perception model splitting system for pipeline parallel training, characterized by: The system comprises: Information Collection Module: This module is used to configure the model structure, parallel training strategy, and hyperparameters, deploy the defined model on a real GPU, run several pre-training iterations, and collect data. This data includes the number of model operators, computation time, output size, and operator dependencies. Parallel training strategies include Tensor parallelism and pipeline parallelism. Model Splitting Module: It is used to provide a recomputation policy-aware model splitting mechanism to achieve uniform workload distribution among multiple devices. The recomputation policy-aware model splitting mechanism is implemented by the Dijkstra algorithm and a two-layer binary search. The first layer of binary search performs a binary search on the approximate solution of the theoretical minimum of the maximum sum of the forward propagation and backpropagation time for training a microbatch in each stage of the pipeline, where the backpropagation time includes the recomputation time. The second layer of binary search is the first layer of binary search. mid As a constraint, calculate the maximum sub-model size that can be supported at each stage. According to the last calculation unit structure of the sub-model, binary search is performed in "with T mid The maximum sub-model size that can be supported at each stage under the constraint condition, where T mid A temporary variable representing the intermediate value of the interval during the first-level binary search process; the Dijkstra algorithm is used to determine whether the intermediate results of the first-level binary search process meet the requirements and to calculate the final result; Model training module: configures the model splitting strategy and the recalculation strategy of each device into the training strategy according to the model splitting mechanism.

2. The recomputation perception model splitting system for pipeline parallel training according to claim 1 is characterized in that: The model splitting module includes a sub-module, which is a memory module containing two tables. One table is used to store the intermediate results of the sub-model scale search ending with the FeedForward unit, and the other table is used to store the intermediate results of the sub-model scale search ending with the Attention unit. The intermediate results include the number of continuous calculation units contained in the sub-model and the recalculation strategy and the estimated calculation time of the sub-model.

3. The recomputation perception model splitting system for pipeline parallel training according to claim 2, characterized in that: The recalculation strategy-aware model splitting mechanism is specifically as follows: S1, the sum of the maximum forward propagation time and the backward propagation time for training a single microbatch in the initialization pipeline stage is the search interval (T low , T high ); the search interval is initially set to (T min , T max ); where T min The time taken for a single microbatch when the current complete model is evenly divided into each stage according to the number of layers and no recalculation strategy is used, T max T is the time it takes to train a single microbatch for the complete model using the full recalculation strategy under the current training and device configuration. The time it takes to train a single microbatch is the sum of the forward propagation and backpropagation time for each stage of the pipeline. low Indicates the lower bound of the search interval in the first-level binary search process, T high Indicates the upper bound of the search interval in the first-level binary search process; S2, use interval (T low , T high )Intermediate value T mid As the constraint on the maximum single microbatch time of the pipeline stage, for each stage i, search for the corresponding constraint T mid The maximum number of consecutive computing units that can be trained under ; S3, for all stages obtained in S2 i，s (i∈{0...p}, s∈{0,1}), establish an edge between the i-1th stage and the i-th stage from l i-1，j to l i，k The edge j, k∈{0, 1}, where i represents the stage number of the pipeline parallelization, s represents which computing unit the stage ends with, 0 represents the end with the FeedForward unit, and 1 represents the end with the Attention unit; l i，s Indicates that the i-th stage obtained by S2 ends with the computation unit represented by s, under the constraint T of the maximum single microbatch time mid The maximum number of consecutive computational units trained under the training phase; j, k indicate which computational unit the corresponding phase ends with, 0 indicates that it ends with the FeedForward unit, and 1 indicates that it ends with the Attention unit; l i-1，j The node number represented is (i-1, j), and l i，k The nodes represented are numbered (i, k) and are eventually connected into a directed acyclic graph: j=0,k=0,edge weight is less than or equal to l i，0 The largest even number of ; j=0,k=1,edge weight is less than or equal to l i，0 The largest odd number of j=1,k=0,edge weight is less than or equal to l i，1 The largest odd number of j=1, k=1, edge weight is less than or equal to l i，1 The largest even number of ; Execute the Dijkstra single-source pathfinding algorithm on the directed acyclic graph to find the longest path from the node ((0, 0) to the node (p, 0). If the length of the longest path is greater than or equal to 2L, then mid Under the constraint, pipeline splitting is feasible, otherwise it is not feasible, where L is the number of layers of the large Transformer model to be trained; if splitting is feasible, the search interval of T is updated (T low , T high ) is (T low , T mid ), if splitting is not feasible, update the search interval of T (T low , T high ) is (T mid , T high ), where T is the theoretical minimum time taken for the longest single microbatch in the pipeline stage; Determine the termination condition T high -T low ≤δ, where δ is the hyperparameter set by the user; if the termination condition is met, return T high As If it is not satisfied, repeat S2-S3 using the updated search interval, and iterate multiple times to obtain the approximate optimal solution of T S4, select the one that satisfies the longest time less than or equal to When the pipeline is in the last iteration where splitting is feasible, any path {(0,0),(1,s1),...(i,s i ),...(p,0)}(s i ∈{0,1}), as the basis for splitting and splitting the model.

4. The recomputation perception model splitting system for pipeline parallel training according to claim 3 is characterized in that: Said S1 comprises: The time t(l q ), and reversely calculate the number of Transformer computational units obtained by splitting each stage under the minimum limit; where the goal and constraints are expressed by the following formula, in which T is optimized as a parameter: t(l i )≤T(i∈{1...p}); By q ) binary search, find t(l q ) Approximate solution of the minimum value T 5. The recomputation perception model splitting system for pipeline parallel training according to claim 3 is characterized in that: The S2 includes: S2.

1. When initializing pipeline stage i, the sum of the maximum forward propagation time and the maximum backward propagation time for training a single microbatch is less than T mid The maximum number of continuous computing units that can be trained in this stage is l i，0 and l i，1 The initial search interval [l i，0 low , l i，0up ) and [l i，1 low , l i，1up ); where l i，0 and l i，1 They represent the maximum number of consecutive computing units in the structure with FeedForward unit and Attention unit as the last computing unit of the sub-model at this stage; Use l′ i Indicates the number of continuous computational units of the sub-model at stage i under the final splitting scheme; during sub-model training, the video memory usage satisfies the requirement of being less than or equal to the training video memory limit and the time consumed is less than or equal to T mid In the case of, if there is stage i, the sub-model size l′ i If the sum satisfies the following equation, the pipeline splitting scheme is feasible, otherwise it is not feasible, where L is the number of layers of the large Transformer model to be trained: By finding l i，0 and l i，1 The value of l′ is used to determine the final splitting solution. i scope; Use binary search to find the i，0 )≤T mid , t(l i，0 )≤T mid The corresponding maximum trainable sub-model size is expressed as follows: argmax l i,0 argmax l i,1 in, and To call the recalculation strategy search method, input the model configuration, training configuration, device information, the size of the sub-model to be searched, and the current stage number i to search for the specific recalculation strategy. and The time taken for a single microbatch under the current recalculation strategy configuration; By determining an interval [l i，0 low , l i，0 high ) and [l i，1 low , l i，1 high ), ensuring the target i，0 and l i，1 Falls within this interval, and then uses binary search to find the maximum value that meets the constraints as l i，0 and l i，1 ; S2.2, the [l i,0 low ,l i,0 high ) The intermediate value is rounded down and used as the number of continuous computing units contained in the sub-model of the pipeline stage l i,0mid , and according to the number of units, the computing units are connected into sub-models M′0 according to the structure of interleaved FeedForward units and Attention units, and the sub-model M′0 corresponding model configuration, model scale and current device, training configuration, and current stage number are input into the recalculation strategy search unit, and the recalculation strategy search unit and its corresponding CostModel are called to obtain the corresponding recalculation strategy And the single microbatch time of the sub-model M′0 calculated according to the Cost Model under the corresponding recalculation strategy S2.3, the [l i，1 low , l i，1 high ) The middle value is rounded down to the number of continuous computing units contained in the sub-model of the pipeline stage l i，1mid , and according to the number of units, the computing units are connected into sub-models M′1 according to the structure of interleaved FeedForward units and Attention units, and the sub-model M′1 corresponding model configuration, model scale and current device, training configuration, and current stage number are input into the recalculation strategy search unit, and the recalculation strategy search unit and its corresponding CostModel are called to obtain the corresponding recalculation strategy And the single microbatch time of the sub-model M′1 calculated according to the Cost Model under the corresponding recalculation strategy 6. The recomputation perception model splitting system for pipeline parallel training according to claim 5, characterized in that: By determining an interval [l i，0 low , l i，0 high ) and [l i，1 low , l i，1 high ), ensuring the target i，0 and l i，1 Falls within this interval, and then uses binary search to find the maximum value that meets the constraints as l i，0 and l i，1 include: Initialize l i，0 and l i，1 The search interval [l i，0 low , l i，0 high ) and [l i，1 low , l i，1 high ), i，0 low and l i，1 low Set to 0, i，0 high and l i，1 high Set to twice the number of model layers L; Through the memory module, we search for the maximum trainable sub-model size l i，0 and l i，1 In the process, the recalculation strategy search unit is called and the temporary variable l′ is input. i，0 or l′ i，1 Get the time t(l′) for a single microbatch i，0 ) or t(l′ i，1 ) and the corresponding recalculation strategy; l′ i，0 and l′ i，1 The corresponding single microbatch time and recalculation strategy are stored in two tables respectively. Each row of the table contains the sub-model size, the sum of the corresponding forward propagation time and backward propagation time, and the corresponding recalculation strategy; The initialization process of the second and subsequent second-level searches is l′ in the memory module i，0 or l′ i，1 Find t(l′ in the corresponding table i，0 ) or t(l′ i，1 ) is less than T mid The number of calculation units of the maximum term is taken as l i，0 low or l i，1 low , and t(l′ i，0 ) or t(l′ i，1 ) is greater than T mid The number of calculation units of the minimum term is the number of calculation units of l i，0 high or l i,1 high If no corresponding item is found in the table, it degenerates into the initialization strategy of the first search.

7. The recomputation perception model splitting system for pipeline parallel training according to claim 5, characterized in that: Said S2.2 includes: exist During the calculation process, before calling the recalculation strategy search unit, the table storing the information of the sub-model ending with the FeedForward unit in the memory module is queried. If the required sub-model size and corresponding strategy are already in the table, the data in the table is used directly. Otherwise, the corresponding information is inserted into the table storing the information of the sub-model ending with the FeedForward unit in the memory module after the search. like Then update the interval [l i，0 low , l i，0 high ) is [l i，0 mid , l i，0 high ), otherwise the update interval is [l i，0 low , l i，0 mid ); Judgment i，0 low , l i，0 high ) Whether the termination condition l is met i，0 low +1=l i，0 high ; If the termination condition is reached, return l i，0 low As l i，0 ; If the termination condition is not met, execute S2.2 again using the updated new interval.

8. The recomputation perception model splitting system for pipeline parallel training according to claim 5, characterized in that: Said S2.3 includes: exist During the calculation process, before calling the recalculation strategy search unit, the table storing the information of the sub-models ending with the Attention unit in the memory module is queried. If the required sub-model size and corresponding strategy are already in the table, the data in the table is used directly. Otherwise, the corresponding information is inserted into the table storing the information of the sub-models ending with the Attention unit in the memory module. like Then update the interval [l i，1 low , l i，1 high ) is [l i，1 mid , l i，1 high ), otherwise update the interval to [l i，1 low , l i，1 mid ); Judgment i，1 low , l i，1 high ) Whether the termination condition l is met i，1 low +1=l i，1 high ; If the termination condition is reached, return l i，1 low As l i，1 ; If the termination condition is not met, execute S2.3 again using the updated interval.

9. The recomputation perception model splitting system for pipeline parallel training according to claim 3, characterized in that: In S4, splitting the model includes: The following constraints are met: The number of units in the i-th stage is less than (i-1, s i-1 ) to (i, s i ) and if k i = 0, then the last unit structure of the sub-model is the FeedForward unit. If k i =1, then the last unit structure of the sub-model is the Attention unit; The total number of units obtained by segmentation at all stages is 2L; The sub-model structure is an interlaced series of Attention units and FeedForward units.

10. A method for splitting a heavy-computation perception model for pipeline parallel training, characterized in that: The method comprises: Collect data for the defined model, including the number of model operators, computation time, output size, and operator dependencies. The collected data is used to formulate subsequent recomputation strategies and model splitting strategies. The workload is evenly distributed across multiple devices based on the recalculation strategy-aware model splitting mechanism. The recalculation strategy-aware model splitting mechanism is implemented using the Dijkstra algorithm and a two-layer binary search. The first layer of binary search performs a binary search on the minimum sum of the forward propagation and backpropagation time for training a microbatch at each stage of the pipeline. The second layer of binary search is the first layer of binary search. mid As a constraint, calculate the maximum sub-model size that can be supported at each stage. According to the last calculation unit structure of the sub-model, binary search is performed in "with T mid The maximum sub-model size that can be supported at each stage under the constraint condition is ""; the Dijkstra algorithm is used to determine whether the intermediate results of the first-level binary search process meet the requirements, and is used to calculate the final result; wherein, T mid A temporary variable representing the middle value of the interval during the first-level binary search process; Configure the model splitting strategy and the recomputation strategy for each device into the training strategy.

Images