Large language model optimization method and system based on variance reduction and momentum acceleration

Abstract

The embodiment of the invention provides a large language model optimization method based on variance reduction and momentum acceleration. The method comprises the following steps: a gradient estimation stage of a large language model: initializing a seed list and a projection list; executing multiple independent query iterations, calling a disturbance subprogram, generating a random seed for the large language model, and storing a table; in the disturbance subprogram, the determined gradient projection value is stored in a projection list, and iteration is carried out for next query; after executing multiple independent query iterations, storing a plurality of random seeds and a plurality of gradient projection values corresponding to the random seeds; in the weight updating stage of the large language model, a gradient norm subprogram is called for each layer of the large language model, and a random seed is obtained to reset a random number generator; and determining variance-reduced gradient estimation according to the gradient projection value extracted from the projection list and the reproduced disturbance vector. According to the embodiment of the invention, gradient information queried for multiple times is aggregated to generate low-noise gradient estimation, and fine tuning of a large language model is completed.

Description

Technical Field

[0001] This invention relates to the fields of deep learning and natural language processing, and in particular to a method and system for optimizing large language models based on variance reduction and momentum acceleration. Background Technology

[0002] In recent years, pre-training LLMs (large language models) on massive datasets and then fine-tuning them for downstream tasks has become the mainstream paradigm for training large models. However, with the continuous expansion of model size, the enormous GPU memory required for fine-tuning large language models has become a major challenge. For example, fully fine-tuning 7B parameters of Llama-2 using the Adam optimizer requires at least 70GB of memory, which easily exceeds the memory limits of consumer-grade GPUs. In other words, traditional fine-tuning methods, even parameter-efficient techniques such as LoRA (Low Rank Adaptation), while significantly reducing the number of parameters to be updated, still rely on first-order gradients and backpropagation for optimization. This means that activation values ​​of intermediate layers must be stored during training to calculate gradients, resulting in peak GPU memory requirements for fine-tuning far exceeding those required for model inference.

[0003] To address the aforementioned memory issues, existing technologies have proposed fine-tuning methods based on ZO (Zeroth-Order Optimization), such as MeZO. These methods estimate gradients using finite differences, completely avoiding backpropagation and thus eliminating the need to store activation values, reducing memory overhead to the inference level.

[0004] In the process of realizing this invention, the inventors discovered at least the following problems in the related technology:

[0005] Existing ZO methods have obvious drawbacks: the MeZO-SGD method based on stochastic gradient descent (SGD) has a slow convergence speed and unstable performance; while ZO-Adam, which attempts to combine momentum optimizers (such as Adam), suffers from severe performance limitations due to the inherent high noise in single query gradient estimation. Summary of the Invention

[0006] In order to at least solve the problems of slow convergence speed and unstable performance caused by the large number of parameters in the existing technology.

[0007] In a first aspect, embodiments of the present invention provide a method for optimizing large language models based on variance reduction and momentum acceleration, comprising:

[0008] In the gradient estimation stage of the large language model, at the beginning of each training step, a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query are initialized;

[0009] A plurality of independent query iterations are performed, in each query iteration, a perturbation subprogram is called to generate random seeds for the large language model and store the random seeds in the seed list, wherein the random seeds are used to reproduce perturbation vectors without consuming GPU memory for storing high-dimensional vectors;

[0010] In the perturbation subprogram, the parameters of the large language model are sequentially subjected to forward perturbation, backward perturbation and the forward perturbation, and gradient projection values are determined and stored in the projection list, and the parameters are restored to the initial state for the next query iteration;

[0011] After performing a plurality of independent query iterations, the seed list and the projection list store a plurality of random seeds and a plurality of gradient projection values corresponding thereto, respectively;

[0012] In the weight updating stage of the large language model, a gradient norm subprogram is called for each layer of the large language model, the gradient norm subprogram obtains random seeds from the seed list to reset the random number generator, and the reproduced perturbation vectors are regenerated without the need for memory caching;

[0013] For each weight matrix, the gradient projection value taken from the projection list and the reproduced perturbation vector are used to determine a variance-reduced gradient estimate, and through the gradient estimate and a selected optimizer, fine-tuning of the large language model is completed without relying on backpropagation.

[0014] In a second aspect, an embodiment of the present application provides a large language model optimization system based on variance reduction and momentum acceleration, comprising:

[0015] The gradient estimation stage module is used in the gradient estimation stage of the large language model, initializes a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query at the beginning of each training step, performs multiple independent query iterations, calls a perturbation subroutine in each query iteration to generate random seeds for the large language model and store them in the seed list, wherein the random seeds are used to reproduce perturbation vectors without consuming GPU memory for storing high-dimensional vectors, in the perturbation subroutine, sequentially performs forward perturbation, backward perturbation and the forward perturbation on the parameters of the large language model, determines the gradient projection values to be stored in the projection list, and restores the parameters to the initial state for the next query iteration, and after performing multiple independent query iterations, the seed list and the projection list store multiple random seeds and multiple gradient projection values corresponding thereto respectively.

[0016] The weight update stage module is used in the weight update stage of the large language model, calls a gradient norm subroutine for each layer of the large language model, the gradient norm subroutine obtains random seeds from the seed list to reset the random number generator, realizes the re-generation of the reproduced perturbation vector without the need for memory caching, for each weight matrix, determines the variance-reduced gradient estimation from the gradient projection value taken out of the projection list and the reproduced perturbation vector, and completes the optimization fine-tuning of the large language model without relying on back propagation through the gradient estimation and the selected optimizer.

[0017] In a third aspect, an electronic device is provided, which includes at least one processor and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to perform the steps of the large language model optimization method based on variance reduction and momentum acceleration according to any one of the embodiments of the present application.

[0018] In a fourth aspect, an embodiment of the present application provides a storage medium having a computer program stored thereon, characterized in that the program, when executed by a processor, implements the steps of the large language model optimization method based on variance reduction and momentum acceleration according to any one of the embodiments of the present application.

