Pipeline parallel partitioning and memory optimization method for large-scale model training

By automatically partitioning large-scale models through dynamic programming algorithms, the problem of uneven memory usage in pipeline parallelism is solved, and the training speed and efficiency are improved. In particular, in the training of large-scale deep learning models, memory consumption and computing overhead are significantly reduced, and the scalability of the model is enhanced.

CN119336489BActive Publication Date: 2025-09-30TSINGHUA SHENZHEN INTERNATIONAL GRADUATE SCHOOL +1
View PDF 2 Cites 0 Cited by

Patent Information

Application Number
CN202411258643.6
Authority / Receiving Office
CN · China
Patent Type
Patents(China)
Current Assignee / Owner
Filing Date
2024-09-09
Publication Date
2025-09-30
Estimated Expiration
2044-09-09

AI Technical Summary

Technical Problem

Existing pipeline parallel methods face challenges in uneven memory usage and improving training efficiency in large-scale model training. Especially when processing large deep learning models containing billions or even trillions of parameters, traditional methods cannot effectively solve the efficiency bottlenecks caused by memory consumption and computational order dependencies.

Method used

A dynamic programming algorithm is used to automatically partition the model. By identifying the operators in the model, a memory usage data model is constructed, and the model partitioning strategy is optimized to reduce the memory usage differences of devices at each stage. The isomorphism between model layers is used to simplify the search space and improve the precision and efficiency of model partitioning.

Benefits of technology

It achieves a balance in memory usage during model training, improves training speed and scalability, reduces the complexity and computing time of manual partitioning, and enhances the efficiency of integration with other distributed parallel technologies.

✦ Generated by Eureka AI based on patent content.

Smart Images

  • Figure CN119336489B_ABST
    Figure CN119336489B_ABST
Patent Text Reader

Abstract

The present invention provides a pipeline parallel partitioning and memory optimization method for large-scale model training, comprising the following steps: S1. Identify and partition model operators; S2. Calculate and number operator memory usage; S3. Construct an objective function to minimize memory differences through dynamic programming; S4. Starting from a single operator, gradually construct and search for an optimal partitioning strategy; S5. Analyze and optimize memory usage to balance device memory at each stage. The present invention uses a dynamic programming algorithm to achieve automatic partitioning of the model, solving the problem of uneven memory usage at each stage in pipeline parallelism. The isomorphism between model layers is utilized to reduce the search space of the dynamic programming algorithm, thereby reducing the time required to search for the optimal solution. The minimum unit of model partitioning is reduced to an operator, which improves the precision of model partitioning and enables the algorithm to search for a better partitioning strategy.
Need to check novelty before this filing date? Find Prior Art

Description

Technical Field

[0001] The present invention relates to machine learning and distributed computing technologies, and in particular to a pipeline parallel partitioning and memory optimization method for large-scale model training. Background Art

[0002] Pipeline parallelism is a key technology for training large-scale models in distributed machine learning. Its key idea is to distribute different layers of a model onto different computing devices, thereby reducing the memory consumption of a single computing device and enabling ultra-large-scale model training. This approach is particularly important when processing large deep learning models containing billions or even trillions of parameters, such as the currently popular Transformer models BERT Large and GPT.

[0003] Pipeline parallelism involves dividing the computational graph into sequential subgraphs and assigning the subgraphs to different stages. In the forward pass, each sample flows from the first stage to the last stage. In the backward pass, the gradient of each sample is backpropagated from the last stage to the first stage. The computation of a single sample exhibits sequential dependencies across the stages, meaning that a sample cannot be processed by two stages simultaneously.

[0004] The main challenges facing pipeline parallelism include improving training speed, optimizing memory allocation, and structural scalability.

[0005] There have been a number of attempts and some progress. Among them, the GPipe pipeline parallel framework proposed by Google in 2018 is one of the most famous works. This work reduces pipeline bubbles caused by the sequential dependencies of calculations by breaking down batch data into smaller micro-batches, thereby improving training efficiency. However, the activation values ​​generated in the forward pass need to be retained during the calculation process until the corresponding backward pass is completed, which results in a large amount of memory consumption. Micro-batches require the early stages to save more activation values, which means that these stages require more memory. This memory imbalance increases with the increase in the number of micro-batches / stages.

