Model distributed training video memory optimization method and system and readable storage medium

Abstract

The invention discloses a model distributed training video memory optimization method and system and a readable storage medium, and the method comprises the steps: obtaining a parameter quantity and an access sequence of each operator in a first training process of a model, carrying out the unique numbering of the operators, generating an execution sequence diagram, and storing the execution sequence diagram; performing uniform segmentation on each operator, and recording a fragment mapping relation; in the training process, according to the execution sequence of the operators, obtaining complete parameters of the initially operated operators through a computing device; calculating a current operator, and pre-fetching parameters of a next operator to be operated through asynchronous communication according to the execution sequence diagram of the operator; after the current operator is executed, only part of original operator parameters are reserved, and the rest of operator parameters are released; and synchronizing operator parameter communication, verifying the integrity of the fragment parameters, then executing training iteration, and circulating the process until the training is completed. By utilizing the method, the efficient training of the model can be ensured while the video memory occupation of a single computing device is reduced.

Description

Technical Field

[0001] The present invention relates to the field of deep learning, and in particular to a method, system and readable storage medium for optimizing video memory for model distributed training. Background Art

[0002] The rapid development of deep learning and AI has achieved remarkable results in computer vision, natural language processing, recommendation systems and other fields. However, as the complexity of tasks increases, the scale of models also expands dramatically, and the number of parameters jumps from millions to trillions, and the demand for video memory surges accordingly. For example, a model with 10 billion parameters requires at least 20GB of video memory at FP16 precision, and the storage requirements of activation functions, gradients, and optimizer states during training may cause the video memory to occupy more than 80GB. Although GPU performance continues to improve, its memory capacity is still limited, becoming the main bottleneck for model size and training speed.

[0003] For video memory optimization, there are the following methods: quantization (compression, mixed precision), recalculation, and swapping in and out. Quantization is the conversion of high-precision values ​​in the model into low-precision values. Recalculation is a strategy that does not store all intermediate activation values ​​during forward propagation, but recalculates these values ​​during backpropagation. Swapping in and out is a way to manage video memory usage by dynamically transferring data between video memory and main memory. For example, the Chinese patent document with publication number CN112329834A discloses a method and device for allocating video memory space during recurrent network model training, which effectively compresses the video memory used in network calculations, thereby increasing the training speed. The Chinese patent document with publication number CN115437795A discloses a method and system for optimizing video memory recalculation with load awareness for heterogeneous GPU clusters, which determines the stage with the highest video memory load among all stages, performs recalculation optimization according to the algorithm for minimizing video memory overhead, and ensures load balance in each stage.

[0004] All of the above methods can reduce video memory usage, but there is a problem of reducing accuracy or training efficiency. At present, there is also a method of using distributed training to optimize video memory. Distributed training uses multi-node parallel acceleration training and is suitable for large-scale data and models, but communication overhead and synchronization problems may become bottlenecks, and the implementation complexity is relatively high.

[0005] There is currently no effective solution to the problems of low video memory utilization and low training efficiency in model training in related technologies. Summary of the invention

[0006] The present invention provides a model distributed training video memory optimization method, system and readable storage medium, which can solve the problems of low video memory utilization and low training efficiency in the prior art.

[0007] A method for optimizing video memory for distributed model training, comprising: (1) Based on the preset network model, obtain the parameter quantity and access order of each operator in the first round of model training, uniquely number the operators, generate and store the topological execution sequence graph; (2) Based on the number of parameters of each operator and the number of computing devices, each operator is evenly divided and the shard mapping relationship is recorded; (3) When the next round of training is carried out, the complete parameters of the initially running operator are obtained through the communication module of the computing device according to the execution order of the operators in the topological execution sequence graph; (4) Perform calculations on the current operator and pre-fetch the parameters of the next operator to be run through asynchronous communication according to the execution sequence diagram of the operator; (5) After the current operator is executed, only some of the original operator parameters are retained according to the shard mapping relationship, and the remaining operator parameters are released; (6) When the next operator runs, synchronize the operator parameter communication first, verify the integrity of the next operator parameter, and then return to step (4). Repeat the above process until the network model is trained on the computing device.

[0008] In step (2), after each operator is evenly divided, each computing device retains part of the operator parameter value after the division, releases the rest of the parameter, and records it as a shard mapping relationship. The shard mapping relationship records which part of the operator parameter is retained by the current computing device.

