An Efficient Large Model Training Method for Federated Learning Based on Dynamic Layer Selection

Abstract

The present invention discloses an efficient large model training method for federated learning based on dynamic layer selection, which relates to the technical field of efficient training of federated learning. In the present invention, the server selects some of the most important layers and hands them over to the client for training, that is, part of the computational workload is allocated to the server side, and the client only needs to update the parameter quantities of very few layers, thus greatly reducing the computational pressure and upload pressure on the client, and significantly reducing the computational pressure and communication pressure of the client; the specific selection of some of the most important layers in the present invention is as follows: the server preliminarily trains the model to be trained based on a small amount of local training data, obtains the gradient value of each linear layer of the model to be trained, and calculates the importance score of each linear layer based on the gradient value. The server selects the linear layers corresponding to the top several highest importance scores and the bottom several lowest importance scores as the current important layers; the client only trains some of the most important layers according to the local training results, thus greatly reducing the computational requirements for training and the parameter quantities of the uploaded data. At the same time, the present invention does not require additional changes to the model, so it has strong scalability.

Description

Technical Field

[0001] The present invention relates to the technical field of efficient training of federated learning, and particularly relates to an efficient large model training method for federated learning based on dynamic layer selection. Background Art

[0002] In recent years, large language models have received attention due to their excellent capabilities in natural language processing tasks. These models rely on deep learning algorithms and have demonstrated superior performance in various applications, including machine translation, text generation, and question answering systems. The effectiveness of large models largely owes to their ability to utilize large amounts of high-quality data, enabling them to capture complex language patterns and semantics. With the improvement of computing power and the increase in data volume, the scale of large models is continuously expanding, allowing them to learn more detailed and rich language representations. This increase in scale further enhances the generalization ability of the models and their performance on multiple tasks, thus driving the progress of the natural language processing field.

[0003] However, in specific domains, due to the critical issue of data privacy, obtaining high-quality data for training large models poses a significant challenge. For example, suppose multiple hospitals wish to jointly train a diagnostic large model. Patient medical records are crucial for training such a model, but these records are highly sensitive and private. Due to privacy and regulatory constraints, traditional centralized training methods, i.e., aggregating data on a single server, are not feasible. Additionally, in industries such as finance, law, and government, data privacy and security also pose major obstacles to training large language models. The data in these domains often contains sensitive information, such as personal identity information, trade secrets, or national security information, and its leakage or misuse may bring serious consequences. Therefore, it has become crucial to develop technologies that can effectively train models while protecting data privacy.

[0004] Federated learning is a distributed machine learning technique where the server only receives the gradients uploaded by users to update the global model. Since the parameter server does not directly access user data, it is generally considered that the privacy of user data is protected. Therefore, federated learning can be used in privacy-protected scenarios such as financial services, medical prediction, speech recognition, user analysis, and keyboard prediction. Recent research has explored combining large models with federated learning, which can both train powerful models and respect data privacy. Additionally, federated learning also provides a method to reduce the data storage and processing requirements of the central server, thereby reducing the dependence on computing resources. In federated learning, the distributed nature of the data also helps to enhance the generalization ability of the model, as it can learn from multiple data sources, thus improving the adaptability and robustness of the model in different environments.

[0005] However, training large models in a federated environment encounters some challenges. First, large models have a large number of parameters and thus require a large amount of computing resources, which may exceed the capabilities of client devices. Second, transmitting gradients from the client to the server imposes a significant burden on the client's upload bandwidth. In a real-world scenario, this imbalance is further exacerbated because servers typically have powerful enterprise-level graphics cards, while clients are equipped with less capable consumer-grade graphics cards. In federated learning, the computational load of gradient calculation mainly falls on the client, and the role of the server is limited to aggregating gradients and updating model parameters, resulting in uneven distribution of computing resources.

[0006] To address the problem of high computational pressure on the client, a feasible solution is to use the Parameter-Efficient Fine-Tuning scheme. This scheme uses the LoRA strategy to parallelly connect a new module with a small number of parameters beside the weights of the original model to simulate changes to the original parameters. Although this scheme can greatly reduce the computational amount required for training, it is still a computationally intensive process for the client, and this scheme cannot solve the problem of uneven computational load between the server and the client. To address the problem of high communication pressure on the client for uploading, a feasible solution is to perform gradient compression on the gradient data. However, considering the large number of parameters of the model, the size of the compressed data is still extremely large. Thus, although some technical solutions touch on related problems to a certain extent, the effects are still not significant, and it is difficult to solve the above problems simultaneously. Summary of the Invention