[0006] To address this memory imbalance, the PipeDream framework proposed by Microsoft and Stanford University in 2019 uses 1F1B (one-forward-one-backward), which uses interleaved forward and backward propagation to release saved activation values ​​in advance. However, this static approach still introduces a certain degree of memory imbalance and does not address the problem of increasing imbalance as the number of stages increases.

[0007] It should be noted that the information disclosed in the above background technology section is only used to understand the background of this application, and therefore may include information that does not constitute prior art known to ordinary technicians in this field. Summary of the Invention

[0008] The main purpose of the present invention is to solve the problems existing in the above-mentioned background technology and provide a pipeline parallel partitioning and memory optimization method for large-scale model training.

[0009] To achieve the above object, the present invention adopts the following technical solutions:

[0010] In a first aspect of the present invention, a pipeline parallel partitioning and memory optimization method for large-scale model training comprises the following steps:

[0011] S1. Identify the operators in the model and divide the model into multiple parts according to the precision of the operators;

[0012] S2. Calculate the memory occupied by each operator in each round of training, number the operators according to the forward propagation order of the model, and build a memory usage data model to provide input for the dynamic programming algorithm;

[0013] S3. Constructing an objective function for model partitioning to minimize the maximum difference in memory usage on devices at different stages through the objective function, and establishing a recursive call relationship for the dynamic programming algorithm to facilitate the search for model partitioning strategies;

[0014] S4. Use a dynamic programming algorithm, starting from the most basic partitioning unit, that is, a single operator, and gradually build more complex model partitioning strategies. Then, through recursive calls and difference minimization strategies, search and determine the optimal model partitioning scheme.

[0015] S5. Based on the optimal model partitioning scheme, analyze the memory usage of each stage, identify the problem of unbalanced memory usage, and adjust the model partitioning strategy to optimize the memory usage so that the memory usage of devices in each stage is as balanced as possible.

[0016] In a second aspect of the present invention, a computer program product includes a computer program, which, when executed by a processor, implements the pipeline parallel partitioning and memory optimization method for large-scale model training.

[0017] The present invention has the following beneficial effects:

[0018] The present invention proposes a pipeline parallel partitioning and memory optimization method for large-scale model training, which innovatively solves the challenges of uneven memory usage and training efficiency improvement faced by pipeline parallelism in large-scale model training. The present invention realizes automatic partitioning of the model through a dynamic programming algorithm, optimizes memory usage, improves training speed, and enhances the scalability of model training. The algorithm can utilize the isomorphism between model substructures to reduce the computing time required to search for the optimal solution, while accurately dividing the smallest unit of the model to the operator level, thereby improving the fineness of the model partitioning. The method of the present invention not only reduces the complexity and cost of manual model partitioning, but also significantly reduces the overall training time by reducing the memory usage of the early stage and the recalculation ratio of the later stage. In addition, the algorithm of the present invention can be effectively combined with other distributed parallel technologies such as data parallelism and tensor parallelism to further improve the efficiency of distributed training, making the present invention particularly important and effective when processing large deep learning models containing billions or even trillions of parameters.

[0019] Compared with traditional methods, the main advantages of the present invention are:

[0020] (1) The present invention can provide an automatic partitioning algorithm for model partitioning in large model parallel training, reducing the cost required for manual partitioning and greatly lowering the entry threshold for model partitioning.

[0021] (2) Use dynamic programming algorithms to automatically partition the model and improve the imbalance of memory usage on devices at different stages under the condition of limited training time.

[0022] (3) By reducing the number of layers saved in the early stages, the memory overhead of these stages in model training can be reduced, and the batch size of hyperparameters in training can be increased, thereby reducing the training time of the model on the full dataset.

[0023] (4) By reducing the proportion of recalculation in the back propagation at the later stages, the additional computation and time overhead required for recalculation at these stages can be reduced, thereby reducing the training time.