[0019] In a fifth aspect, an embodiment of the present application provides a computer program product, including computer programs / instructions, characterized in that the computer programs / instructions, when executed by a processor, implement the steps of the large language model optimization method based on variance reduction and momentum acceleration according to any one of the embodiments of the present application.

[0020] The method has the beneficial effects that the memory occupation can be further compressed by quantizing the optimizer state (MUZO-QAdam). Experiments show that when the OPT-2.7B model is fully parameter fine-tuned, the traditional FO-Adam optimizer requires 90.8 GB of video memory, while the MUZO-QAdam of the method only requires 18.5 GB, and the memory occupation is less than 21% of the former. On the classification, multiple selection and generation tasks of multiple large models, the performance of MUZO-Adam is stable and better than other ZO baseline methods (such as MeZO and LoZO), and reaches the level close to the traditional first-order fine-tuning (FO-Adam). Compared with the baseline method MeZO, the MUZO-Adam of the method can reduce the required GPU hours and the total number of forward propagations by about 70% while achieving similar or better performance. The method can be used as a plug-and-play optimizer and seamlessly combined with full parameter fine-tuning or PEFT methods (such as LoRA). In addition, the method does not depend on a specific model architecture, and the experimental results have been verified on various architectures of LLMs, including masked language models (Masked Language Model, such as RoBERTa-Large) and autoregressive language models (Large Autoregressive Model, such as OPT, Llama-2, Vicuna) and the like. Since it does not depend on back propagation at all, the method provides an innovative technical path for continuous learning and model adaptation in scenarios where back propagation is limited (such as edge devices, quantized models). BRIEF DESCRIPTION OF DRAWINGS

[0021] In order to more clearly illustrate the technical solutions in the embodiments of the present application or the prior art, the following will briefly introduce the drawings needed to be used in the embodiments or prior art description. Obviously, the drawings in the following description are some embodiments of the present application, and those skilled in the art can also obtain other drawings according to these drawings without creative labor.

[0022] Figure 1 is a flowchart of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0023] Figure 2 is a whole algorithm schematic diagram of a MUZO fine-tuning algorithm of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0024] Figure 3 is a MUZO-Adam and MUZO-QAdam fine-tuning algorithm schematic diagram of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0025] Figure 4 is a RoBERTa-Large experimental result schematic diagram of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0026] Figure 5 is a RoBERTa Large training loss curve fine-tuning schematic diagram on MNLI, RTE and TREC data sets of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0027] Figure 6 is a MUZO-Adam best zero-order optimization performance schematic diagram on different tasks of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0028] Figure 7 is a performance schematic diagram on COPA and WinoGrande (Wino.) using different ZO optimizers of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0029] Figure 8 is a memory consumption schematic diagram of MUZO of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0030] Figure 9 is a training loss curve schematic diagram of fine-tuning OPT-13B on SST-2 data sets of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0031] Figure 10 is a computational cost schematic diagram of MUZO, MeZO and FO of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0032] Figure 11 is a schematic diagram of ablation research on the number of queries in the training step of a large language model optimization method based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0033] Figure 12 is a structure schematic diagram of a large language model optimization system based on variance reduction and momentum acceleration provided by an embodiment of the present application;

[0034] Figure 13 is a structure schematic diagram of an embodiment of an electronic device for a large language model optimization based on variance reduction and momentum acceleration provided by an embodiment of the present application. DETAILED DESCRIPTION

[0035] In order to make the purposes, technical solutions and advantages of the embodiments of the present application clearer, the technical solutions in the embodiments of the present application will be described clearly and completely below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are some embodiments but not all of the present application. Based on the embodiments in the present application, all other embodiments obtained by those of ordinary skill in the art without creative work fall within the protection scope of the present application.

[0036] As Figure 1 Fig. 1 shows a flowchart of a large language model optimization method based on variance reduction and momentum acceleration according to an embodiment of the present application, including the following steps:

[0037] S11: In the gradient estimation stage of the large language model, at the beginning of each training step, a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query are initialized;

[0038] S12: Perform multiple independent query iterations. In each query iteration, call a perturbation subroutine to generate random seeds for the large language model and store them in the seed list, wherein the random seeds are used to reproduce perturbation vectors without consuming GPU memory to store high-dimensional vectors;

[0039] S13: In the perturbation subroutine, sequentially perform forward perturbation, backward perturbation and the forward perturbation on the parameters of the large language model, determine the gradient projection value to be stored in the projection list, and restore the parameters to the initial state for the next query iteration;

[0040] S14: After performing multiple independent query iterations, the seed list and the projection list store multiple random seeds and multiple gradient projection values corresponding thereto, respectively;

[0041] S15: In the weight update stage of the large language model, call a gradient norm subroutine for each layer of the large language model, the gradient norm subroutine obtains random seeds from the seed list to reset the random number generator, and re-generates the reproduced perturbation vector without the need for memory caching;

[0042] S16: For each weight matrix, determine the variance-reduced gradient estimate from the gradient projection value taken out of the projection list and the reproduced perturbation vector, and complete the optimization fine-tuning of the large language model without relying on backpropagation through the gradient estimate and the selected optimizer.

[0043] In this embodiment, considering the memory problem in fine-tuning downstream tasks after pre-training LLM, the prior art uses memory-efficient zeroth-order optimization (abbreviated as MeZO). MeZO estimates the gradient through two forward passes, avoiding backpropagation, thus eliminating the need to store activation values. When used in conjunction with the SGD optimizer, MeZO SGD achieves the fine-tuning of LLMs using only inference memory. However, ZO-SGD directly updates model parameters using gradient estimates of the current step, leading to issues such as unstable performance and slow convergence. ZO-Adam smooths gradient estimates and update steps using historical first and second moments, and has shown better performance than ZO-SGD in some studies. However, limited by significant noise in gradient estimates, the moments in Adam also exhibit considerable noise, which limits the performance of ZO-Adam. To solve the above problems, the present method proposes MUZO (MUltiple-query Memory Efficient Zeroth-Order).

