Memory management method and system for large language model training

Abstract

The invention belongs to the technical field of computers, and aims to provide a memory management method and system for large language model training. The method comprises the following steps: constructing a parameter symbol list matched with a specified large language model; according to the parameter symbol list, obtaining a total parameter number when the specified large language model is trained by adopting a preset training mode, and according to the total parameter number and training memory parameters when the specified large language model is trained, calculating to obtain a memory overhead predicted value when the specified large language model is trained by adopting the preset training mode; and obtaining a remaining memory value of the training device, and judging whether the remaining memory value is greater than the memory overhead predicted value, if so, performing memory pre-allocation processing on the training device according to the memory overhead predicted value, so that the training device trains the specified large language model. According to the method, memory use estimation and optimization management in the large language model training process can be realized.

Description

Technical Field

[0001] The present invention belongs to the field of computer technology, and in particular relates to a memory management method and system for large language model training. Background Art

[0002] In recent years, large language models have been widely used in academia and industry. Large language models (LLMs) refer to deep learning models trained with large amounts of text data that can generate natural language text or understand the meaning of language text. Large language models can handle a variety of natural language tasks, such as text classification, question-answering, and dialogue, and are an important path to artificial intelligence.

[0003] In the training process of a large language model, accurately predicting memory usage helps developers understand the internal structure and behavior of the model, while optimizing the system operation status, thereby saving computing resources. In addition, developers usually need to plan hardware resources based on memory prediction results, and reserve too much memory to ensure that the training process will not be interrupted due to insufficient memory. However, in the process of using the prior art, the inventor found that the prior art has at least the following problems: With the development of various structures and training technologies, the complexity of large language model structures has increased. Coupled with the implementation differences of different frameworks, it has become very difficult to predict the memory usage required for these large language models during training or fine-tuning. This not only increases the uncertainty of model development, but may also lead to a waste of resources. Summary of the invention

[0004] The present invention aims to solve the above technical problems at least to a certain extent. The present invention provides a memory management method and system for large language model training.

[0005] In order to achieve the above object, the present invention adopts the following technical solutions: In a first aspect, the present invention provides a memory management method for large language model training, comprising: Perform abstract processing on the specified large language model to obtain an abstracted model; Constructing a parameter symbol list matching the specified large language model according to the abstracted model; Obtaining, according to the parameter symbol list, a total parameter amount of the designated large language model when training in a preset training method, and calculating, according to the total parameter amount and a training memory parameter of the designated large language model when training in the preset training method, a predicted value of memory overhead of the designated large language model when training in the preset training method; Obtain a remaining memory value of the training device, and determine whether the remaining memory value is greater than the memory overhead prediction value. If so, perform memory pre-allocation processing for the training device according to the memory overhead prediction value so that the training device can train the specified large language model.

[0006] In a possible design, the abstracted model includes an input layer, an embedding layer, an intermediate layer, and an output layer connected in sequence; wherein the input dimension of the input layer and the output dimension of the output layer are both b × s , the parameter quantity of the embedding layer is d × v , the intermediate layer includes an attention layer and a feedforward layer connected in sequence, and the number of attention heads in the attention layer is h , the parameters of different linear layers in the attention layer are d 2 , the number of parameters of different linear layers in the feedforward layer is 4 d 2 , b represents the batch size, s represents the sequence length, d represents the hidden layer dimension, v represents the length of the word vector, h is a natural number greater than 1, b , s , d and v are natural numbers greater than 0 respectively; the number of layers of the intermediate layer is L layer, L is a natural number greater than 1.

[0007] In a possible design, the parameter symbol list includes the number of intermediate layers of the abstracted model. L , word vector length v , hidden layer dimension d , batch size b and sequence length s ; The training memory parameters include model weights, gradients, optimizer state information and activation values.

[0008] In a possible design, when the preset training mode is a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, the total parameter amount is: ; The predicted memory cost is: ; In the formula, W , G ,O and A It is calculated based on the total parameter amount and the training memory parameter of the specified large language model during training; wherein, W is the predicted value of the model weight in the training memory parameter, G is the predicted value of the gradient in the training memory parameter, O is the predicted value of the optimizer state information in the training memory parameter, A is the predicted value of the activation value in the training memory parameter.

[0009] In a possible design, when the preset training mode is a single-precision training mode, the ratio of the model weight, the gradient, and the optimizer state information is 1:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0010] In a possible design, when the preset training mode is a half-precision training mode, the ratio of the model weight, the gradient and the optimizer state information is 1:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0011] In a possible design, when the preset training mode is a mixed precision training mode, the ratio of the model weight, the gradient, and the optimizer state information is 1.5:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0012] In one possible design, when the preset training mode is a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, and LoRA is synchronously enabled for the specified large language model, the total parameter amount is: ; In the formula, r The dimension of the trainable linear layer in the specified large language model after LoRA is synchronously enabled for the specified large language model; The predicted memory cost is: ; In the formula, W , G , O and A It is calculated based on the total parameter amount and the training memory parameter of the specified large language model during training; wherein, W is the predicted value of the model weight in the training memory parameter, G is the predicted value of the gradient in the training memory parameter, O is the predicted value of the optimizer state information in the training memory parameter, A is the predicted value of the activation value in the training memory parameter.

