A model training method and device, electronic equipment, storage medium and computer program product

Abstract

The present disclosure relates to a model training method and device, electronic equipment, storage medium and computer program product, the method comprising: determining a checkpoint module in a target model; in a forward propagation stage, performing forward calculation on the checkpoint module, and saving only target parameters in a video memory during the forward calculation of the checkpoint module; in a backward propagation stage, if the last layer included in the checkpoint module is a Flash Attention layer, skipping forward recalculation of the Flash Attention layer, and determining the gradient of the Flash Attention layer according to a preset Flash Attention backward algorithm, and performing forward recalculation on other network layers in the checkpoint module to determine the corresponding gradient. According to the embodiments of the present disclosure, the redundant calculation overhead of forward recalculation of the Flash Attention layer can be saved, and the training efficiency can be effectively improved.

Description

Technical Field

[0001] This disclosure relates to the field of computer technology, and in particular to a model training method and apparatus, electronic equipment, storage medium and computer program product. Background Technology

[0002] In recent years, deep learning technology, especially large-scale model training, has made significant progress in the field of artificial intelligence. However, with the continuous expansion of model size, the issues of memory usage and computational efficiency during training have become increasingly prominent. Memory usage mainly comes from model weights, optimizer states, cached activation values, and temporary tensors, with cached activation values ​​accounting for the majority of memory usage in large-scale model training. To alleviate memory pressure, gradient checkpointing has been proposed in existing technologies, which reduces memory usage by sacrificing some computation. Its core idea is to save only some 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 re-checking of the forward process of each network layer in the module during backpropagation, resulting in additional computational overhead, which significantly reduces training efficiency. Summary of the Invention

[0004] In view of this, this disclosure presents a model training method and apparatus, electronic device, storage medium and computer program product.

[0005] According to one aspect of this disclosure, a model training method is provided, comprising: determining a checkpoint module in a target model; performing forward computation on the checkpoint module during the forward propagation phase, and storing only the target parameters of the checkpoint module during the forward computation process in GPU memory, wherein the target parameters are parameters required during the forward recomputation process of the checkpoint module during the backpropagation phase; during the backpropagation phase, if the last layer included in the checkpoint module is a Flash Attention layer, skipping the forward recomputation of the Flash Attention layer, and determining the gradient of the Flash Attention layer according to a preset Flash Attention backpropagation algorithm, and performing forward recomputation on each other network layer in the checkpoint module to determine the corresponding gradient.

[0006] In one possible implementation, skipping the forward recalculation of the Flash Attention layer and determining the gradient of the Flash Attention layer according to a preset Flash Attention backpropagation algorithm includes: reading the output and output gradient of the Flash Attention layer from the video memory; performing forward recalculation on each network layer in the checkpoint module other than the Flash Attention layer according to the target parameters to determine the input of the Flash Attention layer; and determining the input gradient of the Flash Attention layer using the preset Flash Attention backpropagation algorithm based on the input, output, and output gradient of the Flash Attention layer.

[0007] In one possible implementation, the method further includes: during the forward propagation phase, using the output of the FlashAttention layer as the input of a subsequent network layer, performing forward computation on that network layer to determine its output, and storing the input and computation graph of the network layer in the GPU memory; during the backpropagation phase, determining the input gradient of the network layer based on its input and computation graph, output, and output gradient using an automatic differentiation mechanism, and storing the input gradient of the network layer in the GPU memory, wherein the input gradient of the network layer is the output gradient of the FlashAttention layer.

[0008] In one possible implementation, the method further includes: during the forward propagation phase, after performing forward computation on the FlashAttention layer, storing the intermediate parameters used during the forward computation of the FlashAttention layer in the video memory; and during the backpropagation phase, before executing the preset FlashAttention backpropagation algorithm, reading the intermediate parameters from the video memory.

[0009] In one possible implementation, the checkpoint module includes a linear layer preceding the Flash Attention layer, the input gradient of the Flash Attention layer being the output gradient of the linear layer; the method further includes: after performing forward recomputation on the linear layer, storing the input, output, and computation graph of the linear layer in the video memory.

[0010] In one possible implementation, the step of performing forward recalculation on other network layers in the checkpoint module to determine the corresponding gradient includes: determining the input gradient and weight gradient of the linear layer through an automatic differentiation mechanism based on the input and computation graph, output and output gradient of the linear layer.

[0011] According to another aspect of this disclosure, a model training apparatus is provided, comprising: a determination module for determining checkpoint modules in a target model; a forward propagation module for performing forward computation on the checkpoint modules during the forward propagation phase, and storing only the target parameters of the checkpoint modules during the forward computation process in the video memory, wherein the target parameters are parameters required for performing forward recomputation on the checkpoint modules during the backpropagation phase; and a backpropagation module for skipping the forward recomputation of the Flash Attention layer if the last layer included in the checkpoint module is a Flash Attention layer during the backpropagation phase, and determining the gradient of the Flash Attention layer according to a preset Flash Attention backpropagation algorithm, and performing forward recomputation on other network layers in the checkpoint module to determine the corresponding gradients.

[0012] According to another aspect of this disclosure, an electronic device is provided, 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-described method.

[0013] According to another aspect of this disclosure, a non-volatile computer-readable storage medium is provided, on which a computer program is stored, which, when executed by a processor, implements the steps of the above-described method.

[0014] According to another aspect of this 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, when executed by a processor, implements the steps of the above-described method.

