Inference operation acceleration method for processing extremely long text for large language model

Abstract

The invention discloses a big language model-oriented inference operation acceleration method for processing an extremely long text, which belongs to the technical field of large-scale pre-training language model inference optimization, and specifically comprises the following steps of: selecting a big language model, starting a sparse attention mode, inputting different texts, recording execution time of each attention head and an attention mode type in different layers, and calculating the inference operation speed of the large language model; establishing a statistical database; solving a multi-head attention load balanced distribution strategy of the large language model under the actual text; establishing a weight index table by splitting the weight matrix; retrieving a weight sub-matrix corresponding to each attention head, and loading the weight sub-matrix to a corresponding GPU device; by asynchronously preloading the MHA calculation weight and the MLP calculation weight of the adjacent layers and combining KV cache management, load balancing in the reasoning process is realized. According to the method, dynamic segmentation, recombination and scheduling of the attention head are carried out in the reasoning stage, so that the problems of multi-GPU load imbalance and resource idling are effectively avoided, and the system throughput of long sequence processing is remarkably improved.

Description

Technical Field

[0001] The present invention belongs to the technical field of large-scale pre-trained language model reasoning optimization, and specifically relates to a reasoning operation acceleration method for processing extremely long texts with a large language model. Background Art

[0002] In recent years, large language models (LLMs) based on the Transformer architecture, such as GPT and LLaMA, have made breakthrough progress in the field of natural language processing (NLP), becoming a key milestone in the development of general artificial intelligence. As its application scenarios continue to expand to areas such as intelligent assistants, long document question and answer, multi-round dialogue, and automatic programming, the length of context sequences that large language models need to process during the inference phase has increased significantly, which has put forward higher requirements on the time efficiency and memory management of computing systems.

[0003] The current computing systems for large language models have the following problems:

[0004] 1) The challenge of reasoning with very long sequences: The computational complexity of the self-attention mechanism in the Transformer model grows quadratically with the length of the context, resulting in a sharp increase in reasoning latency and a significant increase in memory consumption.

[0005] 2) Limitations of existing parallel technologies. Academia and industry have proposed a variety of parallel optimization strategies during the training phase, such as data parallelism, pipeline parallelism, tensor parallelism, and sequence parallelism, to accelerate model training and solve the problem of single GPU (graphics processing unit) memory limitations. However, these parallel optimization strategies are mainly designed for homogeneous attention mechanisms. When processing long texts, if heterogeneous and sparse attention modes are enabled, such as retrieval heads, sparse heads, and stride heads, it is easy to cause load imbalance between multiple GPUs and reduce overall computing efficiency.

[0006] 3) The necessity of parallel optimization for long sequence reasoning. In long text application scenarios, the calculation modes of different attention heads often show heterogeneous characteristics, that is, some heads need to calculate for global positions, such as Full Attention Head, and some heads only need to focus on local or jump information, such as Sparse / Strided Attention Head, which will cause significant differences in calculation overhead. Simply using tensor splitting or other isomorphic parallel strategies may lead to uneven allocation of GPU resources and affect end-to-end latency performance.

[0007] Therefore, there is an urgent need for a parallel optimization method specifically for long text reasoning and supporting load balancing. Summary of the invention

[0008] In response to the problems existing in existing large language models in long text reasoning scenarios, the present invention provides an inference operation acceleration method for large language models processing extremely long texts. By dynamically splitting, reorganizing and scheduling attention heads in the reasoning stage, the problems of multi-GPU load imbalance and resource idling are effectively avoided, and the system throughput of long sequence processing is significantly improved.

[0009] In order to achieve the above purpose, the technical method adopted by the present invention is as follows:

[0010] A method for accelerating the inference operation of a large language model for processing extremely long texts includes the following steps:

[0011] Step 1: Select a large language model, including Transformer layers, each self-attention layer includes attention heads and enable sparse attention mode;

[0012] Input texts with different context lengths into the large language model and record the first The execution time of each attention head of the layer and the type of attention mode it adopts are mapped to the corresponding attention mode type, and then a statistical database is established;