[0013] In one possible design, when the preset training mode is a single-precision training mode and LoRA is synchronously enabled for the specified large language model, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0014] In one possible design, when the preset training mode is a half-precision training mode and LoRA is synchronously enabled for the specified large language model, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0015] In one possible design, when the preset training mode is a mixed precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0016] In a second aspect, the present invention provides a memory management system for large language model training, which is used to implement a memory management method for large language model training as described in any one of the above; the memory management system for large language model training includes: The model abstraction module is used to abstract the specified large language model to obtain the abstracted model; A parameter symbol list construction module, which is in communication with the model abstraction module and is used to construct a parameter symbol list matching the specified large language model according to the abstracted model; a memory overhead prediction module, which is in communication connection with the parameter symbol list construction module, and is used to obtain, according to the parameter symbol list, a total parameter amount of the specified large language model when training in a preset training method, and calculate, according to the total parameter amount and a training memory parameter of the specified large language model when training, a memory overhead prediction value of the specified large language model when training in the preset training method; A memory management module is communicatively connected to the memory overhead prediction module, and is used to obtain the remaining memory value of the training device and determine whether the remaining memory value is greater than the memory overhead prediction value. If so, memory pre-allocation processing is performed for the training device according to the memory overhead prediction value so that the training device can train the specified large language model.

[0017] In a third aspect, the present invention provides an electronic device, comprising: a memory for storing computer program instructions; and, A processor is used to execute the computer program instructions to complete the operation of a memory management method for large language model training as described in any one of the above.

[0018] In a fourth aspect, the present invention provides a computer program product, comprising a computer program or instructions, which, when executed by a computer, implements a memory management method for large language model training as described in any one of the above.

[0019] The beneficial effects of the present invention are: The present invention can realize the estimation and optimization management of memory usage during the training of a large language model. Specifically, during the implementation of the present invention, first, the designated large language model is abstracted to obtain the model after abstraction, and a parameter symbol list matching the designated large language model is constructed based on the model after abstraction; then, the total parameter amount of the designated large language model when trained in a preset training method is obtained based on the parameter symbol list, and the memory overhead prediction value of the designated large language model when trained in the preset training method is calculated based on the total parameter amount and the training memory parameters of the designated large language model when training; finally, the remaining memory value of the training device is obtained, and it is determined whether the remaining memory value is greater than the memory overhead prediction value. If so, memory pre-allocation processing is performed for the training device based on the memory overhead prediction value so that the training device can train the designated large language model. In this process, the present invention abstracts the specified large language model and constructs a parameter symbol list matching the specified large language model. Then, based on the parameter symbol list and a preset training method, the memory overhead of the specified large language model is predicted, and based on the obtained memory overhead prediction value, the memory of the training device is managed. This can achieve accurate memory estimation under different training methods, solve the problem of difficult-to-predict memory status in large language model training, and is conducive to ensuring the smooth development of the large language model, while also helping to reduce unnecessary waste of resources.

[0020] Other beneficial effects of the present invention will be further described in the specific embodiments. BRIEF DESCRIPTION OF THE DRAWINGS

[0021] Figure 1 is a flowchart of a memory management method for large language model training in an embodiment; Figure 2 It is a schematic diagram of the structure of the model after abstract processing; Figure 3 It is a simplified diagram of the model linear layer and the trainable linear layers LoRA A and LoRA B; Figure 4 It is a block diagram of a memory management system for large language model training in an embodiment; Figure 5 is a module block diagram of the electronic device in the embodiment. DETAILED DESCRIPTION

[0022] In order to more clearly illustrate the technical solutions in the embodiments of the present invention or the prior art, the present invention will be briefly introduced below in combination with the drawings and the description of the embodiments or the prior art. Obviously, the following description of the structure of the drawings is only some embodiments of the present invention. For ordinary technicians in this field, other drawings can be obtained based on these drawings without creative work. It should be noted that the description of these embodiments is used to help understand the present invention, but does not constitute a limitation of the present invention.

[0023] Embodiment 1: This embodiment discloses a memory management method for large language model training, which can be executed by, but is not limited to, a computer device or a virtual machine with certain computing resources, such as a personal computer, a smart phone, a personal digital assistant, or a wearable device, or by a virtual machine.