[0015] In this embodiment, a checkpoint module in the target model is determined. During the forward propagation phase, only the target parameters are stored in the GPU memory when performing forward computation on the checkpoint module. These target parameters are the parameters required during the forward recomputation of the checkpoint module during the backpropagation phase. During the backpropagation phase, if the last layer in the checkpoint module is a Flash Attention layer, the forward recomputation of the Flash Attention layer is skipped. Instead, a preset Flash Attention backpropagation algorithm is used to determine the gradient of the Flash Attention layer. This gradient is then combined with the gradients determined by performing forward recomputation on other network layers in the checkpoint module to quickly complete the backpropagation of the checkpoint module. This approach ensures the same computational accuracy and GPU memory savings as the standard recomputation scheme while saving redundant computational overhead from the forward recomputation of the Flash Attention layer, effectively improving training efficiency.

[0016] Other features and aspects of this disclosure will become clear from the following detailed description of exemplary embodiments with reference to the accompanying drawings. Attached Figure Description

[0017] The accompanying drawings, which are included in and form part of this specification, illustrate exemplary embodiments, features, and aspects of this disclosure together with the specification and serve to explain the principles of this disclosure.

[0018] Figure 1 A flowchart of a model training method according to an embodiment of the present disclosure is shown.

[0019] Figure 2 A block diagram of a model training apparatus according to an embodiment of the present disclosure is shown.

[0020] Figure 3 A block diagram of an electronic device according to an embodiment of the present disclosure is shown. Detailed Implementation

[0021] Various exemplary embodiments, features, and aspects of this disclosure will now be described in detail with reference to the accompanying drawings. The same reference numerals in the drawings denote elements that have the same or similar functions. Although various aspects of the embodiments are shown in the drawings, they are not necessarily drawn to scale unless specifically indicated otherwise.

[0022] As used herein, the terms “comprising,” “including,” “having,” or variations thereof are open-ended and include one or more of the stated features, integrals, elements, steps, components, or functions, but do not exclude the presence or addition of one or more other features, integrals, elements, steps, components, functions, or groups thereof.

[0023] When an element is referred to as “connected,” “coupled,” “responding,” or a variation thereof relative to another element, it may be directly connected, coupled, or responding to another element, or there may be an intermediate element present.

[0024] 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. Therefore, without departing from the teachings of the inventive concept, a first element / operation in some embodiments may be referred to as a second element / operation in other embodiments.

[0025] The term "exemplary" as used herein means "serving as an example, embodiment, or illustration." Any embodiment illustrated herein as "exemplary" is not necessarily to be construed as superior to or better than other embodiments.

[0026] Furthermore, to better illustrate this disclosure, numerous specific details are set forth in the following detailed description. Those skilled in the art will understand that this disclosure can be practiced without certain specific details. In some instances, methods, means, components, and circuits well known to those skilled in the art have not been described in detail in order to highlight the main points of this disclosure.

[0027] The key technical concepts involved in the embodiments of this disclosure are introduced below.

[0028] Automatic Differentiation (Autograd) is a core mechanism in modern deep learning frameworks (e.g., PyTorch, TensorFlow). During forward propagation, a dynamic computation graph records the forward computation process and the tensors needed for backward propagation. During backpropagation, it automatically calculates the gradient using the chain rule. Taking a linear layer y = Wx as an example, its gradient calculation during backpropagation includes: weight gradients. Input gradient Therefore, during forward propagation, it is necessary to cache the input x and the computation graph so that the gradient can be automatically calculated based on the computation graph and the chain rule during backward propagation.

[0029] GPU memory usage. During the training of deep learning models, GPU memory usage mainly includes: model weights, optimizer state, cached activation values, and temporary tensors. Among these, model weights and optimizer state are relatively fixed in terms of memory usage, and are related to factors such as model structure, data type, distributed strategy, and optimization algorithm. Temporary tensors have relatively small memory usage during the training of large models. The memory usage of cached activation values ​​is positively correlated with model size and batch size. During the training of large models, the variables required for backpropagation of each module (e.g., the input x of a linear layer) will not be released until the backpropagation of that module is completed. For example, after the forward propagation of an MLP is completed, the inputs of both linear layers will be cached, and depending on the type of the intermediate activation function, that activation function may also need to be cached. Therefore, cached activation values ​​are the main source of GPU memory usage.

[0030] Gradient checkpointing is a memory optimization technique that selectively discards some intermediate activation values ​​during forward propagation and recalculates the discarded intermediate activation values ​​during backward propagation, thereby reducing memory usage.

[0031] Transformer layers are the basic building blocks of the Transformer model. A Transformer layer consists of a multi-head self-attention (MHSA) mechanism and a feed-forward network (FFN), and is typically optimized using residual connections and layer normalization.

[0032] Self-attention is a core component of the Transformer architecture. It focuses on information at different positions in the input sequence, capturing long-range dependencies to determine the correlation between positions and generate dynamically weighted representations. The computation process of self-attention is as follows:

[0033] q,k,v=Wq×x1,Wk×x1,Wv×x1;

[0034] o=Attention(q,k,v)=softmax(q×k T / sqrt(d))×v;

[0035] y = Wo × o.

[0036] Flash Attention (FA) is an optimized attention calculation algorithm that reduces the number of GPU memory read / write operations and improves computational efficiency through block-based computation (Tiling) and efficient memory access strategies.

