A Method, System and Readable Storage Medium for Optimizing Video Memory in Distributed Training of Models

Abstract

The present invention discloses a method, system and readable storage medium for optimizing video memory in distributed training of a model, including: obtaining the parameter quantity and access order of each operator during the first round of training of the model, generating an execution order graph after uniquely numbering the operators and storing it; evenly splitting each operator and recording the shard mapping relationship; during the training process, according to the execution order of the operators, obtaining the complete parameters of the initially running operator through a computing device; performing the calculation of the current operator, and prefetching the parameters of the next operator to be run through asynchronous communication according to the operator execution order graph; after the current operator is executed, only retaining the original partial operator parameters and releasing the remaining partial operator parameters; synchronizing the operator parameter communication, verifying the integrity of the shard parameters and then performing the training iteration, and looping the above process until the training is completed. By using the present invention, while reducing the video memory occupation of a single computing device, the efficient training of the model can be ensured.

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 in distributed training of a model. Background Art

[0002] The rapid development of deep learning and AI has achieved remarkable results in the fields of computer vision, natural language processing, recommendation systems, etc. However, with the increase in task complexity, the scale of the model has expanded rapidly, and the number of parameters has jumped from the million level to the trillion level, resulting in a sharp increase in video memory requirements. For example, a model with 10 billion parameters requires at least 20GB of video memory at FP16 precision, and the storage requirements for activation functions, gradients, and optimizer states during training may cause the video memory occupancy to exceed 80GB. Although the performance of GPUs has been continuously improved, their memory capacity is still limited, becoming the main bottleneck for model scale and training speed.

[0003] For video memory optimization, there are the following methods: quantization (compression, mixed precision), recomputation, swapping in and out. Quantization is to convert high-precision numerical values in the model into low-precision numerical values. Recomputation is a strategy that does not store all intermediate activation values during forward propagation but recomputes 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 the publication number CN112329834A discloses a method and device for allocating video memory space during the training of a recurrent network model, effectively compressing the video memory used in network calculations, thereby improving the training speed. The Chinese patent document with the publication number CN115437795A discloses a method and system for optimizing video memory recomputation with load awareness in a heterogeneous GPU cluster, determining the stage with the highest video memory load in all stages, and performing recomputation optimization according to the algorithm for minimizing video memory overhead to ensure load balance in each stage.

[0004] All of the above methods can reduce video memory occupancy, but there are problems of reducing precision or training efficiency. Currently, there are also methods of using distributed training for video memory optimization. Distributed training uses multi-node parallelism to accelerate training, which 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] In view of the problem of low video memory utilization and low training efficiency in model training in the related art, no effective solution has been proposed yet. Summary of the Invention

[0006] The present invention provides a method, system and readable storage medium for optimizing video memory in distributed training of a model, 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 in distributed training of a model, comprising:

[0008] (1) Based on a preset network model, obtain the parameter quantity and access order of each operator in the first round of model training, generate a topologized execution order graph after uniquely numbering the operators, and store it;

[0009] (2) Based on the parameter quantity of each operator and the number of computing devices, evenly divide each operator respectively, and record the shard mapping relationship;

[0010] (3) When conducting the next round of training, according to the execution order of the operators in the topologized execution order graph, obtain the complete parameters of the initially running operator through the communication module of the computing device;

[0011] (4) Conduct the calculation of the current operator, and prefetch the parameters of the next operator to be run through asynchronous communication according to the operator execution order graph;

[0012] (5) After the current operator finishes execution, only retain the original partial operator parameters according to the shard mapping relationship, and release the remaining partial operator parameters;

[0013] (6) When the next operator runs, first synchronize the operator parameter communication, verify the integrity of the next operator parameters, and then return to step (4), and loop the above process until the network model completes training on the computing device.

[0014] In step (2), after evenly dividing each operator, each computing device retains a partial value of the cut operator parameters, releases the other parts of the parameters, and records it as the shard mapping relationship. Among them, the shard mapping relationship records which part of the operator parameters is retained by the current computing device.

[0015] In step (2), it also includes adding attribute information to each operator. The operator attribute information includes: operator number, operator shard parameters, parameter integrity status, and the number of the next operator;

[0016] Among them, the operator shard parameters are the cut operator parameters saved by the current computing device; the parameter integrity status indicates whether the parameters of the current operator are complete; the number of the next operator is used to prefetch the parameters of the next operator during the training process.

