GPU-based high-IO perceptibility long-sequence text injection method, medium and device

Abstract

The invention relates to a GPU-based high-IO perceptibility long-sequence text injection method, medium and equipment, and the method comprises the steps: obtaining and coding a long-sequence text input, processing a coded matrix vector through a multi-head attention mechanism, and realizing text injection, the matrix vector comprises a query Q, a key K and a value V. The method is characterized in that the value V is equal to the query Q, the key K and the value V. In the multi-head attention mechanism, the similarity is calculated by using an L2 distance, the matrix vector is segmented into blocks to realize thread block parallel processing during forward propagation, and the similarity-based result matrix and the gradients of the query Q, the key K and the value V are recalculated by using a recalculation strategy during back propagation. Compared with the prior art, the method has the advantages of reducing the overhead of attention calculation memory access, improving the operation efficiency and the like.

Description

Technical Field

[0001] The present invention relates to the field of artificial intelligence technology, and in particular to a GPU-based high IO-aware long-sequence text injection method, medium and device. Background Art

[0002] In recent years, large-scale Transformer-based language models such as ChatGPT have been widely used in natural language processing tasks. However, since the computational complexity of such models increases at a square rate with the sequence length n, that is, O(n 2 ), so there are some challenges when processing longer texts. The input and output sequence lengths of large language models are both in tokens. Different large models have different token length limits. For example, models such as BART and GPT2 usually only support up to 1024 tokens.

[0003] In more and more scenarios, it is necessary to process longer documents, historical conversations, and code bases in order to capture long-distance dependencies through longer context learning, and use this historical information to improve the semantic understanding and generalization capabilities of the model. In particular, in many AIGC (artificial intelligence content generation) scenarios, text-to-speech, intelligent digital content editing, intelligent digital content generation, and other references require long-sequence text injection. However, the processing cost of long-sequence text is very high, and providing a long context background for the model is still a huge challenge. The attention mechanism occupies a considerable proportion of the model training and inference time.

[0004] Therefore, how to effectively perform long-sequence text injection based on the attention module is an urgent problem to be solved. Since its running time and memory requirements are in a square relationship with the length of the input sequence, it faces dual challenges in terms of memory and speed.

[0005] The existing technology mainly implements long sequence text injection based on attention module through the following methods:

[0006] Flash attention: Some people have proposed approximate attention methods that aim to reduce attention calculations and memory requirements. These methods include sparse approximation, low-rank approximation, and their combination. Although these methods can reduce the calculation to linear or near-linear, they focus too much on reducing the number of floating-point operations per second (FLops) and tend to ignore the overhead from memory access (IO). By reducing GPU memory read / write, Flash Attention runs 2-4 times faster than PyTorch standard attention and requires 5-20 times less memory. Flash Attention uses two methods to reduce the overhead of video memory. One is to split the input into blocks and make multiple passes on the input blocks to perform softmax reduction in an incremental manner; the other is that Flash Attention does not use intermediate attention matrices and reduces HBM memory consumption by storing normalization factors. The Flash Attention algorithm does not write S and P to HBM, but writes them to HBM in blocks, stores the softmax normalization factors of the forward pass, and quickly recalculates the on-chip attention in the backward propagation, which is faster than the standard method of reading intermediate attention matrices from HBM. Even though FLOPS increases due to recomputation, it runs faster and uses less memory (linear in sequence length), mainly because the number of HBM accesses is greatly reduced.

[0007] Flash attention2: Flash Attention achieves linear memory savings and 2-4x runtime speedup by exploiting the asymmetry of the GPU memory hierarchy, while using no approximation methods. However, Flash Attention is still much slower than optimized matrix multiplication operations, achieving only 25-40% of the theoretical maximum FLOPs / s. FlashAttention-2 addresses the inefficiency of Flash Attention by optimizing the work partitioning. Specifically, FlashAttention-2 adjusts the algorithm to reduce the number of non-matrix multiplication FLOPs; parallelizes attention, even for a single head, across different thread blocks to increase occupancy; and within each thread block, assigns work to thread warps to reduce communication through shared memory. These improvements make FlashAttention-2 about 2x faster than Flash Attention, reaching 50-73% of the theoretical maximum FLOPs / s on A100, close to the efficiency of GEMM operations. When used for end-to-end training of GPT-style models, FlashAttention-2 can train at up to 225TFLOPs / s per A100 GPU.