[0013] Step 2: For the actual text , solve The multi-head attention load balancing allocation strategy for the large language model is as follows:

[0014] right GPU devices, Next Each attention head of the layer is assigned to different GPU devices for execution as the first Multi-head attention load distribution scheme for layers ;

[0015] Use the nearest neighbor matching method to find the context length closest to the , and The execution time of each attention head in the layer is Next The corresponding attention head execution time of the layer is calculated to obtain the solution Next Multi-head attention execution time of the layer ;

[0016] by is the load balancing objective function. Through the search algorithm, we can obtain Optimal multi-head attention load distribution scheme for layers , and then obtain Multi-head attention load balancing allocation strategy for large language models ;

[0017] Step 3: The query weight matrix, key weight matrix, and value weight matrix of the self-attention layer are split into The sub-matrices are spliced ​​in the column direction to obtain the query weight sub-matrix, key weight sub-matrix and value weight sub-matrix corresponding to each attention head;

[0018] According to the strategy , determine the attention heads executed on the same GPU device, concatenate the corresponding query weight submatrix, key weight submatrix, and value weight submatrix together to form a reorganized matrix, and get The reorganized matrix is ​​stored in blocks in the CPU memory; and a weight index table is established based on the divided reorganized matrix;

[0019] Step 4: When processing Before the self-attention calculation (MHA calculation) of the layer, according to the scheme and weight index table, retrieve the query weight sub-matrix, key weight sub-matrix and value weight sub-matrix corresponding to each attention head, and load them into the corresponding GPU device;

[0020] Step 5: The calculation of the layer includes MHA calculation and multi-layer perception module calculation (MLP calculation). The load balancing of the reasoning process is achieved by asynchronously preloading the MHA calculation weights and MLP calculation weights of adjacent layers and the Key-Value (KV) cache management used in conjunction.

[0021] Furthermore, the sparse attention mode described in step 1 includes multiple of a local attention mode, a block attention mode, a jumping attention mode and a retrieval attention mode.

[0022] Furthermore, in step 2, The formula is:

[0023] ;

[0024] In the formula, Indicates the GPU device number; Indicates The set of attention head numbers assigned to the GPU devices; express Next Layer The execution time of an attention head.

[0025] Furthermore, the search algorithm in step 2 is a genetic algorithm, a dynamic programming algorithm or a neural network algorithm.

[0026] Furthermore, when the search algorithm is a genetic algorithm, a multi-head attention load balancing allocation strategy is solved. The specific process is:

[0027] By plan For individuals, The layer initializes the population according to the load balancing objective function , calculate individual fitness; iteratively generate candidate solutions through crossover and mutation operations; and update the population through survival competition to obtain the first The solution corresponding to the highest individual fitness of the layer , as a solution ;exist Repeat the iterative process in the layer to obtain a multi-head attention load balancing distribution strategy .

[0028] Furthermore, the specific process of step 5 is:

[0029] In the When calculating the Transformer layer, first perform the The MHA calculation of the layer is performed by using tensor parallel technology to asynchronously preload the first The MLP calculation weights of the layer are transferred to the corresponding GPU device, and the generated KV cache is unloaded to the CPU memory; After the MHA calculation of the layer is completed, the The MLP calculation of the layer is performed, and the first The MHA calculation weights of the layer are transferred to the corresponding GPU device; After the MLP calculation of the layer is completed, the first The MHA calculation of the layer is used to realize the pipeline parallel computing of the reasoning operation.

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

[0031] 1. The present invention proposes an inference operation acceleration method for large language models to process extremely long texts, which can effectively reduce the performance bottleneck caused by the imbalance of attention calculation load in the extremely long sequence reasoning scenario; specifically, by splitting the model weight, each attention head can be independently scheduled and executed, and the heterogeneous load balancing and parallel scheduling of multi-head attention modules can be achieved, thereby effectively overcoming the problem of unbalanced calculation load of long text reasoning, greatly improving the throughput and stability of large language models in extremely long text reasoning tasks, and without changing the model structure. This is the first systematic solution to the problem of calculation imbalance in sparse attention mechanism in long text processing scenarios;

