Assembly line parallel training-oriented re-calculation perception model splitting method and assembly line parallel training-oriented re-calculation perception model splitting system

Abstract

The invention discloses a recalculation perception model splitting method and system oriented to assembly line parallel training, and belongs to the field of artificial intelligence of computer science. The method specifically comprises the steps that a dijkstra algorithm and two-layer binary search are used, and the minimum value of the sum of forward propagation time and reverse propagation time of assembly line microbatch under re-calculation perception model splitting is calculated; and an intermediate result obtained by a dijkstra algorithm in a search process is used to calculate and output an approximately optimal model splitting strategy for re-calculation perception. According to the method, the model splitting method and system are designed by utilizing the model characteristics of a Tasformer large model and the characteristics of an assembly line parallel mode, the model is non-uniformly split to each stage of the assembly line, so that the calculation load of each stage of the assembly line is balanced when re-calculation is started, and the calculation efficiency is improved. Therefore, the overall performance of large model pipeline parallel training and the utilization efficiency of cluster equipment are improved.

Description

Technical Field

[0001] The present invention relates to the field of computer science and artificial intelligence, and in particular to a method and system for splitting a heavy computing perception model for pipeline parallel training. Background Art

[0002] With the development of artificial intelligence technology, the size and complexity of models continue to grow, and more powerful performance and more accurate prediction results are obtained. These large models, such as Transformer 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 places a huge demand on the capacity of video memory, making it very difficult to train these models on a single GPU or CPU. In order to overcome the limitation of video memory capacity, researchers and engineers have developed model splitting technology. The core idea of ​​this technology is to split a large model into multiple smaller sub-models, each of which can be calculated independently on different hardware devices, forming a pipeline with each other to efficiently exert the potential of the device. However, model splitting needs to meet the balance of computing load on each device, otherwise the synchronization operation between devices will reduce the overall training efficiency.

[0003] In a pipeline-parallel distributed training environment, the computational load on different devices may become uneven due to the application of heavy computing techniques.

[0004] In the case of uniform slicing, the computing load on each device in the parallel pipeline is equal, and no device idle time (bubble) will be generated. However, due to the different pipeline stages in which the devices are located, the load of their video memory is also different. In order to prevent video memory overflow, the latest training frameworks (such as Megatron, Deepspeed, etc.) allow users to use recalculation technology. Some new technologies (such as Lynx, AdaPipe, etc.) also propose to use corresponding recalculation strategies according to different pipeline stages. However, directly using different recalculation strategies under the uniformly split pipeline will result in different recalculation time overheads at different stages, which will further cause a lot of idle time for the devices in the pipeline.

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

[0006] In order to solve the problem of low model training hardware performance utilization caused by inefficient segmentation of existing Transformer large model training, the purpose of the present invention is to provide a heavy computing perception model splitting method and system for pipeline parallel training, which balances the time overhead of each stage of pipeline parallel training to reduce the idle time of the device when using pipeline parallel training and improve the overall device utilization.

[0007] The object of the present invention is achieved through the following technical solution: a recalculation perception model splitting system for pipeline parallel training, the system comprising: Information collection module: used to configure the model structure, parallel training strategy and hyperparameters (Global BatchSize, Micro Batch Size), deploy the defined model on a real GPU, pre-run several rounds of training iterations and collect data; the data includes the number of model operators, calculation time, output size and operator dependency; the parallel training strategy includes: Tensor parallelism and pipeline parallelism settings; Model splitting module: used to provide a recomputation strategy-aware model splitting mechanism to achieve uniform distribution of workloads among multiple devices; the recomputation strategy-aware model splitting mechanism is implemented through the Dijkstra algorithm and two-layer binary search, wherein the first-layer binary search performs a binary search on the approximate solution of the theoretical minimum value of the sum of the forward propagation and backpropagation time (including recomputation time) of training a microbatch in each stage of the pipeline; the second-layer binary search is the "with" in the first-layer binary search 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 The maximum sub-model size that can be supported at each stage under the "constraint condition", where is a temporary variable representing the intermediate value of the interval in the first-level binary search process; the Dijkstra algorithm is used to determine whether the intermediate result of the first-level binary search process meets the requirements, and is used 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.

[0008] Furthermore, 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 a FeedForward unit, and the other table is used to store the intermediate results of the sub-model scale search ending with an Attention unit, and 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.