[0024] like Figure 1 As shown, a memory management method for large language model training may include but is not limited to the following steps: S1. Perform abstract processing on the designated large language model to obtain the model after abstract processing; it should be understood that the large language model is a type of machine learning model. During the training process, the input data is calculated inside the large language model to achieve parameter update of the large language model. In this embodiment, the designated large language model is also the large language model to be trained. The designated large language model is abstracted, that is, the basic modules in the designated large language model are abstractly represented, thereby simplifying the model structure, reducing the difficulty of subsequent data processing, and improving the efficiency of subsequent memory overhead prediction.

[0025] Specifically, in step S1 of this embodiment, Figure 2As shown, the abstracted model includes an input layer, an embedding layer, an intermediate layer, and an output layer connected in sequence; wherein the input dimension of the input layer and the output dimension of the output layer are both b × s , the parameter quantity of the embedding layer is d × v , the intermediate layer includes an attention layer and a feedforward layer connected in sequence, and the number of attention heads in the attention layer is h , the parameters of different linear layers in the attention layer are d 2 , the number of parameters of different linear layers in the feedforward layer is 4 d 2 , b Indicates the batch size, that is, the number of samples processed at a time. s represents the sequence length, which is the number of time steps or elements in each sample. d represents the hidden layer dimension, that is, the dimension of each word embedding, v Represents the length of the word vector (Vocabulary Size), that is, the number of different words or symbols in the vocabulary, h is a natural number greater than 1, b , s , d and v are natural numbers greater than 0 respectively; the number of layers of the intermediate layer is L layer, L is a natural number greater than 1. It should be noted that the embedding layer is responsible for converting the input discrete symbols (such as words or characters) into continuous vector representations, namely, embedding vectors. The output of the embedding layer will be used as the input of the subsequent intermediate layer; the intermediate layer consists of L layers, each of which is a composite structure, including two main parts: the attention layer and the feedforward layer. The attention layer is used to capture the relationship between different positions in the sequence, and the feedforward layer is a fully connected neural network, usually including two layers of linear transformations with an activation function in between. To simplify the description, Figure 2 In the figure, the linear layers in the attention layer are represented by Q, K, V and E respectively. The linear layers Q, K and V represent three matrices of query, key and value respectively, which are used to calculate the attention weights in the self-attention mechanism. The linear layer E is the output linear layer of the attention layer. The linear layers in the feedforward layer are represented by UP and DOWN respectively. The linear layers UP and DOWN represent the up-sampling (UpProjection) and down-sampling (DownProjection) linear layers in the feedforward layer respectively, which are used to expand and reduce the feature dimensions respectively.

[0026] S2. constructing a parameter symbol list that matches the specified large language model based on the abstracted model; specifically, constructing the parameter symbol list can help to characterize in detail the state of the specified large language model during runtime.

[0027] In step S2 of this embodiment, the parameter symbol list includes the number of intermediate layers of the model after the abstraction processing. L , word vector length v , hidden layer dimension d , batch size b and sequence length s .

[0028] S3. According to the parameter symbol list, the total parameter amount of the specified large language model when trained by adopting the preset training method is obtained, and according to the total parameter amount and the training memory parameter of the specified large language model when trained by adopting the preset training method, the memory overhead prediction value of the specified large language model when trained by adopting the preset training method is calculated; it should be noted that the training memory parameter can also be referred to as the memory-related parameter or training state parameter in the training process, which describes the complete state of the specified large language model in the training process and constitutes the main memory resource required by the specified large language model during training.

[0029] In step S3 of this embodiment, the training memory parameters include model weights, gradients, optimizer state information, and activation values.

[0030] In step S3 of this embodiment, when the preset training mode is a single precision (FP32) training mode, a half precision (FP16) training mode, or a mixed precision (FP32+FP16) training mode, the total parameter amount is: ; Specifically, in this embodiment, the total parameter amount includes the embedding layer and L All parameters of the intermediate layers.

[0031] In step S3, the memory overhead prediction value is: ; In the formula, W , G , O and A It is calculated based on the total parameter amount and the training memory parameter of the specified large language model during training; wherein, W is the predicted value of the model weight in the training memory parameter, G is the predicted value of the gradient in the training memory parameter, Ois the predicted value of the optimizer state information in the training memory parameter, A is the predicted value of the activation value in the training memory parameter.

[0032] It should be noted that the model weight W It includes the set of all parameters in the specified large language model. These parameters determine the learning ability and performance of the model, are the core part of the model, and directly affect the prediction ability of the model; gradient G Includes the partial derivative of the loss function with respect to each trainable parameter (that is, the gradient of the trainable parameter), which is used to indicate how to adjust the parameters to minimize the loss during backpropagation. It is the information necessary to update the parameters during the optimization process; optimizer state information O Contains optimizer state information of trainable parameters accumulated by optimization algorithms such as Adam (Adaptive Moment Estimation) and SGD (Stochastic Gradient Descent) during training, such as momentum and historical gradients, which can help optimization algorithms better adjust parameters, accelerate convergence or improve stability; weights W ,gradient G and optimizer status information O Together they constitute the core state of the model during training, determining the model's learning process and final performance; activation value A It is the output calculated by each layer and passed to the next layer during the forward propagation process, including the activation value required for training, so as to calculate the gradient in the back propagation and ensure the accuracy of the gradient calculation.