[0044] Overall, the MUZO of the present method consists of two stages: (1) Gradient estimation stage: the model undergoes random perturbation and is passed forward in a cyclic order in q queries, caching random seeds and gradient projection values. (2) Weight update stage: the model reconstructs the gradient estimates of each weight matrix in a cyclic order corresponding to the weight matrix, and updates the weights accordingly.

[0045] For step S11, in the gradient estimation stage (Stage 1: Gradient Estimation), the goal of this stage is to efficiently calculate q independent gradient projection values and cache the minimum information needed to reconstruct these calculations. Its specific execution process can be realized through the first for loop of the main algorithm (Main Algorithm) as shown in Figure 2

[0046] First, the cache needs to be initialized. At the beginning of each training step (t), two empty lists are initialized: seeds is a two-dimensional list used to store the random seeds of each layer of the large language model, and seed is used to initialize random number generation, which plays a key role in ensuring the reproducibility of results in large language models. When a large language model generates text, many operations involve random selection, such as sampling from multiple possible next words. The seed provides a fixed starting state for these random operations. As long as the seed value is fixed, the results generated under the same input and model settings will be the same, facilitating the research and debugging of the model.

[0047] ​proj_grads is a one-dimensional list used to store the gradient projection value of each query. The core of gradient projection determines the feasible descent direction by constructing the projection matrix, and maps the unconstrained optimization direction into the feasible region for iterative optimization. In the specific implementation, the random seed of each query layer pair is sampled and stored as shown in Algorithm 2 (Perturb) of Figure 2 .

[0048] For step S12, multiple independent query iterations are performed, for example, q queries are performed, in order to obtain a more accurate estimate of the true gradient and reduce the estimation variance, the algorithm will perform q independent query iterations. For each query:

[0049] Sampling and caching random seeds: call the perturbation Perturb subroutine to generate a random seed ss for each layer of the large language model and store it in the seeds list. This seed is the key to subsequent reproduction of the perturbation vector, avoiding the GPU memory consumption required by directly storing high-dimensional vectors.

[0050] For step S13, in the perturbation subroutine, the parameters of the large language model are sequentially subjected to forward perturbation, backward perturbation, and the forward perturbation, to determine the gradient projection value stored in the projection list, and to restore the parameters to the initial state for the next query iteration;

[0051] In the perturbation subroutine, a forward perturbation is applied to the parameters of the large language model, where the forward perturbation is θ+∈z, where θ is the model weight, ∈ is a settable value, and z is a perturbation vector generated by sampling from a standard normal distribution N(0,1). Forward propagation is performed on the large language model after the forward perturbation to obtain a first loss;

[0052] The perturbation subroutine is called again to apply a backward perturbation to the parameters of the large language model, so that θ+∈z changes to θ-∈z. Forward propagation is performed on the large language model after the backward perturbation to obtain a second loss;

[0053] Based on the first loss and the second loss, a gradient projection value containing gradient direction information is determined, and the gradient projection value is stored in the projection list;

[0054] The parameters of the large language model are again subjected to forward perturbation to restore the parameters of the large language model to the original state.

[0055] Specifically, forward perturbation and forward propagation: in the perturbation Perturb subroutine, a small forward perturbation (θ+∈z) is applied to each layer of the model parameters θ using a random number generator, where z is a perturbation vector generated by sampling from the standard normal distribution N(0, 1). Then, a complete forward propagation is performed on the perturbed model to calculate the first loss value l + .

[0056] After the above processing, backward perturbation and forward propagation are performed: the perturbation Perturb subroutine is called again, but this time a backward perturbation with a magnitude of twice (the model weight is changed from θ+∈z to θ-∈z) is applied. Then, the second forward propagation is performed to obtain the second loss value l - .

[0057] Gradient projection is calculated and cached: according to the central difference formula, the projection value proj_grad of the gradient in the random direction z is calculated, i.e. (l + -l - ) / 2∈. This projection value is a scalar that contains key information about the gradient direction. After calculation, the scalar is stored in the proj_grads projection list.

[0058] Restore model parameters: finally, the model parameters are restored to the original state θ by a forward perturbation, preparing for the next query.

[0059] For step S14, the above steps describe the steps of each query, and after q queries are completed in the above manner, the seeds seed list and the proj_grads projection list store q*d random seeds and q gradient projection scalar values corresponding to them respectively, and phase one ends.

[0060] For step S15, the goal of this phase is to update the model weights using the information collected in the gradient estimation phase. Specifically, it can be implemented by the second for loop in the main algorithm, and the specific execution process is as follows: Figure 2

[0061] Regarding iterative accumulation update, for each layer of the model, the Getgrad subroutine is called to obtain the average value of the q gradient queries, and the specific function of the Getgrad subroutine is:

[0062] Reset the seed and reproduce the perturbation vector: the jth[j][i]th seed is taken from the seeds seed list (where j is for the number of queries, and i is for the ith weight matrix of the model), and the random number generator of the system is reset using it. This key step can regenerate the same perturbation vector z used in the gradient estimation phase without the need for caching in memory. By, for example, Figure 2 ​The algorithm 3 (Getgrad) implementation resamples the perturbation vector, and the time complexity is O (1). Therefore, the method MUZO can effectively obtain the complete q query gradient estimation of each weight matrix.

[0063] As an implementation, for the i-th layer weight matrix, the gradient projection value taken out from the projection list is multiplied by the reproduced perturbation vector, and the gradient estimation of the i-th layer is obtained by the accumulation of q times of independent query iterations. Divide the gradient estimation by q to obtain the gradient average value of q times of independent query iterations.