[0032] 2. The present invention fully combines the existing sparse attention mechanism and model parallel technology. It can be used alone or in conjunction with the existing distributed training and reasoning framework, and has significant versatility and scalability.

[0033] 3. Compared with the prior art, the multi-head attention load balancing, weight splitting and parallel scheduling scheme proposed in the present invention has the advantages of rapid deployment, automated decision-making and orthogonality with other parallel mechanisms. It can be widely used in programming assistance, legal text processing, academic document comprehension and other fields that require processing of ultra-long texts. BRIEF DESCRIPTION OF THE DRAWINGS

[0034] In order to more clearly illustrate the embodiments of the present invention or the technical solutions in the prior art, the drawings required for use in the embodiments or the description of the prior art will be briefly introduced below. Obviously, the drawings described below are only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without paying creative work.

[0035] Figure 1 It is a flowchart of the method for accelerating the inference operation of a large language model for processing extremely long texts proposed in Example 1;

[0036] Figure 2 is a schematic diagram of weight segmentation and storage in Example 1;

[0037] Figure 3 is a schematic diagram of the weight index table in Example 1;

[0038] Figure 4 It is a schematic diagram of parallel scheduling and weight preloading in Example 1. DETAILED DESCRIPTION

[0039] In order to further understand the present invention, preferred embodiments of the present invention are described below in conjunction with examples, but it should be understood that these descriptions are only for further illustrating the features and advantages of the present invention, rather than limiting the claims of the invention.

[0040] Embodiment 1:

[0041] This embodiment proposes a method for accelerating the inference operation of a large language model for processing extremely long texts. The process is as follows: Figure 1 As shown, the following steps are included:

[0042] Step 1: Select a large language model, including Transformer layers, each self-attention layer includes attention heads, and enable sparse attention modes, including multiple of local attention mode, block attention mode, jumping attention mode and retrieval attention mode; the large language model may be GPT or LLaMA;

[0043] Input texts with different context lengths into the large language model and record the first The execution time of each attention head of the layer and the type of attention mode it adopts (such as local attention mode, block attention mode, jumping attention mode or retrieval attention mode) are mapped to the corresponding attention mode type to establish a statistical database.

[0044] Assume that the total input texts of different context lengths, record the The text is , then the recorded Next Layer The execution time of an attention head is , the corresponding attention mode type is , the mapping relationship formed is .

[0045] Step 2: For the actual text , solve The multi-head attention load balancing allocation strategy for large language models has the core goal of "minimizing the computing time of the slowest device" to improve the overall throughput of the long text reasoning process. The specific process is as follows:

[0046] Step 2.1: Assume that the total number of GPU devices is ,Will Next Layer The attention heads are assigned to GPU devices, denoted as , as the first Multi-head attention load distribution scheme for layers ;

[0047] Use the nearest neighbor matching method to find the context length closest to the The text of ,but Next Layer Execution time of an attention head Approximately equal to The execution time of ,Right now:

[0048] ;

[0049] make For the The set of attention head numbers allocated on GPU devices, then the solution Next Multi-head attention execution time of the layer for:

[0050] ;

[0051] That is, in parallel computing, the execution time of multi-head attention is determined by the slowest GPU device;

[0052] Step 2.2, in In the Transformer layer, the self-attention layer, the multi-layer perception module layer and the RMS Norm normalization layer are combined together. If only the execution time of the multi-head attention is optimized, the solution is obtained. Next Total delay of the layer :

[0053] ;

[0054] In the formula, is the execution time of the MLP layer; is the execution time of the RMSNorm normalization layer;

[0055] Get the solution Total latency for downloading large language models :

[0056] ;

[0057] because can be regarded as a constant term, so determining the multi-head attention load balancing allocation strategy only requires optimizing ;

[0058] Step 2.3: is the load balancing objective function. Through the search algorithm, we can obtain Optimal multi-head attention load distribution scheme for layers , and then obtain Multi-head attention load balancing allocation strategy for large language models ;