[0009] Step (2) also includes adding attribute information for each operator, where the operator attribute information includes: operator number, operator sharding parameters, parameter integrity status, and the number of the next operator; Among them, the operator sharding parameters are the operator parameters after sharding saved by the current computing device; the parameter integrity status indicates whether the parameters of the current operator are complete; and the number of the next operator is used to pre-fetch the parameters of the next operator during the training process.

[0010] In step (4), when the parameter integrity status indicates that the parameters of the current operator are complete, the computing device is started to perform forward calculation and back propagation operations.

[0011] In step (4), asynchronous communication does not block the calculation of the current operator. The parameters of the next operator are pre-fetched while the current operator is calculating, so as to reduce the communication overhead.

[0012] In step (6), when the number of times the verification operator parameters are incomplete exceeds a preset threshold, the training is terminated.

[0013] A model distributed training video memory optimization system includes a memory and one or more processors, wherein the memory stores executable code, and when the one or more processors execute the executable code, they are used to implement the above-mentioned model distributed training video memory optimization method.

[0014] A computer-readable storage medium stores a program, which, when executed by a processor, implements the above-mentioned model distributed training video memory optimization method.

[0015] Compared with the prior art, the present invention has the following beneficial effects: The present invention implements a distributed storage model parameter strategy with low metadata collection and parameter segmentation overhead, effectively reducing video memory occupancy while ensuring training efficiency as much as possible. BRIEF DESCRIPTION OF THE DRAWINGS

[0016] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the accompanying drawings required for use in the description of the embodiments will be briefly introduced below. Obviously, the accompanying drawings in the following description are only some embodiments of the present invention. For ordinary technicians in this field, other accompanying drawings can be obtained based on these accompanying drawings without paying creative labor.

[0017] Figure 1 The present invention is a flowchart of a method for optimizing video memory for model distributed training according to an embodiment of the present invention.

[0018] Figure 2 This is a specific strategy selection flowchart of a model distributed training video memory optimization method according to an embodiment of the present invention. DETAILED DESCRIPTION

[0019] The present invention is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be pointed out that the embodiments described below are intended to facilitate the understanding of the present invention and do not have any limiting effect on the present invention.

[0020] Before further describing the embodiments of the present application in detail, the nouns and terms involved in the embodiments of the present application are explained. The nouns and terms involved in the embodiments of the present application are subject to the following interpretations.

[0021] (1) Operator refers to a function or process that performs a specific mathematical operation or operation. An operator can be a simple mathematical operation (such as addition or multiplication) or a complex neural network layer (such as a convolutional layer or a fully connected layer). In a computational graph, operators are usually represented as edges that connect different nodes (variables) and define the computational relationship between these nodes. The selection and combination of operators determines the structure and function of the model.

[0022] (2) Video memory, also called frame buffer, is used to store data processed by the graphics chip or rendering data to be extracted. Like the computer's memory, video memory is a component used to store graphics information to be processed.

[0023] (3) Model parameters, which are used to control the transmission and transformation of input data in the model. In a neural network, the input data is mathematically operated on the model parameters in the neural network operators to perform forward propagation of the network. The model parameters are adjusted during the model training process through optimization algorithms (such as gradient descent) to minimize the loss function and improve the predictive ability of the model.

[0024] like Figure 1 As shown, a method for optimizing video memory for model distributed training includes the following steps: Step S101, based on the preset network model, obtain the parameter quantity and access order of each operator module in the first round of model training, uniquely number the operators, generate and store a topological execution sequence diagram.

[0025] In some embodiments, the preset network model can be any type of network model to be trained, such as a convolutional neural network model to be trained, a residual network model to be trained, or a Transformer architecture neural network model to be trained. In the first round of the preset network training, the parameter quantity and access order of each operator module are recorded, and the operators are uniquely numbered and then a topological execution sequence diagram is generated and stored.

[0026] In some possible implementations, obtaining the access order and parameter quantity of each operator requires the use of methods provided by a specific deep learning framework. Taking PyTorch as an example, PyTorch provides hook functions. By registering forward and backward propagation hooks, the access time and the parameter quantity of the operator are recorded in the hook function.

[0027] Step S102: based on the parameter quantity of each operator and the number of computing devices, each operator is evenly divided and the shard mapping relationship is recorded.

[0028] In some embodiments, the operator parameters are evenly cut according to the video memory size of the operator parameters recorded above and the number of computing devices, and each computing device retains part of the values ​​of the operator parameters after cutting and releases the other parts of the parameters.

[0029] In some possible implementations, the segmentation of operator parameters requires the use of methods provided by the deep learning framework. Taking PyTorch as an example, PyTorch provides functions such as view and copy. The copy function can back up part of the parameters, and then release the parameters of the operator by releasing the parameters.