[0024] (5) The algorithm is scalable and can be effectively combined with other technologies to further improve the efficiency of distributed training.

[0025] Other beneficial effects of the embodiments of the present invention will be further described below. BRIEF DESCRIPTION OF THE DRAWINGS

[0026] Figure 1 Schematic diagram of the pipeline parallel model division according to an embodiment of the present invention.

[0027] Figure 2This is a simplified flowchart of the pipeline parallel partitioning and memory optimization method for distributed parallel training of a model according to an embodiment of the present invention. DETAILED DESCRIPTION

[0028] The following is a detailed description of the embodiments of the present invention. It should be emphasized that the following description is only exemplary and is not intended to limit the scope of the present invention and its application.

[0029] Furthermore, the terms "first" and "second" are used for descriptive purposes only and should not be understood to indicate or imply relative importance or implicitly specify the number of technical features indicated. Thus, a feature defined as "first" or "second" may explicitly or implicitly include one or more of such features. In the description of the embodiments of the present invention, "plurality" means two or more, unless otherwise specifically defined.

[0030] See Figure 1 and Figure 2 , an embodiment of the present invention provides a pipeline parallel partitioning and memory optimization method for large-scale model training, comprising the following steps:

[0031] S1. Identify the operators in the model, including basic computational units such as matrix multiplication and addition, and divide the model into multiple parts based on the operator's precision to facilitate subsequent memory usage calculation and model partitioning strategy design;

[0032] S2. Calculate the memory usage of each operator in each training round, including parameters, intermediate data, and activation values. Operators are numbered according to the model's forward propagation order and a memory usage data model is constructed to provide input for the dynamic programming algorithm.

[0033] S3. Constructing an objective function for model partitioning to minimize the maximum difference in memory usage on devices at different stages through the objective function, and establishing a recursive call relationship for the dynamic programming algorithm to facilitate the search for model partitioning strategies;

[0034] S4. Use a dynamic programming algorithm, starting from the most basic partitioning unit, that is, a single operator, and gradually build more complex model partitioning strategies. Then, through recursive calls and difference minimization strategies, search and determine the optimal model partitioning scheme.

[0035] S5. Based on the optimal model partitioning scheme, analyze the memory usage of each stage, identify the problem of unbalanced memory usage, and adjust the model partitioning strategy to optimize the memory usage so that the memory usage of devices in each stage is as balanced as possible.

[0036] This invention uses a dynamic programming algorithm to achieve automatic model partitioning. It solves the problem of uneven memory usage across pipeline stages. By leveraging the isomorphism between model layers, the dynamic programming algorithm's search space is reduced, thereby reducing the time required to search for the optimal solution. By reducing the minimum unit of model partitioning to operators, the model partitioning is refined, enabling the algorithm to search for a more optimal partitioning strategy.

[0037] In some embodiments, in step S3, the number of stages p used in the training process is determined, where each stage corresponds to a computing stage in the pipeline parallelism; the number of groups of parameters, memory, and activation values ​​that the device in the sth stage needs to save is analyzed, which is ps, to identify the memory usage requirements of each stage; in step S4, based on the 1F1B (one-forward-one-backward) pipeline parallel scheduling strategy, the model is sequentially divided using a dynamic programming algorithm so that the device in each stage can effectively process the part of the model it is responsible for under the 1F1B pipeline parallel scheduling strategy; under the guidance of this strategy, the model partitioning strategy is optimized so that the stage with a smaller sequence number, i.e., s, can efficiently process more forward computing tasks while reducing the memory usage of these stages; in step S5, under the 1F1B pipeline parallel scheduling strategy, the memory usage of each stage is dynamically adjusted to adapt to the changes in memory requirements of different stages during forward and backward propagation; especially for stages with smaller sequence numbers, the memory overhead of these stages in model training is reduced by optimizing the model partitioning strategy, so as to achieve a balance in the memory usage of devices in each stage.

