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

Abstract

The invention relates to a model training method and device, electronic equipment, a storage medium and a computer program product. The method comprises the steps that a check point module in a target model is determined; in the forward propagation stage, forward calculation is executed on the check point module, and only target parameters in the process that the check point module executes the forward calculation are stored in the video memory; and in the back propagation stage, if the last layer included in the check point module is the Flash Attention layer, the forward re-calculation of the Flash Attention layer is skipped, the gradient of the Flash Attention layer is determined according to a preset Flash Attention reverse algorithm, and the forward re-calculation is executed on other network layers in the check point module to determine the corresponding gradient. According to the embodiment of the invention, the redundant calculation overhead of the forward re-calculation of the Flash Attention layer can be saved, and the training efficiency is effectively improved.

Description

TECHNICAL FIELD

[0001] The present disclosure relates to the technical field of computer, and particularly relates to a model training method and device, electronic equipment, storage medium and computer program product. BACKGROUND

[0002] In recent years, deep learning technology, especially large model training, has made significant progress in the field of artificial intelligence. However, with the continuous expansion of model size, the problems of memory occupation and computing efficiency in the training process are increasingly prominent. The memory occupation mainly comes from model weights, optimizer states, cached activation values and temporary tensors, among which the cached activation values occupy the main part of the memory in large-scale model training. In order to alleviate the memory pressure, the Gradient Checkpointing technology is proposed in the prior art, which reduces the memory occupation by sacrificing part of the computation. The core idea is to save only part of the intermediate results (i.e. Checkpoints) during forward propagation, and to recalculate the discarded intermediate variables during backward propagation.

[0003] However, the existing Gradient Checkpointing technology needs to complete the forward process of each network layer in the checkpoint module during backward propagation, resulting in additional computational overhead, which will significantly reduce the training efficiency. SUMMARY

[0004] Therefore, the present disclosure provides a model training method and device, electronic equipment, storage medium and computer program product.

[0005] According to an aspect of the present disclosure, a model training method is provided, comprising: determining a checkpoint module in a target model; performing forward calculation on the checkpoint module in a forward propagation stage, and saving only target parameters in the memory during the forward calculation of the checkpoint module, wherein the target parameters are parameters required in the forward recalculation of the checkpoint module in a backward propagation stage; in the backward propagation stage, if the last layer included in the checkpoint module is a Flash Attention layer, skipping the 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.

[0006] In a possible implementation, the forward re-computation of the Flash Attention layer is skipped, and the gradient of the Flash Attention layer is determined according to a preset Flash Attention back-propagation algorithm, including: reading the output and the output gradient of the Flash Attention layer from the memory; performing forward re-computation on each network layer in the checkpoint module except 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 by using the preset Flash Attention back-propagation algorithm according to the input, the output and the output gradient of the Flash Attention layer.

[0007] In a possible implementation, the method further includes: in the forward propagation stage, taking the output of the Flash Attention layer as the input of a network layer following the Flash Attention layer, performing forward computation on the network layer, determining the output of the network layer, and saving the input and the computation graph of the network layer in the memory; and in the backward propagation stage, determining the input gradient of the network layer by using an automatic differentiation mechanism according to the input and the computation graph, the output and the output gradient of the network layer, and saving the input gradient of the network layer in the memory, where the input gradient of the network layer is the output gradient of the Flash Attention layer.

[0008] In a possible implementation, the method further includes: in the forward propagation stage, saving, in the memory, intermediate parameters used in the process of performing forward computation on the Flash Attention layer after performing forward computation on the Flash Attention layer; and in the backward propagation stage, reading the intermediate parameters from the memory before performing the preset Flash Attention back-propagation algorithm.

[0009] In a possible implementation, the checkpoint module includes a linear layer before the Flash Attention layer, and the input gradient of the Flash Attention layer is the output gradient of the linear layer; and the method further includes: saving, in the memory, the input, the output and the computation graph of the linear layer after performing forward re-computation on the linear layer.

[0010] In a possible implementation, the determining the gradients of the other network layers in the checkpoint module includes: determining the input gradient and the weight gradient of the linear layer by an automatic differentiation mechanism according to the input and the computation graph of the linear layer, the output and the output gradient.

[0011] According to another aspect of the present disclosure, a model training apparatus is provided, including: a determining module configured to determine a checkpoint module in a target model; a forward propagation module configured to perform forward calculation on the checkpoint module in a forward propagation stage, and save only target parameters in a GPU during the forward calculation of the checkpoint module, wherein the target parameters are parameters required in a backward propagation stage; and a backward propagation module configured to, in the backward propagation stage, skip forward re-calculation of a last layer included in the checkpoint module if the last layer is a Flash Attention layer, and determine a gradient of the Flash Attention layer according to a preset Flash Attention backward algorithm, and determine the gradients of the other network layers in the checkpoint module by performing forward re-calculation.