[0059] Specifically, when the search algorithm is a genetic algorithm, a multi-head attention load balancing allocation strategy is solved The specific process is:

[0060] By plan For individuals, The layer initializes the population according to the load balancing objective function , calculate individual fitness; iteratively generate candidate solutions through crossover and mutation operations; and update the population through survival competition to obtain the first The solution corresponding to the highest individual fitness of the layer , as a solution ;exist Repeat the iterative process in the layer to obtain a multi-head attention load balancing distribution strategy .

[0061] Step 3: To achieve heterogeneous parallelism, The self-attention layer performs weight segmentation and storage. The principle is as follows Figure 2 As shown, Figure 2 In , “head” is used to represent the attention head;

[0062] The specific process of weight segmentation and storage is as follows:

[0063] The said The linear transformation matrix of the self-attention layer includes the query weight matrix , key weight matrix Sum weight matrix , ,in, is the hidden layer dimension of the large language model, is a single attention head dimension, , and Separately divided into Sub-matrices spliced ​​in the column direction, that is:

[0064] W l,Q =[ W l,Q (0) , W l,Q (1) ,…, W l,Q (I-1) ] ;

[0065] W l,K =[ W l,K (0) , W l,K (1) ,…, W l,K (I-1) ] ;

[0066] W l,V =[ W l,V (0) , W l,V (1) ,…, W l,V (I-1) ] ;

[0067] In the formula, For the The query weight sub-matrix of the attention heads; For the The key weight sub-matrix of the attention heads; For the The value weight sub-matrix of the attention head;

[0068] Based on the multi-head attention load balancing distribution strategy , determine the attention heads executed on the same GPU device, concatenate the corresponding query weight submatrix, key weight submatrix, and value weight submatrix together to form a reorganized matrix, and get A reorganized matrix to reduce fragmented access and frequent scheduling during inference; The reorganized matrix is ​​stored in blocks in the CPU memory to ensure that the weights of the corresponding attention heads can be quickly found during the loading phase, avoiding frequent indexing or slicing operations on the original large matrix;

[0069] Based on the divided reorganization matrix, a weight index table is established, such as Figure 3 As shown, , For example, each layer corresponds to 4 GPU devices, and each GPU device corresponds to a two-dimensional list weight [ l,p ] , which contains the query weight submatrix, key weight submatrix, and value weight submatrix of the attention head executed on the GPU device, and has the index function of loading weights to the corresponding GPU device when running at different layers.

[0070] Step 4: When processing Before the self-attention calculation of the layer, according to the scheme The query weight sub-matrix, key weight sub-matrix and value weight sub-matrix corresponding to each attention head are retrieved and loaded into the corresponding GPU device. Each GPU device only needs to read its own weight combination to reduce redundant loading time.

[0071] Step 5: In the long text scenario, the computational overhead exceeds the I / O overhead, and the zigzag pre-filling method has no practical advantage. Therefore, this embodiment proposes a parallel scheduling and weight pre-loading method, namely:

[0072] No. The calculation of the layer includes MHA calculation and multi-layer perception module calculation. By asynchronously preloading the MHA calculation weights and MLP calculation weights of adjacent layers and cooperating with KV cache management, the pipeline parallel calculation of reasoning operation is realized to significantly reduce the GPU memory pressure. The specific process is as follows:

[0073] like Figure 4 As shown in When calculating the Transformer layer, first perform the The MHA calculation of the layer is performed by using tensor parallel technology to asynchronously preload the first The MLP calculation weights of the layer are transferred to the corresponding GPU device, and the generated KV cache is unloaded to the CPU memory. Although the KV cache management spans the two stages of MHA calculation and MLP calculation, the full-duplex communication architecture ensures that the data transmission between storage and loading operations will not be blocked; in the After the MHA calculation of the layer is completed, the The MLP calculation of the layer is performed, and the first The MHA calculation weights of the layer are transferred to the corresponding GPU device; After the MLP calculation of the layer is completed, the first The MHA calculation of the layer is used to parallelize the reasoning and loading processes, so that the calculation process is not interrupted by loading, thereby improving the reasoning efficiency.