[0037] The Flash Attention forward propagation algorithm during the forward propagation of the Flash Attention layer is as follows.

[0038] Input: matrices Q, K, V ∈ R N×d Stored in high-bandwidth memory (HBM), block size parameter B c B r .

[0039] Step 11, divide into blocks.

[0040] Divide Q into T r =[N / B r [a number of blocks] Each sub-block is B in size. r ×d.

[0041] Divide K and V into T c =[N / B c [a number of blocks] and Each sub-block is B in size. c ×d.

[0042] Step 12, output initialization.

[0043] The output matrix O∈R N×d Divided into T r Sub-block Each sub-block is B in size. r ×d.

[0044] The logarithmic summation exponent vector L is divided into T. r Sub-block Each size is B r .

[0045] Step 13, calculate block by block (outer loop i = 1 → T) r ).

[0046] Q i Load from video memory to on-chip memory (SRAM).

[0047] On-chip initialization: Sub-block output initialized to a zero matrix The logarithmic summation exponent vector L sub-block is initialized as a zero matrix. Maximum value cache

[0048] Step 14, Inner loop (j = 1 → T) c ).

[0049] K j V j Load from video memory to on-chip storage.

[0050] Calculate attention score on-chip: Sub-block product renew Calculate normalized score renew

[0051] Update output subblock

[0052] Step 15, Normalization and Writing.

[0053] Calculate the final output

[0054] Calculate the logarithmic summation exponent

[0055] O i and L i Write back to video memory.

[0056] Step 16: Output matrix O and vector L.

[0057] The Flash Attention backpropagation algorithm during backpropagation of the Flash Attention layer is as follows.

[0058] Input: Matrices Q, K, V, O, dO ∈ R in High Bandwidth Memory (HBM) N×d And the logarithmic summation exponential vector L∈R N Block size parameter B c B r .

[0059] Step 21, divide into blocks.

[0060] Divide Q into T r =[N / B r [a number of blocks] Each sub-block is B in size. r ×d.

[0061] Divide K and V into T c =[N / B c [a number of blocks] and Each sub-block is B in size. c×d.

[0062] Step 22, output and output gradient block.

[0063] Divide the output matrix O into T r Sub-block Each sub-block is B in size. r ×d.

[0064] The output gradient matrix dO is divided into T. r Sub-block Each sub-block is B in size. r ×d.

[0065] The logarithmic summation exponential vector L is divided into T. r Sub-block Each size is B r .

[0066] Step 23, Input initialization.

[0067] Initialize dQ to a zero matrix in video memory: dQ = (0) N×d And divided into T r Sub-block Each sub-block is B in size. r ×d;

[0068] Divide dK and dV into T c Sub-block and Each sub-block is B in size. c ×d.

[0069] Step 24: Calculate the auxiliary vector D.

[0070] Calculate the pointwise product D = rowsum(dO·O)∈R N×d Write D to video memory, and divide D into T. r Sub-block Each sub-block is B in size. r .

[0071] Step 25, outer loop (by K, V sub-blocks, j=1→T) c ).

[0072] K j V j Load from video memory to on-chip memory (SRAM).

[0073] On-chip initialization of sub-block gradients as zero matrices

[0074] Step 26, Inner Loop (by Q sub-block, i = 1 → T) r ).

[0075] Q i O i dO i dQ i L i D i Load from video memory to on-chip storage.

[0076] Calculate attention score on-chip: Sub-block product Calculate normalized score

[0077] Gradient update:

[0078] Step 27, convert the gradient dQ i dK j dV j Write back to video memory.

[0079] Step 28, output matrices dQ, dK, and dV.

[0080] The following section introduces the standard gradient checkpointing technique 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 intermediate computations are discarded. Taking the Transformer model as an example, feedforward neural networks (FFNs) and single-layer Transformer network layers can be set as checkpoint modules.

[0082] During the backpropagation phase, for each checkpoint module, its complete forward pass needs to be re-executed to recover intermediate activation values, and then the gradients of all parameters within the module (including linear and nonlinear layers) are calculated through an automatic differentiation mechanism.

[0083] Taking the self-attention mechanism of a Transformer layer as an example, the standard forward propagation process is as follows:

[0084] Input: x1;

[0085] (1.1) q,k,v=Wq×x1,Wk×x1,Wv×x1; x1 is stored in video memory;

[0086] (1.2) o = FA(q,k,v); q,k,v are stored in video memory;

[0087] (1.3)y = Wo × o; o is stored in video memory.

[0088] The standard backpropagation process is as follows:

[0089] enter:

[0090] (1.4) Delete y from video memory;

[0091] (1.5) Delete the 'o' from the video memory;

[0092] (1.6) Delete q, k, v from video memory.

[0093] Here, backward indicates that the automatic differentiation mechanism is invoked.

[0094] By identifying only the first two steps of the Self-Attention mechanism as checkpoint modules and adopting standard recalculation techniques, the computation process is adjusted as follows.

[0095] The forward propagation process is as follows:

[0096] Input: x1;

[0097] (2.1) q,k,v=Wq×x1,Wk×x1,Wv×x1; x1 is stored in video memory;

[0098] (2.2)o=FA(q,k,v);

[0099] (2.3)y = Wo × o; o is stored in video memory.

[0100] The backpropagation process is as follows:

[0101] (2.4) Input:

[0102] (2.5) Delete the 'o' from the video memory;

[0103] (2.6) q,k,v=Wq×x1,Wk×x1,Wv×x1; x1 is stored in video memory;