[0007] The present invention provides an efficient large model training method for federated learning based on dynamic layer selection, which distributes part of the computational pressure to the server to significantly reduce the computational pressure and communication pressure of the client.

[0008] The technical solution adopted by the present invention is as follows:

[0009] An efficient large model training method for federated learning based on dynamic layer selection performs the following steps in a federated learning system including a server and a plurality of clients:

[0010] Step 1, the server determines the model to be trained, and the model to be trained is a large language model;

[0011] The server quantizes the model parameters of the model to be trained (the specific quantization bit can be set based on actual application requirements, and 4-bit quantization is recommended to save video memory to the greatest extent) to significantly reduce the data size of the model and optimize the storage and processing efficiency of the model; and freezes all linear layers (i.e., the fully connected layers of the model) of the model to be trained to ensure that these layers will not be updated during the subsequent training process;

[0012] The server fine-tunes the model parameters of the linear layer of the model to be trained based on the LoRA strategy, that is, in the subsequent federated model training, the network model parameters fine-tuned by the LoRA strategy are updated;

[0013] The server then broadcasts the fine-tuned model to be trained to all participating clients to ensure that each client can start training under the same initial conditions;

[0014] Step 2, the server performs dynamic layer selection based on the preliminary training results of the current model to be trained, and sends the layer numbers of the selected network layers of the model to be trained to all participating clients;

[0015] Among them, the server's dynamic layer selection based on the preliminary training results of the current model to be trained specifically includes:

[0016] The server performs preliminary training on the model to be trained based on local training data and obtains the gradient value of each network layer of the model to be trained;

[0017] For the linear layers in the model to be trained, calculate the importance score I of each layer based on the gradient value l , this importance score I l is used to characterize the importance of each layer in the preliminary training of the model; among them, the amount of local data participating in the preliminary training does not exceed the specified value;

[0018] The server selects the linear layers corresponding to the top α highest importance scores I l and the linear layers corresponding to the bottom β lowest importance scores I l as the currently selected dynamic layers, where the values of α and β are preset values, and α and β are less than or equal to the total number of network layers of the model to be trained;

[0019] Step 3, after receiving the layer numbers specified by the server, the client performs local training on the current model to be trained based on the local training dataset of the client, calculates the gradient of the specified network layer as the update gradient of the client and sends it back to the server;

[0020] Step 4, the server aggregates the update gradients sent by the clients to obtain the global gradients of each specified network layer after aggregation, and then obtains the global model of the current model to be trained in the current round of federated learning based on the gradient values of the non-specified network layers of the current model to be trained obtained locally by the server and sends it to all participating clients to obtain the model to be trained in the next round of federated learning;

[0021] Repeat steps 2 to 4 for multiple rounds of federated learning, and stop when the preset training convergence conditions are met (such as the number of federated learning rounds reaches the preset upper limit or the loss value of the model converges), so as to obtain the target model that meets the expected training objectives and performance standards.

[0022] Further, in step 2, the importance score I of each layer l is specifically calculated as follows:

[0023] Before the initial training, the server takes the derivative of all LoRA parameters of the model to be trained, obtaining the gradients ΔB and ΔA of the two LoRA parameters A and B of the LoRA strategy; wherein, the output of each linear layer of the model to be trained represented by the LoRA parameters A and B is: h = W0x + BAx, where x represents the input data of the current linear layer, and W0 is the LoRA weight of the current linear layer; after the initial training is completed, the importance score I is calculated according to the update degree of each linear layer l ;

[0024] Among them,

[0025]

[0026]

[0027] In the above formula, r i represents the update degree of the i-th network layer using the LoRA strategy, that is, the update degree of the i-th linear layer, and P l represents the set of all linear layers using the LoRA strategy in the model to be trained, and |P l | represents the number of linear layers included in the set P l , and τ is a preset constant used to avoid division by zero.

[0028] Further, in step 4, when obtaining the global model of the current round of federated learning, the gradient values of each non-specified network layer of the current model to be trained obtained by the server in the initial training are respectively multiplied by a scaling factor γ, whose value range is (0, 1), and it is recommended to be set to 0.1, and then combined with the aggregated global gradients of each specified network layer to obtain the global model of the current round of federated learning, so as to avoid the problem that the model may have serious overfitting when directly using the gradient values obtained by the server's local initial training.