[0030] Step S103, during the training process, according to the execution order of the operators, the complete parameters of the initially running operators are obtained through the communication module of the computing device.

[0031] In some embodiments, after the operator parameters are segmented, the parameters of each computing device are incomplete and the operator calculation cannot be performed immediately. Therefore, when the preset network is trained for the next round, it is necessary to obtain the parameters of the initially running operator through the communication module of the computing device according to the execution order of the operator.

[0032] In some possible implementations, different computing devices use corresponding communication modules to communicate and obtain complete operator parameters. Taking NVIDIA GPU as an example, the nccl communication library can be used to communicate operator parameters.

[0033] Step S104, performing calculation of the current operator, and pre-fetching parameters of the next operator to be run through asynchronous communication according to the execution order of the operators.

[0034] In some embodiments, through the above steps, each computing device obtains the complete parameters of the current operator, so that the operation of the operator can be performed. At the same time, according to the execution order of the operator, the communication module of the computing device can be used to pre-fetch the parameters of the next operator to be run through asynchronous communication. Communication is performed while the operator is being calculated, thereby reducing the communication overhead.

[0035] Step S105: After the current operator is executed, only part of the original operator parameters are retained according to the shard mapping relationship, and the remaining operator parameters are released.

[0036] In some embodiments, the operator parameters of the executed operator will not be used for a period of time, so some of the operator parameters can be retained according to the initial partitioning scheme, and the rest of the operator parameters can be released. By releasing some parameters, the video memory of the computing device can be saved more effectively.

[0037] In some possible implementations, retaining part of the original operator parameters and releasing other parts of the operator parameters after the operator is executed requires the use of methods provided by a specific deep learning framework. Taking PyTorch as an example, PyTorch provides hook, view, copy and other functions. By registering the hook, after the operator is executed, you can immediately use view, copy and other functions to split and release the complete parameters.

[0038] Step S106, synchronize operator parameter communication, verify the integrity of the sharding parameters, and then perform training iterations, and repeat the above process until the training is completed.

[0039] In some embodiments, step S104 pre-fetches parameters of the operator to be executed through asynchronous communication. Therefore, when the next operator is executed, it is necessary to ensure that the communication has been completed through communication synchronization to ensure that each computing device has the complete parameters of the operator to be executed at the same time when the operator is calculated.

[0040] In some possible implementations, communication synchronization requires selecting different synchronization primitives according to different computing devices. Taking NVIDIA GPU as an example, cudaStreamSynchronize, cudaEventSynchronize, etc. can be used for communication synchronization.

[0041] like Figure 2 As shown in the figure, the specific strategy selection of the model distributed training memory optimization method includes the following steps: Step S201, run the training program of the preset model, record the execution order of the operators and their parameter quantities, uniquely number the operators, generate a topological execution sequence diagram and store it.

[0042] Step S202: based on the parameter quantity of the operator and the number of computing devices, the operator is segmented, and attribute information is added to each operator.

[0043] In some embodiments, the operator attribute information should include: operator number, operator sharding parameters, parameter integrity status (indicating whether the parameters are complete), and the number of the next operator, etc. The operator sharding parameters refer to the parameters saved by the current computing device; the parameter integrity status indicates whether the parameters of the current operator are complete so that training can be performed; the number of the next operator is used to pre-fetch the parameters of the next operator during training.

[0044] Step S203, obtaining a complete parameter set of the first operator, marking the parameter status as complete, and starting the computing device to perform forward computing and reverse propagation operations.

[0045] In some embodiments, after obtaining the first operator parameter, the operator parameter integrity state needs to be set to complete, and at the same time, the operator is calculated using a computing device to train the neural network.

[0046] Step S204: while the current operator is being calculated, the parameters of the next operator are pre-fetched using asynchronous communication according to the operator execution sequence diagram.

[0047] In some embodiments, asynchronous communication does not block the calculation of the current operator, so the parameters of the next operator can be pre-fetched while the current operator is calculating, reducing communication overhead.

[0048] Step S205: after the current operator is executed, the parameters of the operator are segmented according to a preset segmentation scheme to reduce the usage of video memory.

[0049] Step S206, synchronizing the operator parameter communication in step S204, and setting the parameter integrity status of the operator that has completed the communication to "complete".

[0050] Step S207, continue to execute the neural network training, and determine whether the parameter integrity status of the operator to be executed is "complete".