[0104] (2.7) o = FA(q,k,v); q,k,v are stored in video memory;

[0105] (2.8) Delete the 'o' from the video memory;

[0106] (2.9) Delete q, k, v from video memory.

[0107] After standard recalculation, during backpropagation, (2) will delete the o stored in the video memory after execution, and o will be recalculated in (4). Through the two steps (4) and (5) to perform forward recalculation of the Flash Attention layer, after (4) is executed, q, k, v are stored in the video memory. Then, when the automatic differentiation mechanism is called in (5), o and the corresponding derivative are read from the video memory. And q, k, v, and then, based on the computation graph, the inverse operator of the Flash Attention layer is called to obtain

[0108] The main advantage of standard recalculation is that it effectively saves the memory usage of q, k, v after forward propagation and before back propagation begins (usually during the forward and back propagation of other network layers besides self-attention). However, the corresponding cost is that forward recalculation of [Wq,Wk,Wv]×x1 and FA(q,k,v) needs to be performed during back propagation.

[0109] The standard recomputation in existing technologies completely redoes the forward recomputation of each network layer in the checkpoint module, and then uses an automatic differentiation mechanism to find the derivative. This repeated forward computation results in additional computational overhead.

[0110] To address the aforementioned technical problems, this disclosure provides a model training method that, while maintaining the same computational accuracy and memory savings as the standard recalculation scheme, saves redundant computational overhead in the forward recalculation of the Flash Attention layer, effectively improving training efficiency. The model training method provided by this disclosure is described in detail below.

[0111] Figure 1 A flowchart illustrating a model training method according to an embodiment of this disclosure is shown. The method can be executed by a model training apparatus, such as a terminal device, server, or other electronic device. The terminal device can be a user equipment (UE), mobile device, user terminal, terminal, handheld device, or computing device, etc. In some possible implementations, the model training method can be implemented by a processor calling computer-readable instructions stored in memory. Figure 1 As shown, the method includes:

[0112] In step S11, the checkpoint module in the target model is determined.

[0113] In step S12, during the forward propagation phase, forward computation is performed on the checkpoint module, and only the target parameters during the forward computation of the checkpoint module are saved in the video memory. The target parameters are the parameters required during the forward recomputation of the checkpoint module during the backpropagation phase.

[0114] In step S13, during the backpropagation phase, if the last layer included in the checkpoint module is a FlashAttention layer, the forward recalculation of the FlashAttention layer is skipped, and the gradient of the FlashAttention layer is determined according to the preset FlashAttention backpropagation algorithm. Forward recalculation is then performed on the other network layers in the checkpoint module to determine the corresponding gradients.

[0115] In this embodiment, a checkpoint module in the target model is determined. During the forward propagation phase, only the target parameters are stored in the GPU memory when performing forward computation on the checkpoint module. These target parameters are the parameters required during the forward recomputation of the checkpoint module during the backpropagation phase. During the backpropagation phase, if the last layer in the checkpoint module is a Flash Attention layer, the forward recomputation of the Flash Attention layer is skipped. Instead, a preset Flash Attention backpropagation algorithm is used to determine the gradient of the Flash Attention layer. This gradient is then combined with the gradients determined by performing forward recomputation on other network layers in the checkpoint module to quickly complete the backpropagation of the checkpoint module. This approach ensures the same computational accuracy and GPU memory savings as the standard recomputation scheme while saving redundant computational overhead from the forward recomputation of the Flash Attention layer, effectively improving training efficiency.

[0116] In this embodiment of the disclosure, the target model is any neural network model including a Flash Attention layer, and the specific network structure of the target model is not limited.

[0117] In model training, a checkpoint module can refer to a sub-network or combination of network layers that applies gradient checkpointing techniques. Checkpoint modules can be continuous network layers in the target model that consume significant GPU memory and whose intermediate activation values ​​can be recalculated, such as feedforward networks, residual blocks, and single-layer Transformer layers. The core principle is to save only the input and output of the checkpoint module during forward propagation, discarding intermediate activation values ​​to save GPU memory; during backpropagation, the intermediate activation values ​​are recovered by recalculation.

[0118] In this embodiment of the disclosure, when the checkpoint module ends with the Flash Attention layer, the forward recalculation of the Flash Attention layer can be skipped, and the gradient of the Flash Attention layer can be directly derived using the preset Flash Attention backpropagation algorithm, thereby significantly reducing the computational overhead while maintaining the same memory saving effect.

[0119] In this embodiment, during the forward propagation phase, the inputs, outputs, and all intermediate activation values ​​of each non-checkpoint module can be fully retained in the video memory for direct use in the subsequent backpropagation phase. During the forward propagation phase, only some parameters of the checkpoint modules are retained, namely the target parameters required for the forward recalculation of the checkpoint modules during the backpropagation phase.

[0120] During the backpropagation phase, the handling of the checkpoint module can be divided into the following two cases:

[0121] If the last layer of the checkpoint module is not a Flash Attention layer, a standard recalculation scheme can be adopted, that is, to re-execute the complete forward recalculation of the checkpoint module to recover all intermediate activation values, and then perform gradient calculation according to the automatic differentiation mechanism.

[0122] When the last layer of the checkpoint module is a Flash Attention layer, an optimized computation scheme can be implemented. That is, skip the forward recalculation of the Flash Attention layer of the checkpoint module, determine the gradient of the Flash Attention layer according to the preset Flash Attention backward algorithm, and still perform the standard recalculation process for other network layers in the checkpoint module to determine the corresponding gradient.