[0008] Flash-Decoding: Flash-Decoding draws on the advantages of Flash Attention and extends the parallelization dimension to the length of the key / value sequence. This method is almost unaffected by the length of the h sequence (which is very important for the LLM model capability), can make full use of the GPU, and can greatly improve the encoding speed even when the batch size is small (inference characteristics).

[0009] Although existing methods have made many operator-level optimizations for traditional attention modules, and have been verified in many scenarios and achieved good results, the traditional attention module uses the general Scaled Dot-ProductAttention. In order to have better stability when training large models, the attention module should use the L2 norm instead of the original similarity calculation. This makes it impossible to directly apply the optimization of Flash Attention to the attention module with added L2 norm stability. In long sequence scenarios, it encounters the problem of video memory wall and it is difficult to achieve optimal performance. Therefore, accelerating the L2 norm attention module is a very important issue, and using the distance matrix method to solve the stability problem of long text injection.

[0010] Therefore, the existing technology has disadvantages such as large video memory overhead and no way to accelerate the attention module that uses the L2 norm to calculate similarity. Summary of the invention

[0011] The purpose of the present invention is to overcome the defects of the above-mentioned prior art and the memory wall problem of long sequence text injection based on the attention module, and to provide a high IO-aware long sequence text injection method, medium and device based on GPU by using IO-aware technologies such as blocking, recalculation, parallel computing, and thread block allocation, so as to reduce the overhead of attention calculation memory access.

[0012] The purpose of the present invention can be achieved by the following technical solutions:

[0013] A GPU-based high IO-aware long-sequence text injection method, the method comprising: obtaining a long-sequence text input and encoding it, processing the encoded matrix vector through a multi-head attention mechanism to achieve text injection, the matrix vector comprising a query Q, a key K and a value V, in the multi-head attention mechanism, using L2 distance to calculate similarity, and when performing forward propagation, cutting the matrix vector into blocks to achieve thread block parallel processing, and when performing back propagation, using a recalculation strategy to recalculate a result matrix based on similarity and the gradients of the query Q, key K and value V.

[0014] Furthermore, the calculation formula for calculating the similarity using the L2 distance is:

[0015] S=-L2(Q,K)=-Q 2 +2QK T -K 2

[0016] Among them, S is the similarity, Q is the query in the attention mechanism, and K is the key in the attention mechanism.

[0017] Furthermore, the specific process of the forward propagation includes:

[0018] The query Q, key K and value V in HBM are divided into blocks, and the corresponding block data is loaded into SRAM to build parallel thread blocks;

[0019] The number of blocks for key K and value V is traversed in parallel, similarity is calculated using L2 distance, and the result is updated;

[0020] After the traversal is completed, the final result is returned to the HBM.

[0021] Furthermore, the specific process of the back propagation includes:

[0022] The query Q, key K and value V in HBM are divided into blocks, and the corresponding block data is loaded into SRAM to build parallel thread blocks;

[0023] According to the number of blocks of query Q, traverse each thread block in parallel, recalculate the similarity using L2 distance, and calculate the gradient of the corresponding block data;

[0024] After the traversal is completed, the obtained gradient is returned to the HBM.

[0025] Furthermore, parallelization can be performed across sequence lengths, batches, and / or attention heads.

[0026] Furthermore, if the size of each thread block is b m ×b n , then the number of thread blocks is The number of traversals is Among them, M is the sequence length of query Q, and N is the sequence length of key K and value V.

[0027] Furthermore, the gradient of the query Q is specifically:

[0028] For a single element Q of the matrix Q it , whose gradient dQ it It is expressed as:

[0029] dQ it =S′[i,:]·K[:,t]-Q ij *rowsum(S′[i,:])

[0030] in, S is the similarity, dS is the gradient of S, S′[i,:] is the i-th row element in S′, K[:,t] is the t-th column element of key K, and rowsum() means summing the rows.

[0031] Furthermore, the gradient of the key K is specifically:

[0032] For a single element Q of the matrix K it , whose gradient dK jt It is expressed as:

[0033] dK jt =S′[:,j]Q[:,t]-K jt *colsum(S′[:,j])

[0034] in, S is the similarity, dS is the gradient of S, S′[:,j] is the j-th column element in S′, and Q[:,t] is the t-th column element of the query Q.

[0035] The present invention also provides a computer-readable storage medium, comprising one or more programs for execution by one or more processors of an electronic device, wherein the one or more programs include instructions for executing the GPU-based high IO-aware long sequence text injection method as described in any one of claims 1-8.

[0036] The present invention also provides an electronic device, comprising one or more processors, a memory and one or more programs stored in the memory, wherein the one or more programs include instructions for executing the GPU-based high IO-aware long sequence text injection method as described in any one of claims 1-8.

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

[0038] 1. Under the condition of calculating similarity based on the L2 distance formula, the IO perception level is studied to reduce video memory consumption and quickly calculate the forward and reverse attention. Specifically, under the condition of replacing the conventional attention calculation similarity with the L2 distance formula, the present invention derives the forward and reverse gradient formulas, and then designs and optimizes the parallel strategies according to the characteristics of the forward and reverse calculations. In the forward calculation, the block strategy is used to improve the IO perception, and the thread block parallelism is designed to improve the operation speed; in the reverse propagation, the optimization technique of heavy calculation is used to break the video memory bottleneck.

[0039] 2. The present invention aims at the memory wall problem of the attention module of long-sequence texts. It effectively solves the memory bottleneck through IO-aware methods such as parallel computing, block computing, and recalculation. It effectively combines the multi-level storage structure of hardware to greatly reduce the memory consumption of long sequences, improve the computational efficiency of key steps such as similarity calculation and softmax, and greatly bring into play the performance of hardware equipment.

[0040] 3. Experiments show that compared with the implementation based on pytorch, the video memory of the present invention can be reduced by 2 to 3 orders of magnitude, and the core steps can be accelerated by 5 times. BRIEF DESCRIPTION OF THE DRAWINGS

[0041] Figure 1 A flowchart of a long sequence text injection method applied by the present invention;

[0042] Figure 2 It is a schematic diagram of the forward calculation process of the present invention;

[0043] Figure 3 It is a schematic diagram of the reverse calculation process of the present invention;

[0044] Figure 4 It is a schematic diagram of data block and thread block parallelism of the present invention. DETAILED DESCRIPTION

[0045] The present invention is described in detail below in conjunction with the accompanying drawings and specific embodiments. This embodiment is implemented based on the technical solution of the present invention, and provides a detailed implementation method and specific operation process, but the protection scope of the present invention is not limited to the following embodiments.

[0046] Terminology

[0047] Attention mechanism: This is a method to calculate the output based on the query and key-value pairs, where the query, key, value, and output are all vectors. The output is the weighted sum of the values, and the weight is the similarity between the query and the corresponding key. The attention mechanism allows the model to focus on the important parts of the input or output, improving expressiveness and efficiency.

[0048] Large-scale pre-trained language model: This is a model that uses a large amount of unlabeled text data to learn language knowledge, such as GPT-33, BERT, XLNet, etc. They can be fine-tuned in different natural language processing tasks to improve performance and generalization ability.

[0049] Contextual learning: This is a method that uses historical information in the input or output to influence the current state or behavior. Contextual learning allows the model to capture long-range dependencies and improve semantic understanding and generation capabilities.

[0050] IO perception: The read and write (I / O) speed of GPU SRAM is 19TB / s, which is more than ten times higher than the read and write (I / O) speed of GPU HBM of 1.5TB / s, and the storage capacity is also several orders of magnitude different. Improving IO perception is essentially to design more efficient reading and writing rules for different levels of fast and slow memory.

[0051] Abbreviations

[0052]

[0053]

[0054] Explanation of symbols

[0055]

[0056] This embodiment provides a high IO-aware long-sequence text injection method based on GPU, which can effectively reduce the video memory consumption and training time of the text injection module. The method mainly includes parallel computing, block computing and recalculation.