[0009] Furthermore, the recalculation strategy-aware model splitting mechanism is specifically: S1, the search interval for the sum of the maximum forward propagation time and the backward propagation time (including recalculation time) for training a single microbatch in the initialization pipeline stage ; Initially set the search interval to ,in The time it takes to divide the current complete model into each stage evenly according to the number of layers and without using any recalculation strategy for a single microbatch. It is the time taken for a single microbatch of the current complete model under the current training and device configuration using the full recalculation strategy. The time taken for a single microbatch represents the sum of the forward propagation and backpropagation time for training a microbatch at each stage of the pipeline. Indicates the lower bound of the search interval in the first-level binary search process, Indicates the upper bound of the search interval in the first-level binary search process; S2. Use interval Median As the constraint on the maximum single microbatch time of the pipeline stage, for each stage i, search for the corresponding constraint The maximum number of consecutive computing units that can be trained under ; S3, to verify Is it feasible to split the entire pipeline under the restriction? For all stages obtained in S2 , according to the following edge weight rule, a link is established between the i-1th stage and the i-th stage. arrive Edge , where i represents the stage number of the pipeline parallelism, and 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). Represents the constraint on the maximum single microbatch time consumed by the computation unit represented by s at the end of the i-th stage obtained by S2 The maximum number of consecutive computational units that can be trained under , j, k indicate which computational unit the corresponding stage ends with (0 means it ends with FeedForward unit, 1 means it ends with Attention unit), The node number represented is (i-1, j), 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 The largest even number of ; j=0, k=1, edge weight is less than or equal to The largest odd number of ; j=1, k=0, edge weight is less than or equal to The largest odd number of ; j=1, k=1, edge weight is less than or equal to 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 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 length of the longest path is greater than or equal to , then it is believed that Under the restriction, the pipeline split is feasible, but not otherwise. The number of layers of the Transformer model to be trained; if splitting is feasible, update Search interval for , if splitting is not feasible, update Search interval for ,in The goal is: the theoretical minimum time taken by the longest single microbatch in the pipeline stage; Determine the termination condition , where δ is the hyperparameter set by the user; if the termination condition is met, it returns As If it is not satisfied, repeat S2-S3 with the updated search interval, and iterate multiple times to obtain The approximate optimal solution ; S4, select the longest time less than or equal to In the last iteration of the pipeline splitting under the constraint, the length of any line in the directed acyclic graph is greater than or equal to Path As the basis for splitting and splitting the model.

[0010] Further, the S1 includes: The longest phase time required to search for the training memory limit Minimum value of , and reversely calculate the number of Transformer computing units obtained by splitting each stage under the minimum limit; where the goal and constraint are expressed by the following formula, As a parameter to participate in optimization: ; Through Binary search finds Minimum value of The approximate solution of .

[0011] Further, the S2 includes: S2.1, Initialization pipeline stage The sum of the maximum forward propagation time and the backward propagation time in training a single microbatch is less than The maximum number of continuous computing units that can be trained in this stage and The initial search interval and ,in, and 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 in this stage respectively; use Indicates the next stage of the final split plan The sub-model size (the number of continuous computing units of the sub-model) is such that when the sub-model training is satisfied, the video memory usage is less than or equal to the training video memory limit and the training time is less than or equal to In the case of , sub-model size If the sum satisfies the following equation, the pipeline splitting scheme is feasible, otherwise it is not feasible, where The number of layers of the Transformer model to be trained: ; By looking for and The value of scope; Use binary search to find the , The corresponding maximum trainable sub-model size is expressed by the following formula: ; in, To call the recalculation strategy search method, enter the model configuration, training configuration, device information, sub-model size to be searched, and current stage number The specific recalculation strategy obtained by information search, It is the time taken for a single microbatch under the current recalculation strategy configuration; By determining an interval and , to ensure the target and Falls within this interval, and then uses binary search to find the maximum value that satisfies the constraints as and ; S2.2. Use the interval initialized in S2.1 The intermediate value is rounded down as the number of consecutive computing units contained in the pipeline stage sub-model , and according to the number of units, the computational units are connected into sub-models according to the structure of interleaved FeedForward units and Attention units. , the sub-model Corresponding to the model configuration, model scale and current device, training configuration, and current stage number, input the recalculation strategy search unit, call the recalculation strategy search unit and its corresponding Cost Model to obtain the corresponding recalculation strategy And the sub-model calculated according to the Cost Model under the corresponding recalculation strategy A single microbatch takes ; S2.3. Use the interval initialized in S2.1 The intermediate value is rounded down as the number of consecutive computing units contained in the pipeline stage sub-model , and according to the number of units, the computational units are connected into sub-models according to the structure of interleaved FeedForward units and Attention units. , and the submodel Corresponding to the model configuration, model scale and current device, training configuration, and current stage number, input the recalculation strategy search unit, call the recalculation strategy search unit and its corresponding Cost Model to obtain the corresponding recalculation strategy And the sub-model calculated according to the Cost Model under the corresponding recalculation strategy A single microbatch takes .