[0038] In some embodiments, step S2 specifically includes: dividing the model into basic computing units, such as operators such as matrix multiplication and addition operations; numbering the operators according to the forward propagation order of the model; calculating the memory occupied by each operator during the training process, and constructing a memory occupancy data model. Step S3 specifically includes: determining the goal of model partitioning, that is, minimizing the maximum difference in memory occupancy on devices at different stages; designing a dynamic programming algorithm to search for model partitioning strategies. Step S4 specifically includes: using a dynamic programming algorithm, starting from a single operator, and gradually constructing a model partitioning strategy; searching and determining the optimal model partitioning scheme through recursive calls and difference minimization strategies. Step S5 specifically includes: using the output of the dynamic programming algorithm to determine the device to which each model layer should be assigned, and the memory occupancy that should be allocated on each device; analyzing the memory occupancy of devices at each stage under the optimal model partitioning scheme, and identifying the problem of unbalanced memory occupancy; adjusting the model partitioning strategy for the identified unbalanced memory occupancy problem, and optimizing the memory occupancy of devices at each stage to achieve a balance in overall memory occupancy.

[0039] In some embodiments, in step S3, the optimization objective function F is defined as minimizing the maximum difference D(G) in memory usage on devices at different stages, where G represents the computational graph partitioning of a model; a recursive relationship is set, and the model partitioning is gradually optimized by comparing the memory usage at different stages until the global optimal solution is reached; through recursive calls, starting from the most basic partitioning unit, more complex model partitioning strategies are gradually constructed until the entire model is covered.

[0040] Furthermore, the method also includes: identifying substructures in the model, such as attention layers and feedforward layers, and evaluating their isomorphism; utilizing the isomorphism of the substructures to simplify the search space of the dynamic programming algorithm, thereby reducing computation time and improving algorithm efficiency. The model may include a large Transformer model with attention layers and feedforward layers.

[0041] In some embodiments, the method further includes: analyzing and adjusting the batch size according to memory usage and model partitioning strategy to improve training efficiency; reducing the recalculation ratio in the back propagation in the later stages by optimizing the recalculation strategy, thereby reducing additional calculation and time overhead.

[0042] In some embodiments, the method further includes: verifying the effectiveness of the algorithm in an actual pipeline parallel environment so that the model partitioning strategy can actually improve training efficiency and memory utilization; and tuning the algorithm based on the verification results to adapt to different model sizes and training requirements.

[0043] The method is scalable and can be combined with other distributed parallel technologies such as data parallelism (DP) and tensor parallelism (TP); and allows differentiated recomputation strategies to be set between different layers of the model based on specific performance and resource requirements.

[0044] Specific embodiments of the present invention are further described below.

[0045] A pipeline parallel partitioning and memory optimization method for large-scale model training includes an algorithm for automatically partitioning models in pipeline parallelism, contributing to addressing the imbalance in device memory usage at different stages of pipeline parallelism. This algorithm, preferably based on a 1F1B pipeline parallel scheduling strategy, uses a dynamic programming algorithm to search for the model partitioning strategy and find a near-optimal solution, completing the sequential partitioning of the model.

[0046] Figure 1The pipeline parallel model partitioning of an embodiment of the present invention is shown as an example. Assuming that there are p stages in training, in the 1F1B strategy, the device in the s stage needs to save (ps) sets of parameters, memory and activation values. The above figure shows a schematic diagram with 4 stages. The model is divided in sequence according to a specific strategy and deployed on stage0, stage1, stage2, stage3 respectively. Figure 1 It can be seen that the stage with a smaller sequence number needs to be forwarded more times continuously, which shows that the smaller s is, the more values ​​​​need to be saved and the more memory is occupied.

[0047] In the designed algorithm, the model is first divided into operators with a fine degree of precision, down to operators such as matrix multiplication and addition. These operators are numbered [i] according to the order of calculation in the network forward propagation, and an array M[i] is calculated. M[i] represents the memory occupied by operator i in each round of training. The goal is to balance the memory usage of the device at each stage as much as possible, so the optimization goal is: in Represents the computational graph partitioning of a model, Indicates that the division The maximum difference in memory usage on the device at different stages is obviously In order to solve the minimum problem, the dynamic programming algorithm is used as follows: represents the difference in the optimal partitioning strategy when the first i operators of the model are deployed on j stages, then:

[0048]

[0049] in Indicates the maximum difference between the memory usage of different devices and M under the condition F(i,j).

[0050] In the dynamic programming knapsack problem, the recursive call F(i,1) is known, and the other terms are predictable. Assuming there are p stages and n operators, the target value to be calculated is F(n,p). The above equation is calculated recursively down to F(i,1), iin[1,n]. By recursively calling the above equation, the optimal strategy for deploying the entire model across P stages can be calculated. The algorithm outputs the device to which each model layer should be assigned, as well as the memory usage on each device.

[0051] Taking into account the large number of attention layer + feedforward layer substructures that exist in the current Transformer large model, the present invention uses the isomorphism between model substructures to reduce the complexity of the dynamic programming algorithm, thereby reducing the computing time required to search for the optimal solution.

[0052] The process of the designed algorithm is as follows Figure 2 The designed algorithm is scalable and can be combined with other distributed parallel technologies such as data parallelism (DP) and tensor parallelism (TP). Different recomputation strategies can also be set between layers as needed.

[0053] In summary, this paper proposes a pipeline parallel partitioning and memory optimization method for large-scale model training. This method uses a dynamic programming algorithm to automatically partition the model, resolving the issue of uneven memory usage across pipeline parallelization stages. By leveraging the isomorphism between model layers, the dynamic programming algorithm's search space is reduced, thereby shortening the search time for the optimal solution. By reducing the minimum unit of model partitioning to an operator, the precision of the model partitioning is increased, enabling the algorithm to search for a more optimal partitioning strategy.

[0054] Compared with traditional methods, the main advantages of the present invention are:

[0055] (1) The present invention can provide an automatic partitioning algorithm for model partitioning in large model parallel training, reducing the cost required for manual partitioning and greatly lowering the entry threshold for model partitioning.

[0056] (2) Use dynamic programming algorithms to automatically partition the model and improve the imbalance of memory usage on devices at different stages under the condition of limited training time.

[0057] (3) By reducing the number of layers saved in the early stages, the memory overhead of these stages in model training can be reduced, and the batch size hyperparameter in training can be increased, thereby reducing the training time of the model on the full dataset.

[0058] (4) By reducing the proportion of recalculation in the back propagation at the later stages, the additional computation and time overhead required for recalculation at these stages can be reduced, thereby reducing the training time.

[0059] (5) The algorithm is scalable and can be effectively combined with other technologies to further improve the efficiency of distributed training.

[0060] An embodiment of the present invention further provides a storage medium for storing a computer program, which at least performs the above method when executed.

[0061] An embodiment of the present invention further provides a control device, comprising a processor and a storage medium for storing a computer program; wherein the processor is configured to execute at least the method described above when executing the computer program.

[0062] An embodiment of the present invention further provides a processor, which executes a computer program and at least performs the method described above.

[0063] The storage medium can be implemented by any type of non-volatile storage device, or a combination thereof. Among them, the non-volatile memory can be a read-only memory (ROM), a programmable read-only memory (PROM), an erasable programmable read-only memory (EPROM), an electrically erasable programmable read-only memory (EEPROM), a magnetic random access memory (FRAM), a flash memory, a magnetic surface memory, an optical disc, or a compact disc read-only memory (CD-ROM); the magnetic surface memory can be a magnetic disk memory or a magnetic tape memory. The storage medium described in the embodiments of the present invention is intended to include, but is not limited to, these and any other suitable types of memory.

[0064] In the several embodiments provided by the present invention, it should be understood that the disclosed systems and methods can be implemented in other ways. The device embodiments described above are merely schematic. For example, the division of the units is merely a logical function division. In actual implementation, there may be other division methods, such as: multiple units or components can be combined, or can be integrated into another system, or some features can be ignored or not executed. In addition, the coupling, direct coupling, or communication connection between the components shown or discussed can be through some interfaces, and the indirect coupling or communication connection of the devices or units can be electrical, mechanical or other forms.