[0012] According to another aspect of the present 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 method.

[0013] According to another aspect of the present disclosure, a non-volatile computer readable storage medium is provided, which stores a computer program, wherein the computer program is executed by a processor to implement the steps of the above method.

[0014] 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 is executed by a processor to implement the steps of the above method.

[0015] In the embodiments of the present disclosure, a checkpoint module in the target model is determined, and only target parameters needed in a forward re-computation process of the checkpoint module in a backward propagation phase are saved in the display memory when forward calculation of the checkpoint module is performed in the forward propagation phase, wherein the target parameters are parameters needed in the forward re-computation process of the checkpoint module in the backward propagation phase; in the backward propagation phase, if the last layer included in the checkpoint module is a Flash Attention layer, forward re-computation of the Flash Attention layer is skipped, and a preset Flash Attention backward algorithm is directly used to determine the gradient of the Flash Attention layer, and then the backward propagation of the checkpoint module is quickly completed in combination with the gradients determined by performing forward re-computation on other network layers in the checkpoint module, so as to save the redundant computation overhead of the forward re-computation of the Flash Attention layer while ensuring the same calculation accuracy and display memory saving amount as the standard re-computation scheme, thereby effectively improving the training efficiency.

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

[0017] The accompanying drawings, which are incorporated in and constitute a part of the specification, illustrate exemplary embodiments, features, and aspects of the present disclosure and serve to explain the principles of the present 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 DESCRIPTION

[0021] Various exemplary embodiments, features, and aspects of the present disclosure will be explained in detail below with reference to the accompanying drawings. Like reference numerals in the drawings denote like functional elements throughout. Although various aspects of the embodiments are illustrated in the drawings, the drawings are not necessarily drawn to scale unless specifically indicated.

[0022] As used herein, the terms "comprise", "comprising", "have", "having", "include", "including", "contain", "containing", or variants thereof, are open-ended and include one or more stated features, integers, elements, steps, components, or functions but do not preclude the presence or addition of one or more other features, integers, elements, steps, components, functions, or groups thereof.

[0023] When an element is referred to as being “connected”, “coupled”, “responsive”, or “in communication” with, to or from another element, it can be directly connected, coupled, responsive, or in communication with the other element or intervening elements can be present.

[0024] Although the terms first, second, third, etc. can 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. Thus, a first element / operation in some embodiments could be termed a second element / operation in other embodiments without departing from the teachings of the present inventive concept.

[0025] “Exemplary” is herein used as an example, embodiment, or illustration. Any embodiment described herein as “exemplary” is not necessarily to be construed as preferred or advantageous over other embodiments.

[0026] In addition, for the purpose of convenience and brevity, detailed descriptions of well-known functions, procedures, components, and circuits can not be described in detail herein. It should be appreciated that those specific details are not intended to limit the scope of the present inventive concept, but rather, the aim is to convey the concepts of the present inventive concept to those skilled in the art with a reasonable degree of particularity.

[0027] The following introduces key technical concepts related to embodiments of the present disclosure.

[0028] Automatic Differentiation (Autograd). Automatic differentiation is the core mechanism of modern deep learning frameworks (e.g., PyTorch, TensorFlow, etc.). During forward propagation, the forward computation process and some tensors needed in the backward computation process are recorded through a dynamic computation graph (Dynamic Computation Graph), and the gradients are automatically calculated using the chain rule during backward propagation. Taking a linear layer y = Wx as an example, the gradient calculation during backward propagation includes: Input gradient Therefore, the input x and the computation graph need to be cached during forward propagation to automatically calculate the gradient according to the computation graph and the chain rule during backward propagation.

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

[0030] Gradient Checkpointing: 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 layer. The Transformer layer is the basic building block of the Transformer model. It consists of a multi-head self-attention mechanism (MHSA) and a feed-forward network (FFN), and is typically optimized using residual connections and layer normalization.

[0032] Self-Attention Mechanism. Self-Attention is a core component of the Transformer architecture. It is used to focus on information at different positions in the input sequence, capture long-range dependencies, determine the correlation between positions, and generate dynamic weighted representations. The calculation process of Self-Attention is:

[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). Flash Attention is an optimized attention computation algorithm that reduces the number of memory read and write by using tiling and memory efficient access strategy, improving the computation efficiency.

[0037] The forward algorithm of Flash Attention layer when forward propagating 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, block division.

[0040] Divide Q into T r = [N / B r ] subblocks Each subblock size is B r xd.

[0041] Divide K, V into T c = [N / B c ] subblocks and Each subblock size is B c xd.

[0042] Step 12, output initialization.

[0043] Divide the output matrix O ∈ R N×d into T r subblocks Each subblock size is B r xd.

[0044] Divide the logsumexp vector L into T r subblocks Each size is B r .

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

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