[0012] Furthermore, by determining an interval and , to ensure the target and Falls within this interval, and then uses binary search to find the maximum value that satisfies the constraints as and include: initialization and Search interval and ,Will and Set to 0, and Set to the number of model layers twice as much; In order to reduce the search space during the repeated search process, the memory module is used to search for the maximum trainable sub-model size. and The recalculation strategy search unit will be called each time during the process, and the temporary variable will be entered or The time taken for a single microbatch is or And after the corresponding recalculation strategy, and The corresponding single microbatch time consumption and recalculation strategy are stored in two tables (one table stores Corresponding information, another save Corresponding information), each row of the table contains the sub-model size, the sum of the corresponding forward propagation time and backpropagation time (including recalculation time), and the corresponding recalculation strategy; The initialization process for the second and subsequent second-level searches is to or Find the corresponding table or Less than The number of computational units of the item with the largest number of computational units is or ,as well as or Greater than The number of computational units of the item with the smallest number of computational units is or If the corresponding item is not found in the table, it degenerates to the initialization strategy during the first search.

[0013] Furthermore, the S2.2 includes: Considering the overhead of recalculating the strategy search, 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 scale and the 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 Update interval for , otherwise the update interval is ; judge Whether the termination condition is met ; If the termination condition is reached, return As ; If the termination condition is not met, execute S2.2 again using the updated new interval.

[0014] Furthermore, the S2.3 includes: Considering the overhead of recalculating the strategy search, During the calculation process, before calling the recalculation strategy search unit, the table storing the information of the sub-model ending with the Attention unit in the memory module is queried. If the required sub-model scale 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 Attention unit in the memory module after the search. like Update interval for , otherwise the update interval is ; judge Whether the termination condition is met ; If the termination condition is reached, return As ; If the termination condition is not met, execute S2.3 again using the updated new interval.

[0015] Further, in S4, splitting the model includes: The following constraints are met: The number of units in the i-th stage is less than arrive The length of the side, and if , then the last unit structure of the sub-model is the FeedForward unit. , then the last unit structure of the sub-model is the Attention unit; The total number of units obtained from all stages is ; The sub-model structure is an alternating series of Attention units and FeedForward units.

[0016] The present invention also provides a recalculation perception model splitting method for pipeline parallel training, the method comprising: Collect data for the defined model, including the number of model operators, calculation time, output size, and operator dependency; the collected data is used for the formulation of subsequent recalculation strategies and model splitting strategies; According to the recomputation strategy-aware model splitting mechanism, the workload is evenly distributed among multiple devices; the recomputation strategy-aware model splitting mechanism is implemented by the Dijkstra algorithm and two-layer binary search, wherein the first-layer binary search performs a binary search on the minimum sum of the forward propagation and backpropagation time of training a microbatch at each stage of the pipeline; the second-layer binary search is the first-layer binary search "based on 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 The maximum sub-model size that can be supported at each stage under the "constraint condition"; 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; Configure the model splitting strategy and the recalculation strategy for each device into the training strategy.

[0017] The beneficial effect of the present invention is that by using mathematical formulas and graph structures to establish a mathematical model for the large model pipeline training process, the system can reasonably configure the number of computing units in each stage of the pipeline in a computationally aware manner, ultimately balancing the load of each stage of the training pipeline, increasing the system throughput by 24%-38% compared to the average split, and increasing the system throughput by 3%-6% compared to the split strategy of Lynx, one of the most advanced technologies currently. BRIEF DESCRIPTION OF THE DRAWINGS

[0018] Figure 1 is a system architecture diagram of the present invention; Figure 2 This is the structure diagram of the Attention unit; Figure 3 This is the structure diagram of the FeedForward unit; Figure 4 Search the flowchart for the minimum time spent in the longest single microbatch phase; Figure 5 for Flowchart for searching the maximum model size of the corresponding pipeline stage i under the restriction; Figure 6 This is a test result diagram of the present invention. DETAILED DESCRIPTION

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