[0064] In this embodiment, the query gradient accumulation and average: for the i-th layer weight matrix, the algorithm will multiply the gradient projection scalar proj_grads[j] taken out with the reproduced perturbation vector z, and the gradient estimation value of the layer after q times of accumulation, and finally divide by q to obtain the average value of q times of queries.

[0065] Through the stages of the above steps, the method efficiently utilizes the multi-query mechanism of layer-by-layer perturbation to obtain variance-reduced gradient estimation without relying on back propagation, and can flexibly combine various optimizers to complete the fine-tuning of large language models.

[0066] Among them, the Optimizer optimizer can select SGD-random gradient descent optimizer, SGD-Momentum-random gradient descent optimizer with momentum.

[0067] As can be seen from this embodiment, the method proposes a large language model zero-order optimization fine-tuning method named MUZO based on variance reduction and momentum acceleration. The method generates low-noise gradient estimation by efficiently aggregating gradient information from multiple queries, realizes flexible combination of various optimizers, and completes fine-tuning of large language models.

[0068] As an implementation, the optimizer includes an Adam optimizer;

[0069] The optimization fine-tuning realized by the gradient estimation and the selected optimizer includes:

[0070] The first-order momentum and second-order moment information stored by the Adam optimizer in the optimization fine-tuning process are used to smooth and accelerate the convergence process of optimization.

[0071] On the basis of the method MUZO, in order to achieve better results, the method is further expanded to MUZO-Adam.

[0072] In the present embodiment, the MUZO of the present method is combined with the Adam optimizer, unlike the ZO-Adam in the prior art, the MUZO-Adam of the present method applies the average gradient on multiple queries to estimate the gradient, instead of using a single query. That is, the zero-order gradient method of multiple queries is retained, and the historical first-order momentum and second-order moment information stored by the Adam optimizer is applied when updating the weights to smooth and accelerate the convergence process of optimization.

[0073] Specifically, when each training step t starts iteration, the first-order and second-order estimation equations in the gradient estimation equation of MUZO-Adam are replaced by:

[0074] The first-order matrix smooth gradient is:

[0075]

[0076] The second-order matrix explains the variance of the gradient in the parameter dimension:

[0077]

[0078] Where β1, β2 are the exponential decay rates of the moment estimation.

[0079] The update is replaced by:

[0080]

[0081] Where q is the number of query-independent estimated gradients.

[0082] Further, the coordinates in the matrix are updated as:

[0083]

[0084] Where, is the normalized matrix at iteration t.

[0085] As an embodiment, on the basis of the Adam optimizer, the storage state of the Adam optimizer is quantized, including:

[0086] The storage state of the Adam optimizer is dequantized to restore it to the precision of FP32, FP16 or BF16;

[0087] After the optimization and update of the parameter weights in the large language model, the storage state of the Adam optimizer is quantized to FP8 again, so as to reduce the GPU memory required for storing the Adam optimizer.

[0088] Further, to reduce the extra memory overhead caused by storing Adam optimizer states, the method adds an optimizer state FP8 quantization process to the MUZO-Adam, realizing the MUZO-QAdam algorithm as shown in Figure 3 The purpose of the corresponding steps is:

[0089] For each layer of weight matrix corresponding to the Adam optimizer state, the first step of updating is to perform DeQuantize to restore the precision of FP32, FP16 or BF16, which is used for subsequent update of the first order momentum m i , the second order moment v i and the model weight θ i required tensor calculation.

[0090] After the weight update is completed, the optimizer state is quantized to FP8 format again to reduce the GPU memory required for storing the optimizer.

[0091] Since the above-mentioned dequantization-weight update-quantization process is performed layer by layer, the peak memory of the GPU will only be affected by a small fluctuation (i.e. the parameter amount of a single layer of weight matrix). Overall, the GPU memory required by MUZO-QAdam is significantly less than that required by MUZO-Adam, which is about half of that required by MUZO-Adam under the default FP32 calculation precision.

[0092] Specifically, the above method involves quantizing the optimizer states required by Adam (i.e. the first order momentum and the second order moment) to FP8 precision, while the model weight perturbation and forward propagation remain in FP16 / FP32 format. After calculating the high-precision moments m t and v t , they are converted to FP8 representation ( and ) for storage. This process explicitly includes clipping the values to the representable range of the FP8 data type: where clamp(·) limits its input to the range defined by and , and FP8Quantiize(·) performs the conversion to the FP8 data type. Then, in the subsequent calculation steps, these FP8-stored moments are dequantized back to FP16 / FP32 precision.

[0093] The potential insight is that since the zeroth-order gradient estimation inherently introduces extra noise into the optimizer states, the first and second moments only provide an approximate direction for the weight and step size updates, thus eliminating the need for high precision of these states. Moreover, when applying MUZO-Adam to fine-tune the low-rank parameters introduced by methods like LoRA, the memory footprint of the optimizer states is significantly reduced, less than 3% of the memory required for inference.

[0094] The method discloses the idea, code and steps of the above-mentioned MUZO, MUZO-Adam and MUZO-QAdam. Further theoretical analysis is performed on the above-mentioned steps of the method. The q-query-independent gradient estimation method based on RGE for estimating the gradient is:

[0095]

[0096] Where L is the loss function of the large model fine-tuning, which can be optionally determined as the cross-entropy loss, ∈ is the perturbation scale, which can be understood as a very small quantity, and is generally set to 0.001 in experiments. θ is the model weight of fine-tuning, represents the set of model weight matrices, θ i represents the weight matrix of the i-th layer. z i is a random vector of perturbation generated by a random number generator, and the dimension is equal to the model weight of perturbation. q is the number of gradient estimates used for single-step update, and the gradient estimate value is the average of q times of gradient update.

[0097] In order to determine that the gradient estimation of the method is unbiased and the variance is reduced. In summary: for the MUZO of the method, let be the gradient estimate of the above q-RGE equation:

[0098]

[0099] By the above, it is ensured that is an unbiased estimate of the gradient at iteration t.

[0100] The method assumes the following conditions:

[0101] A1: The function

[0102] is applicable to all θ,

[0103] A2: For a small batch minibatch B: the variance of the stochastic gradient estimator is limited to

[0104] For the MUZO-Adam of the method, by averaging q queries, the variance of the gradient estimate can be reduced variance (σ 2 ). Specifically, for each iteration t:

[0105]

[0106] For the MUZO-QAdam algorithm of the present method, assuming A1, A2 holds, let be the sequence of q queries generated by MUZO-Adam at each iteration. Set θ R is randomly and uniformly chosen from {θ1,..., θ T}, thus we have:

[0107]

[0108] where d is the dimension of the parameter space and T is the number of iterations.

[0109] Finally, in summary, the optimization of the present method includes the following steps:

[0110] Multi-query gradient estimation: by performing multiple independent layer-wise random perturbations on the model weights, and performing 2 x q forward propagation to calculate the change of the loss function value.

[0111] Low-variance gradient generation: averaging the q gradient estimates to generate an aggregated gradient with lower variance and more stable direction. This mechanism significantly reduces the noise of gradient estimation, laying the foundation for the effective application of momentum optimizers.

[0112] Momentum accelerated weight update (MUZO-Adam and MUZO-QAdam): input the above generated low-noise aggregated gradient into the Adam momentum optimizer. Adam optimizer uses the first moment (momentum) and second moment (variance adaptation) of the gradient to accelerate convergence and stabilize the training process, forming an efficient synergy with the low-variance gradient of MUZO. MUZO-QAdam uses quantized optimizer states to reduce the GPU memory required to store the Adam optimizer.

[0113] Vector regeneration based on hierarchical seeds: to efficiently implement multi-query, the present method innovatively uses the "hierarchical random seed reproduction" technique. It only needs to cache the random seed of each layer (i.e. each weight matrix of the model) in each query (rather than the high-dimensional perturbation vector itself), which can accurately reproduce the required perturbation vector at a very small cost during the weight update phase, thereby supporting complex multi-query and momentum calculation without increasing the memory burden.

[0114] As can be seen from the above embodiments, the method can further compress the memory occupation by quantizing the optimizer state (MUZO-QAdam). Experiments show that when the OPT-2.7B model is fully parameter fine-tuned, the traditional FO-Adam optimizer requires 90.8 GB of video memory, while the MUZO-QAdam of the application only requires 18.5 GB, and the memory occupation is less than 21% of the former.

[0115] On the classification, multiple-choice and generation tasks of multiple large models, the performance of MUZO-Adam is stably better than that of other ZO baseline methods (such as MeZO and LoZO), and reaches a level close to that of the traditional first-order fine-tuning (FO-Adam).

[0116] Compared with the baseline method MeZO, the MUZO-Adam of the application can reduce the required GPU hours and total forward propagation times by about 70% while achieving similar or better performance.

[0117] The method can be used as a plug-and-play optimizer, which is seamlessly combined with full-parameter fine-tuning or PEFT methods (such as LoRA). In addition, the method does not depend on a specific model architecture, and the experimental results have been verified on multiple architectures of LLMs, including masked language models (Masked Language Model, such as RoBERTa-Large) and autoregressive language models (Large Autoregressive Model, such as OPT, Llama-2, Vicuna) and the like.

[0118] Since it does not depend on backpropagation at all, the method provides an innovative technical path for continuous learning and model adaptation in scenarios where backpropagation is limited (such as edge devices, quantized models).

[0119] To illustrate the method, the performance of MUZO is evaluated by comparing it with other state-of-the-art ZO methods MeZO and LoZO, as well as various optimizers, including SGD, SGD momentum, and Adam. The method demonstrates through experiments that the MUZO of the method performs well in downstream classification, multiple-choice, and generation tasks. The experiments evaluate the performance of MUZO on RoBERTa Large and autoregressive models of different sizes, including OPT-13B, Llama-2-7B, and Vicuna-V1.5-7B masked language models. The experimental datasets of the method cover various LLM fine-tuning tasks, including most tasks in the GLUE and SuperGLUE benchmarks. The following experimental settings are adopted: for RoBERTa-Large, both few-query (16 examples) and many-query (512 examples) settings are studied. For large autoregressive models, 1000, 500, and 1000 samples are randomly drawn for the training, validation, and test sets of each task.

[0120] The method sets the same total number of forward passes for all experiments involving MUZO, MeZO, and LoZO, calculated by 2 x q x number of training steps. For example, in the experiments of RoBERTa Large, the number of training steps for MeZO is set to 100000, while for MUZO, the number of steps is set to 100000 / q. Due to faster convergence, the FT steps are set to 1000. All experiments are conducted on Nvidia Geforce RTX 3090 and A800 GPUs.

[0121] For masked language models, the performance of the method MUZO is always better than that of MeZO.

[0122] In the experiments of Figure 4 , the combination of MUZO with different optimizers achieves higher accuracy than single-query MeZO on most tasks and models, and this performance is more prominent in the few-query learning setting. However, in the many-query learning scenario, MeZO SGD outperforms MUZO-SGD in some cases (e.g., the SST-5 task). One possible reason is that although a relatively high learning rate is used to alleviate the under-training caused by the reduction of training steps, MUZO-SGD still requires more forward passes to converge, where FT in the figure refers to first-order fine-tuning using the Adam optimizer, and the best results are shown in bold and the second-best results are shown in underlined.

[0123] Unlike the fluctuating performance of MUZO-SGD, the method MUZO-Adam is always better than the 1-query method in both few-query and many-query settings. Figure 2They are shown their loss curves, where it is clear that MUZO Adam converges better, mitigating the MUZO's reliance on overpassing and improving its performance.