[0047] On-chip initialization: subblock output is initialized to zero matrix Logsumexp vector L subblock is initialized to zero matrix Maximum cache

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

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

[0050] Compute attention scores on-chip: sub-block product Update Compute normalization scores Update

[0051] Update output sub-block

[0052] Step 15, normalization and write.

[0053] Compute final output

[0054] Compute log-sum-exp

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

[0056] Step 16, output matrix O and vector L.

[0057] The Flash Attention backpropagation algorithm when the Flash Attention layer is backpropagated is as follows.

[0058] Input: matrix Q, K, V, O, dO∈R N×d , and log-sum-exp vector L∈R N in high-bandwidth memory (HBM), block size parameter B c , B r .

[0059] Step 21, block partitioning.

[0060] Partition Q into T r = [N / B r ] sub-blocks Each sub-block size is B r ×d.

[0061] Partition K, V into T c = [N / B c ] sub-blocks and Each sub-block size is B c×d.

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

[0063] The output matrix O is divided into T r sub-blocks Each sub-block size is B r ×d.

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

[0065] The log sum exponent vector L is divided into T r sub-blocks Each size is B r .

[0066] Step 23, input initialization.

[0067] Initialize dQ to zero matrix in the video memory dQ = (0) N×d , and divide it into T r sub-blocks Each sub-block size is B r ×d;

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

[0069] Step 24, calculate auxiliary vector D.

[0070] Calculate point-by-point product D = rowsum(dO·O) ∈ R N×d , write D into the video memory, and divide D into T r sub-blocks Each sub-block size is B r .

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

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

[0073] Initialize the sub-block gradient to zero matrix on-chip

