Large language model memory scheduling management method, system and storage medium
Through unified paging scheduling, the memory of key-value cache and activation value storage area is managed, and the memory waste of activation value storage area in the large language model is solved, which improves memory utilization and service system processing capabilities.
Patent Information
- Application Number
- CN202410484218.2
- Authority / Receiving Office
- CN · China
- Patent Type
- Patents(China)
- Current Assignee / Owner
- Filing Date
- 2024-04-22
- Publication Date
- 2024-12-10
- Estimated Expiration
- 2044-04-22
AI Technical Summary
The memory management solution of the existing large language model ignores the activation value storage area, resulting in memory waste during the decoding stage, especially in cloud service systems, which are prone to blockage and cannot effectively handle a large number of concurrent requests.
Adopt a unified paging scheduling management method to share memory for key-value cache and activation value storage area. The memory page is allocated in the pre-filling stage to activate the value storage area. When the demand is reduced in the decoding stage, it is reassigned for key-value cache, and improves memory utilization. Rate.
Improves memory utilization, improves the processing capability of the service system under high load conditions, and enhances the load limit of non-blocking processing, especially when handling a large number of concurrent requests.
Smart Images
Figure CN118331734B_ABST
Abstract
Description
Technical Field
[0001] The present invention relates to the field of artificial intelligence technology, and in particular to a large language model memory scheduling management method, system and storage medium. Background Art
[0002] At present, the "emergence" phenomenon of existing large language models (LLMs) after they have grown to a certain size has demonstrated amazing application potential in various fields, such as text generation, machine translation, programming assistance, etc.
[0003] Most existing large language models run on memory, such as graphics processing unit (GPU) memory, to handle their huge number of parameters and complex computing tasks. However, large language models require much more memory resources than traditional convolutional neural networks (CNNs). Commonly used CNN parameters are generally in the millions, but commonly used large language model parameters are generally in the billions, such as multiple versions of Llama2 (7 billion, 13 billion, 70 billion). Even GPT4 may have hundreds or even thousands of billion parameters. The ultra-large-scale parameters put great pressure on memory.
[0004] The current problem is that the existing memory management scheme is limited to the design of key-value cache, ignoring the activation value storage area. Since the memory requirement of the activation value storage area in the decoding stage of large language model inference is much smaller than that in the pre-filling stage, the memory used for the activation value storage area in the existing scheme is actually wasted in the decoding stage. In addition, in the large language model service system in the cloud, the memory is under greater pressure, and memory management is particularly important. Existing memory management schemes, such as the simplest LLM service dynamic batch processing scheme, set the maximum batch size (max batch size) according to the maximum memory and the maximum number of tokens (tokens) per request preset by the service system. The batch size supported by the system during the service process cannot exceed this value, even if the number of tokens actually generated for each request does not reach or is far less than the preset maximum number of tokens. This causes the key value cache actually used by the system to be extremely small, and it is easy to get blocked when the request load is large.
[0005] In addition, some other background technologies (or concepts) involved in the present invention are further described as follows:
[0006] The general processing flow for a large language model usually involves the following steps:
[0007] 1. Input processing: receiving user input, which may be a piece of text, a query or a specific instruction.
[0008] 2. Tokenization: Convert the input text into a format that the model can understand, usually involving breaking the text into a series of tokens.
[0009] 3. Embedding: Convert the segmented word units into numerical embedding vectors that the model can handle.
[0010] 4. Encoding: The input embedding vector is processed by the encoder part of the model to extract contextual features of the text.
[0011] 5. Attention Mechanism: During the encoding process, the model uses a self-attention mechanism to strengthen the connection between different parts of the text.
[0012] 6. Prefill: Some models perform a prefill phase before generating a response to prepare the contextual information required for generation.
[0013] 7. Decoding: The model uses the decoder part to gradually generate responses based on the output of the encoder and the information from the pre-filling stage.
[0014] 8. Post-processing: Perform necessary formatting and adjustments on the generated response to ensure that it conforms to the expected output format.
[0015] 9. Output: Output the response generated by the model to the user.
[0016] 10. Memory management: Throughout the entire processing flow, memory is used to store the model's weights (weight storage areas), key-value caches, and activations. Memory management is critical to ensuring that the model can run efficiently.
[0017] In cloud service systems, in order to improve the output quality of the model and handle more user requests, models with larger parameter sizes are usually used and dynamic batch processing is performed to maximize memory utilization. However, memory resources are limited, so an effective memory management strategy is needed to optimize the operation of the model. Summary of the invention
[0018] The purpose of the present invention is to provide a large language model memory scheduling management method, a large language model system and a computer storage medium, which improve the utilization of memory and increase the upper limit of the load that a service system can process without blocking, especially when processing a large number of concurrent requests, thereby enhancing the processing capability of the system under high load conditions.
[0019] In order to achieve the above technical objectives, the present invention adopts the following technical solutions:
[0020] A large language model memory scheduling management method, wherein the memory includes a weight storage area, a key value cache, and an activation value storage area;
[0021] The memory scheduling management method comprises:
[0022] Allocate memory in memory for the weight storage area;
[0023] Allocate the remaining memory after removing the weight storage area as key value cache and activation value storage area;
[0024] Implement unified paging scheduling management for key-value cache and activation value storage area.
[0025] Furthermore, the unified paging scheduling management is implemented for the key value cache and the activation value storage area, which specifically includes:
[0026] Allocate a specific number of memory pages in the shared memory of the key-value cache and the activation value storage area;
[0027] In the pre-filling phase, memory pages are allocated as activation value storage areas according to demand;
[0028] During the decoding phase, when the memory requirement of the activation value storage area decreases, some memory pages originally used as the activation value storage area are reallocated as key-value cache.
[0029] Furthermore, for the shared memory of the key-value cache and the activation value storage area, the memory page is divided according to 128 words as one memory page unit.
[0030] Furthermore, in the process of implementing unified paging scheduling management on the key value cache and the activation value storage area, memory pages are dynamically allocated and scheduled according to the number of word units.
[0031] A large language model system adopts the above-mentioned memory scheduling management method to schedule and manage memory.
[0032] A computer storage medium stores program instructions; when the program instructions are executed by a processor, the processor executes the memory scheduling management method as described above.
[0033] Compared with the prior art, the memory scheduling management method, large language model system and computer storage medium of the present invention have the following beneficial effects: the remaining part of the memory after removing the weight storage area part is uniformly managed by paging, and in the decoding stage, when the memory demand of the activation value storage area is reduced, part of the memory pages originally used as the activation value storage area are reallocated as key value cache, thereby avoiding the waste of memory when the memory demand of the activation value storage area is reduced in the decoding stage, and by uniformly managing the key value cache and the activation value storage area memory, the memory utilization rate is improved, and the load upper limit that the service system can process without blocking is improved, especially when processing a large number of concurrent requests, the processing capability of the system under high load conditions is enhanced. BRIEF DESCRIPTION OF THE DRAWINGS
[0034] Figure 1 A schematic diagram of key-value cache paging management during the pre-filling and decoding phases for an existing large language model;
[0035] Figure 2 The diagram is a schematic diagram of key-value cache paging management in the pre-filling stage and the decoding stage of a large language model after the memory scheduling management method of the present invention is adopted. DETAILED DESCRIPTION
[0036] The present invention will be further described below with specific embodiments:
[0037] This embodiment provides a memory scheduling management method for a large language model, which improves memory utilization and increases the load limit that the service system can process without blocking, especially when processing a large number of concurrent requests, thereby enhancing the system's processing capabilities under high load conditions.
[0038] It should be noted that the memory described in this embodiment is GPU memory, which is the main memory form currently used by large language models.
[0039] In the processing flow of the large language model, there are two key stages: the pre-filling stage and the decoding stage.
[0040] The pre-population phase mainly refers to a series of initialization and preparation work performed by the model before receiving input. The main purpose of this phase is to provide the model with a good starting state so that it can better process subsequent input data. Specifically, it may include operations such as loading pre-trained parameters, setting the internal state of the model, and initializing the encoder. Through these preparations, the model can quickly enter the state when receiving the input, providing a solid foundation for the subsequent decoding process.
[0041] The decoding phase is the core process of the model generating output based on input. In this phase, the model will gradually generate an output sequence until a certain termination condition is met. Specifically, the model will predict the next most likely output word or symbol based on the input data and the current internal state, and add it to the output sequence. Then, the model will continue to predict the next output based on the new output sequence and internal state until the entire output sequence is generated.
[0042] These two stages are interrelated and interdependent in the processing of large language models. The pre-filling stage provides the necessary preparation and initialization for the decoding stage, ensuring that the model can maintain the best state when processing input. The decoding stage is the key link for the model to achieve natural language processing tasks. It generates high-quality output based on the input and completes the conversion process from input to output.
[0043] Through the collaborative work of these two stages, the large language model can achieve efficient processing and understanding of natural language, providing strong support for various natural language processing tasks. Whether it is text generation, question-answering system or language translation, the large language model can generate accurate and fluent natural language output through the fine processing of these two stages.
[0044] The memory used by large language models during reasoning and serving is mainly divided into three key parts: weight storage area, key-value cache, and activation value storage area.
[0045] First, the weight storage area is used to store the set of parameters learned during model training. It represents the features and knowledge extracted by the model from a large amount of data. In large language models, the weight storage area usually occupies a considerable portion of the memory because it contains all the information required for the model to understand and generate natural language. These weights are used to calculate the output of the model during inference and are the basis for realizing the model's functions.
[0046] Secondly, the key-value cache is mainly used to store the intermediate results generated during the model inference process. In large language models, especially those using the attention mechanism, key-value cache plays a key role. It can help the model quickly retrieve relevant information during inference and improve inference efficiency. By properly managing the key-value cache, repeated calculations can be effectively reduced, thereby saving memory and increasing inference speed.
[0047] Finally, the activation value storage area is used to store the activation values generated by the model during the inference process, which represents the processing results of each layer of the model on the input data. In a large language model, as the input sequence is gradually processed, a large number of activation values will be generated, which need to be stored in memory for subsequent calculations. The size and number of activation values depend on the complexity of the model and the size of the input data, and are an important part of memory consumption.
[0048] In general, these three parts of memory play their respective roles in the reasoning and serving process of the large language model, and jointly support the operation of the model. Their collaboration enables the model to efficiently process natural language input and generate accurate output. By optimizing the memory usage of these parts, the performance and efficiency of the large language model can be further improved.
[0049] It should be noted that during the inference and service process of the large language model, the size of the weight storage area is fixed and is continuously used throughout the inference service process, so there is no problem of memory waste. Even if technologies such as quantizing the weight storage area or offloading are used, the memory size of the weight storage area remains unchanged.
[0050] The core concept of the memory scheduling management method of this embodiment is:
[0051] Taking advantage of the fact that the activation value storage area required in the pre-filling stage is much higher than that in the decoding stage during the inference process of large language models, the actual memory usage of the service system is improved by "allowing the decoding stage to use the memory saved in the decoding stage as a key-value cache". Based on the most efficient "Paged KVcache cache (paged key-value cache)" solution of the current large language model service system, a "unified memory paging management solution" is proposed on the basis of minimizing the impact on the use of existing systems.
[0052] The important innovation of the memory scheduling management method of this embodiment is that after subtracting the portion required for the weight storage area from the total memory (for example, the GPU's memory), all the remaining memory is included in the paging management. This is fundamentally different from the existing method of only performing paging management on the key-value cache. This means that both the key-value cache and the activation value storage area will adopt a unified paging management strategy, so that the good effect of "increasing the scheduling freedom of memory pages" can be achieved.
[0053] It is worth noting that the memory scheduling management method of this embodiment is completely independent of the method specifically used to reduce the memory consumption of part of the activation value storage area (such as the three-cache strategy), and can be used in parallel to further improve memory usage efficiency.
[0054] In order to deeply understand the memory scheduling management method of this embodiment, it is first necessary to understand the paging management logic of the key-value cache in the existing large language model service method.
[0055] Specifically,
[0056] Assuming that the size of the key-value cache is fixed, if the maximum number of tokens generated by each request (for example, 4096 tokens for the llama2 model) is used as the minimum granularity of the key-value cache, the maximum number of batches that the key-value cache can accommodate (i.e., max batch size) can be calculated. Even if the actual number of tokens in some requests is much smaller than 4096, such as only 128 tokens, they will still occupy a key-value cache equivalent to 4096 tokens. The paging management of the key-value cache further refines this granularity. For example, the memory corresponding to 128 tokens is used as the minimum unit of paging, so that the original memory requirement of 4096 tokens can be divided into smaller pages, which can theoretically increase the maximum batch size by 32 times. In this way, the key-value cache occupied by smaller requests will be reduced, thereby increasing the number of batches and accelerating those inference processes that are severely limited by memory access.
[0057] The existing key-value cache paging management scheme mainly focuses on the division and scheduling method inside the key-value cache, but it cannot change the overall size of the key-value cache. In addition, the existing scheme ignores the activation value storage area, which is an important part of the memory besides the weight storage area.
[0058] In this implementation, both the key value cache and the activation value storage area are considered, because they occupy almost all the memory except the weight storage area. In this implementation, these two parts of memory are uniformly managed by paging, and no clear distinction is made as to whether a memory page belongs to the key value cache or the activation value storage area, thereby improving the scheduling flexibility of the memory page.
[0059] Assuming that the size of each memory page is the same, in the memory scheduling management method of this embodiment, more memory pages can be divided, and these memory pages can be used as activation value storage areas in the pre-filling stage of model reasoning, and can also be used as key value caches in the decoding stage when the memory demand for the activation value storage area is reduced. Since the number of memory pages that the scheduler can schedule has increased, this is equivalent to increasing the size of the key value cache, allowing the service system to accommodate more requests, thereby providing better performance.
[0060] The following example illustrates how the memory scheduling management method of this embodiment can improve the performance and effectiveness of a large language model service system.
[0061] See also Figure 1 and Figure 2, assuming that the service system's memory can accommodate a key-value cache of 5120 tokens. In the pre-filling phase of model inference, the activation value storage area requires a space equivalent to a key-value cache of 1024 tokens, which makes the actual number of tokens that the key-value cache can accommodate 4096. With 128 tokens as a memory page unit, this part of the key-value cache can be divided into 32 memory pages.
[0062] See also Figure 1 , imagine that the service system receives and processes 32 requests, and each request allocates a memory page for its key-value cache. When all requests complete the pre-filling phase and enter the decoding phase, if the number of tokens processed by a request reaches 128, the service system may be blocked because there are no additional memory pages to allocate. In this case, the existing processing method is to remove the key-value cache data of a request from the memory, thereby freeing up the memory page to allocate to the request that caused the blocking. This practice is very time-consuming for the service system because it may require data transfer between the memory and the CPU, or the request needs to be reprocessed. This operation is unfair to the user because it may affect other requests that do not cause blocking. If this blocking occurs frequently, it will be based on the worst-case scenario of the key-value cache paging management method, resulting in a serious degradation of service performance.
[0063] See also Figure 2 , if the memory scheduling management method of this embodiment is adopted, under the same conditions, the service system can divide 40 memory pages. For 32 requests, each request is allocated a memory page for the key value cache, which is the same as the previous situation. However, after the pre-filling stage, when the service system enters the decoding stage, since the memory demand of the activation value storage area is greatly reduced, only 1 memory page is required. At this time, 7 of the remaining 8 memory pages are free and can be allocated to requests that need to process 128 word elements. In this way, the service system can continue to execute normally without blocking, which is equivalent to the available key value cache increasing to 39 memory pages, which increases the number of available memory pages compared to the previous method.
[0064] The benefit of this implementation is that it increases the upper limit of the load that the service system can handle without blocking, thereby enhancing the overall performance of the service system.
[0065] The memory scheduling management method of this embodiment can be summarized as follows:
[0066] The service system of the large language model receives input requests from users, which may be a piece of text, a query, or a specific instruction. Then, the service system performs a series of preprocessing on the user's input request, such as word segmentation, embedding, etc., and then generates a response corresponding to the input request and outputs it to the user. The large language model runs based on memory, which includes a weight storage area, a key value cache, and an activation value storage area. The process of the large language model processing "user input request" includes a pre-filling stage and a decoding stage.
[0067] The memory scheduling management method of this embodiment includes:
[0068] During the operation of the large language model, memory is allocated in the memory for use as a weight storage area. It should be noted that the memory space allocated to the weight storage area is usually fixed, and generally speaking, the allocated memory space is tens of GB or even hundreds of GB.
[0069] All remaining memory after removing the part required by the weight storage area is included in the unified paging management. It should be noted that the remaining memory after removing the part required by the weight storage area is the memory space used by the key value cache and the activation value storage area. For the convenience of description, this memory included in the unified paging management is called the key value cache and activation value storage area shared memory.
[0070] That is to say, a unified paging scheduling management strategy is used for the shared memory of the key value cache and the activation value storage area. For the convenience of description, this paging scheduling management strategy is called "unified paging scheduling management strategy for key value cache and activation value storage area".
[0071] Specifically, a specific number of memory pages are divided in the shared memory of the key-value cache and the activation value storage area; in the pre-filling stage, memory pages are allocated as the activation value storage area according to demand; in the decoding stage, when the memory demand of the activation value storage area decreases, some of the memory pages originally used as the activation value storage area are reallocated as the key-value cache.
[0072] The memory scheduling management method of this embodiment can be used for most of the existing large language models to schedule and manage memory, for example, GPU memory.
[0073] The memory scheduling management method of this embodiment has the following advantages:
[0074] In the memory scheduling management method of the present embodiment, the remaining part of the memory after removing the weight storage area is uniformly managed by paging. In the decoding stage, when the memory demand of the activation value storage area decreases, some memory pages originally used as the activation value storage area are reallocated for use as key value cache, thereby avoiding memory waste when the memory demand for the activation value storage area decreases in the decoding stage. By uniformly managing the key value cache and the activation value storage area memory, the memory utilization rate is improved, and the load upper limit that the service system can process without blocking is improved, especially when processing a large number of concurrent requests, the system's processing capability under high load conditions is enhanced.
[0075] The above are only preferred embodiments of the present invention and are not intended to limit the protection scope of the present invention. Therefore, any modifications, equivalent substitutions, improvements, etc. made within the spirit and principles of the present invention should be included in the protection scope of the present invention.
Claims
1. A memory scheduling management method for a large language model, wherein the memory includes a weight storage area, a key value cache, and an activation value storage area; Features: The memory scheduling management method comprises: Allocate memory in memory for the weight storage area; Allocate the remaining memory after removing the weight storage area as key value cache and activation value storage area; Implement unified paging scheduling management for key-value cache and activation value storage areas; The unified paging scheduling management for the key value cache and the activation value storage area specifically includes: Allocate a specific number of memory pages in the shared memory of the key-value cache and the activation value storage area; In the pre-filling phase, memory pages are allocated as activation value storage areas according to demand; During the decoding phase, when the memory requirement of the activation value storage area decreases, some memory pages originally used as the activation value storage area are reallocated as key-value cache.
2. According to the large language model memory scheduling management method of claim 1, Features: For the shared memory of the key-value cache and the activation value storage area, the memory pages are divided according to 128 words as one memory page unit.
3. According to the large language model memory scheduling management method of claim 1, Features: In the process of implementing unified paging scheduling management on the key value cache and the activation value storage area, memory pages are dynamically allocated and scheduled according to the number of word units.
4. A large language model system, Features: The large language model system uses the memory scheduling management method as described in any one of claims 1 to 3 to schedule and manage the memory.
5. A computer storage medium having program instructions stored thereon; Features: When the program instructions are executed by a processor, the processor executes the memory scheduling management method according to any one of claims 1 to 3.