[0124] As Figure 5 shown, in the early stage of training, MUZO-Adam and MUZO-QAdam do not show obvious advantages over ZO-SGD, and may even lag behind ZO-SGD (e.g. RTE task). However, after the number of training steps accumulates, ZO-Adam finds the optimization direction, and the convergence speed significantly accelerates, while ZO-SGD, which relies on the current step gradient estimate, shows slow convergence. In addition, the convergence speed of the MUZO QAdam curve is almost the same as that of the MUZO Adam curve, and due to the low accuracy, it only converges prematurely when the loss is less than 1e-5. That is, the MUZO-Adam of the present method needs a slow preheating process to obtain an approximate optimization direction.

[0125] For autoregressive language models, such as Figure 6 and as Figure 7 shown, in experiments on large autoregressive models, the present method tests different models and datasets, and draws similar conclusions as RoBERTa-large. Specifically, MUZO-Adam outperforms MeZO and LoZO on most tasks. In addition, due to the adoption of the LoRA method, MUZO-Adam requires almost no additional memory compared to the ZO-SGD method. As Figure 8 shows the GPU peak memory required by different optimization algorithms, indicating that MUZO Adam maintains its memory advantage.

[0126] As Figure 9 shows the loss curves under different query numbers. When the number of queries increases, the loss curve drops more smoothly, and the drop speed will change. The present method demonstrates the performance under different query numbers in the ablation study. Increasing the number of queries may improve performance under the same number of forward steps. However, when the number of queries is too high, the number of iterations of the model in the dataset iteration will decrease, resulting in a decrease in performance.

[0127] As Figure 8 shows the memory requirements of MUZO Adam when choosing LoRA fine-tuning or optimizer quantization. Compared with inference memory, directly using MUZO-Adam for full-parameter fine-tuning requires more than 150% additional memory. However, after applying the MUZO-QAdam optimizer with quantization, this overhead can be reduced to less than 50%. The most memory-saving MUZO-Adam fine-tuning still needs to be used in combination with LoRA. In this setting, MUZO-Adam can achieve optimization of inference memory while obtaining performance superior to ZO-SGD.

[0128] As Figure 10 The GPU hours and forward pass iterations before convergence of fine-tuning RoBERTa on SST2 task using MUZO are analyzed. Compared with MeZO, MUZO-Adam reduces 70% GPU time and total forward pass time, but still needs three times GPU time compared with FO-Adam.

[0129] The ablation study on the number of queries is performed to investigate whether directly increasing the number of queries to q and scaling the learning rate by q times can improve the performance of MUZO-Adam. The experimental results are shown in Figure 11 and Figure 9 When q = 8 or q = 16, scaling the learning rate can directly obtain good performance, but with further increase of q, the performance may deteriorate due to too high learning rate and fewer iterations on the training set.

[0130] In summary, the present method discloses a new zero-order optimization method MUZO, which effectively calculates the average value of multiple query gradients, thereby reducing the variance of gradient estimation. The present method proposes a MUZO-Adam optimizer, which uses the Adam optimizer to determine the optimization direction, thereby accelerating the convergence. In order to alleviate the memory overhead related to the Adam optimizer, the present method further introduces a MUZO-QAdam optimizer using FP8 quantization. The above experiments show that in the full parameter and LoRA fine-tuning settings, both MUZO-Adam and MUZO-QAdam show excellent performance and convergence speed compared with other zero-order optimization baselines.

[0131] As Figure 12 The system can execute the large language model optimization method based on variance reduction and momentum acceleration provided in any of the above embodiments, and is configured in a terminal.

[0132] The large language model optimization system 10 based on variance reduction and momentum acceleration provided in the present embodiment comprises a gradient estimation stage module 11 and a weight update stage module 12.

[0133] The gradient estimation stage module 11 is used in the gradient estimation stage of the large language model, and at the beginning of each training step, a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query are initialized, a plurality of independent query iterations are performed, in each query iteration, a perturbation subroutine is called, random seeds are generated for the large language model and stored in the seed list, wherein the random seeds are used to reproduce perturbation vectors without the need for GPU memory consumption to store high-dimensional vectors, in the perturbation subroutine, the parameters of the large language model are sequentially subjected to forward perturbation, reverse perturbation and the forward perturbation, gradient projection values are determined and stored in the projection list, and the parameters are restored to the initial state for the next query iteration, after a plurality of independent query iterations are performed, the seed list and the projection list store a plurality of random seeds and a plurality of gradient projection values corresponding thereto respectively; the weight update stage module 12 is used in the weight update stage of the large language model, and the gradient norm subroutine is called for each layer of the large language model, the gradient norm subroutine obtains random seeds from the seed list to reset the random number generator, and the reproduced perturbation vectors are regenerated without the need for memory caching, for each weight matrix, the gradient projection values taken from the projection list and the reproduced perturbation vectors are used to determine variance-reduced gradient estimates, and through the gradient estimates and the selected optimizer, the optimization fine-tuning of the large language model is completed without relying on back propagation.

[0134] The embodiment of the present application also provides a non-volatile computer storage medium, and the computer storage medium stores computer executable instructions, and the computer executable instructions can execute the large language model optimization method based on variance reduction and momentum acceleration in any method embodiment.

[0135] As an implementation manner, the non-volatile computer storage medium of the present application stores computer executable instructions, and the computer executable instructions are configured to:

[0136] In the gradient estimation stage of the large language model, at the beginning of each training step, a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query are initialized;

[0137] A plurality of independent query iterations are performed, in each query iteration, a perturbation subroutine is called, random seeds are generated for the large language model and stored in the seed list, wherein the random seeds are used to reproduce perturbation vectors without the need for GPU memory consumption to store high-dimensional vectors;

[0138] In the perturbation subprogram, the parameters of the large language model are sequentially subjected to forward perturbation, reverse perturbation and the forward perturbation, and the gradient projection value is determined and stored in the projection list, and the parameters are restored to the initial state for the next query iteration;