[0029] The technical solution provided by the present invention at least brings the following beneficial effects:

[0030] In the present invention, the server selects some of the most important layers and hands them over to the client for training, that is, distributes part of the computational workload to the server side, while the client only needs to update the parameter quantities of a very small number of layers, thus greatly reducing the computational pressure and upload pressure on the client, and significantly reducing the computational pressure and communication pressure of the client; in the present invention, the server uses a small part of auxiliary data to pre-compute the gradients of all parameters of the model, and the client only trains some of the most important layers according to the computed results, thus greatly reducing the computational requirements for training and the parameter quantities of the uploaded data. At the same time, the present invention does not require additional changes to the model, so it has strong scalability. BRIEF DESCRIPTION OF THE DRAWINGS

[0031] To more clearly illustrate the technical solutions in the embodiments of the present invention, the following will briefly introduce the drawings required for the description of the embodiments. Obviously, the drawings in the following description are only some embodiments of the present invention. For those of ordinary skill in the art, without creative efforts, other drawings can be obtained based on these drawings.

[0032] Figure 1 is a schematic diagram of the overall process of an efficient large model training method for federated learning based on dynamic layer selection provided by an example of the present invention;

[0033] Figure 2 is a schematic diagram of the principle of the method for calculating the importance score of each fully connected layer adopted in an example of the present invention;

[0034] Figure 3 is a flowchart of an efficient large model training method for federated learning based on dynamic layer selection provided by an example of the present invention. DETAILED DESCRIPTION OF THE EMBODIMENTS

[0035] To make the objectives, technical solutions, and advantages of the embodiments of the present invention clearer, the following will describe the technical solutions in the embodiments of the present invention in detail and completely with reference to the drawings in the embodiments of the present invention. Obviously, the described embodiments are only some embodiments of the present application, rather than all embodiments. Usually, the components of the embodiments of the present invention described and shown in the drawings can be arranged and designed in different configurations. Therefore, the following detailed description of the embodiments of the present invention provided in the drawings is not only to limit the scope of the present application claimed, but only represents the selected embodiments of the present invention.

[0036] The embodiments of the present invention propose an efficient large model training method for federated learning based on dynamic layer selection, aiming to construct a large model federated learning training framework with low computational pressure and communication pressure on the client, and further achieve the training of large language models without infringing on user privacy.

[0037] As a possible implementation, an efficient federated learning large model training method based on dynamic layer selection provided in an embodiment of the present invention includes the following steps:

[0038] Step 1: Prepare the model to be trained:

[0039] In the initial stage of federated learning, the server determines a large language benchmark model, which will be used as the starting model of the federated learning process and will be used in all subsequent training stages; in order to optimize the storage and processing efficiency of the model, the server quantizes the model to 4 bits, which significantly reduces the data size of the model. In addition, all fully connected layers will be frozen to ensure that these layers will not be updated during subsequent training; the server also integrates the LoRA module into the model, which is designed to reduce the number of parameters in the model training process, which means that only the LoRA module will be updated in subsequent training; finally, the model after these preprocessing steps will be broadcast by the server to all clients participating in the training, ensuring that each client can start training under the same initial conditions.

[0040] Step 2: The server performs dynamic layer selection:

[0041] The server uses a small amount of local data (local auxiliary data) to perform preliminary training on the model. The main purpose of this stage is to evaluate the performance of each layer in the model, especially to record and analyze the gradient value of each layer; based on the gradient value of each layer, the server calculates the importance score of each layer I l , these scores reflect the importance of each layer in model training and are ranked accordingly;

[0042] After sorting, the server will select the first α (experience value) layer with the highest importance score and the last β (experience value) layer with the lowest score, considering these layers to be the most in need of training or the most likely to need further optimization. The layer numbers of these layers are then sent to each client to guide them to conduct targeted training.

[0043] Step 3: Train the model locally on the client:

[0044] After receiving the layer number specified by the server, the client will use its own local data set to train these specified layers based on this information. That is, the model to be trained is forward propagated on the local training data set, but only the gradient of the specified layer is calculated when calculating the gradient in the back propagation. The client can ignore the unspecified layers during the back propagation and only calculate the gradient of the specified layer; after the training is completed, the client calculates the difference between the weights before and after training. This difference represents the update amount of the model and is a direct product of local training; then, the client sends these updates back to the server as key data in the global model update process, that is, the client update gradient.

