Model training method and device, storage medium and program product

Abstract

The invention relates to a model training method and device, a storage medium and a program product. The method comprises the following steps: in a forward propagation stage, for any check point module, storing input and output of the check point module in a video memory, and releasing an intermediate activation value of the check point module in the video memory; the input of the check point module is used for recalculating an intermediate activation value, and the output of the check point module is used for forward calculation of subsequent modules of the check point module; and in a back propagation stage, for any check point module, responding to the condition that the last layer of the check point module is a linear layer, skipping forward calculation of the last layer, calculating a gradient according to a gradient formula corresponding to the last layer, and completing back propagation of the check point module according to the gradient of each layer in the check point module. According to the method and the device, the calculation cost can be remarkably reduced while the calculation precision and the video memory saving amount are the same as those of a standard re-calculation scheme.

Description

Technical Field

[0001] The present disclosure relates to the field of computer technology, and in particular to a model training method, a model training device, a non-volatile computer-readable storage medium, and a computer program product. Background Art

[0002] In recent years, deep learning technology, especially large-model training, has made significant progress in the field of artificial intelligence. However, as the scale of models continues to expand, the problems of video memory usage and computational efficiency during training have become increasingly prominent. Video memory usage mainly comes from model weights, optimizer states, cached activation values, and temporary tensors. Among them, cached activation values ​​occupy a major part of video memory in large-scale model training. In order to alleviate video memory pressure, gradient checkpointing technology has been proposed in the existing technology to reduce video memory usage by sacrificing part of the computational effort. Its core idea is to save only part of the intermediate results (i.e., checkpoints) during forward propagation, and recalculate the discarded intermediate variables during backpropagation.

[0003] However, existing gradient checkpointing techniques require a complete recalculation of the forward pass of the linear layer during backpropagation (e.g., re-performing the matrix multiplication operation y = Wx), resulting in additional computational overhead. This significantly reduces training efficiency, especially for computationally intensive operations like fully connected layers in large models. Furthermore, due to the high computational cost of recalculation, existing methods are typically limited to specific modules (e.g., attention layers) and lack the flexibility to scale to other architectures, limiting the overall effectiveness of video memory optimization. Summary of the Invention

[0004] In view of this, the present disclosure provides a model training technology solution.

[0005] According to one aspect of the present disclosure, a model training method is provided, comprising:

[0006] Identify the checkpoint modules in the target model;

[0007] During the forward propagation phase, for any checkpoint module, the input and output of the checkpoint module are saved in the video memory, and the intermediate activation value of the checkpoint module in the video memory is released; wherein the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module during the backward propagation phase, and the output of the checkpoint module is used for the forward calculation of the subsequent modules of the checkpoint module;

[0008] In the back-propagation phase, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, the forward calculation of the last layer is skipped, the gradient of the last layer is calculated according to the gradient formula corresponding to the last layer, and the back-propagation of the checkpoint module is completed according to the gradients of each layer in the checkpoint module.

[0009] In a possible implementation, in response to the last layer of the checkpoint module being a linear layer, skipping the forward calculation of the last layer, and calculating the gradient of the last layer according to the gradient formula corresponding to the last layer, includes:

[0010] In response to the last layer of the checkpoint module including only one linear layer, skipping the forward calculation of the linear layer of the last layer, and calculating the gradient of the linear layer of the last layer according to the gradient formula corresponding to the linear layer of the last layer;

[0011] or,

[0012] In response to the last layer of the checkpoint module including at least two parallel linear layers, the forward calculation of each linear layer of the last layer is skipped, and the gradients of each linear layer of the last layer are calculated respectively according to the gradient formulas corresponding to the each linear layer of the last layer.

[0013] In a possible implementation, calculating the gradient of the last layer according to the gradient formula corresponding to the last layer includes:

[0014] Reading the output gradient of the last layer and the input of the last layer from the video memory, wherein the output gradient of the last layer is obtained by passing the gradient back through the subsequent module of the checkpoint module, and the input of the last layer is obtained by re-executing the forward calculation of the network layers other than the last layer in the checkpoint module;

[0015] Calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer;

[0016] Calculate the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer.

[0017] In one possible implementation,

[0018] Calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer includes: multiplying the transpose of the weight matrix of the last layer by the output gradient of the last layer to obtain the input gradient of the last layer;

[0019] Calculating the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer includes: multiplying the output gradient of the last layer by the transpose of the input of the last layer to obtain the weight gradient of the last layer.

[0020] In a possible implementation, the method further includes:

[0021] In the backpropagation stage, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, for any network layer other than the last layer in the checkpoint module, the forward calculation of the network layer is re-executed to obtain the intermediate activation value of the network layer, and the gradient of the network layer is calculated based on the intermediate activation value of the network layer.

[0022] In a possible implementation, calculating the gradient of the network layer according to the intermediate activation value of the network layer includes:

[0023] Through the automatic differentiation mechanism, back propagation calculation is performed based on the intermediate activation values ​​of the network layer to obtain the gradient of the network layer.

[0024] In one possible implementation, the checkpoint module includes at least one of the following types:

[0025] Feedforward network module;

[0026] Multilayer Perceptron module;

[0027] A module that combines layer normalization and linear layers.

[0028] According to another aspect of the present disclosure, there is provided a model training device, comprising:

[0029] A determination module, used to determine a checkpoint module in a target model;

[0030] A forward propagation module, configured to, during the forward propagation phase, store the input and output of any checkpoint module in a video memory and release the intermediate activation value of the checkpoint module in the video memory; wherein the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module during the backward propagation phase, and the output of the checkpoint module is used for forward computation of subsequent modules of the checkpoint module;

[0031] A back propagation module is configured to, in the back propagation phase, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, skip the forward calculation of the last layer, calculate the gradient of the last layer according to the gradient formula corresponding to the last layer, and complete the back propagation of the checkpoint module according to the gradients of each layer in the checkpoint module.

[0032] In one possible implementation, the back propagation module is used to:

[0033] In response to the last layer of the checkpoint module including only one linear layer, skipping the forward calculation of the linear layer of the last layer, and calculating the gradient of the linear layer of the last layer according to the gradient formula corresponding to the linear layer of the last layer;

[0034] or,

[0035] In response to the last layer of the checkpoint module including at least two parallel linear layers, the forward calculation of each linear layer of the last layer is skipped, and the gradients of each linear layer of the last layer are calculated respectively according to the gradient formulas corresponding to the each linear layer of the last layer.

[0036] In one possible implementation, the back propagation module is used to:

[0037] Reading the output gradient of the last layer and the input of the last layer from the video memory, wherein the output gradient of the last layer is obtained by passing the gradient back through the subsequent module of the checkpoint module, and the input of the last layer is obtained by re-executing the forward calculation of the network layers other than the last layer in the checkpoint module;