[0139] After performing multiple independent query iterations, the seed list and the projection list store multiple random seeds and multiple gradient projection values corresponding thereto, respectively;

[0140] In the weight updating stage of the large language model, the gradient norm subprogram is called for each layer of the large language model, the random seed is obtained from the seed list to reset the random number generator, and the reproducible perturbation vector is regenerated without the need for memory caching;

[0141] For each layer weight matrix, the gradient projection value taken from the projection list and the reproducible perturbation vector are used to determine the variance-reduced gradient estimate, and the optimization fine-tuning of the large language model is completed without relying on backpropagation through the gradient estimate and the selected optimizer.

[0142] As a non-volatile computer readable storage medium, it can be used to store non-volatile software programs, non-volatile computer executable programs and modules, such as program instructions / modules corresponding to the method in the embodiment of the present application. One or more program instructions are stored in the non-volatile computer readable storage medium, and when executed by the processor, the above-mentioned any method embodiment based on variance reduction and momentum acceleration large language model optimization method is executed.

[0143] Figure 13 is the hardware structure schematic diagram of the electronic device provided by another embodiment of the present application based on the variance reduction and momentum acceleration large language model optimization method, as shown in Figure 13 The device comprises:

[0144] One or more processors 1310 and memories 1320, Figure 13 In the embodiment, the processor 1310 is taken as an example. The device based on the variance reduction and momentum acceleration large language model optimization method can further comprise an input device 1330 and an output device 1340.

[0145] The processor 1310, the memory 1320, the input device 1330 and the output device 1340 can be connected through a bus or other means, Figure 13 In the embodiment, the connection through the bus is taken as an example.

[0146] The memory 1320, as a non-volatile computer readable storage medium, can be used to store non-volatile software programs, non-volatile computer executable programs and modules, such as program instructions / modules corresponding to the large language model optimization method based on variance reduction and momentum acceleration in the embodiments of the present application. The processor 1310 executes various function applications and data processing of the server by running the non-volatile software programs, instructions and modules stored in the memory 1320, that is, implements the large language model optimization method based on variance reduction and momentum acceleration of the above method embodiments.

[0147] The memory 1320 can include a program storage area and a data storage area, wherein the program storage area can store an operating system, application programs required by at least one function; and the data storage area can store data and the like. In addition, the memory 1320 can include a high-speed random access memory, and can also include a non-volatile memory, such as at least one disk storage device, a flash memory device, or other non-volatile solid-state memory device. In some embodiments, the memory 1320 can optionally include a memory remotely arranged with respect to the processor 1310, and these remote memories can be connected to the mobile device through a network. Examples of the above network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0148] The input device 1330 can receive input digital or character information. The output device 1340 can include a display device such as a display screen.

[0149] The one or more modules are stored in the memory 1320, and when executed by the one or more processors 1310, perform the large language model optimization method based on variance reduction and momentum acceleration in any of the above method embodiments.

[0150] The above product can perform the method provided by the embodiments of the present application, and has the corresponding function modules and beneficial effects of performing the method. Technical details not described in detail in the embodiments can be referred to the method provided by the embodiments of the present application.

[0151] The non-volatile computer readable storage medium can include a program storage area and a data storage area, wherein the program storage area can store an operating system, application programs required by at least one function; and the data storage area can store data created according to the use of the device and the like. In addition, the non-volatile computer readable storage medium can include a high-speed random access memory, and can also include a non-volatile memory, such as at least one disk storage device, a flash memory device, or other non-volatile solid-state memory device. In some embodiments, the non-volatile computer readable storage medium can optionally include a memory remotely arranged with respect to the processor, and these remote memories can be connected to the device through a network. Examples of the above network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0152] This invention also provides an electronic device comprising: at least one processor and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, the instructions being executed by the at least one processor to enable the at least one processor to perform the steps of the large language model optimization method based on variance reduction and momentum acceleration according to any embodiment of this invention.

[0153] The electronic devices described in this application exist in various forms, including but not limited to:

[0154] (1) Mobile communication devices: These devices are characterized by their mobile communication capabilities and primarily aim to provide voice and data communication. These terminals include smartphones, multimedia phones, feature phones, and low-end phones.

[0155] (2) Ultra-mobile personal computer devices: These devices fall under the category of personal computers, possessing computing and processing capabilities, and generally also have mobile internet access features. These terminals include PDAs, MIDs, and UMPCs, such as tablet computers.

[0156] (3) Portable entertainment devices: These devices can display and play multimedia content. This category includes audio and video players, handheld game consoles, e-book readers, as well as smart toys and portable car navigation devices.

[0157] (4) Other electronic devices with data processing functions.

[0158] In this document, relational terms such as "first" and "second" are used merely to distinguish one entity or operation from another, without necessarily requiring or implying any such actual relationship or order between these entities or operations. Furthermore, the terms "comprising" or "including" include not only those elements but also other elements not expressly listed, or elements inherent to such a process, method, article, or apparatus. Unless otherwise specified, an element defined by the phrase "comprising..." does not exclude the presence of additional identical elements in the process, method, article, or apparatus that includes said element.

[0159] The device embodiments described above are merely illustrative, wherein the units described as separate components can or can not be physically separate, and the components displayed as units can or can not be physical units, i.e., can be located in one place, or can be distributed to multiple network units. Part or all of the modules can be selected to achieve the purposes of the embodiments according to actual needs. Those skilled in the art can understand and implement without creative labor.

[0160] Through the description of the above embodiments, those skilled in the art can clearly understand that the embodiments can be realized by means of software and the necessary general hardware platform, and of course can also be realized by hardware. Based on such understanding, the above technical solutions can be embodied in the form of a software product, which can be stored in a computer readable storage medium, such as a ROM / RAM, a magnetic disk, an optical disk, etc., and includes a number of instructions to make a computer device (which can be a personal computer, a server, or a network device, etc.) execute the methods described in each embodiment or some parts of the embodiments.