[0020] The Transformer model is structurally composed of a multi-layer Transformer Layer structure, and each layer of the Transformer structure is composed of two units, Attention and FeedForward. Considering that the pipeline parallel delay is affected by the communication time between pipeline stages, the Transformer model is often split according to the Transformer Layer structure or the Attention / FeedForward computing unit in the past pipeline splitting method. 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 according to the Attention / FeedForward computing unit.

[0021] The scenario targeted by the present invention is to use distributed GPU for training, and use pipeline parallelism and recalculation methods at the same time during training. The overall structure is designed as follows Figure 1 As shown, it specifically includes the following three modules.

[0022] Information collection module: used to configure the model structure, parallel training strategy (including Tensor parallelism and pipeline parallelism settings) and hyperparameters (Global Batch Size, Micro Batch Size), deploy the defined model on a real GPU, pre-run several rounds of training iterations and collect data; the data includes the number of model operators, computing time, output size and operator dependencies.

[0023] Model splitting module: used to provide a recomputation strategy-aware model splitting mechanism to achieve uniform distribution of workloads among multiple devices; the recomputation strategy-aware model splitting mechanism is implemented through the Dijkstra algorithm and two-layer binary search, wherein the first-layer binary search performs a binary search on the approximate solution of the theoretical minimum value of the sum of the forward propagation and backpropagation time (including recomputation time) of training a microbatch in each stage of the pipeline; the second-layer binary search is the "with" in the first-layer binary search 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 The maximum sub-model size that can be supported at each stage under the "constraint condition", where It is a temporary variable representing the intermediate value of the interval in the first-level binary search process; the Dijkstra algorithm is used to determine whether the intermediate result of the first-level binary search process meets the requirements, and is used to calculate the final result.

[0024] Model training module: The model splitting strategy and the recalculation strategy of each device are configured into the training strategy, and the system framework deploys the model to the GPU device for training.

[0025] The recomputation strategy-aware model splitting mechanism is implemented through the Dijkstra algorithm and 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 (including recomputation time) of training a microbatch at each stage of the pipeline. The process of the first layer is as follows: Figure 4 As shown; the second level search is the first level "with As a constraint, calculate the maximum sub-model size that can be supported at each stage. This layer of search is based on the last calculation unit structure of the sub-model. The binary search is performed in the " Under the "constraint condition", the maximum sub-model size that can be supported at each stage, the second-level process is as follows Figure 5 As shown in the figure, the Dijkstra algorithm is mainly used in the recalculation strategy-aware model splitting mechanism to determine whether the intermediate results of the first-level binary search process meet the requirements, and to calculate the final results.

[0026] The following is a detailed description of the recalculation strategy-aware model splitting mechanism: Step 1: Initialize the pipeline stage to train a single microbatch. The maximum forward propagation time and the sum of the backward propagation time (including recalculation time) search interval .

[0027] Specifically, under pipeline parallelism, the optimal splitting scheme of the model requires that the sum of the forward propagation and backward propagation time (including recalculation time) of training a microbatch at each stage of the pipeline (hereinafter referred to as the single microbatch time) is equal, that is: .

[0028] in For stage The scale of the sub-model is The time taken for a single microbatch is is the total number of parallel stages in the pipeline.

[0029] If a splitting scheme does not satisfy this equation, the stage with a shorter time for a single microbatch needs to wait for the stage with a longer time to complete the calculation task during the pipeline parallel training process. 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 by algorithm until the pipeline parallelization reaches the expected efficiency.

[0030] This patent observed that the overall training time and the longest phase of a single microbatch A single microbatch takes Positive correlation. Thus, the longest phase can be minimized by Starting from the beginning, we optimize the overall efficiency of the pipeline. The difficulty here is the longest stage under different splitting schemes. It is not a fixed stage. The longest stage under different pipeline splitting strategies is also different. Therefore, it is necessary to optimize the overall splitting plan from a global perspective to obtain the sub-model scale of each stage.

[0031] Globally, the time taken for a single microbatch in each stage must satisfy the following inequality group: .