[0038] Calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer;

[0039] Calculate the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer.

[0040] In one possible implementation, the back propagation module is used to:

[0041] Multiplying the transpose of the weight matrix of the last layer by the output gradient of the last layer to obtain the input gradient of the last layer;

[0042] The output gradient of the last layer is multiplied by the transpose of the input of the last layer to obtain the weight gradient of the last layer.

[0043] In a possible implementation, the back propagation module is further configured to:

[0044] In the backpropagation stage, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, for any network layer other than the last layer in the checkpoint module, the forward calculation of the network layer is re-executed to obtain the intermediate activation value of the network layer, and the gradient of the network layer is calculated based on the intermediate activation value of the network layer.

[0045] In one possible implementation, the back propagation module is used to:

[0046] Through the automatic differentiation mechanism, back propagation calculation is performed based on the intermediate activation values ​​of the network layer to obtain the gradient of the network layer.

[0047] In one possible implementation, the checkpoint module includes at least one of the following types:

[0048] Feedforward network module;

[0049] Multilayer Perceptron module;

[0050] A combination module of layer normalization layers and linear layers.

[0051] According to another aspect of the present disclosure, a model training device is provided, comprising a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of the above method.

[0052] According to another aspect of the present disclosure, a non-volatile computer-readable storage medium is provided, on which a computer program is stored. When the computer program is executed by a processor, the steps of the above method are implemented.

[0053] According to another aspect of the present disclosure, a computer program product is provided, including a computer program, or a non-volatile computer-readable storage medium carrying the computer program, wherein the computer program implements the steps of the above method when executed by a processor.

[0054] In an embodiment of the present disclosure, by determining a checkpoint module in a target model, in the forward propagation phase, for any checkpoint module, the input and output of the checkpoint module are saved in a video memory, and the intermediate activation value of the checkpoint module in the video memory is released, wherein the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module in the backward propagation phase, and the output of the checkpoint module is used for the forward calculation of subsequent modules of the checkpoint module. In the backward propagation phase, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, the forward calculation of the last layer is skipped, the gradient of the last layer is calculated according to the gradient formula corresponding to the last layer, and the backward propagation of the checkpoint module is completed according to the gradients of each layer in the checkpoint module. Therefore, by skipping the forward calculation of the last linear layer of the checkpoint module and directly calculating its gradient using a mathematical formula, the computational overhead can be significantly reduced while ensuring the same computational accuracy and video memory savings as the standard recalculation scheme. For example, in a multi-layer perceptron module, traditional recalculation requires recalculating the forward process of all linear layers, while the disclosed embodiment can reduce the computational overhead to 2 / 3 of the original by optimizing the back propagation process of the last linear layer.

[0055] In addition, the disclosed embodiments make the configuration of recomputation more flexible, and provide more fine-grained configuration possibilities for video memory optimization in model training. Due to computational overhead limitations, traditional solutions usually only apply recomputation to specific modules (such as core-attention modules) or the entire Transformer layer, while the disclosed embodiments can optimize any module in the model that ends with a linear layer, thereby achieving a better balance between video memory savings and computational efficiency. For example, for a combination module of a layer normalization layer and a linear layer, the additional overhead caused by recomputation is almost negligible because the forward calculation of the linear layer is skipped. This allows small modules that were originally difficult to apply recomputation due to excessive computational costs to be flexibly incorporated into the video memory optimization strategy, expanding the scope of application and configuration flexibility of the video memory optimization strategy.

[0056] Further features and aspects of the present disclosure will become apparent from the following detailed description of exemplary embodiments with reference to the attached drawings. BRIEF DESCRIPTION OF THE DRAWINGS

[0057] The accompanying drawings, which are incorporated in and constitute a part of the specification, illustrate exemplary embodiments, features, and aspects of the disclosure and, together with the description, serve to explain the principles of the disclosure.

[0058] Figure 1 Schematic diagram showing the mechanism for preserving the input x during the forward computation of the linear layer y=Wx.

[0059] Figure 2 Figure 2 shows the memory usage of a typical multilayer perceptron (MLP) module during training.

[0060] Figure 3 A schematic diagram showing the structure of the multilayer perceptron module and its connection relationship with other modules in a neural network.

[0061] Figure 4 A flowchart of the model training method provided by an embodiment of the present disclosure is shown.

[0062] Figure 5 A schematic diagram showing the data flow of the model training method provided by an embodiment of the present disclosure.

[0063] Figure 6 A block diagram of a model training device provided by an embodiment of the present disclosure is shown.

[0064] Figure 7 is a block diagram of a model training device 1900 according to an exemplary embodiment. DETAILED DESCRIPTION

[0065] Various exemplary embodiments, features, and aspects of the present disclosure will be described in detail below with reference to the accompanying drawings. The same reference numerals in the accompanying drawings represent elements with the same or similar functions. Although various aspects of the embodiments are shown in the accompanying drawings, the drawings are not necessarily drawn to scale unless otherwise indicated.

[0066] As used herein, the terms "comprises," "comprising," "having," or variations thereof are open ended and include one or more stated features, integers, elements, steps, parts, or functions, but do not preclude the presence or addition of one or more other features, integers, elements, steps, parts, functions, or groups thereof.

[0067] When an element is referred to as being "connected," "coupled," "responsive" or variations thereof to another element, it can be directly connected, coupled or responsive to the other element or intervening elements may be present.

[0068] Although the terms first, second, third, etc. may be used herein to describe various elements / operations, these elements / operations should not be limited by these terms. These terms are only used to distinguish one element / operation from another element / operation. Therefore, without departing from the teachings of the present invention, the first element / operation in some embodiments may be referred to as the second element / operation in other embodiments.

[0069] The word “exemplary” is used exclusively herein to mean “serving as an example, example, or illustration.” Any embodiment described herein as “exemplary” is not necessarily to be construed as preferred or advantageous over other embodiments.

[0070] In addition, numerous specific details are provided in the following detailed description to better illustrate the present disclosure. Those skilled in the art will appreciate that the present disclosure can be practiced without certain specific details. In some instances, methods, means, components, and circuits well known to those skilled in the art are not described in detail in order to highlight the main points of the present disclosure.

[0071] The following is an introduction to the key technical concepts involved in the embodiments of the present disclosure:

[0072] 1. Automatic Differentiation (Autograd)

[0073] Automatic differentiation is the core mechanism of modern deep learning frameworks (such as PyTorch and TensorFlow). It records the forward computation process through a dynamic computation graph and automatically calculates the gradient using the chain rule during backpropagation.