[0057] refer to Figure 1 The overall process of the GPU-based high IO-aware long-sequence text injection method of this embodiment is shown. The method includes: obtaining long-sequence text input and encoding, processing the encoded matrix vector through a multi-head attention mechanism to achieve text injection, the matrix vector includes a query Q, a key K and a value V, and in the multi-head attention mechanism, the L2 distance is used to calculate the similarity, and when forward propagation is performed, the matrix vector is cut into blocks to achieve parallel processing of thread blocks, and when backpropagation is performed, the recalculation strategy is used to recalculate the result matrix based on similarity and the gradients of the query Q, key K and value V. The main work of this method is to target Figure 1 Improvements made by multi-head attention in .

[0058] Different from the traditional attention mechanism, this embodiment uses L2 distance to calculate the similarity to improve the stability of the generated content after long text injection. The similarity calculation formula is:

[0059] S=-L2(Q,K)=-Q 2 +2QK T -K 2 (1)

[0060] In addition, the traditional attention mechanism will generate a squared matrix P = softmax(S)∈R after the softmax calculation of similarity. M×N .

[0061] In back propagation, the traditional Q and K gradients are In this embodiment, we first calculate the matrix P, and reversely deduce the gradient of S based on P, dS = dsoftmax(dP)∈R M×N , and then re-derive the gradient formula of Q and K based on the L2 distance similarity.

[0062] During back propagation, the gradient of the known matrix S is According to dS and The matrix Q gradient dQ∈R can be obtained M×d .because More complicated, we take a single element Q of the matrix Q it Conduct analysis. it In the forward calculation, only the i-th row element of S, that is, S i1 ,S i2 ,…,S iN So we only need to consider S[i,:] for Q it The derivative of :

[0063]

[0064] Q it The derivative of depends on the i-th row element S[i,:] of S, so the calculation of Q it The gradient dQ it When it depends only on the i-th row dS[i,:] of dS, we define Then dQ it It can be expressed as:

[0065]

[0066] in, S is the similarity, dS is the gradient of S, calculated based on the matrix P, S′[i,:] is the i-th row element in S′, K[:,t] is the t-th column element of key K, and rowsum() means summing the rows.

[0067] Similarly, we can calculate the gradient of matrix K. Take a single element K of matrix K jt For analysis, K jt In the forward calculation, it is only related to the j-th column element of S, so only the j-th column element of S needs to be considered for K jt The influence of K jt The gradient dK jt :

[0068] dK jt =S′[:,j]Q[:,t]-K jt *colsum(S′[:,j]) (4)

[0069] Among them, S′[:,j] is the j-th column element in S′, and Q[:,t] is the t-th column element of query Q.

[0070] In the reverse calculation, the video memory of S′ grows quadratically with the length of the sequence, and the row and column dimensions of S′ grow linearly with the length of the sequence. Next, we will introduce how to use IO-aware technologies such as blocking and recalculation to solve the GPU computing QK T , S′’s IO bottleneck.

[0071] 1. Forward Propagation

[0072] The forward propagation is to get the result and loss. Figure 2 The specific process of forward propagation includes: dividing the query Q, key K and value V in HBM into blocks, loading the corresponding block data into SRAM, and building parallel thread blocks; according to the number of blocks of key K and value V, traversing each thread block in parallel, using L2 distance to calculate the similarity, and updating the result; after the traversal is completed, the final result is returned to HBM.

[0073] In the forward propagation, we divide Q, K, V, etc. into blocks and load them into HBM. Each thread block calculates a result block. Since the read and write speed of the GPU's global memory HBM is slower than the on-chip cache SRAM, we put the divided blocks from HBM into SRAM with more efficient reading, and then calculate the corresponding results of these blocks. We set the size of each thread block to b m ×b n , each thread block corresponds to the matrix multiplication of the left input b m b of the row and matrix multiplication result n rows, and the matrix multiplication right input is traversed, the number of traversals is Each traversal will update the softmax result. In this way, we do not need to write the M×N intermediate calculation results in HBM, but can m ×b n Stored in shared memory, to achieve the effect of reducing video memory.