[0123] In one possible implementation, skipping the forward recalculation of the Flash Attention layer and determining the gradient of the Flash Attention layer according to a preset Flash Attention backpropagation algorithm includes: reading the output and output gradient of the Flash Attention layer from GPU memory; performing forward recalculation on each network layer in the checkpoint module other than the Flash Attention layer according to the target parameters to determine the input of the Flash Attention layer; and determining the input gradient of the Flash Attention layer using a preset Flash Attention backpropagation algorithm based on the input, output, and output gradient of the Flash Attention layer.

[0124] Based on the above description of Flash Attention, the inputs to the preset Flash Attention inverse algorithm include: Q, K, V, O, and dO, where Q, K, and V are the inputs to the Flash Attention layer, and O is the output of the Flash Attention layer. This is the output gradient of the Flash Attention layer.

[0125] The output of the Flash Attention layer is the input of the next network layer. Therefore, when the next network layer performs forward propagation, its input is stored in video memory. Thus, when performing backward propagation on the Flash Attention layer, the output of the Flash Attention layer can be read from video memory.

[0126] The output gradient of the Flash Attention layer is the input gradient of the next network layer after the Flash Attention layer. Therefore, after the backpropagation of the next network layer after the Flash Attention layer is completed, its determined input gradient will be saved in the GPU memory. Thus, when performing backpropagation on the Flash Attention layer, the output gradient of the Flash Attention layer can be read from the GPU memory.

[0127] The input to the Flash Attention layer is the output of the network layer preceding the Flash Attention layer. Therefore, by performing forward recalculation on each network layer in the checkpoint module other than the Flash Attention layer according to the target parameters, the input to the Flash Attention layer can be determined.

[0128] Then, based on the input, output, and output gradient of the Flash Attention layer, the input gradient of the Flash Attention layer is determined using a pre-defined Flash Attention backpropagation algorithm.

[0129] In one possible implementation, the method further includes: during the forward propagation phase, using the output of the Flash Attention layer as the input of the next network layer after the Flash Attention layer, performing forward computation on the network layer to determine the output of the network layer, and saving the input and computation graph of the network layer in the GPU memory; during the backpropagation phase, based on the input and computation graph, output and output gradient of the network layer, determining the input gradient of the network layer through an automatic differentiation mechanism, and saving the input gradient of the network layer in the GPU memory, wherein the input gradient of the network layer is the output gradient of the Flash Attention layer.

[0130] After performing forward computation on the Flash Attention layer during the forward propagation phase, the output of the Flash Attention layer is used as the input of the next network layer. Forward computation is then performed on the next network layer to determine its output, and the input and computation graph of the next network layer are saved in memory.

[0131] When the network layer following the Flash Attention layer is not a layer in the checkpointing module, the input gradient of that network layer can be determined during the backpropagation phase using an automatic differentiation mechanism based on its input and computation graph, output and output gradient. This input gradient is also stored in GPU memory, and the output gradient is also the output gradient of the Flash Attention layer. Therefore, when performing backpropagation on the Flash Attention layer, its output gradient can be retrieved from GPU memory.

[0132] Taking the Self-Attention mechanism of a Transformer layer as an example, only the first two steps of the Self-Attention mechanism are still identified as checkpoint modules. Using the recomputation technique of this disclosure, the computation process is adjusted as follows.

[0133] The forward propagation process is as follows:

[0134] (3.1) q,k,v=Wq×x1,Wk×x1,Wv×x1; x1 is stored in video memory;

[0135] (3.2)o=FA(q,k,v);

[0136] (3.3)y = Wo × o; o is stored in video memory.

[0137] It is evident that the forward propagation process remains unchanged from the standard recalculation scheme.

[0138] Compared to the standard recalculation technique described above, the backpropagation process is adjusted as follows:

[0139] (3.4) Input:

[0140] (3.5)

[0141] (3.6) q,k,v=Wq×x1,Wk×x1,Wv×x1; x1 is stored in video memory;

[0142] (3.7)

[0143] (3.8)

[0144] Among them, (3.5) is the process of performing reverse computation on the network layer after the Flash Attention layer during the backpropagation phase, which is determined by... (3.3) is the process of performing forward computation on the network layer after the Flash Attention layer during the forward propagation stage. The output o of the Flash Attention layer is used as the input of the network layer and is stored in the video memory. (3.6) is the process of performing forward recomputation on the network layers other than the Flash Attention layer in the checkpoint module. The output q, k, v is the input of the Flash Attention layer.

[0145] Based on the input q,k,v, output o, and output gradient of the Flash Attention layer Using the pre-defined FlashAttention inverse algorithm, the input gradient of the FlashAttention layer can be directly determined. Compared to the backpropagation process of the standard recalculation technique described above, the backpropagation process of this embodiment can maintain the same gradient calculation results as (2.7) and (2.8) through (3.7), while saving the computational overhead of performing forward recalculation on the Flash Attention layer in (2.7).

[0146] In the back propagation process (3.5) of this embodiment, compared with the back propagation process (2.5) of the standard recomputation technique described above, after performing back propagation on a network layer after the Flash Attention layer, o in the video memory is no longer deleted, so that when the preset Flash Attention back propagation algorithm is executed in the subsequent (3.7), o can be directly read from the video memory.