[0032] Therefore, we can find the longest phase time that satisfies the training memory limit by searching Minimum value of And reversely calculate the number of Transformer computing units obtained by splitting each stage under the minimum limit, that is, the size of the sub-model at each stage. So far, the optimization method starts from first determining the size of the sub-model at each stage and then making The minimum becomes to find first The minimum value of determines the scale of each sub-model. The specific objectives and constraints are expressed by the following formula: Also participates in optimization as a parameter: .

[0033] Through direct Binary search can find Minimum value of The approximate solution of .

[0034] The first step is to determine a search range , the present invention initially sets this interval to ,in The time it takes to divide the current complete model into each stage evenly according to the number of layers and without using any recalculation strategy for a single microbatch. The time it takes to run a single microbatch for the current complete model using the full recalculation strategy for the current training and device configuration.

[0035] Step 2: Use interval Median As the constraint on the maximum single microbatch time of the pipeline stage, for each stage i, search for the corresponding constraint The maximum number of consecutive computing units that can be trained under this condition is , and the search process includes the following sub-steps.

[0036] Step 2.1: Initialize pipeline stages The sum of the maximum forward propagation time and the maximum backward propagation time (including recalculation time) for training a single microbatch is less than The maximum number of continuous computing units that can be trained in this stage and The initial search interval and ,in, and Respectively represent the FeedForward unit (structure such as Figure 3 As shown in the figure) and the Attention unit (structure as Figure 2 As shown in Figure 2, the maximum number of consecutive computing units under the last computing unit structure of the sub-model of this stage.

[0037] use Refers to the next stage of the final split plan The sub-model scale (the number of continuous computing units of the sub-model scale) is such that when the sub-model training is satisfied, the video memory usage is less than or equal to the training video memory limit and the time consumption is less than or equal to In the case of Submodel size If the sum satisfies the following equation, the pipeline splitting scheme is feasible, otherwise it is not feasible, where The number of layers of the Transformer model to be trained: .

[0038] Therefore, by finding and The value of can determine the final splitting solution range.

[0039] Consider the recalculation strategy affected by the current pipeline stage The model size is generated by the training features and the size of the sub-model in the current pipeline stage and Time consumption for a single microbatch and It is difficult to directly predict the specific mapping relationship between Calculated and .

[0040] Although the above effects cannot be directly predicted and , but due to and Time consumption for a single microbatch and There is a positive correlation between them, so binary search can be used to find , The corresponding maximum trainable sub-model size, the specific problem and goal are expressed by the following formula, where To call the recalculation strategy search method, enter the model configuration, training configuration, device information, sub-model size to be searched, and current stage number The specific recalculation strategy obtained by searching for information such as The time taken for a single microbatch under the current recalculation strategy configuration is: .

[0041] Specifically, by determining an interval and , to ensure the target and Falls within this interval, and then uses binary search to find the maximum value that satisfies the constraints as and .

[0042] The first step is initialization and Search interval and , this patent will be executed when this step is first performed and Set to 0, and Set to the number of model layers twice as much.

[0043] In order to reduce the search space when repeating the search process, a memory module is added to search for the maximum trainable sub-model size. and The recalculation strategy search unit will be called each time during the process, and the temporary variable will be entered or The time taken for a single microbatch is or And after the corresponding recalculation strategy, and The corresponding single microbatch time consumption and recalculation strategy are stored in two tables (one table stores Corresponding information, another save Corresponding information), each row of the table contains the sub-model size, the sum of the corresponding forward propagation time and backpropagation time (including recalculation time), and the corresponding recalculation strategy.

[0044] The initialization process of the second and subsequent second-level searches becomes in the memory module or Find the corresponding table or Less than The number of computational units of the item with the largest number of computational units is or ,as well as or Greater than The number of computational units of the item with the smallest number of computational units is or If the corresponding item is not found in the table, it degenerates to the initialization strategy during the first search.

[0045] Step 2.2: Use the interval initialized in 2.1 The intermediate value is rounded down as the number of consecutive computing units contained in the pipeline stage sub-model , and according to the number of units, the computational units are connected into sub-models according to the structure of interleaved FeedForward units and Attention units. , and the submodel The corresponding model configuration, model scale, current device, training configuration, current stage number and other information are input into the recalculation strategy search unit, and the recalculation strategy search unit and its corresponding Cost Model are called to obtain the corresponding recalculation strategy. And the sub-model calculated according to the Cost Model under the corresponding recalculation strategy A single microbatch takes .

[0046] Considering the overhead of recalculating the strategy search, 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 scale and the 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.

[0047] like Update interval for , otherwise the update interval is .

[0048] judge Whether the termination condition is met . If the termination condition is met, return As ; If the termination condition is not met, execute step 2.2 again using the updated new interval.

[0049] Step 2.3: Use the interval initialized in 2.1 The intermediate value is rounded down as the number of consecutive computing units contained in the pipeline stage sub-model , and according to the number of units, the computational units are connected into sub-models according to the structure of interleaved FeedForward units and Attention units. , and the submodel The corresponding model configuration, model scale, current device, training configuration, current stage number and other information are input into the recalculation strategy search unit, and the recalculation strategy search unit and its corresponding Cost Model are called to obtain the corresponding recalculation strategy. And the sub-model calculated according to the Cost Model under the corresponding recalculation strategy A single microbatch takes .

[0050] Considering the overhead of recalculating the strategy search, During the calculation process, before calling the recalculation strategy search unit, the table storing the information of the sub-model ending with the Attention unit in the memory module is queried. If the required sub-model scale 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 Attention unit in the memory module after the search.

[0051] like Update interval for , otherwise the update interval is .

[0052] judge Whether the termination condition is met . If the termination condition is met, return As ; If the termination condition is not met, execute step 2.3 again using the updated new interval.

[0053] Step 3: After step 2, each stage will obtain the solutions of the two maximum trainable sub-model sizes generated by the FeedForward unit and the Attention unit as the last layer structure: and ending with an Attention unit , special, the last stage Not calculated This patent also adds a blank stage 0, which is not calculated. ,and Fixed to 0.

[0054] To verify Is it feasible to split the entire pipeline under the constraints? , according to the following edge weight rule, a link is established between the i-1th stage and the i-th stage. (numbered (i-1, j)) to The edge (numbered (i, k)) ), where i represents the stage number of the pipeline parallelism, and 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). Represents the constraint on the maximum single microbatch time consumed by the computation unit represented by s at the end of the i-th stage obtained by S2 The maximum number of consecutive computational units that can be trained under , j, k indicate which computational unit the corresponding stage ends with (0 means it ends with FeedForward unit, 1 means it ends with Attention unit), The node number represented is (i-1, j), The nodes represented are numbered (i, k) and connected into a directed acyclic graph: j=0, k=0, edge weight is less than or equal to The largest even number of ; j=0, k=1, edge weight is less than or equal to The largest odd number of ; j=1, k=0, edge weight is less than or equal to The largest odd number of ; j=1, k=1, edge weight is less than or equal to The largest even number of .