[0017] In step (4), when the parameter integrity status indicates that the parameters of the current operator are complete, start the computing device to execute the forward calculation and backpropagation operations.

[0018] In step (4), the asynchronous communication does not block the calculation of the current operator, and prefetches the parameters of the next operator while the current operator is being calculated to reduce the communication overhead.

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

[0020] A system for optimizing video memory in distributed training of a model includes a memory and one or more processors. Executable code is stored in the memory. When the one or more processors execute the executable code, it is used to implement the above method for optimizing video memory in distributed training of a model.

[0021] A computer-readable storage medium has a program stored thereon. When the program is executed by a processor, it implements the above method for optimizing video memory in distributed training of a model.

[0022] Compared with the prior art, the present invention has the following beneficial effects:

[0023] The present invention realizes a strategy for decentralized storage of model parameters with relatively low metadata collection and parameter splitting overheads, while effectively reducing video memory occupancy and ensuring training efficiency as much as possible. BRIEF DESCRIPTION OF THE DRAWINGS

[0024] In order to more clearly illustrate the technical solutions in the embodiments of the present invention, the drawings required for description in the embodiments will be briefly introduced below. Obviously, the drawings in the following description are only some embodiments of the present invention. For those of ordinary skill in the art, other drawings can be obtained based on these drawings without creative efforts.

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

[0026] Figure 2 It is a flowchart of specific strategy selection for a method for optimizing video memory in distributed training of a model according to an embodiment of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0027] The present invention will be further described in detail below with reference to the drawings and embodiments. It should be noted that the following embodiments are intended to facilitate the understanding of the present invention and do not limit it in any way.

[0028] Before further elaborating on the embodiments of the present application, the nouns and terms involved in the embodiments of the present application are described. The nouns and terms involved in the embodiments of the present application are applicable to the following explanations.