[0147] In one possible implementation, the method further includes: during the forward propagation phase, after performing forward computation on the Flash Attention layer, storing the intermediate parameters used during the forward computation of the Flash Attention layer in the video memory; and during the backpropagation phase, before executing the preset Flash Attention backpropagation algorithm, reading the intermediate parameters from the video memory.

[0148] Based on the above description of Flash Attention, it can be seen that the input to the preset Flash Attention backpropagation algorithm also includes intermediate parameters: the logarithmic summation exponent vector L and the block size parameter B. c B r Therefore, during the forward propagation phase, after performing forward computation on the FlashAttention layer, the logarithmic summation exponent vector L and the block size parameter B are stored in GPU memory. c B r This allows the logarithmic summation exponent vector L and the block size parameter B to be read from GPU memory before executing the preset Flash Attention backpropagation algorithm during the backpropagation phase. c B r .

[0149] In one possible implementation, the checkpoint module includes a linear layer preceding the Flash Attention layer, with the input gradient of the Flash Attention layer being the output gradient of the linear layer; the method also includes: after performing forward recomputation on the linear layer, saving the input, output, and computation graph of the linear layer in GPU memory.

[0150] Taking the Self-Attention example above, (3.6) is the process of performing forward computation on a linear layer before the Flash Attention layer in the checkpoint module. After the forward recomputation is completed, the input x1, output q, k, v and computation graph are saved.

[0151] In one possible implementation, forward recalculation is performed on each other network layer in the checkpoint module to determine the corresponding gradient, including: determining the input gradient and weight gradient of the linear layer through an automatic differentiation mechanism based on the input and computation graph, output and output gradient of the linear layer.

[0152] Taking the Self-Attention algorithm mentioned above as an example, the input gradient of the FlashAttention layer, determined by the pre-defined FlashAttention backpropagation algorithm, is the output gradient of the linear layer preceding the FlashAttention layer. Therefore, based on the input and computation graph, output and output gradient of the linear layer, the input gradient and weight gradient of the linear layer are determined through an automatic differentiation mechanism.

[0153] (3.8) describes the process of performing reverse computation on a linear layer preceding the Flash Attention layer in the checkpoint module, based on the input x1 and computation graph, output q,k,v and output gradient. The input gradient of the linear layer is determined through an automatic differentiation mechanism. and weight gradient

[0154] In this embodiment, a checkpoint module in the target model is determined. During the forward propagation phase, only the target parameters are stored in the GPU memory when performing forward computation on the checkpoint module. These target parameters are the parameters required during the forward recomputation of the checkpoint module during the backpropagation phase. During the backpropagation phase, if the last layer in the checkpoint module is a Flash Attention layer, the forward recomputation of the Flash Attention layer is skipped. Instead, a preset Flash Attention backpropagation algorithm is used to determine the gradient of the Flash Attention layer. This gradient is then combined with the gradients determined by performing forward recomputation on other network layers in the checkpoint module to quickly complete the backpropagation of the checkpoint module. This approach ensures the same computational accuracy and GPU memory savings as the standard recomputation scheme while saving redundant computational overhead from the forward recomputation of the Flash Attention layer, effectively improving training efficiency.

[0155] In one example, the target model of this embodiment can be a MoE (Mixture of Experts) model, in which the MLA part is defined as a checkpoint module. Before the checkpoint module is a first network layer (q / k_down_project, which performs dimensionality reduction on q and k); the checkpoint module includes: a second network layer (q / k_layernorm, which performs layer normalization on q and k), a third network layer (q / k_up_projectcct, which performs dimensionality increase on q and k), and a Flash Attention layer. After the checkpoint module is a fourth network layer (o_project, which outputs o). In this case, during the backpropagation phase, only the second network layer (q / k_layernorm) and the third network layer (q / k_up_projectcct) in the checkpoint module need to be recomputed forward, skipping the forward recomputed Flash Attention layer.

[0156] Assuming the above MoE model includes 60 checkpoint modules, and each checkpoint module skips the forward recalculation of the Flash Attention layer, it can save 528MB of video memory. In this case, a total of 30GB of video memory can be saved.

[0157] In the standard recalculation scheme, forward recalculation is required for each network layer in the checkpoint module. This means additional computational overhead is incurred for forward recalculation of the second network layer (q / k_layernorm), the third network layer (q / k_up_projecct), and the FlashAttention layer. The computational overhead for forward recalculation of the second network layer (q / k_layernorm) is relatively small and can be ignored; the computational overhead for forward recalculation of the third network layer (q / k_up_projecct) is 0.44T Flops; and the computational overhead for forward recalculation of the FlashAttention layer is 1.25T Flops. Therefore, the standard recalculation scheme incurs an additional computational cost of 1.69T Flops.

[0158] By employing the recalculation scheme of this embodiment, skipping the forward recalculation of the Flash Attention layer, the computational overhead can be reduced to 0.44T Flops, thereby effectively improving training efficiency.

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

[0160] In addition, this disclosure also provides a model training device, electronic device, computer-readable storage medium, and program, all of which can be used to implement any of the model training methods provided in this disclosure. The corresponding technical solutions and descriptions are described in the corresponding section of the method and will not be repeated here.

[0161] Figure 2 A block diagram of a model training apparatus according to an embodiment of the present disclosure is shown. Figure 2 As shown, device 20 includes:

[0162] Module 21 is used to determine the checkpoint module in the target model;