[0055] 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 It is believed that Under the restriction, the pipeline split is feasible, but otherwise it is not feasible. The number of layers of the Transformer model to be trained. Update if splitting is feasible Search interval for , if splitting is not feasible, update Search interval for ,in The goal is to achieve the theoretical minimum time required for the longest single microbatch in the pipeline stage.

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

[0057] Step 4: After iterating through the above steps, you can find The approximate optimal solution , 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, the length of any line in the directed acyclic graph is greater than or equal to Path As the basis for splitting, 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 arrive The length of the side and if Then the last unit structure of the sub-model is the FeedForward unit. If Then the last unit structure of the sub-model is the Attention unit; 2. The total number of units obtained by segmentation in all stages is ; 3. The sub-model structure is an alternating series of Attention units and FeedForward units.

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

[0059] The specific experiments are as follows: Experimental configuration: (1) Operating system: Ubuntu 22.04 LTS; (2) CPU: 2.25 GHz AMD(R) EPYC(R) 7742 with 755GB DRAM; (3) GPU: 8*NVIDIA A800 40GB.

[0060] Model configuration: (1) Model: Llama2-13B; (2) Dataset: Wikitext2-v1, containing 36,718 text data.

[0061] The final result is Figure 6 As 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%.

[0062] The above description is only a preferred embodiment of the present invention and is not intended to limit the present invention. Any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the scope of protection of the present invention.

Claims