[0033] In this embodiment, by analyzing in detail the memory usage of the four parameters in the training memory parameters, the memory requirements of the specified large language model under different training methods can be more accurately estimated, thereby optimizing resource allocation and improving training efficiency.

[0034] In this embodiment, a baseline is provided to help understand the memory resources required for model training if no efficient parameter fine-tuning method is used.

[0035] Specifically, in this embodiment, the preset training mode can be a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, that is, this embodiment supports memory overhead prediction under the above three training modes. In addition, in training modes with different types of precision, data is stored in a corresponding format. For example, in a single-precision training mode, a single data occupies 4 bytes, and in a half-precision training mode, a single data occupies 2 bytes. Based on this, memory overhead prediction can be achieved.

[0036] Therefore, this embodiment can realize memory prediction of a large language model under three training methods, so that the versatility of this embodiment is improved.

[0037] Specifically, in step S3, when the preset training mode is the single-precision training mode, the ratio of the model weight, the gradient and the optimizer state information is 1:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0038] It should be noted that when the preset training mode is single-precision training mode, each data occupies 4 bytes. , so the predicted value of the model weight in the training memory parameter ; Then, according to the ratio of the model weight, the gradient and the optimizer state information in the training memory parameter of 1:1:2, the predicted value of the gradient in the training memory parameter can be obtained in turn , the predicted value of the optimizer state information in the training memory parameter .

[0039] The activation values ​​mainly exist before and after some specific modules of the large language model, and are used to calculate the gradients of the parameters of the large language model during the training process. When pre-calculating the activation values, if the training method of the large language model or its internal operating mechanism is unclear, it usually leads to calculation errors. To this end, the present embodiment distinguishes different training methods in advance, and predicts the activation values ​​according to the different behaviors within the model under different training methods. Specifically, the present embodiment first analyzes the specific positions where the activation values ​​are generated in the attention layer and the feedforward layer, and then calculates the activation value overhead of the single intermediate layer of the specified large language model during training according to the parameters in the parameter symbol list. When the preset training method is a single-precision training method, the predicted values ​​of the activation values ​​in the attention layer and the feedforward layer are shown in Tables 1 and 2 below, respectively. Tables 1 and 2 both show the distribution positions of the activation values ​​and the predicted results of the activation values ​​at different positions: Table 1 The distribution of activation values ​​in the attention layer and the prediction results of activation values ​​at different positions

[0040] Table 2 Distribution of activation values ​​in the feedforward layer and prediction results of activation values ​​at different positions

[0041] The predicted results of the activation values ​​at the above different positions are summed up to obtain the predicted value of the activation value in the training memory parameter when the single-precision training method is adopted. .

[0042] It should be understood that by summarizing the predicted values ​​of the above four parameters, it can be seen that when the preset training mode is the single-precision training mode, the memory overhead prediction value .

[0043] In step S3, when the preset training mode is the half-precision training mode, the ratio of the model weight, the gradient and the optimizer state information is 1:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0044] It should be noted that when the preset training mode is the half-precision training mode, each data occupies 2 bytes. , so the predicted value of the model weight in the training memory parameter ; Then, according to the ratio of the model weight, the gradient and the optimizer state information in the training memory parameter of 1:1:2, the predicted value of the gradient in the training memory parameter can be obtained in turn , the predicted value of the optimizer state information in the training memory parameter .

[0045] In this embodiment, when the preset training mode is the half-precision training mode, the predicted values ​​of the activation values ​​in the attention layer and the feedforward layer are shown in Tables 3 and 4, respectively. Tables 3 and 4 both show the distribution positions of the activation values ​​and the predicted results of the activation values ​​at different positions: Table 3 Distribution of activation values ​​in the attention layer and prediction results of activation values ​​at different positions

[0046] Table 4 Distribution of activation values ​​in the feedforward layer and prediction results of activation values ​​at different positions

[0047] The predicted results of the activation values ​​at the above different positions are summed up to obtain the predicted values ​​of the activation values ​​in the training memory parameters when the half-precision training method is adopted. .

[0048] It should be understood that by summarizing the predicted values ​​of the above four parameters, it can be seen that when the preset training method is the half-precision training method, the memory overhead prediction value .

[0049] In step S3, when the preset training mode is the mixed precision training mode, the ratio of the model weight, the gradient and the optimizer state information is 1.5:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0050] It should be noted that when the preset training mode is the mixed precision training mode, the specified large language model is trained by combining single precision training with half precision training. At this time, the predicted value of the model weight in the training memory parameter is ; Then, according to the ratio of the model weight, the gradient and the optimizer state information in the training memory parameter of 1.5:1:2, the predicted value of the gradient in the training memory parameter can be obtained in turn , the predicted value of the optimizer state information in the training memory parameter .