[0074] Since softmax is row-dependent, the calculations of each thread block do not interfere with each other, so we have The parallel operation of multiple thread blocks can fully utilize the GPU's stream multiprocessor computing and improve the GPU occupancy rate. Figure 4As shown in the figure. The blue box is the input matrix placed in HBM, the yellow box indicates the part where the input matrix is ​​divided into blocks and loaded into SRAM for calculation, and the dotted box indicates the super large square matrix that is not placed in HBM but is calculated in blocks in SRAM. In addition to parallelization on the sequence length M, we also parallelize on batches and attention heads. Batches and attention heads can be used as a dimension of cuda kernel gridDim (CUDA kernel grid dimension) separately.

[0075] 2. Back Propagation

[0076] Back propagation calculates and updates the gradient of each layer based on the results of the forward calculation and the loss. In order to calculate the gradient in the backward pass, some intermediate results usually need to be written back to the HBM during the forward calculation, which will generate additional HBM reads and writes and slow down the running time.

[0077] In this embodiment, reference Figure 3 As shown in FIG. 1 , the specific process of back propagation includes: dividing the query Q, key K, and value V in HBM into blocks, loading the corresponding block data into SRAM, and building parallel thread blocks; traversing each thread block in parallel according to the number of blocks of query Q, recalculating the similarity using L2 distance, and calculating the gradient of the corresponding block data; after the traversal is completed, returning the obtained gradient to HBM.

[0078] Because the intermediate results of attention need to occupy M*N video memory, we do not want to store this type of intermediate data for reverse gradient calculation. Backward transfer usually requires matrices S, P∈R M×N to compute the gradient relative to QKV. However, by storing logsumexpL (j) =m (j) +log(l (j) ), where m (j) is the maximum value of the row, l (j) is the denominator of softmax, we can recalculate S,P∈R during back propagation M×N By recomputing in blocks, we achieve a speedup of the reverse computation and reduce the memory requirements for intermediate results.

[0079] If the above method is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present invention, or the part that contributes to the prior art or the part of the technical solution, can be embodied in the form of a software product. The computer software product is stored in a storage medium, including several instructions for a computer device (which can be a personal computer, a server, or a network device, etc.) to perform all or part of the steps of the method described in each embodiment of the present invention. The aforementioned storage medium includes: U disk, mobile hard disk, read-only memory (ROM, Read-Only Memory), random access memory (RAM, Random Access Memory), disk or optical disk and other media that can store program codes.

[0080] This embodiment also provides an electronic device, including one or more processors, a memory and one or more programs stored in the memory, wherein the one or more programs include instructions for executing the GPU-based high IO-aware long sequence text injection method as described above.

[0081] It will be appreciated by those skilled in the art that embodiments of the present invention may be provided as methods, systems, or computer program products. Therefore, the present invention may take the form of a complete hardware embodiment, a complete software embodiment, or an embodiment combining software and hardware. Moreover, the present invention may take the form of a computer program product implemented on one or more computer-usable storage media (including but not limited to disk storage, CD-ROM, optical storage, etc.) containing computer-usable program codes. The schemes in the embodiments of the present invention may be implemented in various computer languages, for example, object-oriented programming language Java and literal scripting language JavaScript, etc.

[0082] The present invention is described with reference to flowcharts and / or block diagrams of methods, devices (systems), and computer program products according to embodiments of the present invention. It should be understood that each process and / or block in the flowchart and / or block diagram, as well as the combination of processes and / or blocks in the flowchart and / or block diagram, can be implemented by computer program instructions. These computer program instructions can be provided to a processor of a general-purpose computer, a special-purpose computer, an embedded processor, or other programmable data processing device to produce a machine, so that the instructions executed by the processor of the computer or other programmable data processing device generate instructions for implementing the processes in the flowchart and / or block diagram. Figure 1 A process or multiple processes and / or boxes Figure 1 A device that provides the functions specified in a block or multiple blocks.

[0083] These computer program instructions may also be stored in a computer-readable memory capable of directing a computer or other programmable data processing device to operate in a specific manner, so that the instructions stored in the computer-readable memory produce an article of manufacture comprising an instruction device, which implements the process Figure 1 A process or multiple processes and / or boxes Figure 1 A function specified in one or more boxes.