1. A recomputation perception model splitting system for pipeline parallel training, characterized in that: The system comprises: Information collection module: used to configure the model structure, parallel training strategy and hyperparameters, deploy the defined model on a real GPU, pre-run several rounds of training iterations and collect data; the data includes the number of model operators, calculation time, output size and operator dependency; the parallel training strategy includes: Tensor parallelism and pipeline parallelism settings; Model splitting module: used to provide a recomputation strategy-aware model splitting mechanism to achieve uniform distribution of workloads among multiple devices; the recomputation strategy-aware model splitting mechanism is implemented by the Dijkstra algorithm and two-layer binary search, wherein the first-layer binary search performs a binary search on the approximate solution of the theoretical minimum value of the maximum value of the sum of the forward propagation and backpropagation time for training a microbatch in each stage of the pipeline, and the backpropagation time includes the recomputation time; the second-layer binary search is the first-layer binary search "with 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 The maximum sub-model size that can be supported at each stage under the "constraint condition", where is a temporary variable representing the intermediate value of the interval in the first-level binary search process; the Dijkstra algorithm is used to determine whether the intermediate result of the first-level binary search process meets the requirements, and is used 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. According to the recomputation perception model splitting system for pipeline parallel training according to claim 1, it is characterized in that: The model splitting module includes a submodule, 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 model splitting mechanism of the recalculation strategy awareness is specifically as follows: S1, the sum of the maximum forward propagation time and the reverse propagation time for training a single microbatch in the initialization pipeline stage is the search interval ; Initially set the search interval to ;in, The time it takes to divide the current complete model into each stage evenly according to the number of layers and without using any recalculation strategy for a single microbatch. It is the time taken for a single microbatch of the current complete model under the current training and device configuration using the full recalculation strategy. The time taken for a single microbatch represents the sum of the forward propagation and backpropagation time for training a microbatch at each stage of the pipeline. Indicates the lower bound of the search interval in the first-level binary search process, Indicates the upper bound of the search interval in the first-level binary search process; S2. Use interval Median As the constraint on the maximum single microbatch time of the pipeline stage, for each stage i, search for the corresponding constraint The maximum number of consecutive computing units that can be trained under ; S3, for all stages obtained in S2 , according to the following edge weight rule, a link is established between the i-1th stage and the i-th stage. arrive Edge , where i represents the stage number of the pipeline parallelism, 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; Indicates that the i-th stage obtained by S2 ends with the computational unit represented by s, within the constraint of the maximum single microbatch time The maximum number of consecutive computational units trained under ; j, k indicate which computational unit the corresponding stage ends with, 0 means it ends with FeedForward unit, 1 means it ends with Attention unit; The node number represented is (i-1, j), 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 The largest even number of ; j=0, k=1, edge weight is less than or equal to The largest odd number of ; j=1, k=0, edge weight is less than or equal to The largest odd number of ; j=1, k=1, edge weight is less than or equal to 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 numbered (0, 0) to the node numbered (p, 0). If the length of the longest path is greater than or equal to , then in Under the restriction, the pipeline split is feasible, but not otherwise. The number of layers of the Transformer model to be trained; if splitting is feasible, update Search interval for , if splitting is not feasible, update Search interval for ,in The theoretical minimum time taken for the longest single microbatch in the pipeline stage; Determine the termination condition , where δ is the hyperparameter set by the user; if the termination condition is met, it returns As If it is not satisfied, repeat S2-S3 with the updated search interval, and iterate multiple times to obtain The approximate optimal solution ; S4, select the longest time less than or equal to When the pipeline is in the last iteration where splitting is feasible, any line in the directed acyclic graph has a length greater than or equal to Path , as the basis for splitting and splitting the model.

4. The recomputation perception model splitting system for pipeline parallel training according to claim 3, characterized in that: The S1 includes: Search for the longest time in the phase that meets the training memory limit Minimum value of , and reversely calculate the number of Transformer computing units obtained by splitting each stage under the minimum limit; where the goal and constraint are expressed by the following formula, As a parameter to participate in optimization: ； Through Binary search, find Minimum The approximate solution of .

5. The recomputation perception model splitting system for pipeline parallel training according to claim 3, characterized in that: The S2 includes: S2.1, Initialization pipeline stage When training a single microbatch, the sum of the maximum forward propagation time and the maximum backpropagation time is less than The maximum number of continuous computing units that can be trained in this stage and The initial search interval and ;in, and 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 in this stage respectively; use Indicates the next stage of the final split plan The number of continuous computing units of the sub-model; when training the sub-model, the video memory usage is less than or equal to the training video memory limit and the time consumed is less than or equal to In the case of , sub-model size If the sum satisfies the following equation, the pipeline splitting scheme is feasible, otherwise it is not feasible, where The number of layers of the Transformer model to be trained: ； By looking for and The value of scope; Use binary search to find the , The corresponding maximum trainable sub-model size is expressed by the following formula: ； in, To call the recalculation strategy search method, enter the model configuration, training configuration, device information, sub-model size to be searched, and current stage number The specific recalculation strategy obtained by information search, It is the time taken for a single microbatch under the current recalculation strategy configuration; By determining an interval and , to ensure the target and Falls within this interval, and then uses binary search to find the maximum value that satisfies the constraints as and ; S2.