[0051] In this embodiment, when the preset training mode is a mixed precision training mode, the predicted values ​​of the activation values ​​in the attention layer and the feedforward layer are shown in Table 5 and Table 6, respectively. Table 5 and Table 6 both show the distribution positions of the activation values ​​and the predicted results of the activation values ​​at different positions: Table 5 Distribution of activation values ​​in the attention layer and prediction results of activation values ​​at different positions

[0052] Table 6 Distribution of activation values ​​in the feedforward layer and prediction results of activation values ​​at different positions

[0053] The predicted results of the activation values ​​at the above different positions are summed up to obtain the predicted values ​​of the activation values ​​in the training memory parameters when the mixed precision training method is adopted. .

[0054] It should be understood that by summarizing the predicted values ​​of the above four parameters, when the preset training mode is the mixed precision training mode, the predicted value of the memory overhead is .

[0055] In step S3 of this embodiment, when the preset training mode is a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, and LoRA (Low-Rank Adaptation) is synchronously enabled for the specified large language model, the total parameter amount is: ; In the formula, r The dimension of the trainable linear layer in the specified large language model after LoRA is synchronously enabled for the specified large language model; In step S3, the memory overhead prediction value is: ; In the formula, W , G , O and A It is calculated based on the total parameter amount and the training memory parameter of the specified large language model during training; wherein, W is the predicted value of the model weight in the training memory parameter, G is the predicted value of the gradient in the training memory parameter, O is the predicted value of the optimizer state information in the training memory parameter, A is the predicted value of the activation value in the training memory parameter.

[0056] It should be noted that after LoRA is synchronously enabled for the specified large language model, only a portion of the additional parameters in the specified large language model are trained, which can greatly reduce the cost of model training and reduce the computing power required during model training. Specifically, during implementation, it can be decided whether to enable LoRA for the specified large language model based on user selection, task requirements or hardware limitations.

[0057] In the implementation process, if LoRA is enabled for the specified large language model, a pair of parameter quantities will be d ×r and r × d The trainable linear layers LoRA A and LoRA B are attached to the linear layer of the specified large language model. The simplified schematic diagram of the model linear layer and the trainable linear layers LoRA A and LoRA B is as shown in Figure 3 As shown, r It is one of the dimensions of the trainable linear layers LoRA A and LoRA B, which directly affects the parameter quantity of the trainable linear layers LoRA A and LoRA B. In the subsequent training process, only the trainable linear layers LoRA A and LoRA B are trained, involving the memory usage of the relevant gradients and optimizer state information, while the linear layer of the specified large language model does not participate in the training.

[0058] Based on this, each parameter in the specified large language model can be d × d The number of parameters added by LoRA is ; For each parameter in the specified large language model d ×4 d The number of parameters added when LoRA is enabled synchronously for the specified large language model is , then the total number of parameters added in each intermediate layer is . It can be calculated that when LoRA is enabled synchronously for the specified large language model, the total number of parameters of the specified large language model is: , where the number of trainable parameters is: .

[0059] Therefore, this embodiment integrates the latest large language model training technology, which can predict the memory overhead of supporting LoRA fine-tuning to support large language model training in LoRA scenarios.

[0060] In step S3, when the preset training mode is the single-precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0061] It should be noted that when the preset training mode is the single-precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is , however, since the trainable parameters in the total parameter amount are , so the predicted value of the gradient in the training memory parameter , the predicted value of the optimizer state information in the training memory parameter .

[0062] In addition, since the trainable linear layers LoRA A and LoRA B are trainable, the dimensions of the two layers need to be and The activation value is used to calculate the gradient of the two. When the preset training mode is the single-precision training mode and LoRA is synchronously enabled for the specified large language model, it can be analyzed that the activation value has the following behaviors: the activation value of LoRA is the same as the activation value of the linear layer of the model during non-LoRA training; the synchronous enabling of LoRA for the specified large language model does not affect the distribution of other activation values. Based on the above two behaviors, when the activation value is actually calculated, compared with non-LoRA training, for each linear layer of the model, LoRA introduces an activation value of dimension size, thereby incrementally calculating the overall activation value. In addition, since each intermediate layer in the specified large language model has 6 linear layers, when using single-precision training, LoRA introduces a dimension of size in each intermediate layer. Activation value memory; when using half-precision training and mixed-precision training, LoRA introduces a size of Activation value memory.

[0063] It should be understood that by summarizing the predicted values ​​of the above four parameters, it can be seen that when the preset training mode is the single-precision training mode and LoRA is enabled synchronously for the specified large language model, the memory overhead prediction value .