[0074] Taking the linear layer y=Wx as an example, during the forward calculation process, the framework will record the operational relationship between the input x and the weight W, and save the key intermediate variable x. When performing backpropagation, the system automatically derives the gradient calculation formula according to the chain rule, including the weight gradient and input gradient Figure 1 Schematic diagram showing the mechanism for saving input x during the forward calculation of the linear layer y=Wx. Figure 1 As shown in Figure 1, the linear layer saves the input x during forward propagation so that the gradient can be accurately calculated during backpropagation.

[0075] 2. Video memory usage

[0076] During the training of deep learning models, video memory usage mainly consists of four parts: model weights, optimizer states, cached activation values, and temporary tensors. Among them, the video memory usage of model weights and optimizer states is relatively fixed, and is determined by the model structure, data type, distribution strategy, and optimization algorithm. Temporary tensors account for a relatively small proportion in large model training. Cache activation values ​​are the main source of video memory usage, and the size of the video memory space occupied by these saved intermediate results is directly related to the model depth and batch size. Specifically, the deeper the model (the more layers) and the larger the batch size, the more activation values ​​need to be cached, and the more video memory capacity is consumed. These saved intermediate results will be used for gradient calculation during backpropagation, so they must be fully retained during the forward propagation process.

[0077] Figure 2 Figure 2 shows the memory usage of a typical multilayer perceptron (MLP) module during training. Figure 2 As shown, the input of the linear layer must be preserved (corresponding to Figure 2 The activation function input is saved or not according to the specific function type (corresponding to Figure 2 In large model training scenarios, the pressure on video memory caused by caching activation values ​​is particularly prominent. Figure 2 As shown in the figure, even a simple multi-layer perceptron module needs to store intermediate results for multiple linear layers and activation functions. This memory usage often means that the memory requirements for model training often far exceed the storage requirements of the model parameters themselves, becoming a key bottleneck limiting the training efficiency of large models.

[0078] 3. Gradient Checkpointing

[0079] Gradient checkpointing is a memory optimization technique that reduces memory usage by selectively discarding some intermediate activation values ​​and recalculating them during backpropagation.

[0080] The following introduces the standard gradient checkpoint technology and the standard recomputation scheme (Vanilla Recomputation).

[0081] In standard gradient checkpointing techniques, during the forward propagation phase, the model is first divided into several modules, and some modules are selected as checkpoint modules. For these checkpoint modules, only their input and output data are saved, while activation values ​​generated during the intermediate calculations are discarded. Taking the Transformer model as an example, the feedforward neural network (FFN) is usually set as the checkpoint module.

[0082] In the backpropagation phase, for each checkpoint module, its complete forward calculation process (Forward Pass) needs to be re-executed to restore the intermediate activation value, and then the gradients of all parameters in the module (including linear layers and nonlinear layers) are calculated through the automatic differentiation mechanism.

[0083] Figure 3 A schematic diagram showing the structure of the multilayer perceptron module and its connection relationship with other modules in the neural network. Figure 3 In the example shown, the neural network model includes module 1, multi-layer perceptron module, module 2 and module 3. Figure 3 In the example, data flows from left to right and the loss is calculated at the end.

[0084] Module 1 can receive input data and process it according to y=F1(x) to generate output x1. The output x1 of module 1 will serve as the input of the multilayer perceptron module.

[0085] The multilayer perceptron module consists of a linear layer 1, an activation function, and a linear layer 2. Linear layer 1 performs a linear transformation on input x1 according to y = W1*x, producing an intermediate output y1. Subsequently, the activation function (y = ρ(x)) performs a nonlinear processing on y1, generating an intermediate output x2 = ρ(y1). Finally, linear layer 2 performs a second linear transformation on x2 according to y = W2*x, producing the output y2.

[0086] The output y2 of the multilayer perceptron module will be transmitted to module 2. Modules 2 and 3 will process according to y=F2(x) and y=F3(x) respectively, and finally calculate the loss.

[0087] Taking the multi-layer perceptron module as an example, the standard forward propagation process is:

[0088] Input: x1;

[0089] y1=W1*x1; save x1;

[0090] x2 = act(y1); save y1;

[0091] y2=W2*x2; save x2.

[0092] Among them, act() represents the activation function.

[0093] The standard backpropagation process is:

[0094] enter:

[0095]

[0096] Among them, backward() calls the automatic differentiation mechanism.

[0097] After adopting the standard recalculation technology, the calculation process is adjusted to:

[0098] Forward propagation:

[0099] Input: x1;

[0100] y1=W1*x1; save x1;

[0101] x2=act(y1);

[0102] y2=W2*x2.

[0103] Backward Propagation:

[0104] enter:

[0105] Recalculate: y1 = W1 * x1; save x1;

[0106] Recalculate: x2 = act(y1); save y1;

[0107] Recalculate: y2=W2*x2; save x2;

[0108]

[0109] The main advantage of the standard recalculation technique is that it can save memory usage for y1 and x2 after the forward propagation is completed and before the backpropagation begins (usually during the forward and backpropagation of other modules outside the multilayer perceptron). However, the corresponding cost is that the entire forward propagation process of the multilayer perceptron must be recalculated.

[0110] Taking a multilayer perceptron as an example, standard recalculation requires recalculating the forward pass of all linear layers (e.g., calculating y = Wx) during backpropagation, and then performing automatic differentiation. This approach incurs additional computational overhead due to repeated forward computations of the linear layers, especially for large matrix multiplications (e.g., fully connected layers in large language models (LLMs)).

[0111] Furthermore, when applying gradient checkpoints, related technologies typically recalculate only specific modules (such as the core attention module) or the entire Transformer layer, and rarely use recalculation techniques on small internal modules. This is because the computational overhead of traditional recalculation on these small modules is too high to be acceptable. As a result, related technologies have a clear trade-off between video memory optimization and computational efficiency, which limits the flexibility of recalculation strategies.

[0112] To address technical problems similar to those described above, embodiments of the present disclosure provide a model training method. A checkpoint module in a target model is determined. During the forward propagation phase, for any checkpoint module, the input and output of the checkpoint module are saved in video memory, and the intermediate activation value of the checkpoint module in the video memory is released. The input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module during the backward propagation phase, and the output of the checkpoint module is used for forward computation of subsequent modules of the checkpoint module. During the backward propagation phase, for any checkpoint module, if the last layer of the checkpoint module is a linear layer, the forward computation of the last layer is skipped, and the gradient of the last layer is calculated according to the gradient formula corresponding to the last layer. Backward propagation of the checkpoint module is then completed based on the gradients of each layer in the checkpoint module. Thus, by skipping the forward computation of the last linear layer of the checkpoint module and directly calculating its gradient using a mathematical formula, computational overhead can be significantly reduced while maintaining the same computational accuracy and video memory savings as a standard recalculation solution. For example, in a multi-layer perceptron module, traditional recalculation requires recalculating the forward process of all linear layers, while the disclosed embodiment can reduce the computational overhead to 2 / 3 of the original by optimizing the back propagation process of the last linear layer.

[0113] In addition, the disclosed embodiments make the configuration of recomputation more flexible, and provide more fine-grained configuration possibilities for video memory optimization in model training. Due to computational overhead limitations, traditional solutions usually only apply recomputation to specific modules (such as core-attention modules) or the entire Transformer layer, while the disclosed embodiments can optimize any module in the model that ends with a linear layer, thereby achieving a better balance between video memory savings and computational efficiency. For example, for a combination module of a layer normalization layer and a linear layer, the additional overhead caused by recomputation is almost negligible because the forward calculation of the linear layer is skipped. This allows small modules that were originally difficult to apply recomputation due to excessive computational costs to be flexibly incorporated into the video memory optimization strategy, expanding the scope of application and configuration flexibility of the video memory optimization strategy.

[0114] The model training method provided by the embodiment of the present disclosure is described in detail below with reference to the accompanying drawings.

[0115] Figure 4 A flow chart of the model training method provided by an embodiment of the present disclosure is shown. In one possible implementation, the execution subject of the model training method may be a model training device. For example, the model training method may be executed by a terminal device or a server or other electronic device. The terminal device may be a user equipment (UE), a mobile device, a user terminal, a terminal, a handheld device or a computing device. In some possible implementations, the model training method may be implemented by a processor calling computer-readable instructions stored in a memory. Figure 4 As shown, the model training method includes steps S41 to S43.

[0116] In step S41 , a checkpoint module in the target model is determined.

[0117] In step S42, in the forward propagation phase, for any checkpoint module, the input and output of the checkpoint module are saved in the video memory, and the intermediate activation value of the checkpoint module in the video memory is released; wherein, the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module in the backward propagation phase, and the output of the checkpoint module is used for the forward calculation of the subsequent modules of the checkpoint module.

[0118] In step S43, in the back propagation stage, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, the forward calculation of the last layer is skipped, the gradient of the last layer is calculated according to the gradient formula corresponding to the last layer, and the back propagation of the checkpoint module is completed according to the gradients of each layer in the checkpoint module.

[0119] In the embodiments of the present disclosure, the target model may refer to any neural network model to be trained.

[0120] In model training, a checkpoint module can refer to a subnetwork or layer combination that applies gradient checkpoint technology. A checkpoint module can be a continuous network layer in the target model that has a high video memory usage and whose intermediate activation values ​​can be recalculated, such as a feedforward network, a residual block, etc. Its core principle is to save only the input and output of the module during forward propagation, discarding the intermediate activation values ​​to save video memory; and restore the intermediate results by recalculation during backward propagation. In the embodiment of the present disclosure, when the checkpoint module ends with a linear layer, the forward recalculation of the linear layer can be skipped, and the gradient can be directly derived using a mathematical formula, thereby significantly reducing the computational overhead while maintaining the same video memory saving effect.

[0121] In a possible implementation, the checkpoint module includes at least one of the following types: a feedforward network module; a multi-layer perceptron module; and a combination module of a layer normalization layer and a linear layer.

[0122] As an example of this implementation, the checkpoint module can include a feedforward neural network (FFN) module. A feedforward neural network module typically contains two linear transformation layers and an intermediate nonlinear activation function, which generates a large number of intermediate activation values ​​that need to be cached. Taking the feedforward neural network module of the Transformer as an example, its calculation process involves dimensional expansion (such as from 512 dimensions to 2048 dimensions) and contraction. The output activation values ​​of the expansion phase will occupy a significant amount of video memory, but these values ​​can be completely recalculated based on the input.

[0123] As an example of this implementation, the checkpoint module may include a multi-layer perceptron module. As a general fully connected network structure, the multi-layer perceptron module has obvious video memory occupation characteristics. A typical multi-layer perceptron module may contain a stack of multiple linear layers and activation functions. If the gradient checkpoint technology is not used, the output of each layer needs to be cached for back propagation. When the multi-layer perceptron module ends with a linear layer, the embodiment of the present disclosure can only save the input of the first layer and recalculate the intermediate results during back propagation, and can skip the forward calculation of the last linear layer.

[0124] As an example of this implementation, the checkpoint module may include a combination module of a layer normalization (LayerNorm) layer and a linear layer. In traditional solutions, the combination of layer normalization and subsequent linear layers is often difficult to optimize due to the high recalculation cost. However, the method provided by the embodiment of the present disclosure only requires the layer normalization part to be recalculated (the computational cost is extremely small), and the gradient of the linear layer is directly derived through mathematical formulas, making the recalculation overhead of this type of module almost negligible.

[0125] In the disclosed embodiment, when the target model performs forward computation, the input and output of each checkpoint module can be saved in the video memory, and the intermediate activation values ​​of each checkpoint module stored in the video memory can be discarded (i.e., released). The input of any checkpoint module is used to recalculate the intermediate activation values ​​of the checkpoint module during the backpropagation phase, and the output of any checkpoint module is used for the forward computation of subsequent modules of the checkpoint module.

[0126] Taking the feedforward network module of Transformer as an example, the module usually contains two linear layers and an activation function (such as GeLU). The traditional training method requires caching all intermediate activation values ​​between linear layers (such as the output of the first linear layer, the output of the activation function, etc.), and these temporary data will take up a large amount of video memory. However, using the embodiment of the present disclosure, the system can only retain the initial input tensor and final output tensor of the feedforward network module, and all temporary data generated in the intermediate process (such as the output tensor of the first linear layer, the output of the activation function, etc.) are temporarily stored in the video memory and released immediately. For example, in the checkpoint module, after a network layer outputs an intermediate activation value, the intermediate activation value can be temporarily stored in the video memory. When the next network layer of the network layer has used up the intermediate activation value, the intermediate activation value can be released.

[0127] Although intermediate activation values ​​are discarded, the checkpoint module's inputs and outputs are preserved, allowing necessary intermediate results to be restored during backpropagation by re-executing the forward computation. Combined with the optimization strategy for linear layers in the disclosed embodiments, when the checkpoint module ends with a linear layer, some redundant recalculation operations can be avoided, further improving computational efficiency while maintaining memory savings.

[0128] In addition, in the embodiment of the present disclosure, in the forward propagation stage, the input, output and all intermediate activation values ​​of each non-checkpoint module can be completely retained in the video memory for direct use in the subsequent backward propagation stage.

[0129] In the back-propagation phase, the following two situations can be distinguished for the processing of the checkpoint module:

[0130] When the last layer of the checkpoint module is not a linear layer, a standard recalculation scheme can be used, that is, re-execute the full forward calculation of the checkpoint module to restore all intermediate activation values, and then perform gradient calculation;

[0131] When the last layer of the checkpoint module is a linear layer, an optimized calculation scheme can be implemented, that is, skipping the forward calculation recalculation step of the last linear layer of the checkpoint module, calculating the gradient according to the gradient calculation formula of the last linear layer, and still executing the standard recalculation process for other network layers in the checkpoint module.

[0132] In one possible implementation, in response to the last layer of the checkpoint module being a linear layer, skipping the forward calculation of the last layer, and calculating the gradient of the last layer according to the gradient formula corresponding to the last layer, includes: in response to the last layer of the checkpoint module including only one linear layer, skipping the forward calculation of the linear layer of the last layer, and calculating the gradient of the linear layer of the last layer according to the gradient formula corresponding to the linear layer of the last layer; or, in response to the last layer of the checkpoint module including at least two parallel linear layers, skipping the forward calculation of each linear layer of the last layer, and calculating the gradient of each linear layer of the last layer according to the gradient formula corresponding to each linear layer of the last layer.

[0133] As an example of this implementation, in response to the last layer of the checkpoint module including only one linear layer, the forward calculation of the linear layer of the last layer is skipped, and the gradient of the linear layer of the last layer is calculated according to the gradient formula corresponding to the linear layer of the last layer.

[0134] In this example, for a checkpoint module that ends with a single linear layer (such as a standard fully connected layer), during the backpropagation phase, the forward calculation of the last linear layer can be skipped and the mathematical formula can be directly applied to calculate the gradient.

[0135] As another example of this implementation, in response to the last layer of the checkpoint module including at least two parallel linear layers, the forward calculation of each linear layer of the last layer is skipped, and the gradients of each linear layer of the last layer are calculated separately according to the gradient formulas corresponding to the each linear layer of the last layer.

[0136] In this example, for the checkpoint module at the end of the parallel linear layer, during the backpropagation phase, the forward calculation of each linear layer in the last layer can be skipped, and the gradient is calculated by applying a mathematical formula to each linear layer separately. There is no data dependency between these parallel linear layers. For example, in the MoE (Mixture of Experts) model, assuming that the expert layer contains four parallel linear projection layers, the traditional method needs to recalculate the forward propagation of these four linear projection layers, while this solution can skip the forward calculation of these four linear projection layers during the backpropagation phase.

[0137] In a possible implementation, calculating the gradient of the last layer according to the gradient formula corresponding to the last layer includes: reading the output gradient of the last layer and the input of the last layer from the video memory, wherein the output gradient of the last layer is obtained by passing the gradient back through the subsequent module of the checkpoint module, and the input of the last layer is obtained by re-executing the forward calculation of the network layers other than the last layer in the checkpoint module; calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer; calculating the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer, wherein the input of the last layer is restored by re-executing the forward calculation of the network layers other than the last layer in the checkpoint module.

[0138] Among them, the output gradient of the last layer of the checkpoint module is the output gradient of the checkpoint module.

[0139] During the back-propagation process, the output gradient of the last layer of the checkpoint module It is obtained by the chain rule from the downstream network layer or loss function. When the back propagation calculation is propagated to the checkpoint module, the subsequent layers (such as the loss function or higher-level network) first calculate the gradient of the checkpoint module output (y) (Right now ), the gradient value can be stored in the video memory by the automatic differentiation framework (such as PyTorch / TensorFlow) and used as the input for the gradient calculation of the checkpoint module.

[0140] If a checkpoint module ends with a linear layer, the input value of the last linear layer can be recovered by recalculating all predecessor layers in the checkpoint module except the last linear layer, and temporarily storing the input value of the last linear layer in video memory. Then, based on the mathematical properties of linear transformations, the input gradient and weight gradient can be directly calculated using the output gradient and the recovered input value, respectively.

[0141] This implementation avoids the need to fully recalculate the forward process of the entire checkpoint module in traditional solutions. It can significantly reduce the recalculation burden of the backpropagation stage while maintaining the mathematical equivalence of gradient calculations. This can save considerable computing resources, especially when the last linear layer of the checkpoint module is large (such as the high-dimensional projection layer in large language models).

[0142] In a possible implementation, calculating the input gradient of the last layer based on the weight matrix of the last layer and the output gradient of the last layer includes: multiplying the transpose of the weight matrix of the last layer by the output gradient of the last layer to obtain the input gradient of the last layer; calculating the weight gradient of the last layer based on the output gradient of the last layer and the input of the last layer includes: multiplying the output gradient of the last layer by the transpose of the input of the last layer to obtain the weight gradient of the last layer.

[0143] For example, the forward computation of the last linear layer in the checkpoint module is y = Wx + b, where W represents the weight matrix of the last linear layer, with dimensions m × n; x represents the input vector of the last linear layer, with dimensions n × 1; and y represents the output vector of the last linear layer, with dimensions m × 1.

[0144] Then, the input gradient of the last linear layer in, Represents the output gradient of the last linear layer (dimension m×1). By transposing the weight matrix With the output gradient Multiply to get the input gradient (dimension n×1).

[0145] The weight gradient of the last linear layer That is, the output gradient (dimension m×1) and the transpose of the input vector Multiply to get the weight gradient (dimension m×n).

[0146] This implementation can significantly improve the training efficiency of linear layers (such as fully connected layers in large models) while maintaining mathematical equivalence.

[0147] In one possible implementation, the method further includes: in the back-propagation stage, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, for any network layer other than the last layer in the checkpoint module, re-executing the forward calculation of the network layer to obtain the intermediate activation value of the network layer, and calculating the gradient of the network layer based on the intermediate activation value of the network layer.

[0148] In this implementation, if the last layer of a checkpoint module is a linear layer, then the forward computations of all layers except the last layer in the checkpoint module can be re-executed to recover their intermediate activation values. Subsequently, the gradients of the corresponding layers can be calculated based on these recovered intermediate activation values.

[0149] In a possible implementation, calculating the gradient of the network layer according to the intermediate activation value of the network layer includes: performing backpropagation calculation based on the intermediate activation value of the network layer through an automatic differentiation mechanism to obtain the gradient of the network layer.

[0150] In this implementation, during the backward propagation of model training, automatic differentiation (Autograd) is used to calculate gradients for all layers other than the final linear layer within the checkpoint module. This mechanism dynamically tracks the computational relationships in the forward computation graph to automatically construct the backward propagation path. In gradient checkpointing, automatic differentiation works in conjunction with the recalculation mechanism. The system only needs to temporarily restore the necessary intermediate activation values ​​to automatically complete the gradient propagation from the output layer to the target layer using the chain rule.

[0151] In this implementation, the model training method can be used in a framework with an automatic differentiation mechanism, such as Torch, to achieve a recalculation method with the same memory saving effect and accuracy as the standard recalculation method, but with higher computational efficiency.

[0152] Corresponding to the standard recalculation process described above, the improved recalculation process in the embodiment of the present disclosure is as follows:

[0153] Forward propagation:

[0154] Input: x1;

[0155] y1=W1*x1; save x1;

[0156] x2=act(y1);

[0157] y2=W2*x2.

[0158] Backward Propagation:

[0159] enter:

[0160] Recalculate: y1 = W1 * x1; save x1;

[0161] Recalculate: x2 = act(y1); save y1;

[0162] Skip the forward calculation of y2 and use the gradient formula to calculate the gradient of linear layer 2:

[0163]

[0164] The model training method provided in the embodiments of the present disclosure can be applied to technical fields such as artificial intelligence (AI), artificial intelligence infrastructure (AI-Infra), large model training framework, and training acceleration, and is not limited here.

[0165] The following describes the model training method provided by the embodiment of the present disclosure through a specific application scenario.

[0166] Figure 5 A schematic diagram showing the data flow of the model training method provided by an embodiment of the present disclosure. Figure 5 Contains the processing logic of the two stages of forward propagation and backpropagation.

[0167] 1. Forward propagation stage

[0168] x1=F1(x0), non-checkpoint module 1, saves all activation values;

[0169] y1=W1*x1, the linear layer 1 in the checkpoint module, saves the input x1 (i.e., the input of the checkpoint module);

[0170] x2 = ρ(y1), the activation function in the checkpoint module, y1 is not saved;

[0171] y2=W2*x2, linear layer 2 in the checkpoint module, x2 is not saved;

[0172] y3 = F2(y2), non-checkpoint module 2, saves all activation values;

[0173] Loss=F3(y3), loss function, saves all activation values.

[0174] 2. Backpropagation phase

[0175] By automatic differentiation, the output gradient of non-checkpoint module 2 is calculated:

[0176] By automatic differentiation, the output gradient of linear layer 2 is calculated:

[0177] Recalculate y1=W1*x1 and temporarily store x1 in the video memory;

[0178] Recalculate x2 = ρ(y1) and temporarily store y1 in the video memory;

[0179] Use the input gradient formula corresponding to linear layer 2 to calculate the input gradient of linear layer 2:

[0180] The parameter gradient formula corresponding to linear layer 2 is used to calculate the parameter gradient of linear layer 2:

[0181] By automatic differentiation, the gradient of the activation function is calculated:

[0182] By automatic differentiation, the input gradient and parameter gradient of linear layer 1 are calculated:

[0183] By automatic differentiation, the input gradient of module 1 is calculated:

[0184] exist Figure 5 In the example shown, during the forward propagation phase, only the checkpoint module inputs and outputs (x1, y2) are saved, and the intermediate activation values ​​(y1, x2) are discarded. During the backward propagation phase, the partial activation values ​​(y1, x2) are selectively recalculated, and the recalculation of the linear layer 2 is skipped, directly calculating the gradient using the formula. This reduces the need to save intermediate activation values ​​during the forward propagation phase, saving video memory. Skipping the recalculation of the final linear layer during the backward propagation phase reduces the amount of computation.

[0185] In one example, the multilayer perceptron module in the MoE model can be used as a checkpoint module. The multilayer perceptron module in the MoE model consists of a grouped linear layer 1, a SiLU activation function, and a grouped linear layer 2. Grouped linear layer 2 is the last layer of the multilayer perceptron module and includes multiple linear layers that have no data dependencies. By applying the optimized recalculation technology provided by the embodiments of the present disclosure, only the forward process of grouped linear layer 1 and the SiLU activation function needs to be recalculated, without recalculating the forward process of grouped linear layer 2.

[0186] When applying the gradient checkpoint technology to the multi-layer perceptron module, 288MB of video memory can be saved for each multi-layer perceptron module. If the entire model has 60 multi-layer perceptron modules, a total of 17GB of video memory can be saved.

[0187] If the standard recomputation solution is used, the additional computational overhead includes 1 TFlops of forward computation overhead for grouped linear layer 1 and 0.5 TFlops of forward computation overhead for grouped linear layer 2. In other words, the standard recomputation solution requires an additional 1.5 TFlops of computational cost. However, through optimization, the disclosed embodiments reduce this overhead to 1 TFlops, significantly improving training efficiency.

[0188] It is understood that the above-mentioned various method embodiments mentioned in this disclosure can be combined with each other to form combined embodiments without violating the principle logic. Due to space limitations, this disclosure will not go into details. It is understood by those skilled in the art that in the above-mentioned methods of specific implementation, the specific execution order of each step should be determined by its function and possible internal logic.

[0189] In addition, the present disclosure also provides a model training device, a non-volatile computer-readable storage medium, and a computer program product, all of which can be used to implement any model training method provided by the present disclosure. The corresponding technical solutions and technical effects can be found in the corresponding records in the method section and will not be repeated here.

[0190] Figure 6 FIG. 1 is a block diagram of a model training device provided by an embodiment of the present disclosure. Figure 6 As shown, the model training device includes:

[0191] A determination module 61 is used to determine a checkpoint module in a target model;

[0192] A forward propagation module 62 is configured to, during the forward propagation phase, store the input and output of any checkpoint module in a video memory and release the intermediate activation value of the checkpoint module from the video memory; the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module during the backward propagation phase, and the output of the checkpoint module is used for forward computation of subsequent modules of the checkpoint module;

[0193] The back propagation module 63 is used to, in the back propagation phase, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, skip the forward calculation of the last layer, calculate the gradient of the last layer according to the gradient formula corresponding to the last layer, and complete the back propagation of the checkpoint module according to the gradients of each layer in the checkpoint module.

[0194] In a possible implementation, the back propagation module 63 is used to:

[0195] In response to the last layer of the checkpoint module including only one linear layer, skipping the forward calculation of the linear layer of the last layer, and calculating the gradient of the linear layer of the last layer according to the gradient formula corresponding to the linear layer of the last layer;

[0196] or,

[0197] In response to the last layer of the checkpoint module including at least two parallel linear layers, the forward calculation of each linear layer of the last layer is skipped, and the gradients of each linear layer of the last layer are calculated respectively according to the gradient formulas corresponding to the each linear layer of the last layer.

[0198] In a possible implementation, the back propagation module 63 is used to:

[0199] Reading the output gradient of the last layer and the input of the last layer from the video memory, wherein the output gradient of the last layer is obtained by passing the gradient back through the subsequent module of the checkpoint module, and the input of the last layer is obtained by re-executing the forward calculation of the network layers other than the last layer in the checkpoint module;

[0200] Calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer;

[0201] Calculate the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer.

[0202] In a possible implementation, the back propagation module 63 is used to:

[0203] Multiplying the transpose of the weight matrix of the last layer by the output gradient of the last layer to obtain the input gradient of the last layer;

[0204] The output gradient of the last layer is multiplied by the transpose of the input of the last layer to obtain the weight gradient of the last layer.

[0205] In a possible implementation, the back propagation module 63 is further configured to:

[0206] In the backpropagation stage, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, for any network layer other than the last layer in the checkpoint module, the forward calculation of the network layer is re-executed to obtain the intermediate activation value of the network layer, and the gradient of the network layer is calculated based on the intermediate activation value of the network layer.

[0207] In a possible implementation, the back propagation module 63 is used to:

[0208] Through the automatic differentiation mechanism, back propagation calculation is performed based on the intermediate activation values ​​of the network layer to obtain the gradient of the network layer.

[0209] In one possible implementation, the checkpoint module includes at least one of the following types:

[0210] Feedforward network module;

[0211] Multilayer Perceptron module;

[0212] A combination module of layer normalization layers and linear layers.

[0213] In some embodiments, the functions or modules included in the device provided by the embodiments of the present disclosure can be used to execute the method described in the above method embodiments. Its specific implementation and technical effects can refer to the description of the above method embodiments. For the sake of brevity, they will not be repeated here.

[0214] An embodiment of the present disclosure also provides a model training device, including a memory, a processor, and a computer program stored in the memory, wherein the processor executes the computer program to implement the steps of the above method.

[0215] An embodiment of the present disclosure further provides a non-volatile computer-readable storage medium having a computer program stored thereon, wherein the computer program implements the steps of the above method when executed by a processor.

[0216] An embodiment of the present disclosure further provides a computer program product, including a computer program, or a non-volatile computer-readable storage medium carrying the computer program, wherein the computer program implements the steps of the above method when executed by a processor.

[0217] Figure 7 1 is a block diagram of a model training apparatus 1900 according to an exemplary embodiment. For example, the apparatus 1900 can be provided as a server or a terminal device. Figure 7 The apparatus 1900 includes a processing component 1922, which further includes one or more processors, and a memory resource represented by a memory 1932 for storing instructions, such as an application, that can be executed by the processing component 1922. The application stored in the memory 1932 may include one or more modules, each corresponding to a set of instructions. In addition, the processing component 1922 is configured to execute the instructions to perform the above-described method.

[0218] The device 1900 may also include a power supply component 1926 configured to perform power management of the device 1900, a wired or wireless network interface 1950 configured to connect the device 1900 to a network, and an input / output interface 1958 (I / O interface). The device 1900 may operate based on an operating system stored in the memory 1932, such as Windows Server 2003. TM , MacOS X TM , Unix TM ,Linux TM , FreeBSD TM or similar.

[0219] In an exemplary embodiment, a non-volatile computer-readable storage medium is also provided, such as a memory 1932 including computer program instructions that can be executed by the processing component 1922 of the apparatus 1900 to perform the above-described method.

[0220] A computer-readable storage medium can be a tangible device that can hold and store programs / instructions used by an instruction execution device. A computer-readable storage medium can be, for example, but not limited to, an electrical storage device, a magnetic storage device, an optical storage device, an electromagnetic storage device, a semiconductor storage device, or any suitable combination thereof. More specific examples (a non-exhaustive list) of computer-readable storage media include: a portable computer disk, a hard disk, a random access memory (RAM), a read-only memory (ROM), an erasable programmable read-only memory (EPROM or flash memory), a static random access memory (SRAM), a portable compact disc read-only memory (CD-ROM), a digital versatile disk (DVD), a memory stick, a floppy disk, a mechanical encoding device, such as a punch card or a raised structure in a groove on which instructions are stored, and any suitable combination thereof. As used herein, a computer-readable storage medium is not to be construed as a transient signal per se, such as a radio wave or other freely propagating electromagnetic wave, an electromagnetic wave propagating through a waveguide or other transmission medium (e.g., a light pulse through a fiber optic cable), or an electrical signal transmitted through an electrical wire.

[0221] The computer programs (or computer-readable program instructions) described herein can be downloaded from a computer-readable storage medium to each computing / processing device, or downloaded to an external computer or external storage device via a network, such as the Internet, a local area network, a wide area network, and / or a wireless network. The network can include copper transmission cables, optical fiber transmission, wireless transmission, routers, firewalls, switches, gateway computers, and / or edge servers. The network adapter card or network interface in each computing / processing device receives the computer-readable program instructions from the network and forwards the computer-readable program instructions to be stored in the computer-readable storage medium in each computing / processing device.

[0222] The computer program (or computer program instructions) for performing the operations of the present disclosure may be assembly instructions, instruction set architecture (ISA) instructions, machine instructions, machine-dependent instructions, microcode, firmware instructions, state setting data, or source code or object code written in any combination of one or more programming languages, including object-oriented programming languages ​​such as Smalltalk, C++, and conventional procedural programming languages ​​such as "C" or similar programming languages. The computer readable program instructions may be executed entirely on the user's computer, partially on the user's computer, as a separate software package, partially on the user's computer and partially on a remote computer, or entirely on a remote computer or server. In the case of a remote computer, the remote computer may be connected to the user's computer via any type of network, including a local area network (LAN) or a wide area network (WAN), or may be connected to an external computer (e.g., via the Internet using an Internet service provider). In some embodiments, by utilizing state information of computer-readable program instructions to personalize and customize an electronic circuit, such as a programmable logic circuit, a field programmable gate array (FPGA), or a programmable logic array (PLA), the electronic circuit can execute the computer-readable program instructions to implement various aspects of the present disclosure.

[0223] Various aspects of the present disclosure are described herein with reference to flowcharts and / or block diagrams of methods, apparatus (systems), and computer program products according to embodiments of the present disclosure. It should be understood that each block of the flowcharts and / or block diagrams, and combinations of blocks in the flowcharts and / or block diagrams, can be implemented by computer-readable program instructions.

[0224] These computer-readable program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, or other programmable data processing device, thereby producing a machine, so that when these instructions are executed by the processor of the computer or other programmable data processing device, a device is generated that implements the functions / actions specified in one or more blocks in the flowchart and / or block diagram. These computer-readable program instructions can also be stored in a computer-readable storage medium, where these instructions cause the computer, programmable data processing device, and / or other device to operate in a specific manner. Thus, the computer-readable medium storing the instructions comprises an article of manufacture that includes instructions for implementing various aspects of the functions / actions specified in one or more blocks in the flowchart and / or block diagram.

[0225] Computer-readable program instructions may also be loaded onto a computer, other programmable data processing apparatus, or other device so that a series of operational steps are performed on the computer, other programmable data processing apparatus, or other device to produce a computer-implemented process, thereby causing the instructions executed on the computer, other programmable data processing apparatus, or other device to implement the functions / actions specified in one or more blocks in the flowchart and / or block diagram.

[0226] The flow charts and block diagrams in the accompanying drawings show the possible architecture, functions and operations of the systems, methods and computer program products according to multiple embodiments of the present disclosure. In this regard, each box in the flow chart or block diagram can represent a part of a module, program segment or instruction, and the part of the module, program segment or instruction contains one or more executable instructions for realizing the prescribed logical function. In some alternative implementations, the functions marked in the box can also occur in a sequence different from that marked in the accompanying drawings. For example, two consecutive boxes can actually be executed substantially in parallel, and they can sometimes be executed in the opposite order, depending on the functions involved. It should also be noted that each box in the block diagram and / or flow chart, and the combination of the boxes in the block diagram and / or flow chart can be implemented by a dedicated hardware-based system that performs the prescribed function or action, or can be implemented by a combination of dedicated hardware and computer instructions.

[0227] The computer program product may be implemented in hardware, software, or a combination thereof. In one embodiment, the computer program product is implemented as a computer storage medium. In another embodiment, the computer program product is implemented as a software product, such as a software development kit (SDK).

[0228] The above description of the various embodiments tends to emphasize the differences between the various embodiments. The same or similar aspects can be referenced with each other and will not be repeated herein for the sake of brevity.

[0229] If the technical solutions of the embodiments of the present disclosure involve personal information, the products applying the technical solutions of the embodiments of the present disclosure have clearly informed the personal information processing rules and obtained the individual's voluntary consent before processing the personal information. If the technical solutions of the embodiments of the present disclosure involve sensitive personal information, the products applying the technical solutions of the embodiments of the present disclosure have obtained the individual's separate consent before processing the sensitive personal information, and at the same time meet the "explicit consent" requirement. For example, on personal information collection devices such as cameras, clear and prominent signs are set to inform that the personal information collection scope has been entered and personal information will be collected. If the individual voluntarily enters the collection scope, it is deemed that they agree to the collection of their personal information; or on the personal information processing device, when the personal information processing rules are notified by obvious signs / information, the individual's authorization is obtained through pop-up information or by asking the individual to upload their personal information. The personal information processing rules may include information such as the personal information processor, the purpose of personal information processing, the processing method, and the type of personal information processed.

[0230] While various embodiments of the present disclosure have been described above, the above descriptions are illustrative, non-exhaustive, and not intended to be limiting of the disclosed embodiments. Many modifications and variations will be apparent to those skilled in the art without departing from the scope and spirit of the described embodiments. The terminology used herein is selected to best explain the principles of the embodiments, their practical applications, or improvements to existing technologies, or to enable others skilled in the art to understand the embodiments disclosed herein.

Claims

1. A model training method, characterized in that: include: Identify the checkpoint modules in the target model; During the forward propagation phase, for any checkpoint module, the input and output of the checkpoint module are saved in the video memory, and the intermediate activation value of the checkpoint module in the video memory is released; wherein the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module during the backward propagation phase, and the output of the checkpoint module is used for the forward calculation of the subsequent modules of the checkpoint module; In the back-propagation phase, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, the forward calculation of the last layer is skipped, the gradient of the last layer is calculated according to the gradient formula corresponding to the last layer, and the back-propagation of the checkpoint module is completed according to the gradients of each layer in the checkpoint module.

2. The method according to claim 1, characterized in that In response to the last layer of the checkpoint module being a linear layer, skipping the forward calculation of the last layer, and calculating the gradient of the last layer according to a gradient formula corresponding to the last layer, including: In response to the last layer of the checkpoint module including only one linear layer, skipping the forward calculation of the linear layer of the last layer, and calculating the gradient of the linear layer of the last layer according to the gradient formula corresponding to the linear layer of the last layer; or, In response to the last layer of the checkpoint module including at least two parallel linear layers, the forward calculation of each linear layer of the last layer is skipped, and the gradients of each linear layer of the last layer are calculated respectively according to the gradient formulas corresponding to the each linear layer of the last layer.

3. The method according to claim 1, characterized in that Calculating the gradient of the last layer according to the gradient formula corresponding to the last layer includes: Reading the output gradient of the last layer and the input of the last layer from the video memory, wherein the output gradient of the last layer is obtained by passing the gradient back through the subsequent module of the checkpoint module, and the input of the last layer is obtained by re-executing the forward calculation of the network layers other than the last layer in the checkpoint module; Calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer; Calculate the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer.

4. The method according to claim 3, characterized in that Calculating the input gradient of the last layer according to the weight matrix of the last layer and the output gradient of the last layer includes: multiplying the transpose of the weight matrix of the last layer by the output gradient of the last layer to obtain the input gradient of the last layer; Calculating the weight gradient of the last layer according to the output gradient of the last layer and the input of the last layer includes: multiplying the output gradient of the last layer by the transpose of the input of the last layer to obtain the weight gradient of the last layer.

5. The method according to any one of claims 1 to 4, characterized in that The method further comprises: In the backpropagation stage, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, for any network layer other than the last layer in the checkpoint module, the forward calculation of the network layer is re-executed to obtain the intermediate activation value of the network layer, and the gradient of the network layer is calculated based on the intermediate activation value of the network layer.

6. The method according to claim 5, characterized in that Calculating the gradient of the network layer according to the intermediate activation value of the network layer includes: Through the automatic differentiation mechanism, back propagation calculation is performed based on the intermediate activation values ​​of the network layer to obtain the gradient of the network layer.

7. The method according to any one of claims 1 to 4, characterized in that The checkpoint module includes at least one of the following types: Feedforward network module; Multilayer Perceptron module; A combination module of layer normalization layers and linear layers.

8. A model training device, characterized in that: include: A determination module, used to determine a checkpoint module in a target model; A forward propagation module, configured to, during the forward propagation phase, store the input and output of any checkpoint module in a video memory and release the intermediate activation value of the checkpoint module in the video memory; wherein the input of the checkpoint module is used to recalculate the intermediate activation value of the checkpoint module during the backward propagation phase, and the output of the checkpoint module is used for forward computation of subsequent modules of the checkpoint module; A back propagation module is configured to, in the back propagation phase, for any checkpoint module, in response to the last layer of the checkpoint module being a linear layer, skip the forward calculation of the last layer, calculate the gradient of the last layer according to the gradient formula corresponding to the last layer, and complete the back propagation of the checkpoint module according to the gradients of each layer in the checkpoint module.

9. A model training device comprising a memory, a processor, and a computer program stored in the memory, characterized in that: The processor executes the computer program to implement the steps of the method according to any one of claims 1 to 7.

10. A non-volatile computer-readable storage medium having a computer program stored thereon, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

11. A computer program product comprising a computer program, or a non-volatile computer-readable storage medium carrying a computer program, characterized in that: When the computer program is executed by a processor, the steps of the method according to any one of claims 1 to 7 are implemented.

Images