2. The intermediate value is rounded down to the nearest integer and used as the number of consecutive computing units contained in the sub-model of the pipeline stage. , and according to the number of units, the computational units are connected into sub-models according to the structure of interleaved FeedForward units and Attention units. , the sub-model Corresponding to the model configuration, model scale and current device, training configuration, and current stage number, input the recalculation strategy search unit, call the recalculation strategy search unit and its corresponding CostModel to obtain the corresponding recalculation strategy And the sub-model calculated according to the Cost Model under the corresponding recalculation strategy A single microbatch takes ; S2.3, the The intermediate value is rounded down to the number of continuous computing units contained in the sub-model of the pipeline stage. , and according to the number of units, the computational units are connected into sub-models according to the structure of interleaved FeedForward units and Attention units. , and the submodel Corresponding to the model configuration, model scale and current device, training configuration, and current stage number, input the recalculation strategy search unit, call the recalculation strategy search unit and its corresponding Cost Model to obtain the corresponding recalculation strategy And the sub-model calculated according to the Cost Model under the corresponding recalculation strategy A single microbatch takes .

6. The recomputation perception model splitting system for pipeline parallel training according to claim 5, characterized in that: By determining an interval and , to ensure the target and Falls within this interval, and then uses binary search to find the maximum value that satisfies the constraints as and include: initialization and Search interval and ,Will and Set to 0, and Set to the number of model layers twice as much; Through the memory module, we search for the maximum trainable sub-model size and In the process, the recalculation strategy search unit is called and the temporary variable is entered or Get the time taken for a single microbatch or and the corresponding recalculation strategy; and The corresponding single microbatch time consumption and recalculation strategy are stored in two tables. Each row of the table contains the sub-model size, the sum of the corresponding forward propagation time and the backpropagation time, and the corresponding recalculation strategy. The initialization process of the second and subsequent second-level searches is to or Find the corresponding table or Less than The number of calculation units of the maximum term is the number of calculation units of or ,as well as or Greater than The number of computational units of the minimum term is the number of computational units of or If the corresponding item is not found in the table, it degenerates to the initialization strategy during the first search.

7. The recomputation perception model splitting system for pipeline parallel training according to claim 5, characterized in that: The 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 scale and the 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 Update interval for , otherwise the update interval is ; judge Whether the termination condition is met ; If the termination condition is reached, return As ; 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: The S2.3 includes: exist During the calculation process, before calling the recalculation strategy search unit, query the table in the memory module that stores the information of the sub-model ending with the Attention unit. If the required sub-model scale and corresponding strategy are already in the table, use the data in the table directly. Otherwise, insert the corresponding information into the table in the memory module that stores the information of the sub-model ending with the Attention unit. like , then update the interval for , otherwise update the interval to ; judge Whether the termination condition is met ; If the termination condition is reached, return As ; 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 arrive The length of the side, and if , then the last unit structure of the sub-model is the FeedForward unit. , then the last unit structure of the sub-model is the Attention unit; The total number of units obtained from all stages is ; The sub-model structure is an alternating series of Attention units and FeedForward units.

10. A recomputation perception model splitting method for pipeline parallel training, characterized in that: The method comprises: Collect data for the defined model, including the number of model operators, calculation time, output size, and operator dependency; the collected data is used for the formulation of subsequent recalculation strategies and model splitting strategies; According to the recomputation strategy-aware model splitting mechanism, the workload is evenly distributed among multiple devices; the recomputation strategy-aware model splitting mechanism is implemented by the dijkstra algorithm and two-layer binary search, wherein the first layer of binary search performs a binary search on the minimum value of the sum of the forward propagation and backpropagation time of training a microbatch at each stage of the pipeline; the second layer of binary search is the first layer of binary search "with 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 The maximum sub-model size that can be supported at each stage under the "constraint condition"; 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; Configure the model splitting strategy and the recalculation strategy for each device into the training strategy.

Images