[0161] Finally, it should be noted that: the above embodiments are only used to illustrate the technical solutions of the present application, and not to limit them; although the present application has been described in detail with reference to the foregoing embodiments, those skilled in the art should understand that: it can still modify the technical solutions recorded in the foregoing embodiments, or make equivalent replacement to part of the technical features; and these modifications or replacements do not make the essence of the corresponding technical solutions deviate from the spirit and scope of the technical solutions of the embodiments of the present application.

Claims

1. A large language model optimization method based on variance reduction and momentum acceleration, comprising: in the gradient estimation stage of the large language model, at the beginning of each training step, initializing a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query; performing multiple independent query iterations, in each query iteration, calling a perturbation subroutine to generate random seeds for the large language model and store them in the seed list, wherein the random seeds are used to reproduce perturbation vectors without consuming GPU memory for storing high-dimensional vectors; in the perturbation subroutine, sequentially performing forward perturbation, backward perturbation and the forward perturbation on the parameters of the large language model, determining the gradient projection values and storing them in the projection list, and restoring the parameters to the initial state for the next query iteration; after performing multiple independent query iterations, the seed list and the projection list store multiple random seeds and multiple corresponding gradient projection values, respectively; in the weight update stage of the large language model, calling a gradient norm subroutine for each layer of the large language model, the gradient norm subroutine obtains random seeds from the seed list to reset the random number generator, and re-generates the reproduced perturbation vector without the need for memory caching; for each weight matrix, the gradient projection value taken from the projection list and the reproduced perturbation vector are used to determine the variance-reduced gradient estimate, and the selected optimizer is used to complete the optimization fine-tuning of the large language model without relying on backpropagation.

2. The method of claim 1, wherein, the sequentially performing forward perturbation, backward perturbation and the forward perturbation on the parameters of the large language model, determining the gradient projection values and storing them in the projection list, and restoring the parameters to the initial state for the next query iteration comprises: in the perturbation subroutine, applying forward perturbation to the parameters of the large language model, wherein the forward perturbation is θ+∈z, wherein θ is the model weight, ∈ is a settable value, and z is a perturbation vector generated by sampling from a standard normal distribution N(0, 1), performing forward propagation on the large language model after the forward perturbation to obtain a first loss; calling the perturbation subroutine again to apply backward perturbation to the parameters of the large language model, changing θ+∈z to θ-∈z, and performing forward propagation on the large language model after the backward perturbation to obtain a second loss; determining the gradient projection value containing gradient direction information based on the first loss and the second loss, and storing the gradient projection value in the projection list; restoring the parameters of the large language model to the original state by performing forward perturbation on the parameters of the large language model.

3. The method of claim 1, wherein, the optimizer comprises an Adam optimizer; the optimization fine-tuning of the large language model without relying on backpropagation comprises: storing the first-order momentum and second-order moment information in the optimization fine-tuning stored by the Adam optimizer to smooth and accelerate the convergence process of the optimization.

4. The method of claim 1, wherein, The variance-reduced gradient estimate for each layer weight matrix is determined by multiplying the gradient projection value taken from the projection list with the replicated perturbation vector. For the i-th layer weight matrix, the gradient projection value taken from the projection list is multiplied with the replicated perturbation vector, and the gradient estimate for the i-th layer is obtained by accumulating the q independent query iterations, and the gradient estimate is divided by q to obtain the gradient average value of q independent query iterations.

5. The method of claim 3, wherein, On the basis of the optimizer being an Adam optimizer, the storage state of the Adam optimizer is quantized, including: The storage state of the Adam optimizer is dequantized to restore it to the precision of FP32, FP16 or BF16. After the optimization and update of the parameter weight in the large language model, the storage state of the Adam optimizer is quantized to FP8 again to reduce the GPU memory required for storing the Adam optimizer.

6. The method of claim 1, wherein, The optimizer further includes a stochastic gradient descent optimizer and a stochastic gradient descent optimizer with momentum.

7. A large language model optimization system based on variance reduction and momentum acceleration, comprising: A gradient estimation stage module for the gradient estimation stage of a large language model, at the beginning of each training step, initializing a seed list for storing random seeds of each layer of the large language model and a projection list for storing gradient projection values of each query, performing multiple independent query iterations, in each query iteration, calling a perturbation subroutine to generate random seeds for the large language model and store them in the seed list, wherein the random seeds are used to reproduce perturbation vectors without consuming GPU memory for storing high-dimensional vectors, in the perturbation subroutine, sequentially performing forward perturbation, backward perturbation and the forward perturbation on the parameters of the large language model, determining the gradient projection value and storing it in the projection list, and restoring the parameters to the initial state for the next query iteration, after performing multiple independent query iterations, the seed list and the projection list store multiple random seeds and multiple gradient projection values corresponding thereto respectively; A weight update stage module for calling a gradient norm subroutine for each layer of the large language model in the weight update stage of the large language model, the gradient norm subroutine obtains random seeds from the seed list to reset the random number generator, re-generates the replicated perturbation vector without the need for memory caching, for each layer weight matrix, determines the variance-reduced gradient estimate by multiplying the gradient projection value taken from the projection list with the replicated perturbation vector, and completes the optimization and fine-tuning of the large language model without relying on backpropagation through the gradient estimate and the selected optimizer.

8. A storage medium having stored thereon a computer program product, characterized in that, The program is executed by the processor to implement the steps of the method of any one of claims 1-6.

9. A computer program product embedded with instructions on a storage medium, the instructions implementing the steps of the method of any one of claims 1-6.

10. An electronic device comprising: at least one processor, and a memory communicatively connected to the at least one processor, wherein the memory stores instructions executable by the at least one processor, and the instructions are executed by the at least one processor to enable the at least one processor to perform the steps of the method of any one of claims 1-6.

Images