[0064] In step S3, when the preset training mode is the half-precision training mode and LoRA is synchronously enabled for the specified large language model, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0065] It should be noted that when the preset training mode is the half-precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter , however, since the trainable parameters in the total parameter amount are , so the predicted value of the gradient in the training memory parameter , the predicted value of the optimizer state information in the training memory parameter .

[0066] It should be understood that by summarizing the predicted values ​​of the above four parameters, it can be seen that when the preset training mode is the half-precision training mode and LoRA is enabled synchronously for the specified large language model, the memory overhead prediction value .

[0067] In step S3, when the preset training mode is a mixed precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is: ; The predicted value of the gradient in the training memory parameter is: ; The predicted value of the optimizer state information in the training memory parameter is: ; The predicted value of the activation value in the training memory parameter is: .

[0068] It should be noted that when the preset training mode is the mixed precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter , the predicted value of the gradient in the training memory parameter , the predicted value of the optimizer state information in the training memory parameter .

[0069] It should be understood that by summarizing the predicted values ​​of the above four parameters, when the preset training mode is the mixed precision training mode and LoRA is enabled synchronously for the specified large language model, the memory overhead prediction value is .

[0070] It should also be noted that in this embodiment, when the preset training mode is a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, the memory overhead prediction value is recorded as , when the preset training mode is single-precision training mode, half-precision training mode or mixed-precision training mode, and the memory overhead prediction value when LoRA is enabled synchronously for the specified large language model is recorded as , only in Greater than In the case of , only by enabling LoRA for the specified large language model can the memory overhead of large model training be effectively reduced. At this time, the optimization problem and constraints are as follows. By solving this problem, the invention obtains b and s Value makes LoRA effectively reduce memory overhead r Value boundaries: .

[0071] S4. Obtain the remaining memory value of the training device, and determine whether the remaining memory value is greater than the memory overhead prediction value. If so, perform memory pre-allocation processing for the training device according to the memory overhead prediction value so that the training device can train the specified large language model; if not, output memory warning information, such as memory optimization suggestions, to avoid training failures due to insufficient memory during subsequent training. It should be noted that in this embodiment, the training device is also a computing device used to train the specified large language model, such as a specified GPU (Graphics Processing Unit), TPU (tensor processing unit) and other devices, and the remaining memory value is also the remaining video memory (VRAM) space of the training device.

[0072] This embodiment can realize the estimation and optimization management of memory usage during the training of a large language model. Specifically, during the implementation of this embodiment, first, the designated large language model is abstracted to obtain the model after abstraction, and a parameter symbol list matching the designated large language model is constructed based on the model after abstraction; then, the total parameter amount of the designated large language model when trained in a preset training method is obtained based on the parameter symbol list, and the memory overhead prediction value of the designated large language model when trained in the preset training method is calculated based on the total parameter amount and the training memory parameter of the designated large language model when training; finally, the remaining memory value of the training device is obtained, and it is determined whether the remaining memory value is greater than the memory overhead prediction value. If so, memory pre-allocation processing is performed for the training device based on the memory overhead prediction value so that the training device can train the designated large language model. In this process, this embodiment abstracts the specified large language model and constructs a parameter symbol list matching the specified large language model. Then, based on the parameter symbol list and a preset training method, the memory overhead of the specified large language model is predicted, and the memory of the training device is managed based on the obtained memory overhead prediction value. This can achieve accurate memory estimation under different training methods, solve the problem of difficult-to-predict memory status in large language model training, and is conducive to ensuring the smooth development of the large language model, while also helping to reduce unnecessary waste of resources.

[0073] Embodiment 2: This embodiment discloses a memory management system for large language model training, which is used to implement the memory management method for large language model training in Embodiment 1; Figure 4 As shown, the memory management system for large language model training includes: The model abstraction module is used to abstract the specified large language model to obtain the abstracted model; A parameter symbol list construction module, which is in communication with the model abstraction module and is used to construct a parameter symbol list matching the specified large language model according to the abstracted model; a memory overhead prediction module, which is in communication connection with the parameter symbol list construction module, and is used to obtain, according to the parameter symbol list, a total parameter amount of the specified large language model when training in a preset training method, and calculate, according to the total parameter amount and a training memory parameter of the specified large language model when training, a memory overhead prediction value of the specified large language model when training in the preset training method; A memory management module is communicatively connected to the memory overhead prediction module, and is used to obtain the remaining memory value of the training device and determine whether the remaining memory value is greater than the memory overhead prediction value. If so, memory pre-allocation processing is performed for the training device according to the memory overhead prediction value so that the training device can train the specified large language model.

[0074] It should be noted that the working process, working details and technical effects of the memory management system for large language model training provided in this embodiment 2 can be found in embodiment 1 and will not be repeated here.

[0075] Embodiment 3: Based on Embodiment 1 or 2, this embodiment discloses an electronic device, which may be a smart phone, a tablet computer, a laptop computer, or a desktop computer. The electronic device may be referred to as a user terminal, a portable terminal, a desktop terminal, etc. Figure 5 As shown, the electronic equipment includes: a memory for storing computer program instructions; and, A processor is used to execute the computer program instructions to complete the operation of a memory management method for large language model training as described in any one of Embodiment 1.