[0065] The units described above as separate components may or may not be physically separated, and the components displayed as units may or may not be physical units, that is, they may be located in one place or distributed on multiple network units; some or all of the units may be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0066] In addition, all functional units in the embodiments of the present invention may be integrated into one processing unit, or each unit may be separately used as a unit, or two or more units may be integrated into one unit; the above-mentioned integrated units may be implemented in the form of hardware or in the form of hardware plus software functional units.

[0067] Those skilled in the art will appreciate that all or part of the steps of the above-mentioned method embodiments may be implemented by hardware associated with program instructions, and the aforementioned program may be stored in a computer-readable storage medium. When the program is executed, the program executes the steps of the above-mentioned method embodiments. The aforementioned storage medium includes various media that can store program codes, such as mobile storage devices, read-only memories (ROMs), random access memories (RAMs), magnetic disks, or optical disks.

[0068] Alternatively, if the above-mentioned integrated unit of the present invention is implemented in the form of a software functional module and sold or used as an independent product, it can also be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the embodiment of the present invention, or the part that contributes to the prior art, can be embodied in the form of a software product. The computer software product is stored in a storage medium and includes a number of instructions for enabling a computer device (which can be a personal computer, server, or network device, etc.) to execute all or part of the methods described in each embodiment of the present invention. The aforementioned storage medium includes: various media that can store program codes, such as mobile storage devices, ROM, RAM, magnetic disks or optical disks.

[0069] The methods disclosed in the several method embodiments provided by the present invention can be arbitrarily combined without conflict to obtain new method embodiments.

[0070] The features disclosed in several product embodiments provided by the present invention can be arbitrarily combined without conflict to obtain new product embodiments.

[0071] The features disclosed in several method or device embodiments provided by the present invention can be arbitrarily combined without conflict to obtain new method embodiments or device embodiments.

[0072] The above is a further detailed description of the present invention in conjunction with specific preferred embodiments, and the specific implementation of the present invention should not be considered to be limited to these descriptions. Those skilled in the art will recognize that, without departing from the scope of the present invention, several equivalent substitutions or obvious variations can be made, and the performance or use of the same should be considered to fall within the scope of protection of the present invention.

Claims

1. A pipeline parallel partitioning and memory optimization method for large-scale model training, characterized in that: The following steps are involved: S1. Identify the operators in the model and divide the model into multiple parts according to the precision of the operators; S2. Calculate the memory occupied by each operator in each round of training, number the operators according to the model's forward propagation order, and build a memory usage data model to provide input for the dynamic programming algorithm; S3. Constructing an objective function for model partitioning to minimize the maximum difference in memory usage on devices at different stages through the objective function, and establishing a recursive call relationship for the dynamic programming algorithm to facilitate the search for model partitioning strategies; S4. Use a dynamic programming algorithm, starting from the most basic partitioning unit, that is, a single operator, and gradually build more complex model partitioning strategies. Then, through recursive calls and difference minimization strategies, search and determine the optimal model partitioning scheme. S5. Based on the optimal model partitioning scheme, analyze the memory usage of each stage, identify the problem of unbalanced memory usage, and adjust the model partitioning strategy to optimize the memory usage so that the memory usage of devices in each stage is as balanced as possible.

2. The pipeline parallel partitioning and memory optimization method for large-scale model training according to claim 1, characterized in that: In step S3, the number of stages p used in the training process is determined, where each stage corresponds to a computational stage in the pipeline parallelism; the number of groups of parameters, memory, and activation values ​​required to be saved by the device at stage s is analyzed, which is ps, to identify the memory usage requirements of each stage; In step S4, based on the 1F1B pipeline parallel scheduling strategy, a dynamic programming algorithm is used to sequentially partition the model so that the device in each stage can effectively process the part of the model it is responsible for under the 1F1B pipeline parallel scheduling strategy. Under the guidance of this strategy, the model partitioning strategy is optimized so that stages with smaller sequence numbers (i.e., s) can efficiently process more forward computing tasks, while reducing the memory usage of these stages. In step S5, under the 1F1B pipeline parallel scheduling strategy, the memory usage of each stage is dynamically adjusted to adapt to the changes in memory requirements of different stages during forward and backward propagation. In particular, for stages with smaller sequence numbers, the memory overhead of these stages in model training is reduced by optimizing the model partitioning strategy, thereby achieving a balance in the memory usage of devices at each stage.