[0045] Step 4, the server aggregates the gradients and updates the global model:

[0046] After the server receives the gradients submitted by all clients, it aggregates this data to form global gradients. This step is the core model synchronization process in federated learning, ensuring that the training results of all clients are integrated;

[0047] Using the aggregated global gradients, the server can update the model parameters by using the Adam optimizer to generate the updated global model;

[0048] After the update is completed, the server will initiate the next round of federated learning, repeating the above steps until the entire model training process converges and finally reaches the expected training goals and performance standards.

[0049] As Figure 1 shown, the core idea of the method in the embodiment of the present invention is that before the start of each round of federated learning, the server uses a small portion of auxiliary data to pre-compute the gradients of all model parameters. These pre-computed gradients are used to evaluate the relative importance of each layer of the model in the current learning context. Based on this evaluation, when the client calculates the gradients using its local data, it only updates the specified partial layers by the server instead of all parameters of the entire model. Subsequently, the server combines the received client gradients with its own pre-computed gradients to update the global model.

[0050] The large language model of the present invention is divided into several layers according to the Transformer decoder layer. The key problem lies in how to measure the importance of each layer based on the gradients and select the layers to be trained accordingly. In the fine-tuning training of the large model, before the start of each round of federated learning, the server first takes the derivative of all LoRA parameters. That is, for all linear layers h = W0x + BAx, the gradients ΔB and ΔA of the LoRA parameters B and A can be obtained, where W0 is the layer weight parameter (abbreviated as LoRA weight) under the LoRA strategy, x represents the input data of the current layer. Since LoRA simulates the change amount of the original linear layer, therefore, the change amount ΔW of the original parameters (weight parameters) in the current training round t can be estimated in the following way:

[0051] ΔW t = B·ΔB + A·ΔA

[0052] Since the parameter value scales of different layers are not consistent, directly comparing the parameter change amounts is not fair. In the current training round, the original parameter W of the linear layer t can be calculated from W0, B, and A. Therefore, the embodiment of the present invention uses the ratio r of the parameter change amount to the size of the original parameter value to measure the update degree of a linear layer:

[0053]

[0054] The overall principle is as Figure 2 shown, where τ is a minimum value (i.e., a preset constant) used to avoid the problem of division by zero. At the same time, since the original parameters of the model adopt 4-bit quantization, they need to be restored to float-type values before calculation for further calculation. For each Transformer decoder layer of the model (i.e., each layer of the model to be trained), the L1 norm of the r values of all linear layers therein can be calculated, and its average value can be statistically calculated to measure the parameter variation of the entire decoder layer. In the embodiment of the present invention, it is assumed that for the decoder layers with a relatively large update amplitude (r), their contributions in the current training stage are also greater. Therefore, it can be considered that the importance of these layers in the current training stage is higher. Therefore, the importance degree I l of the l-th layer can be calculated using the following formula:

[0055]

[0056] where P l represents the set of all parameters using LoRA in the l-th layer of the large model, and |P l | represents the number of parameters in this set. Through this method, the importance of each Transformer decoder layer can be calculated on the server side. Subsequently, the client will only calculate the gradients for the most important several layers and upload them to the server. It should be emphasized that the small amount of auxiliary data sets used on the server side are only used to evaluate the importance of each layer, and the actual training process still depends on the local data of the client. Through this mechanism, the number of parameters that the client needs to train will be significantly reduced, and the amount of gradient data uploaded to the server will also be greatly reduced.

[0057] However, since this method tends to select the layers with relatively large gradient values, a problem may occur: some layers may fall into local optimal solutions, resulting in their abnormally low gradient values and being unable to jump out of the local optimal points for update. To avoid this situation, when selecting layers, the embodiment of the present invention will select the α layers with the largest importance I and the β layers with the smallest importance I as the layers for which the client needs to calculate gradients. This strategy can effectively ensure that each layer is fully trained.

[0058] Meanwhile, during training, it is desirable to make full use of the gradient values calculated using auxiliary data on the server side instead of wasting them. However, since the scale of the auxiliary dataset is relatively small and its data distribution may be inconsistent with that of real user data, directly using these gradient values may cause serious overfitting problems in the model, thus affecting the performance of the final model. To solve this problem, the embodiments of the present invention only use the gradients calculated on the server side to update the layers that have not been trained by the clients in the current round of federated learning, that is, those relatively less important layers. At the same time, the embodiments of the present invention multiply the gradients calculated on the server side by a scaling factor γ to minimize their impact on the model. In this way, the global model of the embodiments of the present invention can still be updated and can explore more updated directions and possibilities. Finally, the server side will use all the collected gradients to update the global model by using the Adam optimizer, and its specific process is as Figure 3 shown.