[0051] In some embodiments, if the parameter integrity status of the operator is "complete", it indicates that the actual execution order of the operator is consistent with the recorded operator order, and training can continue, entering step S208; otherwise, it indicates that the operator execution order may be deviated, and enters step S209.

[0052] Step S208: The operator parameters are complete and the operator calculation is performed.

[0053] Step S209, determine whether the number of times the operator parameters are incomplete exceeds a preset threshold (for example, three times). If it exceeds the threshold, proceed to step S211.

[0054] In some embodiments, if the parameter integrity status of an operator is "incomplete", it indicates that the actual execution order of the operator is inconsistent with the recorded operator order. If the number of times the operator parameters are incomplete exceeds a preset threshold, it indicates that the operator execution order changes frequently. Frequent changes in the operator execution order indicate that the training process of the neural network is sensitive to the training data and is not suitable for the video memory optimization strategy of the present invention.

[0055] Step S210, determine whether the neural network training is completed. If it is completed, proceed to step S212; otherwise, return to step S204.

[0056] Step S211, terminating the training, indicating that the video memory optimization strategy of the present invention is not suitable for the neural network.

[0057] Step S212: training is completed and a trained neural network model is obtained.

[0058] Based on the same inventive principle, this embodiment also provides a model distributed training memory optimization system, including a memory and one or more processors, wherein the memory stores executable code, and when the one or more processors execute the executable code, they are used to implement the model distributed training memory optimization method mentioned in the above embodiment.

[0059] Based on the same inventive principle, this embodiment also provides a computer-readable storage medium on which a program is stored. When the program is executed by a processor, the model distributed training video memory optimization method mentioned in the above embodiment is implemented.

[0060] The embodiments described above provide a detailed description of the technical solutions and beneficial effects of the present invention. It should be understood that the above are only specific embodiments of the present invention and are not intended to limit the present invention. Any modifications, supplements and equivalent substitutions made within the scope of the principles of the present invention should be included in the protection scope of the present invention.

Claims

1. A method for optimizing video memory for model distributed training, characterized in that: include: (1) Based on the preset network model, obtain the parameter quantity and access order of each operator in the first round of model training, uniquely number the operators, generate a topological execution sequence graph and store it; (2) Based on the number of parameters of each operator and the number of computing devices, each operator is evenly divided and the shard mapping relationship is recorded; (3) When the next round of training is carried out, the complete parameters of the initially running operator are obtained through the communication module of the computing device according to the execution order of the operators in the topological execution sequence graph; (4) Perform calculations on the current operator and pre-fetch the parameters of the next operator to be run through asynchronous communication according to the execution sequence diagram of the operator; (5) After the current operator is executed, only some of the original operator parameters are retained according to the shard mapping relationship, and the remaining operator parameters are released; (6) When the next operator runs, synchronize the operator parameter communication first, verify the integrity of the next operator parameter, and then return to step (4). Repeat the above process until the network model is trained on the computing device.

2. The model distributed training video memory optimization method according to claim 1, characterized in that: In step (2), after each operator is evenly divided, each computing device retains part of the value of the operator parameter after the division, releases the other part of the parameter, and records it as a shard mapping relationship.

3. The model distributed training video memory optimization method according to claim 1, characterized in that: Step (2) also includes adding attribute information for each operator, where the operator attribute information includes: operator number, operator sharding parameters, parameter integrity status, and the number of the next operator; Among them, the operator sharding parameters are the operator parameters after sharding saved by the current computing device; the parameter integrity status indicates whether the parameters of the current operator are complete; and the number of the next operator is used to pre-fetch the parameters of the next operator during the training process.

4. The model distributed training video memory optimization method according to claim 3 is characterized in that: In step (4), when the parameter integrity status indicates that the parameters of the current operator are complete, the computing device is started to perform forward calculation and back propagation operations.

5. The model distributed training video memory optimization method according to claim 3 is characterized in that: In step (4), asynchronous communication does not block the calculation of the current operator. The parameters of the next operator are pre-fetched while the current operator is calculating, so as to reduce the communication overhead.

6. The model distributed training video memory optimization method according to claim 1, characterized in that: In step (6), when the number of times the verification operator parameters are incomplete exceeds a preset threshold, the training is terminated.

7. A model distributed training memory optimization system, characterized in that: It includes a memory and one or more processors, wherein the memory stores executable code, and when the one or more processors execute the executable code, it is used to implement the model distributed training video memory optimization method described in any one of claims 1-6.

8. A computer-readable storage medium, characterized in that: A program is stored thereon, and when the program is executed by a processor, the model distributed training video memory optimization method described in any one of claims 1-6 is implemented.

Images