3. The pipeline parallel partitioning and memory optimization method for large-scale model training according to claim 1 or 2, characterized in that: Step S2 specifically includes: Dividing the model into basic computing units, wherein the basic computing units include matrix multiplication and addition operators; Number the operators according to the forward propagation order of the model; Calculate the memory occupied by each operator during training and build a memory usage data model; Step S3 specifically includes: Determine the goal of model partitioning, which is to minimize the maximum difference in memory usage on devices at different stages; Design a dynamic programming algorithm to search for model partitioning strategies; Step S4 specifically includes: Using dynamic programming algorithms, starting from a single operator, gradually build a model partitioning strategy; Search and determine the optimal model partitioning solution through recursive calling and difference minimization strategy; Step S5 specifically includes: Use the output of the dynamic programming algorithm to determine the device to which each model layer should be assigned and the memory usage that each device should have; Analyze the memory usage of devices at each stage under the optimal model partitioning scheme and identify memory usage imbalance issues; In response to the identified memory usage imbalance, the model partitioning strategy was adjusted to optimize the memory usage of devices at each stage to achieve a balanced overall memory usage.

4. The pipeline parallel partitioning and memory optimization method for large-scale model training according to claim 3, characterized in that: In step S3, the optimization objective function F is defined as minimizing the maximum difference D(G) in memory usage on devices at different stages, where G represents the computational graph partitioning of a model. A recursive relationship is set, and the model partitioning is gradually optimized by comparing the memory usage at different stages until a global optimal solution is reached. Through recursive calls, starting from the most basic partitioning unit, more complex model partitioning strategies are gradually built until the entire model is covered.

5. The pipeline parallel partitioning and memory optimization method for large-scale model training according to any one of claims 1 to 2, characterized in that: The method is scalable and can be combined with data parallelism (DP) and tensor parallelism (TP) technologies; and allows setting differentiated recomputation strategies among different layers of the model according to specific performance and resource requirements.

6. The pipeline parallel partitioning and memory optimization method for large-scale model training according to any one of claims 1 to 2, characterized in that: Also includes: Analyze and adjust batch size based on memory usage and model partitioning strategy to improve training efficiency; By optimizing the recalculation strategy, the recalculation ratio in the back propagation in the later stages is reduced, thus reducing the additional calculation and time overhead.

7. The pipeline parallel partitioning and memory optimization method for large-scale model training according to any one of claims 1 to 2, characterized in that: Also includes: Identify substructures in the model, including attention layers and feed-forward layers, and assess their isomorphism; The isomorphism of substructures is used to simplify the search space of dynamic programming algorithms, reduce computing time and improve algorithm efficiency.

8. The pipeline parallel partitioning and memory optimization method for large-scale model training according to claim 7, characterized in that: The model includes a large Transformer model with attention layer and feed-forward layer substructures.

9. The pipeline parallel partitioning and memory optimization method for large-scale model training according to any one of claims 1 to 2, characterized in that: Also includes: Verify the effectiveness of the algorithm in an actual pipeline parallel environment, so that the model partitioning strategy can actually improve training efficiency and memory utilization; tune the algorithm based on the verification results to adapt to different model sizes and training requirements.

10. A computer program product comprising a computer program, characterized in that When the computer program is executed by a processor, it implements the pipeline parallel partitioning and memory optimization method for large-scale model training as described in any one of claims 1 to 9.

Citation Information

Patent Citations

  • Model training method and device based on data parallel normal form

    CN116862019A

  • Data processing method and device, equipment and medium

    CN117170685A