[0076] Specifically, the processor 301 may include one or more processing cores, such as a 4-core processor, an 8-core processor, etc. The processor 301 may be implemented in at least one hardware form of DSP (Digital Signal Processing), FPGA (Field-Programmable Gate Array), and PLA (Programmable Logic Array). The processor 301 may also include a main processor and a coprocessor. The main processor is a processor for processing data in the awake state, also known as a CPU (Central Processing Unit); the coprocessor is a low-power processor for processing data in the standby state. In some embodiments, the processor 301 may be integrated with a GPU (Graphics Processing Unit), which is responsible for rendering and drawing the content to be displayed on the display screen.

[0077] The memory 302 may include one or more computer-readable storage media, which may be non-transitory. The memory 302 may also include a high-speed random access memory, and a non-volatile memory, such as one or more disk storage devices, flash memory storage devices. In some embodiments, the non-transitory computer-readable storage medium in the memory 302 is used to store at least one instruction, which is used to be executed by the processor 301 to implement the memory management method for large language model training provided in Example 1 of the present application.

[0078] In some embodiments, the terminal may further optionally include: a communication interface 303 and at least one peripheral device. The processor 301, the memory 302 and the communication interface 303 may be connected via a bus or a signal line. Each peripheral device may be connected to the communication interface 303 via a bus, a signal line or a circuit board. Specifically, the peripheral device includes: at least one of a radio frequency circuit 304, a display screen 305 and a power supply 306.

[0079] The communication interface 303 may be used to connect at least one peripheral device related to I / O (Input / Output) to the processor 301 and the memory 302. In some embodiments, the processor 301, the memory 302, and the communication interface 303 are integrated on the same chip or circuit board; in some other embodiments, any one or two of the processor 301, the memory 302, and the communication interface 303 may be implemented on a separate chip or circuit board, which is not limited in this embodiment.

[0080] The radio frequency circuit 304 is used to receive and transmit RF (Radio Frequency) signals, also known as electromagnetic signals. The radio frequency circuit 304 communicates with a communication network and other communication devices through electromagnetic signals.

[0081] The display screen 305 is used to display a UI (User Interface). The UI may include any combination of graphics, text, icons, and videos.

[0082] The power supply 306 is used to supply power to various components in the electronic device.

[0083] Embodiment 4: Based on any one of Embodiments 1 to 3, this embodiment discloses a computer program product, including a computer program or an instruction, which, when executed by a computer, implements a memory management method for large language model training as described in any one of Embodiments 1. The computer may be a general-purpose computer, a special-purpose computer, a computer network, or other programmable devices.

[0084] Obviously, those skilled in the art should understand that the above modules or steps of the present invention can be implemented by a general computing device, they can be concentrated on a single computing device, or distributed on a network composed of multiple computing devices, and optionally, they can be implemented by a program code executable by a computing device, so that they can be stored in a storage device and executed by the computing device, or they can be made into individual integrated circuit modules, or multiple modules or steps therein can be made into a single integrated circuit module for implementation. Thus, the present invention is not limited to any specific combination of hardware and software.

[0085] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention, rather than to limit them; although the present invention has been described in detail with reference to the above embodiments, a person skilled in the art should understand that the technical solutions described in the above embodiments can still be modified, or some of the technical features can be replaced by equivalents. However, these modifications or replacements do not deviate the essence of the corresponding technical solutions from the spirit and scope of the technical solutions of the embodiments of the present invention.

Claims

1. A memory management method for large language model training, characterized in that: include: Perform abstract processing on the specified large language model to obtain an abstracted model; Constructing a parameter symbol list matching the specified large language model according to the abstracted model; Obtaining, according to the parameter symbol list, a total parameter amount of the designated large language model when training in a preset training method, and calculating, according to the total parameter amount and a training memory parameter of the designated large language model when training in the preset training method, a predicted value of memory overhead of the designated large language model when training in the preset training method; Obtain a remaining memory value of the training device, and determine whether the remaining memory value is greater than the memory overhead prediction value. If so, perform memory pre-allocation processing for the training device according to the memory overhead prediction value so that the training device can train the specified large language model.

2. A memory management method for large language model training according to claim 1, characterized in that: The abstracted model includes an input layer, an embedding layer, an intermediate layer, and an output layer connected in sequence; wherein the input dimension of the input layer and the output dimension of the output layer are both b × s , the parameter quantity of the embedding layer is d × v , the intermediate layer includes an attention layer and a feedforward layer connected in sequence, and the number of attention heads in the attention layer is h , the parameters of different linear layers in the attention layer are d 2 , the number of parameters of different linear layers in the feedforward layer is 4 d 2 , b represents the batch size, s represents the sequence length, d represents the hidden layer dimension, v represents the length of the word vector, h is a natural number greater than 1, b , s , d and v are natural numbers greater than 0 respectively; the number of layers of the intermediate layer is L layer, L is a natural number greater than 1.