[0074] Example: Under torch.bfloat16 precision, the LLaMA3-8B model using the inference operation acceleration method for large language models processing extremely long texts described in this embodiment can reduce the weight storage requirement from 13 GB = MHA 2.5 GB + MLP 10.5 GB to 0.33 GB per operation module, a reduction of 40 times. Through this parallel scheduling and weight preloading method, I / O operations are completely overlapped with calculations, effectively converting the main cost factor from memory limitations to pure computing overhead, allowing larger models to be executed with limited hardware resources.

[0075] In order to comprehensively evaluate the performance of the inference operation acceleration method for large language models to process extremely long texts, this embodiment conducts a large number of experiments on a variety of advanced large language models and long text inference algorithms. This experiment mainly conducts application tests on two mainstream long text sparse algorithms: MInference (Minference 1.0: Accelerating Pre-filling for Long-Context LLMs via Dynamic Sparse Attention) and DuoAttention (DuoAttention: Efficient Long-Context LLM Inference with Retrieval and Streaming HeadsChunk), and comprehensively verifies the performance under highly sparse attention. In the MInference algorithm environment, the large language models selected are LLaMA3-8B, Yi-9B, and ChatGLM4-9B. In the DuoAttention algorithm environment, the large language models selected are LLaMA2-7B and LLaMA3-8B. Under different context lengths (768 K and 160 K), the inference efficiency of the inference operation acceleration method for processing extremely long texts for large language models proposed in this embodiment is compared with three existing benchmark methods, and the inference throughput (in tokens / second or its equivalent index) data results are shown in Table 1. The larger the inference throughput value, the higher the inference efficiency. Among them, the three existing benchmark methods are: (1) Random allocation: For each layer of the multi-head attention mechanism, the attention heads are randomly allocated to each GPU device; (2) Uniform partitioning (Megatron): Similar to Megatron-LM (Training Multi-Billion Parameter Language Models Using Model Parallelism), the attention heads are uniformly distributed to each GPU device in sequence; (3) Random uniform allocation: The number of heads on each GPU device is maintained equal through restricted randomization.

[0076] Table 1:

[0077]

[0078] As can be seen from Table 1, under a long context of 768 K, when using MInference, LLaMA3-8B can significantly increase the inference throughput from 1896.93 to 2447.07 under the inference operation acceleration method proposed in this embodiment, significantly alleviating the inference bottleneck caused by long text input, and is also applicable to other large language models. For a length of 160 K, the inference throughput of ChatGLM4-9B under the Random strategy is only 3885.62, while this embodiment can reach 4358.08, greatly improving the downstream inference efficiency. Similarly, in the DuoAttention environment, the inference operation acceleration method proposed in this embodiment can increase the throughput of LLaMA3-8B from 991.13 to 1292.08, highlighting the effectiveness of "multi-head parallel load balancing" in extremely long text scenarios.

[0079] In summary, the inference operation acceleration method for large language models processing extremely long texts proposed in this embodiment can effectively deal with the inference load imbalance problem caused by extremely long sequences in the applications of multiple advanced large language models and long text algorithms (including MInference, DuoAttention, etc.), greatly reducing end-to-end inference latency and improving throughput. Compared with the comparison methods such as uniform partitioning, random allocation, and random uniform allocation, this embodiment relies on the core ideas of offline data collection, attention head balancing strategy solution, weight segmentation and parallel scheduling, and has both performance and versatility. It can be widely used in various large language model reasoning tasks that need to process ultra-long texts, such as programming assistance, legal text processing, and academic literature understanding, and has significant practical and application value.

[0080] The above embodiments are intended to provide a better understanding of the present invention, but are not limited to the optimal implementation mode and do not constitute a limitation on the content and protection scope of the present invention. Any product identical or similar to the present invention obtained by anyone under the inspiration of the present invention or by combining the present invention with the features of other prior arts shall be within the protection scope of the present invention.