[0029] (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.

[0030] (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.

[0031] (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.

[0032] like Figure 1 As shown, a method for optimizing video memory for model distributed training includes the following steps:

[0033] 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.

[0034] 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.

[0035] 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.

[0036] 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.

[0037] 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.

[0038] 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.

[0039] 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.

[0040] 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.

[0041] 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.

[0042] 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.

[0043] 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.

[0044] 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.

[0045] 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.

[0046] In some possible implementation manners, after the operator execution is completed, preserving the original partial operator parameters and releasing the other parts of the operator parameters need to rely on the methods provided by the specific deep learning framework. Taking PyTorch as an example, PyTorch provides functions such as hook, view, and copy. By registering a hook, after the operator execution is completed, functions such as view and copy can be immediately used to split and release the complete parameters.

[0047] Step S106, synchronize the operator parameter communication, verify the integrity of the sharded parameters, and then execute the training iteration. Loop the above process until the training is completed.

[0048] In some embodiments, in step S104, the parameter prefetching for the upcoming operator is performed through asynchronous communication. Therefore, when the next operator is executed, it is necessary to ensure through communication synchronization that the communication has been completed, so as to ensure that each computing device has the complete parameters of the upcoming operator when performing operator calculations.

[0049] In some possible implementation manners, different synchronization primitives need to be selected for communication synchronization according to different computing devices. Taking NVIDIA GPUs as an example, cudaStreamSynchronize, cudaEventSynchronize, etc. can be used for communication synchronization.

[0050] As Figure 2 shown, the specific strategy selection of the model distributed training video memory optimization method includes the following steps:

[0051] Step S201, run the training program of the preset model, record the execution order of the operators and their parameter quantities, generate a topological execution order graph after uniquely numbering the operators, and store it.

[0052] Step S202, based on the parameter quantity of the operator and the number of computing devices, split the operator, and add attribute information to each operator.

[0053] In some embodiments, the operator attribute information should include: operator number, operator sharded parameters, parameter integrity status (indicating whether the parameters are complete), and the number of the next operator, etc. The operator sharded parameters point 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 prefetch the parameters of the next operator during the training process.

[0054] Step S203, obtain the complete parameter set of the first operator, mark the parameter status as complete, and then start the computing device to perform forward calculation and backward propagation operations.

[0055] In some embodiments, after obtaining the first operator parameter, the integrity status of the operator parameter needs to be set to complete, and at the same time, a computing device is used to calculate the operator for neural network training.

[0056] Step S204, while calculating the current operator, according to the operator execution sequence diagram, use asynchronous communication to prefetch the parameters of the next operator.

[0057] In some embodiments, asynchronous communication does not block the calculation of the current operator. Therefore, the parameters of the next operator can be prefetched while calculating the current operator, reducing communication overhead.

[0058] Step S205, after the current operator is executed, the parameters of the operator are sliced according to a preset slicing scheme to reduce video memory occupancy.

[0059] Step S206, synchronize the operator parameter communication in step S204, and set the integrity status of the parameters of the operator that has completed communication to "complete".

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

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

[0062] Step S208, the operator parameters are complete, and the operator is calculated.

[0063] 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, enter step S211.

[0064] In some embodiments, if the integrity status of the operator parameter 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 the preset threshold, it indicates that the operator execution order changes frequently. Frequent changes in the operator execution order indicate that the training process of this neural network is sensitive to training data and is not applicable to the video memory optimization strategy of the present invention.

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

[0066] Step S211, terminate the training, indicating that the video memory optimization strategy of the present invention is not applicable to this neural network.

[0067] Step S212, the training is completed, and a trained neural network model is obtained.

[0068] Based on the same inventive principle, this embodiment also provides a memory optimization system for distributed training of a model, including a memory and one or more processors. Executable code is stored in the memory. When the one or more processors execute the executable code, it is used to implement the method for optimizing the video memory for distributed training of the model mentioned in the above embodiment.

[0069] 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, it implements the method for optimizing the video memory for distributed training of the model mentioned in the above embodiment.

[0070] The above embodiments have described in detail 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 used to limit the present invention. Any modifications, supplements, and equivalent replacements made within the scope of the principles of the present invention shall be included within the protection scope of the present invention.

Claims

1. A method for optimizing video memory in distributed training of a model, characterized in that, Including: (1) Based on a preset network model, obtain the number of parameters and access order of each operator during the first round of model training. After uniquely numbering the operators, generate a topologized execution order graph and store it; (2) Based on the number of parameters of each operator and the number of computing devices, evenly slice each operator's parameters respectively; After evenly slicing each operator's parameters, each computing device retains a partial value of the sliced operator's parameters, releases the other parts of the parameters, and records them as a sharding mapping relationship; (3) When conducting the next round of training, according to the execution order of the operators in the topologized execution order graph, obtain the complete parameters of the initially running operator through the communication module of the computing device; (4) Conduct the calculation of the current operator, and according to the operator's execution order graph, prefetch the parameters of the next operator to be run through asynchronous communication; (5) After the current operator finishes execution, only retain the original partial operator's parameters according to the sharding mapping relationship, and release the remaining operator's parameters; (6) When the next operator runs, first synchronize the operator's parameter communication, verify the integrity of the next operator's parameters, and then return to step (4). Loop the above process until the network model completes training on the computing device.

2. The method for optimizing video memory in distributed training of the model according to claim 1, wherein In step (2), it also includes adding attribute information to each operator. The operator attribute information includes: operator number, operator sharded parameters, parameter integrity status, and the number of the next operator; Among them, the operator sharded parameters are the sliced operator's parameters saved by the current computing device; the parameter integrity status indicates whether the parameters of the current operator are complete; the number of the next operator is used to prefetch the parameters of the next operator during the training process.

3. The method for optimizing video memory in model distributed training according to claim 2, wherein, In step (4), when the parameter integrity status indicates that the parameters of the current operator are complete, start the computing device to execute the forward calculation and backpropagation operations.

4. The method for optimizing video memory in distributed training of the model according to claim 2, wherein In step (4), the asynchronous communication does not block the calculation of the current operator, and prefetches the parameters of the next operator while the current operator is being calculated to reduce the communication overhead.

5. The model distributed training video memory optimization method according to claim 1, characterized in that In step (6), when the number of times of verifying that the operator's parameters are incomplete exceeds a preset threshold, terminate the training.

6. A model distributed training video memory optimization system, characterized in that Including a memory and one or more processors. The memory stores executable code. When the one or more processors execute the executable code, it is used to implement the method for optimizing video memory in distributed training of the model according to any one of claims 1 - 5.

7. A computer-readable storage medium, characterized in that, A program is stored thereon. When the program is executed by a processor, it implements the method for optimizing video memory in distributed training of the model according to any one of claims 1 - 5.

Images