[0084] These computer program instructions can also be loaded onto a computer or other programmable data processing device so that a series of operating steps are executed on the computer or other programmable device to produce a computer-implemented process, thereby providing instructions for implementing the process. Figure 1 A process or multiple processes and / or boxes Figure 1 The steps for the functions specified in one or more boxes.

[0085] The preferred specific embodiments of the present invention are described in detail above. It should be understood that a person skilled in the art can make many modifications and changes based on the concept of the present invention without creative work. Therefore, any technical solution that can be obtained by a person skilled in the art through logical analysis, reasoning or limited experiments based on the concept of the present invention on the basis of the prior art should be within the scope of protection determined by the claims.

Claims

1. A GPU-based high IO-aware long-sequence text injection method, the method comprising: A long sequence text input is obtained and encoded, and the encoded matrix vector is processed through a multi-head attention mechanism to realize text injection, wherein the matrix vector includes a query Q, a key K and a value V. It is characterized in that in the multi-head attention mechanism, the L2 distance is used to calculate the similarity, and when forward propagation is performed, the matrix vector is cut into blocks to realize thread block parallel processing, and when back propagation is performed, a recalculation strategy is used to recalculate the result matrix based on the similarity and the gradients of the query Q, key K and value V.

2. The high IO-aware long sequence text injection method based on GPU according to claim 1 is characterized in that: The calculation formula for calculating the similarity using the L2 distance is: S=-L2(Q,K)=-Q 2 +2QK T -K 2 Among them, S is the similarity, Q is the query in the attention mechanism, and K is the key in the attention mechanism.

3. The high IO-aware long sequence text injection method based on GPU according to claim 1 is characterized in that: The specific process of the forward propagation includes: The query Q, key K and value V in HBM are divided into blocks, and the corresponding block data is loaded into SRAM to build parallel thread blocks; The number of blocks for key K and value V is traversed in parallel, similarity is calculated using L2 distance, and the result is updated; After the traversal is completed, the final result is returned to the HBM.

4. The high IO-aware long sequence text injection method based on GPU according to claim 1 is characterized in that: The specific process of back propagation includes: The query Q, key K and value V in HBM are divided into blocks, and the corresponding block data is loaded into SRAM to build parallel thread blocks; According to the number of blocks of query Q, traverse each thread block in parallel, recalculate the similarity using L2 distance, and calculate the gradient of the corresponding block data; After the traversal is completed, the obtained gradient is returned to the HBM.

5. The method for injecting long sequence text with high IO awareness based on GPU according to claim 3 or 4, characterized in that: Parallelize across sequence length, batches, and / or attention heads.

6. The method for injecting long sequence text with high IO awareness based on GPU according to claim 3 or 4, characterized in that: If the size of each thread block is b m ×b n , then the number of thread blocks is The number of traversals is Among them, M is the sequence length of query Q, and N is the sequence length of key K and value V.

7. The high IO-aware long sequence text injection method based on GPU according to claim 4 is characterized in that: The gradient of the query Q is specifically: For a single element Q of the matrix Q it , whose gradient dQ it It is expressed as: dQ it =S′[i,:]·K[:,t]-Q ij *rowsum(S′[i,:]) in, S is the similarity, dS is the gradient of S, S′[i,:] is the i-th row element in S′, K[:,t] is the t-th column element of key K, and rowsum() means summing the rows.

8. The high IO-aware long sequence text injection method based on GPU according to claim 4 is characterized in that: The gradient of the key K is specifically: For a single element Q of the matrix K it , whose gradient dK jt It is expressed as: dK jt =S′[:,j]Q[:,t]-K jt *colsum(S′[:,j]) in, S is the similarity, dS is the gradient of S, S′[:,j] is the j-th column element in S′, and Q[:,t] is the t-th column element of the query Q.

9. A computer-readable storage medium, characterized in that: It includes one or more programs for execution by one or more processors of an electronic device, wherein the one or more programs include instructions for executing the GPU-based high IO-aware long sequence text injection method as described in any one of claims 1-8.

10. An electronic device, characterized in that: It includes one or more processors, a memory and one or more programs stored in the memory, wherein the one or more programs include instructions for executing the GPU-based high IO-aware long sequence text injection method as described in any one of claims 1-8.

Images