[0059] Finally, it should be noted that the above embodiments are only used to illustrate the technical solutions of the present invention and are not intended to limit them. Although the present invention has been described in detail with reference to the foregoing embodiments, those of ordinary skill in the art should understand that they can still modify the technical solutions described in the foregoing embodiments or equivalently replace some of the technical features. However, such modifications or replacements do not cause the essence of the corresponding technical solutions to deviate from the spirit and scope of the technical solutions of the embodiments of the present invention.

[0060] The above are only some embodiments of the present invention. For those of ordinary skill in the art, without departing from the inventive concept of the present invention, several deformations and improvements can still be made, and these all belong to the protection scope of the present invention.

Claims

1. An efficient federated learning large model training method based on dynamic layer selection, constructing a large model federated learning training framework with low client computing and communication pressure, which is used to train large language models without infringing user privacy, characterized by: The following steps are performed in a federated learning system consisting of a server and several clients: Step 1: The server determines a model to be trained, which is a large diagnostic model of a large language model for medical prediction; The server quantizes the model parameters of the model to be trained and freezes all linear layers of the model to be trained; The server fine-tunes the model parameters of all linear layers of the training model based on the LoRA strategy for efficient parameter fine-tuning; The server then broadcasts the fine-tuned model to be trained to all clients participating in the training; Step 2: The server performs dynamic layer selection based on the preliminary training results of the current model to be trained, and sends the layer number of the network layer of the selected model to be trained to all clients participating in the training; The server performs dynamic layer selection based on the preliminary training results of the current model to be trained, specifically including: The server performs preliminary training on the model to be trained based on local training data, and obtains the gradient value of each network layer of the model to be trained; For each linear layer of the model to be trained, the importance score I of each layer is calculated based on the gradient value l ; The amount of local data involved in the preliminary training does not exceed the specified value; The server selects the first α highest importance scores I l The corresponding linear layer and the last β lowest importance scores I l The corresponding linear layer is used as the currently selected dynamic layer, where the values ​​of α and β are preset values, and α and β are less than or equal to the total number of network layers of the model to be trained; Step 3: After receiving the layer number specified by the server, the client performs local training on the current model to be trained based on the client's local training data set, calculates the gradient of the specified network layer as the client's update gradient and sends it back to the server; Step 4: The server aggregates the updated gradients sent by the client to obtain the aggregated global gradients of each designated network layer, and then obtains the global model of the current model to be trained in the current round of federated learning based on the gradient values ​​of each non-designated network layer of the current model to be trained obtained by the server in the preliminary training, and sends it to all clients participating in the training to obtain the model to be trained in the next round of federated learning; Repeat steps 2 to 4 for multiple rounds of federated learning and stop when the preset training convergence conditions are met.

2. The method according to claim 1, characterized in that In step 1, when the server quantizes the model parameters of the model to be trained, the quantization bit used is 4 bits.

3. The method according to claim 1, characterized in that In step 2, the importance score of each layer is I l The specific calculations include: Before preliminary training, the server derives all LoRA parameters of the model to be trained to obtain the gradients ΔB and ΔA of the two LoRA parameters A and B of the LoRA strategy; the output of each linear layer of the model to be trained represented by parameters A and B is: h = W0x + BAx, x represents the input data of the current linear layer, and W0 is the LoRA weight of the current linear layer; after the preliminary training is completed, the importance score I is calculated according to the update degree of each linear layer l ; in, Among them, r i represents the update degree of the i-th linear layer using the LoRA strategy, P l Represents the set of all linear network layers that use the LoRA strategy in the model to be trained, |P l | represents the set P l The number of linear layers included, τ is a preset constant used to avoid division by zero.

4. The method according to claim 1, characterized in that In step 4, when obtaining the global model of the current round of federated learning, the gradient values ​​of each non-specified network layer of the current model to be trained obtained by the server in the preliminary training are multiplied by a scaling factor γ, where the value range of the scaling factor γ is (0,1); and then the global gradients of each specified network layer after aggregation are combined to obtain the global model of the current round of federated learning.

5. The method according to claim 4, characterized in that: The scaling factor γ is set to 0.1.

Images