[0163] The forward propagation module 22 is used to perform forward calculation on the checkpoint module during the forward propagation phase, and to save only the target parameters in the video memory during the forward calculation of the checkpoint module. The target parameters are the parameters required during the forward recalculation of the checkpoint module during the back propagation phase.

[0164] Backpropagation module 23 is used to skip the forward recalculation of the Flash Attention layer if the last layer included in the checkpoint module is a FlashAttention layer during the backpropagation stage, and to determine the gradient of the Flash Attention layer according to the preset Flash Attention backpropagation algorithm, and to perform forward recalculation on the other network layers in the checkpoint module to determine the corresponding gradient.

[0165] In one possible implementation, the backpropagation module 23 is specifically used for:

[0166] Read the output and output gradient of the Flash Attention layer from the video memory;

[0167] Based on the target parameters, forward recalculation is performed on each network layer in the checkpoint module except for the Flash Attention layer to determine the input of the Flash Attention layer.

[0168] Based on the input, output, and output gradient of the Flash Attention layer, the input gradient of the Flash Attention layer is determined using a pre-defined Flash Attention backpropagation algorithm.

[0169] In one possible implementation, the forward propagation module 22 is specifically used to: during the forward propagation stage, take the output of the FlashAttention layer as the input of a network layer after the FlashAttention layer, perform forward computation on the network layer, determine the output of the network layer, and save the input and computation graph of the network layer in the video memory.

[0170] The backpropagation module 23 is specifically used to: determine the input gradient of the network layer through an automatic differentiation mechanism based on the input and computation graph, output and output gradient of the network layer during the backpropagation stage, and save the input gradient of the network layer in the video memory, wherein the input gradient of the network layer is the output gradient of the Flash Attention layer.

[0171] In one possible implementation, the forward propagation module 22 is specifically used to: after performing forward computation on the FlashAttention layer during the forward propagation phase, save the intermediate parameters used in the forward computation process of the FlashAttention layer in the video memory.

[0172] The backpropagation module 23 is specifically used to read intermediate parameters from the video memory before executing the preset Flash Attention backpropagation algorithm during the backpropagation stage.

[0173] In one possible implementation, the checkpoint module includes a linear layer preceding the Flash Attention layer, where the input gradient of the Flash Attention layer is the output gradient of the linear layer.

[0174] Backpropagation module 23 is specifically used for:

[0175] After performing forward recomputation on the linear layer, the input, output, and computation graph of the linear layer are saved in video memory.

[0176] In one possible implementation, the backpropagation module 23 is specifically used for:

[0177] Based on the input and computation graph, output and output gradient of the linear layer, the input gradient and weight gradient of the linear layer are determined through an automatic differentiation mechanism.

[0178] In some embodiments, the functions or modules of the apparatus provided in this disclosure can be used to perform the methods described in the above method embodiments. The specific implementation can be referred to the description of the above method embodiments, and for the sake of brevity, it will not be repeated here.

[0179] This disclosure also provides an electronic 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.

[0180] This disclosure also provides a non-volatile computer-readable storage medium having a computer program stored thereon, which, when executed by a processor, implements the steps of the above-described method.

[0181] This disclosure also 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, when executed by a processor, implements the steps of the above method.

[0182] Figure 3 A block diagram of an electronic device according to an embodiment of the present disclosure is shown. (Refer to...) Figure 3 Device 1900 can be provided as a server or terminal device. (See reference...) Figure 3 The apparatus 1900 includes a processing component 1922, which further includes one or more processors, and memory resources represented by memory 1932 for storing instructions, such as application programs, that can be executed by the processing component 1922. The application programs stored in memory 1932 may include one or more modules, each corresponding to a set of instructions. Furthermore, the processing component 1922 is configured to execute instructions to perform the methods described above.

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

[0184] 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 a processing component 1922 of the device 1900 to perform the above-described method.

[0185] Computer-readable storage media can be tangible devices capable of holding and storing programs / instructions used by instruction execution devices. Computer-readable storage media can be, for example—but not limited to—electrical storage devices, magnetic storage devices, optical storage devices, electromagnetic storage devices, semiconductor storage devices, or any suitable combination of the foregoing. More specific examples (a non-exhaustive list) of computer-readable storage media include: portable computer disks, hard disks, random access memory (RAM), read-only memory (ROM), erasable programmable read-only memory (EPROM or flash memory), static random access memory (SRAM), portable compact disc read-only memory (CD-ROM), digital multifunction disc (DVD), memory sticks, floppy disks, mechanical encoding devices, such as punch cards or recessed protrusions storing instructions thereon, and any suitable combination of the foregoing. The computer-readable storage media used herein are not to be construed as transient signals themselves, such as radio waves or other freely propagating electromagnetic waves, electromagnetic waves propagating through waveguides or other transmission media (e.g., light pulses through fiber optic cables), or electrical signals transmitted through wires.

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

[0187] The computer program (or computer program instructions) used to perform the operations of this disclosure may be assembly instructions, instruction set architecture (ISA) instructions, machine instructions, machine-dependent instructions, microcode, firmware instructions, status 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++, etc., and conventional procedural programming languages ​​such as the "C" language or similar programming languages. The computer-readable program instructions may execute entirely on the user's computer, partially on the user's computer, as a standalone software package, partially on the user's computer and partially on a remote computer, or entirely on a remote computer or server. In cases involving remote computers, 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, electronic circuits, such as programmable logic circuits, field-programmable gate arrays (FPGAs), or programmable logic arrays (PLAs), are personalized by utilizing state information of computer-readable program instructions. These electronic circuits can execute computer-readable program instructions to implement various aspects of this disclosure.

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