3. A memory management method for large language model training according to claim 1, characterized in that: The parameter symbol list includes the number of intermediate layers of the abstracted model L , word vector length v , hidden layer dimension d , batch size b and sequence length s ; The training memory parameters include model weights, gradients, optimizer state information and activation values.

4. A memory management method for large language model training according to claim 3, characterized in that: When the preset training mode is a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, the total parameter amount is: ； The predicted memory cost is: ； In the formula, W , G , O and A It is calculated based on the total parameter amount and the training memory parameter of the specified large language model during training; wherein, W is the predicted value of the model weight in the training memory parameter, G is the predicted value of the gradient in the training memory parameter, O is the predicted value of the optimizer state information in the training memory parameter, A is the predicted value of the activation value in the training memory parameter.

5. A memory management method for large language model training according to claim 4, characterized in that: When the preset training mode is a single-precision training mode, the ratio of the model weight, the gradient, and the optimizer state information is 1:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ； The predicted value of the gradient in the training memory parameter is: ； The predicted value of the optimizer state information in the training memory parameter is: ； The predicted value of the activation value in the training memory parameter is: 。 6. A memory management method for large language model training according to claim 4, characterized in that: When the preset training mode is the half-precision training mode, the ratio of the model weight, the gradient and the optimizer state information is 1:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ； The predicted value of the gradient in the training memory parameter is: ； The predicted value of the optimizer state information in the training memory parameter is: ； The predicted value of the activation value in the training memory parameter is: 。 7. A memory management method for large language model training according to claim 4, characterized in that: When the preset training mode is the mixed precision training mode, the ratio of the model weight, the gradient and the optimizer state information is 1.5:1:2; correspondingly, the predicted value of the model weight in the training memory parameter is: ； The predicted value of the gradient in the training memory parameter is: ； The predicted value of the optimizer state information in the training memory parameter is: ； The predicted value of the activation value in the training memory parameter is: 。 8. The memory management method for large language model training according to claim 3, characterized in that: When the preset training mode is a single-precision training mode, a half-precision training mode, or a mixed-precision training mode, and LoRA is synchronously enabled for the specified large language model, the total parameter amount is: ； In the formula, r The dimension of the trainable linear layer in the specified large language model after LoRA is synchronously enabled for the specified large language model; The predicted memory cost is: ； In the formula, W , G , O and A It is calculated based on the total parameter amount and the training memory parameter of the specified large language model during training; wherein, W is the predicted value of the model weight in the training memory parameter, G is the predicted value of the gradient in the training memory parameter, O is the predicted value of the optimizer state information in the training memory parameter, A is the predicted value of the activation value in the training memory parameter.

9. A memory management method for large language model training according to claim 8, characterized in that: When the preset training mode is the single-precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is: ； The predicted value of the gradient in the training memory parameter is: ； The predicted value of the optimizer state information in the training memory parameter is: ； The predicted value of the activation value in the training memory parameter is: 。 10. A memory management method for large language model training according to claim 8, characterized in that: When the preset training mode is the half-precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is: ； The predicted value of the gradient in the training memory parameter is: ； The predicted value of the optimizer state information in the training memory parameter is: ； The predicted value of the activation value in the training memory parameter is: 。 11. A memory management method for large language model training according to claim 8, characterized in that: When the preset training mode is a mixed precision training mode and LoRA is enabled synchronously for the specified large language model, the predicted value of the model weight in the training memory parameter is: ； The predicted value of the gradient in the training memory parameter is: ； The predicted value of the optimizer state information in the training memory parameter is: ； The predicted value of the activation value in the training memory parameter is: 。 12. A memory management system for large language model training, characterized in that: Used to implement a memory management method for large language model training as described in any one of claims 1 to 11; The memory management system for large language model training includes: The model abstraction module is used to abstract the specified large language model to obtain the abstracted model; A parameter symbol list construction module, which is in communication with the model abstraction module and is used to construct a parameter symbol list matching the specified large language model according to the abstracted model; a memory overhead prediction module, which is in communication connection with the parameter symbol list construction module, and is used to obtain, according to the parameter symbol list, a total parameter amount of the specified large language model when training in a preset training method, and calculate, according to the total parameter amount and a training memory parameter of the specified large language model when training, a memory overhead prediction value of the specified large language model when training in the preset training method; A memory management module is communicatively connected to the memory overhead prediction module, and is used to obtain the remaining memory value of the training device and determine whether the remaining memory value is greater than the memory overhead prediction value. If so, memory pre-allocation processing is performed for the training device according to the memory overhead prediction value so that the training device can train the specified large language model.

Images