[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 memory.

[0076] Compute attention scores on-chip: sub-block product Compute normalization scores

[0077] Gradient update:

[0078] Step 27, write back to video memory. i , dK j , dV j

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

[0080] The standard gradient checkpointing technique and the standard recomputation scheme (Vanilla Recomputation) are introduced as follows.

[0081] In the standard gradient checkpointing technique, during the forward propagation phase, the model is first divided into several modules, and some of the modules are selected as checkpoint modules. For these checkpoint modules, only their input and output data are saved, and the activation values generated in the intermediate calculation process are discarded. Taking the Transformer model as an example, the Feedforward Neural Network (FFN) and the single-layer Transformer network layer can be set as checkpoint modules.

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

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

[0084] Input: x1;

[0085] (1.1) q, k, v = Wq x x1, Wk x x1, Wv x x1; save x1 in video memory

[0086] ​(1.2) o = FA(q, k, v); save q, k, v in memory;

[0087] (1.3) y = Wo x o; save o in memory.

[0088] The standard back propagation process is:

[0089] Input:

[0090] (1.4) Delete o in memory;

[0091] (1.5) Delete o in memory;

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

[0093] Wherein, backward indicates calling automatic differentiation mechanism.

[0094] Only the first two steps in the self-attention mechanism Self-Attention are determined as checkpoint modules, and after adopting the standard recalculation technology, the calculation process is adjusted as follows.

[0095] The forward propagation process is:

[0096] Input: x1;

[0097] (2.1) q, k, v = Wq x x1, Wk x x1, Wv x x1; save x1 in memory.

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

[0099] (2.3) y = Wo x o; save o in memory.

[0100] The back propagation process is:

[0101] (2.4) Input:

[0102] (2.5) Delete o in memory;

[0103] (2.6) q, k, v = Wq x x1, Wk x x1, Wv x x1; save x1 in memory.

[0104] (2.7) o = FA(q, k, v); save q, k, v in memory.

[0105] (2.8) Delete o in memory;

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

[0107] After the standard recomputation, in the back propagation process, (2) is executed to delete o saved in the video memory, and o is recomputed in (4). The forward recomputation of the Flash Attention layer is executed through (4) and (5). After (4) is executed, q, k, and v are saved in the video memory, and then o, the corresponding derivative and q, k, and v are read from the video memory when the automatic differentiation mechanism is called in (5), and the back operator of the Flash Attention layer is called based on the computation graph to obtain

[0108] The main advantage of the standard recomputation is that the video memory occupation of q, k, and v is effectively saved after the forward propagation is completed and before the back propagation is started (usually in the forward propagation and back propagation process of other network layers except the self-attention mechanism). However, the corresponding cost is that the forward recomputation of [Wq, Wk, Wv] x x1 and FA(q, k, v) needs to be performed in the back propagation process.

[0109] The standard recomputation in the prior art completely reperforms the forward recomputation of each network layer in the checkpoint module and then derives through the automatic differentiation mechanism. Due to the repeated execution of the forward computation, additional computational overhead is caused.

[0110] To solve the above technical problems, the model training method provided in the embodiments of the present disclosure can save the redundant computational overhead of the forward recomputation of the Flash Attention layer while ensuring the same computational accuracy and video memory saving amount as the standard recomputation scheme, thereby effectively improving the training efficiency. The model training method provided in the embodiments of the present disclosure is described in detail below.

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

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

[0113] In step S12, in the forward propagation stage, forward calculation is performed on the checkpoint module, and only target parameters in the checkpoint module during the forward calculation are saved in the display memory, where the target parameters are parameters required in the forward recalculation of the checkpoint module in the backward propagation stage.

[0114] In step S13, in the backward propagation stage, if the last layer included in the checkpoint module is a Flash Attention layer, forward recalculation of the Flash Attention layer is skipped, and the gradient of the Flash Attention layer is determined according to a preset Flash Attention backward algorithm, and forward recalculation is performed on other network layers in the checkpoint module to determine the corresponding gradient.

[0115] In the embodiments of the present disclosure, the checkpoint module in the target model is determined, and then only target parameters required in the forward recalculation of the checkpoint module in the backward propagation stage are saved in the display memory when the forward calculation of the checkpoint module is performed in the forward propagation stage; in the backward propagation stage, if the last layer included in the checkpoint module is a Flash Attention layer, forward recalculation of the Flash Attention layer is skipped, and the gradient of the Flash Attention layer is directly determined by using a preset Flash Attention backward algorithm, and then the backward propagation of the checkpoint module is quickly completed in combination with the gradient determined by performing forward recalculation on other network layers in the checkpoint module, so as to save the redundant calculation overhead of the forward recalculation of the Flash Attention layer while ensuring the same calculation accuracy and display memory saving amount as the standard recalculation scheme, and effectively improve the training efficiency.

[0116] In the embodiments of the present 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, the checkpoint module can refer to a subnetwork or a combination of network layers to which the gradient checkpoint technology is applied. The checkpoint module can be a continuous network layer with higher display memory occupation and intermediate activation values that can be recalculated in the target model, for example, a feedforward network, a residual block, a single-layer Transformer Layer, etc. The core principle is to save only the input and output of the checkpoint module during forward propagation, and discard the intermediate activation values to save the display memory; during backward propagation, the intermediate activation values are recovered by recalculation.

[0118] In the embodiments of the present disclosure, when the checkpoint module ends with a Flash Attention layer, the forward recalculation of the Flash Attention layer can be skipped, and the gradient of the Flash Attention layer can be derived directly by using a preset Flash Attention reverse algorithm, so as to significantly reduce the calculation overhead while maintaining the same memory saving effect.

[0119] In the embodiments 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 memory for direct use in the subsequent backward propagation stage. In the forward propagation stage, only part of the parameters of the checkpoint module are retained, that is, the target parameters required in the forward recalculation process of the checkpoint module in the backward propagation stage.

[0120] In the backward propagation stage, the processing of the checkpoint module can be divided into the following two cases:

[0121] When the last layer of the checkpoint module is not a Flash Attention layer, a standard recalculation scheme can be used, that is, the complete forward recalculation of the checkpoint module is performed again to recover all intermediate activation values, and then the gradient calculation is performed according to the automatic differentiation mechanism;

[0122] When the last layer of the checkpoint module is a Flash Attention layer, an optimized calculation scheme can be implemented, that is, the forward recalculation of the Flash Attention layer of the checkpoint module is skipped, the gradient of the Flash Attention layer is determined according to the preset Flash Attention reverse algorithm, and the standard recalculation process is still performed on other network layers in the checkpoint module to determine the corresponding gradient.

[0123] In a possible implementation, 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 reverse algorithm, including: reading the output and the output gradient of the Flash Attention layer from the memory; performing forward recalculation on each network layer in the checkpoint module except 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 by using the preset Flash Attention reverse algorithm according to the input, the output and the output gradient of the Flash Attention layer.

[0124] Based on the above description of the Flash Attention, it can be known that the input of the preset Flash Attention reverse algorithm comprises Q, K, V, O and dO, wherein Q, K and V are inputs of the Flash Attention layer, O is an output of the Flash Attention layer, and dO is an output gradient of the Flash Attention layer. The output of the Flash Attention layer is an input of a network layer after the Flash Attention layer.

[0125] The output of the Flash Attention layer is an input of a network layer after the Flash Attention layer, and therefore, when the network layer after the Flash Attention layer performs the forward propagation, the input thereof is saved in the video memory, and therefore, when the reverse propagation is performed on the Flash Attention layer, the output of the Flash Attention layer can be read from the video memory.

[0126] The output gradient of the Flash Attention layer is an input gradient of a network layer after the Flash Attention layer, and therefore, after the network layer after the Flash Attention layer performs the reverse propagation, the determined input gradient thereof is saved in the video memory, and therefore, when the reverse propagation is performed on the Flash Attention layer, the output gradient of the Flash Attention layer can be read from the video memory.

[0127] The input of the Flash Attention layer is an output of a network layer before the Flash Attention layer, and therefore, according to the target parameters, the forward recalculation is performed on each network layer except the Flash Attention layer in the checkpoint module, and the input of the Flash Attention layer can be determined.

[0128] Furthermore, according to the input, the output and the output gradient of the Flash Attention layer, the input gradient of the Flash Attention layer is determined by using the preset Flash Attention reverse algorithm.

[0129] In a possible implementation, the method further includes: in the forward propagation stage, taking the output of the Flash Attention layer as the input of a network layer following the Flash Attention layer, performing forward calculation on the network layer, determining the output of the network layer, and saving the input and the computation graph of the network layer in the GPU memory; in the backward propagation stage, determining the input gradient of the network layer according to the input and the computation graph, the output and the output gradient of the network layer by means of the 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 the forward calculation on the Flash Attention layer is completed in the forward propagation stage, the output of the Flash Attention layer is taken as the input of a network layer following the Flash Attention layer, and then forward calculation is performed on the network layer, the output of the network layer is determined, and the input and the computation graph of the network layer are saved in the GPU memory.

[0131] In the case where the network layer following the Flash Attention layer is not the network layer in the checkpoint module, the input gradient of the network layer can be determined according to the input and the computation graph, the output and the output gradient of the network layer by means of the automatic differentiation mechanism in the backward propagation stage, and the input gradient of the network layer is saved in the GPU memory, and the output gradient is also the output gradient of the Flash Attention layer. Then, when the backward propagation is performed on the Flash Attention layer, the output gradient of the Flash Attention layer can be read from the GPU memory.

[0132] Taking the self-attention mechanism Self-Attention of one Transformer network layer (Transformer Layer) as an example, only the first two steps in the self-attention mechanism Self-Attention are determined as the checkpoint module. The calculation process of the recalculation technology of the embodiment of the present disclosure is adjusted as follows.

[0133] The forward propagation process is as follows:

[0134] (3.1) q, k, v = Wq x x1, Wk x x1, Wv x x1; x1 is saved in the GPU memory;

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

[0136] (3.3) y = Wo x o; o is saved in the GPU memory.

[0137] It can be seen that the forward propagation process does not change from the standard recalculation scheme.

[0138] Compared with the above-mentioned standard recomputation technique, the back propagation process is adjusted as follows:

[0139] (3.4) Input:

[0140] (3.5)

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

[0142] (3.7)

[0143] (3.8)

[0144] wherein (3.5) is a process of performing back computation on a network layer after the Flash Attention layer in the back propagation stage, and determines both the input gradient of the network layer and the output gradient of the Flash Attention layer; (3.3) is a process of performing forward computation on a network layer after the Flash Attention layer in the forward propagation stage, and the output o of the Flash Attention layer is saved in the video memory as the input of the network layer; (3.6) is a process of performing forward recomputation on the network layer except the Flash Attention layer in the checkpoint module, and the output q, k, v is the input of the Flash Attention layer.

[0145] According to the input q, k, v, the output o and the output gradient of the Flash Attention layer, the input gradient of the Flash Attention layer can be directly determined by using the preset Flash Attention back algorithm. Compared with the back propagation process of the above-mentioned standard recomputation technique, the back propagation process of the embodiment of the present disclosure can, through (3.7), not only maintain the same gradient calculation result as (2.7) and (2.8), but also save the calculation overhead of (2.7) for performing forward recomputation on the Flash Attention layer.

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

[0147] In a possible implementation, the method further includes: in the forward propagation phase, saving, in the video memory, intermediate parameters used in the forward calculation of the Flash Attention layer after the forward calculation of the Flash Attention layer is performed; and in the backward propagation phase, reading the intermediate parameters from the video memory before the preset Flash Attention backward algorithm is executed.

[0148] Based on the above description of the Flash Attention, it is known that the input of the preset Flash Attention backward algorithm further includes intermediate parameters: the logarithm sum exponent vector L and the block size parameter B c , B r . Therefore, in the forward propagation phase, the logarithm sum exponent vector L and the block size parameter B c , B r are saved in the video memory after the forward calculation of the Flash Attention layer is performed, so that the logarithm sum exponent vector L and the block size parameter B c , B r are read from the video memory before the preset Flash Attention backward algorithm is executed in the backward propagation phase.

[0149] In a possible implementation, the checkpoint module includes a linear layer before the Flash Attention layer, and the input gradient of the Flash Attention layer is the output gradient of the linear layer; the method further includes: saving, in the video memory, the input, output and computation graph of the linear layer after the forward recalculation of the linear layer is performed.

[0150] Taking the Self-Attention as an example, (3.6) is the process of performing the forward calculation of the linear layer before the checkpoint module of the Flash Attention layer, and the input x1, the output q, k, v and the computation graph are saved after the forward recalculation is performed.

[0151] In a possible implementation, the forward recalculation of the other network layers in the checkpoint module to determine the corresponding gradient includes: determining the input gradient and the weight gradient of the linear layer by the automatic differentiation mechanism according to the input and the computation graph of the linear layer, the output and the output gradient of the linear layer.

[0152] Taking the above Self-Attention as an example, the input gradient of the Flash Attention layer determined based on the preset Flash Attention reverse algorithm is the output gradient of a linear layer before the Flash Attention layer. Therefore, according to the input and the computation graph of the linear layer, the output and the output gradient, the input gradient and the weight gradient of the linear layer are determined through the automatic differentiation mechanism.

[0153] (3.8) For the process of performing reverse calculation on a linear layer before the Flash Attention layer in the checkpoint module, according to the input x1 and the computation graph, the output q, k, v and the output gradient the input gradient of the linear layer is determined through the automatic differentiation mechanism and the weight gradient

[0154] In the embodiments of the present disclosure, the checkpoint module in the target model is determined, and only the target parameters needed in the forward re-computation process of the checkpoint module in the backward propagation phase are saved in the video memory when the forward calculation of the checkpoint module is performed in the forward propagation phase, wherein the target parameters are the parameters needed in the forward re-computation process of the checkpoint module in the backward propagation phase; in the backward propagation phase, if the last layer included in the checkpoint module is a Flash Attention layer, the forward re-computation of the Flash Attention layer is skipped, and the gradient of the Flash Attention layer is directly determined by using the preset Flash Attention reverse algorithm, and then the backward propagation of the checkpoint module is quickly completed in combination with the gradients determined by performing the forward re-computation on each network layer in the checkpoint module, so as to save the redundant calculation overhead of the forward re-computation of the Flash Attention layer while ensuring the same calculation accuracy and memory saving amount as the standard re-computation scheme, thereby effectively improving the training efficiency.

[0155] In an example, the target model of the embodiments of the present disclosure can be a MoE (Mixture of Experts) model, and the MLA part therein is determined as a checkpoint module. The checkpoint module is preceded by a first network layer (q / k_down_project, dimension reduction processing of q and k); the checkpoint module includes: a second network layer (q / k_layernorm, layer normalization of q and k), a third network layer (q / k_up_project, dimension increase processing of q and k), and a Flash Attention layer. The checkpoint module is followed by a fourth network layer (o_project, output o). At this time, in the back propagation phase, only the second network layer (q / k_layernorm) and the third network layer (q / k_up_project) in the checkpoint module need to perform forward recalculation, and the forward recalculation of the Flash Attention layer is skipped.

[0156] Suppose that the MoE model includes 60 checkpoint modules, and the forward recalculation of the Flash Attention layer is skipped for each checkpoint module, which can save 528 MB of GPU memory. At this time, a total of 30 GB of GPU memory can be saved.

[0157] In the standard recalculation technical solution, the forward recalculation of each network layer in the checkpoint module is required, that is, the calculation overhead of performing forward recalculation on the second network layer (q / k_layernorm), the third network layer (q / k_up_project), and the Flash Attention layer is required. Among them, the calculation overhead of the forward recalculation of the second network layer (q / k_layernorm) is small and can be ignored; the calculation overhead of the forward recalculation of the third network layer (q / k_up_project) is 0.44T Flops in total; and the calculation overhead of the forward recalculation of the Flash Attention layer is 1.25T Flops. Therefore, the standard recalculation scheme needs to pay an additional calculation cost of 1.69T Flops.

[0158] Using the recalculation calculation scheme of the embodiments of the present disclosure, the forward recalculation of the Flash Attention layer is skipped, which can reduce the calculation overhead to 0.44T Flops, thereby effectively improving the training efficiency.

[0159] It can be understood that the above-mentioned various method embodiments mentioned in the present disclosure can be combined with each other to form combined embodiments without violating the principle logic. Limited by the length, the present disclosure will not be repeated. Those skilled in the art can understand that the specific execution order of each step in the above-mentioned method should be determined according to its function and possible internal logic.

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

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

[0162] A determination module 21 is configured to determine a checkpoint module in a target model.

[0163] A forward propagation module 22 is configured to perform forward calculation on the checkpoint module in a forward propagation stage, and save only target parameters in a video memory during the forward calculation of the checkpoint module, wherein the target parameters are parameters required during forward re-calculation of the checkpoint module in a backward propagation stage.

[0164] A backward propagation module 23 is configured to, in the backward propagation stage, skip forward re-calculation of a Flash Attention layer if the last layer included in the checkpoint module is the Flash Attention layer, and determine a gradient of the Flash Attention layer according to a preset Flash Attention backward algorithm, and perform forward re-calculation on other network layers in the checkpoint module to determine corresponding gradients.

[0165] In a possible implementation, the backward propagation module 23 is specifically configured to:

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

[0167] perform forward re-calculation on other network layers in the checkpoint module except the Flash Attention layer according to the target parameters, to determine an input of the Flash Attention layer;

[0168] determine an input gradient of the Flash Attention layer according to the input, the output, and the output gradient of the Flash Attention layer, by using the preset Flash Attention backward algorithm.

[0169] ​In a possible implementation, the forward propagation module 22 is specifically configured to: in the forward propagation phase, take the output of the Flash Attention layer as the input of a network layer after the Flash Attention layer, perform forward calculation on the network layer, determine the output of the network layer, and save the input and the computation graph of the network layer in the video memory;

[0170] The backward propagation module 23 is specifically configured to: in the backward propagation phase, determine the input gradient of the network layer by means of an automatic differentiation mechanism according to the input and the computation graph of the network layer, the output and the output gradient, 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 a possible implementation, the forward propagation module 22 is specifically configured to: after performing forward calculation on the Flash Attention layer, save the intermediate parameters used in the process of performing forward calculation on the Flash Attention layer in the video memory;

[0172] The backward propagation module 23 is specifically configured to: in the backward propagation phase, read the intermediate parameters from the video memory before performing the preset Flash Attention backward algorithm.

[0173] In a possible implementation, the checkpoint module includes a linear layer before the Flash Attention layer, and the input gradient of the Flash Attention layer is the output gradient of the linear layer;

[0174] The backward propagation module 23 is specifically configured to:

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

[0176] In a possible implementation, the backward propagation module 23 is specifically configured to:

[0177] Determine the input gradient and the weight gradient of the linear layer by means of an automatic differentiation mechanism according to the input and the computation graph of the linear layer, the output and the output gradient.

[0178] In some embodiments, the apparatus provided by the embodiments of the present disclosure has functions or includes modules that can be used to execute the methods described in the above method embodiments, and the specific implementation can refer to the description of the above method embodiments. For brevity, details are not repeated here.

[0179] The embodiments of the present disclosure also provide an electronic device, including a memory, a processor and a computer program stored in the memory, and the processor executes the computer program to implement the steps of the above method.

[0180] The embodiment of the disclosure further provides a nonvolatile computer readable storage medium, which stores a computer program, and the computer program is executed by a processor to realize the steps of the method.

[0181] The embodiment of the disclosure further provides a computer program product, which comprises a computer program or a nonvolatile computer readable storage medium carrying the computer program, and the computer program is executed by a processor to realize the steps of the method.

[0182] Figure 3 A block diagram of an electronic device according to an embodiment of the disclosure is shown. Referring to Figure 3 , the apparatus 1900 can be provided as a server or a terminal device. Referring to Figure 3 , the apparatus 1900 comprises a processing component 1922, which further comprises one or more processors, and a memory resource represented by a memory 1932, for storing instructions executable by the processing component 1922, such as an application program. The application program stored in the memory 1932 can comprise one or more than one module each corresponding to a set of instructions. In addition, the processing component 1922 is configured to execute the instructions to perform the above method.

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

[0184] In an exemplary embodiment, a nonvolatile computer readable storage medium, such as a memory 1932 comprising computer program instructions, is also provided, which can be executed by the processing component 1922 of the apparatus 1900 to complete the above method.

[0185] Computer readable storage media can be any media that can be read by a machine. Such media can include, but is not limited to, optical discs, magnetic discs, magnetic tapes, electronic memories, and / or any combination thereof. Computer readable storage media can be non-transitory, in that it can be a tangible medium. In some embodiments, computer readable storage media can be non-transitory, in that it can not be a signal per se. In other embodiments, computer readable storage media can be a transitory medium, in that it can be a signal. In some embodiments, computer readable storage media can be non-transitory, in that it can not be a signal per se, but can be a tangible medium. In other embodiments, computer readable storage media can be a transitory medium, in that it can be a signal. In some embodiments, computer readable storage media can be non-transitory, in that it can not be a signal per se, but can be a tangible medium. In other embodiments, computer readable storage media can be a transitory medium, in that it can be a signal.

[0186] The computer programs (or computer readable program instructions) described herein can be downloaded from a computer readable storage medium to respective computing / processing devices or external computer or external storage device via a network, for example, the Internet, a local area network, a wide area network, and / or a wireless network. The network can include copper transmission cables, optical transmission fibers, wireless transmission, routers, firewalls, switches, gateway computers, and / or edge servers. A network adapter card or network interface in each computing / processing device receives computer readable program instructions from the network and forwards the computer readable program instructions to storage media in respective computing / processing devices for execution.

[0187] Computer readable program instructions for carrying out operations of the present disclosure can be assembler instructions, instruction-set-architecture (ISA) instructions, machine instructions, machine dependent instructions, microcode, firmware instructions, state-setting data, or either source code or object code written in any combination of one or more programming languages, including an object oriented programming language such as Smalltalk, C++ or the like and conventional procedural programming languages, such as the "C" programming language or similar programming languages. The computer readable program instructions can execute entirely on the user's computer, partly on the user's computer, as a stand-alone software package, partly on the user's computer and partly on a remote computer or entirely on the remote computer or server. In the latter scenario, the remote computer can be connected to the user's computer through any type of network, including a local area network (LAN) or a wide area network (WAN), or the connection can be made to an external computer (for example, through the Internet using an Internet Service Provider). In some embodiments, electronic circuitry including, for example, programmable logic circuitry, field-programmable gate array (FPGA), or programmable logic array (PLA) can execute the computer readable program instructions by utilizing state information of the computer readable program instructions to personalize the electronic circuitry, in order to perform aspects of the present disclosure.

[0188] The computer readable program instructions can 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 apparatus or other device to produce a computer implemented process such that the instructions which execute on the computer or other programmable apparatus provide processes for implementing the functions / acts specified in the flowchart and / or block diagram block or blocks.

[0189] The computer readable program instructions can 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 apparatus or other device to produce a computer implemented process such that the instructions which execute on the computer or other programmable apparatus provide processes for implementing the functions / acts specified in the flowchart and / or block diagram block or blocks.

[0190] The computer readable program instructions can 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 such that the instructions which execute on the computer, other programmable data processing apparatus, or other device implement the functions / acts specified in the flowchart and / or block diagram block or blocks.

[0191] The flow diagrams and the block diagrams in the drawings are presented to 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 the flow diagrams and the block diagrams can represent a module, segment, or portion of instructions, which comprises one or more executable instructions for implementing the specified logic functions. In some alternative implementations, the functions noted in the blocks can occur out of the order noted in the figures. For example, two blocks shown in succession may, in fact, be executed substantially concurrently, or the blocks can sometimes be executed in the reverse order, depending upon the functionality involved. It will also be noted that each block of the block diagrams and / or flow diagrams, and combinations thereof, can be implemented by special purpose hardware-based systems that perform the specified functions or acts, or combinations of special purpose hardware and

[0192] Embodiments of the present disclosure have been described above, and the description is intended to be illustrative of the embodiments and not restrictive. Many modifications and variations of the described embodiments are possible and are within the scope of the disclosure. The selection of terms is intended to best describe the principles of the embodiments, practical application, or technical improvements in the art, 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; In the forward propagation phase, forward calculation is performed on the checkpoint module, and only target parameters of the checkpoint module in the forward calculation process are stored in the video memory, wherein the target parameters are parameters required for the forward recalculation process of the checkpoint module in the backward propagation phase; In the back-propagation phase, if the last layer included in the checkpoint module is the 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 reverse algorithm, and the forward recalculation is 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 a preset Flash Attention reverse algorithm includes: Read the output and output gradient of the Flash Attention layer from the video memory; Perform forward recalculation on each network layer in the checkpoint module except the Flash Attention layer according to the target parameters to determine the input of the Flash Attention layer; According to 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 reverse algorithm.

3. The method according to claim 2, characterized in that The method further comprises: In the forward propagation phase, the output of the Flash Attention layer is used as the input of a network layer following the Flash Attention layer, forward calculation is performed on the network layer to determine the output of the network layer, and the input and calculation graph of the network layer are saved in the video memory; In the back-propagation phase, the input gradient of the network layer is determined through an automatic differentiation mechanism based on the input and computational graph, output, and output gradient of the network layer, and the input gradient of the network layer is saved 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 comprises: In the forward propagation stage, after performing forward calculation on the Flash Attention layer, intermediate parameters used in the forward calculation of the Flash Attention layer are stored in the video memory; In the back propagation phase, before executing the preset Flash Attention reverse algorithm, the intermediate parameters are read from the video memory.

5. The method according to claim 2, characterized in that The checkpoint module includes a linear layer before the FlashAttention layer, and the input gradient of the FlashAttention layer is the output gradient of the linear layer; The method further comprises: 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 performing forward recalculation on other network layers in the checkpoint module to determine corresponding gradients includes: According to the input and calculation graph, output and output gradient of the linear layer, the input gradient and weight gradient of the linear layer are determined through the automatic differentiation mechanism.

7. 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 perform forward calculations on the checkpoint module during a forward propagation phase, and to store in a video memory only target parameters during the forward calculations performed by the checkpoint module, wherein the target parameters are parameters required for performing forward recalculation on the checkpoint module during a backward propagation phase; A back-propagation module is configured to, during a back-propagation phase, skip forward recalculation of the Flash Attention layer if the last layer included in the checkpoint module is the Flash Attention layer, determine the gradient of the Flash Attention layer according to a preset Flash Attention reverse 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, the steps of the method according to any one of claims 1 to 6 are implemented.

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

Images