[0189] 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 apparatus to produce a machine such that, when executed by the processor of the computer or other programmable data processing apparatus, they create means for implementing the functions / actions specified in one or more blocks of the flowchart and / or block diagram. These computer-readable program instructions can also be stored in a computer-readable storage medium that causes a computer, programmable data processing apparatus, and / or other device to operate in a particular manner; thus, the computer-readable medium storing the instructions comprises an article of manufacture that includes instructions for implementing aspects of the functions / actions specified in one or more blocks of the flowchart and / or block diagram.

[0190] Computer-readable program instructions may also be loaded onto a computer, other programmable data processing apparatus, or other device to cause a series of operational steps to be 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 perform the functions / actions specified in one or more boxes of a flowchart and / or block diagram.

[0191] The flowcharts and block diagrams in the accompanying drawings illustrate the architecture, functionality, and operation of possible implementations of systems, methods, and computer program products according to various embodiments of the present disclosure. In this regard, each block in a flowchart or block diagram may represent a module, segment, or portion of an instruction containing one or more executable instructions for implementing a specified logical function. In some alternative implementations, the functions marked in the blocks may occur in a different order than those shown in the drawings. For example, two consecutive blocks may actually be executed substantially in parallel, and they may sometimes be executed in reverse order, depending on the functions involved. It should also be noted that each block in the block diagrams and / or flowcharts, and combinations of blocks in the block diagrams and / or flowcharts, may be implemented using a dedicated hardware-based system that performs the specified function or action, or using a combination of dedicated hardware and computer instructions.

[0192] The various embodiments of this disclosure have been described above. These descriptions are exemplary and not exhaustive, nor are they limited to 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 chosen to best explain the principles, practical application, or technical improvements to the embodiments in the market, 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, forward computation is performed on the checkpoint module, and only the target parameters during the forward computation of the checkpoint module are stored in the video memory. The target parameters are the parameters required during the forward recomputation of the checkpoint module during the backpropagation phase. During the backpropagation phase, if the last layer in the checkpoint module is a Flash Attention layer, the forward recalculation of the Flash Attention layer is skipped, and the gradient of the Flash Attention layer is determined according to the preset Flash Attention backpropagation algorithm. Forward recalculation is then performed on the other network layers in the checkpoint module to determine the corresponding gradients.

2. The method according to claim 1, characterized in that, The step of skipping the forward recalculation of the Flash Attention layer and determining the gradient of the Flash Attention layer according to the preset Flash Attention backpropagation algorithm includes: Read the output and output gradient of the Flash Attention layer from the video memory; Based on the target parameters, forward recalculation is performed on each network layer in the checkpoint module except for the Flash Attention layer to determine the input of the Flash Attention layer; Based on the input, output, and output gradient of the Flash Attention layer, the input gradient of the Flash Attention layer is determined using the preset Flash Attention inverse algorithm.

3. The method according to claim 2, characterized in that, The method further includes: During the forward propagation phase, the output of the Flash Attention layer is used as the input of the next network layer after the Flash Attention layer. Forward computation is performed on the network layer to determine the output of the network layer, and the input and computation graph of the network layer are saved in the video memory. During the backpropagation phase, based on the input and computation graph, output and output gradient of the network layer, the input gradient of the network layer is determined through an automatic differentiation mechanism, and the input gradient of the network layer is stored in the video memory, wherein the input gradient of the network layer is the output gradient of the Flash Attention layer.

4. The method according to claim 2, characterized in that, The method further includes: During the forward propagation phase, after performing forward computation on the Flash Attention layer, the intermediate parameters used in the forward computation process of the Flash Attention layer are stored in the video memory. During the backpropagation phase, the intermediate parameters are read from the video memory before executing the preset Flash Attention backpropagation algorithm.

5. The method according to claim 2, characterized in that, The checkpoint module includes a linear layer preceding the FlashAttention layer, and the input gradient of the FlashAttention layer is the output gradient of the linear layer; The method further includes: After performing forward recalculation on the linear layer, the input, output, and computation graph of the linear layer are saved in the video memory.

6. The method according to claim 5, characterized in that, The step of performing forward recalculation on each of the other network layers in the checkpoint module to determine the corresponding gradient includes: Based on the input and computation graph, output and output gradient of the linear layer, the input gradient and weight gradient of the linear layer are determined through an automatic differentiation mechanism.

7. A model training device, characterized in that, include: The determination module is used to determine the checkpoint module in the target model; The forward propagation module is used to perform forward calculation on the checkpoint module during the forward propagation phase, and to save only the target parameters of the checkpoint module during the forward calculation process in the video memory, wherein the target parameters are the parameters required for the forward recalculation process of the checkpoint module during the back propagation phase. The backpropagation module is used to, during the backpropagation phase, skip the forward recalculation of the Flash Attention layer if the last layer included in the checkpoint module is a FlashAttention layer, and determine the gradient of the Flash Attention layer according to the preset Flash Attention backpropagation algorithm, and perform forward recalculation on the other network layers in the checkpoint module to determine the corresponding gradients.

8. An electronic 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 6.

9. 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, it implements the steps of the method according to any one of claims 1 to 6.

10. 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, it implements the steps of the method according to any one of claims 1 to 6.

Images