Claims

1. A method for accelerating the inference operation of a large language model for processing extremely long texts, characterized in that: The following steps are involved: Step 1: Select a large language model, including Transformer layers, each self-attention layer includes attention heads and enable sparse attention mode; Input texts with different context lengths into the large language model and record the first The execution time of each attention head of the layer and the type of attention mode it adopts are mapped to the corresponding attention mode type, and then a statistical database is established; Step 2: For the actual text , solve The multi-head attention load balancing allocation strategy for the large language model is as follows: right GPU devices, Next Each attention head of the layer is assigned to different GPU devices for execution as the first Multi-head attention load distribution scheme for layers ; Use the nearest neighbor matching method to find the context length closest to , and The execution time of each attention head in the layer is Next The corresponding attention head execution time of the layer is then calculated to obtain the solution Next Multi-head attention execution time of the layer ; by is the load balancing objective function. Through the search algorithm, we can obtain Optimal multi-head attention load distribution scheme for layers , and then obtain Multi-head attention load balancing allocation strategy for large language models ; Step 3: The query weight matrix, key weight matrix, and value weight matrix of the self-attention layer are split into The sub-matrices are spliced ​​in the column direction to obtain the query weight sub-matrix, key weight sub-matrix and value weight sub-matrix corresponding to each attention head; According to the strategy , determine the attention heads executed on the same GPU device, concatenate the corresponding query weight submatrix, key weight submatrix, and value weight submatrix together to form a reorganized matrix, and get The reorganized matrix is ​​stored in blocks in the CPU memory; and a weight index table is established based on the divided reorganized matrix; Step 4: According to the plan and weight index table, retrieve the query weight sub-matrix, key weight sub-matrix and value weight sub-matrix corresponding to each attention head, and load them into the corresponding GPU device; Step 5: The calculation of the layer includes MHA calculation and MLP calculation. The load balancing of the reasoning process is achieved by asynchronously preloading the MHA calculation weights and MLP calculation weights of adjacent layers and cooperating with KV cache management.

2. The method for accelerating the inference operation of a large language model for processing extremely long text according to claim 1, characterized in that: The sparse attention mode described in step 1 includes multiple of a local attention mode, a block attention mode, a jumping attention mode and a retrieval attention mode.

3. The method for accelerating the inference operation of a large language model for processing extremely long text according to claim 1, characterized in that: Calculated in step 2 The formula is: ； In the formula, Indicates the GPU device number; Indicates The set of attention head numbers assigned to the GPU devices; express Next Layer The execution time of an attention head.

4. The method for accelerating the inference operation of a large language model for processing extremely long text according to claim 1, characterized in that: The search algorithm in step 2 is a genetic algorithm, a dynamic programming algorithm or a neural network algorithm.

5. The method for accelerating the inference operation of a large language model for processing extremely long text according to claim 4, characterized in that: When the search algorithm is a genetic algorithm, solve the multi-head attention load balancing allocation strategy The specific process is: By plan For individuals, The layer initializes the population according to the load balancing objective function , calculate individual fitness; iteratively generate candidate solutions through crossover and mutation operations; And update the population through survival competition, and obtain the first The solution corresponding to the highest individual fitness of the layer , as a solution ;exist Repeat the iterative process in the layer to obtain a multi-head attention load balancing distribution strategy .

6. The method for accelerating the inference operation of a large language model for processing extremely long text according to claim 1, characterized in that: The specific process of step 5 is: In the When calculating the Transformer layer, first perform the The MHA calculation of the layer is performed by using tensor parallel technology to asynchronously preload the first The MLP calculation weights of the layer are transferred to the corresponding GPU device, and the generated KV cache is unloaded to the CPU memory; After the MHA calculation of the layer is completed, the The MLP calculation of the layer is performed, and the first The MHA calculation weights of the layer are transferred to the corresponding GPU device; After the MLP calculation of the layer is completed, the first The MHA calculation of the layer is used to realize the pipeline parallel computing of the reasoning operation.

Images