Federated Fine-tuning of Large Language Models under Heterogeneous Tasks and Client Resources

Authors: jiamu  bai
daoyuan  chen
bingchen  qian
liuyi  yao
yaliang  li
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 08:58:24

Abstract: Federated Learning (FL) has recently been applied to the parameter-efficient fine-tuning of Large Language Models (LLMs). While promising, it raises significant challenges due to the heterogeneous resources and data distributions of clients.This study introduces FlexLoRA, a simple yet effective  aggregation scheme for LLM fine-tuning, which mitigates the "buckets effect" in traditional FL that restricts the potential of clients with ample resources by tying them to the capabilities of the least-resourced participants. FlexLoRA allows for dynamic adjustment of local LoRA ranks, fostering the development of a global model imbued with broader, less task-specific knowledge. By synthesizing a full-size LoRA weight from individual client contributions and employing Singular Value Decomposition (SVD) for weight redistribution, FlexLoRA fully leverages heterogeneous client resources. Involving thousands of clients performing heterogeneous NLP tasks and client resources, our experiments validate the efficacy of FlexLoRA, with the federated global model achieving consistently better improvement over SOTA FL methods in downstream NLP task performance across various heterogeneous distributions. FlexLoRA's practicality is further underscored by our theoretical analysis and its seamless integration with existing LoRA-based FL methods, offering a path toward cross-device, privacy-preserving federated tuning for LLMs.

Main paper: 
 introduction: 
 Large Language Models (LLMs) have propelled advancements in natural language processing (NLP), offering breakthroughs in various tasks [45] . Finetuning LLMs on specific datasets enhances their applicability [12, 33] , yet collecting such datasets raises concerns regarding cost and privacy [29, 6] .Researchers have turned to Federated Learning (FL) as a means to fine-tune LLMs using more data across distributed clients without compromising data privacy [28, 2, 43, 32] . In these settings, parameter-efficient fine-tuning techniques [34] , particularly Low-Rank Adaptation (LoRA) [16] , become attractive for reducing computational and communicational burdens [44, 41, 8] . * Equal contribution. Work done during Jiamu Bai's internship at Alibaba Group.38th Conference on Neural Information Processing Systems (NeurIPS 2024). 3 Methodology of FlexLoRA 
intrinsic dimension and generalization: 
 Fine-tuning LLMs to enhance task-specific performance inevitably encounters cost of reduced generalization ability: a trade-off supported by the "no-free-lunch" theorem and empirical evidence usually called "alignment tax" of LLM [40, 31] . The generalization capability of LLMs is influenced by complexity of applied tasks and their solution spaces, which can be characterized by the concept of an intrinsic dimension -typically far smaller than the total number of model parameters [1] .The insight of intrinsic dimension informs the design of LoRA to fine-tune LLMs' pre-trained weights in a parameter-efficient manner, utilizing compact and low-rank matrices. Specifically, matrices A ∈ R r×p , B ∈ R d×r are introduced, where r denotes the rank that encapsulates intrinsic dimension. These matrices form a low-rank approximation for tuning original weights W 0 as h = W 0 x + sBAx, where x is the input of the parameter to be tuned, h is the output, and s is a scaling constant. Previous studies show that different ranks produce weights with attributes particularly tailored to specific downstream tasks [16, 18] . Consequently, the rank value plays a critical role in not only task-specific solution subspaces but also in determining a model's ability to generalize to various tasks.In scenarios where clients have highly heterogeneous task and resource distributions, a uniform LoRA rank usually does not suffice for model performance, especially in its zero-shot generalization ability for unseen clients and tasks. Employing a small LoRA rank potentially leads to under-fitting in a global context by capturing only a subset of task-specific features, while a large rank is usually infeasible due to the "bucket effect" of existing FL solutions constrained by least-resourced clients.FlexLoRA emerges as a solution to this dilemma by dynamically adjusting the rank in response to the variability in local client resources. By increasing the LoRA rank for clients with greater resources to contribute more global knowledge, FlexLoRA enhances the model's ability to generalize across diverse data distributions without sacrificing local performance accuracy. This strategy allows for federated fine-tuning of LLMs to navigate between the extremes of task-specific optimization and generalization to unseen clients and tasks.Traditional FL methods like FedAvg aggregate local LoRA weights by computing a weighted average of the decomposed matrices A and B as B g = (m i=1 n i B i l )/( m i=1 n i ), A g = ( m i=1 n i A i l )/( m i=1 n i ), where B g , A g are the global LoRA decomposed matrices, and B i l , A i l are the local LoRA decomposed matrices of i-th client, n i is the size of the i-th client's local training dataset, m is the number of FL clients. However, this scheme is restricted by the lowest LoRA rank among participating clients for aggregation compatibility, which makes it hard to capture the full diversity of client contributions and fully utilize ample client resources.FlexLoRA takes a different yet simple approach to enable decomposed matrix with different LoRA ranks to be mixed together. Specifically, it first forms a low-rank approximation of the LoRA matrix for each client, W i l , before computing the weighted average:W g = ( n i W i l )/( m i=1 n i ) = ( n i sB i l A i l )/( m i=1 n i ).After the weighted average with heterogeneous LoRA ranks, the resulting global LoRA weight W g is decomposed using SVD. Then the SVD components U, Σ, V are redistributed to clients in a low-rank approximation that preserves as much information of W g as possible meanwhile based on clients' local resources characterized by r i :SVD(W g ) = U ΣV T , W i g = U [:, : r i ]Σ[: r i , : r i ]V [: r i , :] T ≈ W g ,where U , Σ, and V T are the SVD components of W g , the r i within [] indicates the indexing operator of each client to select their singular vectors corresponding to top r i singular values. As a result, client i receives the aggregated knowledge W i g from server and incorporates W i g into its local LoRA weightW i g = sB i g A i g with B i g = U [:, : r i ]Σ[: r i , : r i ]/s, and A i g = V [: r i , :] T .The local training then proceeds as similar to those in standard FL approaches, using W i l as the local weights to be tuned. This aforementioned process is repeated until convergence is achieved or a predetermined number of rounds is completed.To fully utilize the local resource of a local client, we adhere to the principle of allocating the highest feasible rank given a client's resource budget, which is motivated by our empirical finding that larger ranks generally yield better generalization. Figure 1 demonstrates that models trained with uniformly higher ranks outperform those with lower ranks under conventional parameter-average aggregation schemes. For single-client performances, the zero-shot performance is also boosted in the majority of the cases. The Table 13 and Figure 9 from Appendix J show that performance improves with higher LoRA ranks uniformly regardless of the tasks assigned to each client. While there might be an ideal LoRA rank that maximizes a single client's performance-potentially as high as 200-practical resource limitations may necessitate settling for a lower rank, such as 8. Therefore, we adopt the principle of setting the rank to be as large as possible to completely utilize the resources in FlexLoRA, which is easy to implement and under low risk of overfitting as Occam's razor suggests.In Appendix B, the overall procedure of FlexLoRA and its core function of server update are summarized in Algorithm 1 and Algorithm 2 respectively. FlexLoRA optimally leverages the inherent characteristics of LoRA, boosting model generalization effectively by increasing local ranks, while without sacrificing overall training efficiency. Compared to FedAvg and homogeneous rank-based FL methods, FlexLoRA incorporates a lightweight SVD procedure, but the overhead from SVD is negligible compared to the local LLM training procedure. Moreover, the SVD is performed only once per round and is independent of client numbers. Notably, FlexLoRA enables heterogeneous ranks without the need for any additional hyperparameter tuning. As we empirically demonstrate in Table 4 , the improved convergence rate, thanks to larger ranks, more than compensates for the extra overhead introduced by training on more parameters per round, resulting in a net gain in overall efficiency and a reduced overall time to completion. These features enhance its efficiency and scalability in cross-device FL settings where thousands or millions of devices are involved.We analyze the generalization ability of FlexLoRA by extending Baxter's model of learning [3] . Here, h W (•) represents the hypothesis generated by the model with LoRA weights W , and f W ; (x, y) denotes the loss function for a single data point (x, y). The expected loss is denoted as L(W g ) ≜ E (x,y)∼Di f W ; (x, y) . The two key assumptions underpinning our analysis are as follows:Assumption 1. The following Lipschitz conditions hold: |f (W ; x, y)-f (W ′ ; x, y)| ≤ L f ||W -W ′ || and ||h(W ; x) -h(W ′ ; x)|| ≤ L h ||W -W ′ ||.Following current analysis in SOTA methodologies in LLM research, we assume Lipschitz conditions in our analysis for f and h [27, 17, 13] . This assumption indicates that f and h are Lipschitz continuous with respect to the LoRA weights W , ensuring the stability of the loss landscape. For simplicity, we denote SVD(W g , r i ) as using the top r i singular values and the corresponding singular vectors to approximate W g . We denote the parameter solution space of W g to be k. Next, due to the federated average and indexing operation based on the largest singular values, we assume that the dissimilarity between the global model and its rank-constrained approximation can be bounded: Assumption 2. The LoRA weights can be bounded in a ball with radius R, and the error induced by the SVD approximation for each client is bounded by a constantϕ i as ||SVD(W g , r i ) -W g || ≤ ϕ i . Theorem 1.Under Assumptions 1 and 2, with probability at least 1 -δ, there exists a sample sizeÑ = O( k |C|ϵ 2 log( RL f L h ϵ-2ϕ i L f L h ) -log δ |C|ϵ 2 )such that for all W g , the bound ||L(W g ) -L(W ′ g )|| ≤ ϵ holds when the number of local data samples for each client i exceeds Ñ .Detailed proof is in Appendix C. This theorem suggests that the generalization ability of the global model is influenced by the LoRA rank r i chosen by each local client. Specifically, as ϕ i is the error bound of SVD approximation, increasing rank r i makes the approximation more accurate, thus reducing ϕ i . Consequently, this reduction in error bound decreases the requisite number of samples, denoted as Ñ , required for effective generalization. Moreover, an increase in the number of clients |C| also contributes positively to the generalization of the federated model in the order of O( 1 |C| ), a stronger impact than ϕ i whose impact is in a logarithmic fashion O(log( 1 ϕ i )). Collectively, FlexLoRA is effective in cross-device FL settings, where the generalization capability of the global model can be significantly enhanced by the participation of massive clients (larger |C|) with heterogeneous resources (larger r i ). Note that Assumption 1 is standard in FL literature [22] , and our proof do not rely on simplified assumptions that often do not hold in cross-device cases, such as identically distributed data. We provide empirical support for distribution-related generalization ability, the effect of SVD (Assumption 2), and the effect of client numbers in Sections 4.3, 4.5 and 4.6 respectively.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Answer: [Yes]Answer: [Yes] Justification: We have released the code and data used in our paper.Guidelines:• The answer NA means that paper does not include experiments requiring code. • Providing as much information as possible in supplemental material (appended to the paper) is recommended, but including URLs to data and code is permitted. 
setup for cross-device fl environments: 
 Resource Heterogeneity. We make FL clients resource-heterogeneous by crafting four distinct LoRA configurations as listed in Table 1 . Type 1, 2, and 4 assign the same LoRA on all tunable layers, while Type 3 assigns small ranks on attention layers and large ranks on FFN layers, following the design of the MAM adapter [14] . Clients are randomly assigned a configuration type, simulating four types of heterogeneous resource environments similar to [5] . As shown in Figure 3 , we consider uniform resource distribution where each LoRA configuration type is equally likely to be assigned to each client, heavy tail resource distribution where either Type 1 or Type 4 is dominant, and normal distribution where the LoRA configuration types are normal distribution and Type 2 and Type 3 are dominant. A comparison of the active memory cost with different LoRA configurations is shown in Appendix D. Besides, we add LoRA on top of all the linear layers of LLMs based on the empirical results in Appendix E.3.Task Heterogeneity. We further make FL clients task-heterogeneous by utilizing the natural instruction dataset [39] . The dataset consists of over 1600 distinct natural language tasks that come from 76 NLP task types and is split based on its meta-info of the belonging NLP tasks, such that each client holds a unique task to mirror a task heterogeneous environment. Notably, while the FL setup includes over 1600 clients, the distribution of 76 task types across these clients means that some will inherently share similar local data distributions, thereby mirroring the natural variability and overlapping task characteristics often encountered in real-world settings. Unless stated otherwise, we conduct all our FL experiments on this dataset and adopt DataJucier (1.3B) as our foundation models [6] , chosen for its suitability for edge devices with constrained resources. More details about data preparation are included in Appendix E.2. Baselines with Homogeneous Rank. We adopt FedAvg [28] , FedIT [43] , and SLoRA [2] We note that both FlexLoRA and HETLORA are able to be plugged into the above-mentioned FL methods with homogeneous LoRA ranks. In our experiment, for each baseline with homogeneous rank, we also examine their performance after integration by either FlexLoRA or HETLORA.Evaluation Setup. Our initial examination focuses on the generalization capabilities of global models to unseen clients by deploying the models to clients with unseen data distributions. The unseen clients are newly sampled clients from the next communication round. This assessment allows us to measure zero-shot performance, a key indicator of a model's ability to generalize beyond the data available during the training phase. This approach is to simulate the real FL setting, where well-trained global weights will be deployed to new clients rather than clients participating in the previous round. Specifically, we investigate the performance of global models trained with baseline methods both with and without the integration of FlexLoRA and HETLORA under four distinct resource heterogeneity scenarios.Overview Performance Comparison. Table 2 displays the zero-shot performance under Rouge-L scores on the test set from unseen clients, facilitating a comparison of the generalization capabilities across different federated global models. It is observed from the table that in most of the cases, methods with heterogeneous LoRA ranks have better performance than that of homogeneous ranks, indicating that heterogeneous LoRA ranks enhance the clients with larger ranks to fully exploit their capability. Furthermore, among the heterogeneous LoRA rank methods, our proposed FlexLoRA consistently outperforms HETLORA across all resource distribution settings. This shows that FlexLoRA is able to take advantage of heterogeneous resource distribution, and is more capable of leveraging the general information from heterogeneous LoRA configurations. Effect of Specific Resource Distributions. To gain further insight into the effect of FlexLoRA and the effect of heterogeneity of LoRA ranks, in Table 10 in Appendix F, we list the percentage improvement for each FL methods when incorporating FlexLoRA in comparison with the respective standard homogeneous rank implementations. Notably, after integrating FlexLoRA, the average performance gains are 2.14% for FedAvg, 0.86% for FedIT, and 1.94% for SLoRA. These enhancements lend empirical support to our theoretical generalization analysis that clients utilizing higher LoRA ranks tend to exhibit improved generalization abilities. The most substantial performance improvements are observed in the heavy-tail-strong resource distribution, followed by the normal distribution. This is consistent with our expectations since the heavy-tail-strong distribution predominantly comprises clients with Type 4 LoRA configurations (rank 200). The limited presence of Type 1 clients (rank 8) in the normal distribution minimizes the risk of the global model being excessively influenced by task-specific LoRA weights. Therefore, FlexLoRA is able to leverage the heterogeneous resource distributions to boost the zero-shot generalization, and the gain from integrating FlexLoRA is directly related to the ratio of clients with heavy resources.Evaluation Setup. To assess the natural language understanding capabilities of the FlexLoRAenhanced global model, we evaluate its performance on a range of downstream NLP tasks. Specifically, the model is tested on the English Track of the evaluation tasks from [39] , featuring 12 categories and 119 tasks. For each task, a random sample of 100 data points is chosen for testing. We summarize the average percentage improvement achieved by integrating FlexLoRA across different resource distributions in Table 3 .Effect of Specific Resource Distributions. In the majority of cases, the global models augmented with FlexLoRA demonstrate marked improvements over the vanilla implementations of FedAvg, FedIT, and SLoRA by up to 1.99%, suggesting that the FlexLoRA also improves the natural language analysis capabilities. An exception is noted in the SLoRA on the heavy-tail-light distribution, potentially due to the predominance of the Type 1 LoRA configuration (rank 8), which may limit the overall language processing capabilities when such clients are disproportionately represented. This configuration's minimal rank assignment across all linear layers suggests that the local weight aggregation on the server side might not fully leverage the capabilities of clients with larger resources, potentially detracting from the model's performance on certain tasks.Effect of Specific Language Tasks. We further illustrate the task-specific improvements of integrating FlexLoRA in comparison with the standard FedAvg configuration for various resource distributions in Figure 4 . The task-wise improvement figures for other FL methods are included in Appendix G. We observe that the global models trained using the FlexLoRA aggregation scheme generally outperform others on tasks requiring the parsing of logical relationships between sentences. Particularly, it gains improvements at most 4% in the overlap extraction task, and around 2.5% in the textual entailment, cause-effect classification, and dialogue act recognition task, verifying again the effectiveness of FlexLoRA. SVD on Convergence. FlexLoRA's aggregation scheme constructs full-size LoRA weights before averaging, unlike the conventional FedAvg method's parameter-wise averaging. To understand the effect of this difference on model performance, we assess the performance of FlexLoRA in a controlled environment using homogeneous LoRA ranks and compare it to the standard FL aggregation scheme. Figure 6 (a) presents the test loss trajectories for FedIT with and without the FlexLoRA enhancement, both utilizing a homogeneous rank of 8. The loss curves for the standard FedIT and FedIT with FlexLoRA closely align, suggesting comparable performance. In a nutshell, we empirically demonstrate that under homogeneous conditions, FlexLoRA's aggregation does not negatively impact model performance compared to traditional methods. Besides, it's worth noting that the test loss for FedIT with a homogeneous rank of 200 is significantly lower, underscoring the benefits of higher rank configurations and evidencing our Theorem 1.SVD v.s. LoRA Rank. To gain further insight into the effect of SVD, we calculate and sort the singular values of the global LoRA weights from largest to smallest. We focus on specific layers where LoRA is applied within the transformer blocks 1, 8, and 14 (each block includes one attention layer and one FFN layer). Figure 6 (b) displays the scale of singular values and the error ratio between the global LoRA weights approximated by the top i singular values and the full-rank global LoRA weights in the FedIT setting with a heavy-tail-strong resource distribution. The approximation error is quantified as the norm of the difference between the approximated and full-rank weights. The error curves for q proj layers across all transformer blocks nearly overlap. With the weight approximated from the top 30 ranks, the error ratio is as low as 0.16, suggesting the approximated weights are in close proximity to the actual full-rank weights, lending empirical support to our Assumption 2.Larger Model Size and Lower Degree of Task Heterogeneity. While the aforementioned experiments with 1.3B LLM and meta-task dataset split effectively showcases FlexLoRA's capabilities against baselines in a highly heterogeneous cross-device environment, real-world settings may be relaxed involving fewer clients with a mixture of tasks and larger LLM. To better evaluate FlexLoRA's performance in such scenarios, we expanded our study to include settings where each client manages not just one specific "meta" task but a variety of different tasks. We use Dolly-15K dataset [9] , which supports instruction tuning and includes 8 tasks in total. We distributed this dataset among 200 clients using a Dirichlet distribution with α = 0.5 to simulate the non-IID data distributions. Moreover, we also incorporate the most cutting-edge LLaMA-3 [37] with 8B parameter size as our foundation model to assess FlexLoRA's effectiveness with advanced, larger-scale models.Table 12 in Appendix I illustrates FlexLoRA's efficacy under both smaller and larger foundation models within the mixture of task settings. Compared with LLaMA-3 (8B), finetuning on DataJucier(1.3B) demonstrates more generalization improvement for unseen clients. This enhanced performance when using the smaller DataJuicer model suggests that FlexLoRA is particularly effective for foundation models that are scalable to edge devices. Such ability is instrumental in maximizing the utility and efficiency of smaller models in resource-constrained environments. The "Cost all " indicates total cost calculated as multiply of R and Cost R , and by setting the result of "Homo Rank" as the baseline. From the results, we can see that FlexLoRA achieves faster convergence with slightly increased parameter percentages under heterogeneous resource distributions. For the overall efficiency, "Heavy-Tail-Light" has a total reduction of 49.4%, and "Uniform" has 66.9%, indicating good scalability of FlexLoRA.Effect of Larger Client Number. We designed an experiment to empirically validate Theorem 1 and demonstrate how FlexLoRA performs with varied client numbers. We use subsets of 10, 50, and 100 clients from a pool of 200 clients on Dolly-15K dataset and DataJuicer (1.3B), with FedAvg as baseline FL method and each FL round always sampling 10 participants. From Figure 5 , there is a marked improvement in generalization as the number of participating clients increases. Moreover, from Theorem 1, we can derive the functional relationship between client number |C| and the generalization loss ϵ as |C| = A 1 /ϵ 2 (A 2 -log(ϵ -A 3 )), where A 1,2,3 are constants absorbing the other factors impacted by specific resource distribution in this experiment. We thus fit these coefficients using the derived form and empirical observations, and find that the dotted curves gain good fitness for all the tested distributions. There results affirm the preciseness of our Theorem 1 again, indicating the suitableness of FlexLoRA for cross-device FL scenarios where leveraging a broad client pool can boost the generalization across diverse data distributions. 
conclusion: 
 In this work, we propose a simple yet effective method named FlexLoRA to address the challenges posed by resource and data heterogeneity among clients during the federated fine-tuning of LLMs. By leveraging larger local LoRA ranks, FlexLoRA not only improves the generalization ability of the global model but also ensures that all clients, irrespective of their resource capabilities, can contribute meaningfully. Theoretical analysis and extensive experiments verify the effectiveness and scalability of FlexLoRA. Due to resource limitations, we have not tested the LLaMA-3 model on thousands-client scenarios, which we leave as future work. We hope this study can enlighten more future research and development in data-efficient and privacy-preserving enhancement of LLMs.Compute each client LoRA weight based on their resource limitation:B i g = U [:, : r i ]Σ[: r i , : r i ]/s A i g = V [: r i , :] T 10 return {B i g , A i g } i∈CWith the inequalities introduced by Assumptions 1 and 2, we have:∆(W i g -W i, ′ g ) = 1 n E x,y∼Di (f (W i g ; x, y) - f (W i, ′ g ; x, y) ≤L f ||h W i g -h W i, ′ g || ≤L f L h ||SVD(W g , r i ) -SVD(W ′ g , r i )|| =L f L h ||SVD(W g , r i ) -W g -SVD(W ′ g , r i ) + W ′ g + W g -W ′ g || ≤L f L h ||SVD(W g , r i ) -W g || + L f L h ||SVD(W ′ g , r i ) -W ′ g || + L f L h ||W g -W ′ g || ≤L f L h 2ϕ i + ||W g -W ′ g || . (2)Let the parameter solution space of W g denote as k. We can get an ϵ-covering in metric ∆(W g,i -W ′ g,i ) if we select a covering in the parameter space with ||W g -W′ g || equal to ϵ L f L h -2ϕ i . Therefore, the covering number of H |C| , denoted as B(ϵ, H |C| ) is: log B(ϵ, H |C| ) = O k log( RL f L h ϵ-2ϕ i L f L h ) .According to previous studies [3, 35, 7] , there exists Ñ andÑ = O 1 nϵ 2 log B(ϵ,H |C| ) δ = O( k |C|ϵ 2 log( RL f L h ϵ-2ϕ i L f L h ) -log δ |C|ϵ 2 ). 
experimental setting/details: 
 Question: Does the paper specify all the training and test details (e.g., data splits, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?Answer: [Yes]Justification: Implementation details can be founded in both Section 4 and Appendix E.1.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them. • The full details can be provided either with the code, in appendix, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes]Justification: The error bars are shown in the tables of main results, and significance analysis is in Appendix F.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper. • The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions). • The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.) • The assumptions made should be given (e.g., Normally distributed errors).• It should be clear whether the error bar is the standard deviation or the standard error of the mean. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: We have not found such risks in this paper.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The authors properly credit the assets used in the paper.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: We calculate the % trainable parameters and active memory cost for LoRA under different ranks compared with full finetuning in our experiment setting, summarized in Table 5 . Although the size of LoRA weights is small compared with the pretrained model because LoRA only tunes a small proportion of the trainable parameters, LoRA tuning is still much more efficient than full parameter finetuning, which highlights the necessity to scale LoRA rank size on clients with diverse computation resources. Table 6 illustrates the empirical time consumption for tuning client with LoRA and full finetuning. For the sparse fine-tuning stage of SLoRA, we set its sparsity corresponding with the resource cost of the client's LoRA configuration shown in Table1. For example, if the client is assigned with Type 1 LoRA configuration, we will generate a mask with a sparsity of 0.12%. For HETLORA, following the original paper, we adopt 0.99 as the decay factor for rank pruning and search the regularization factor from a range of {5e-2, 5e-3, 5e-4}. For both the experiments integrated with and without FlexLoRA, we grid search their learning rates from a range of {5e-2, 5e-3, 5e-4} for FedAvg, and {5e-4, 1e-4, 5e-5, 1e-5} for FedIT and SLoRA, both accompanied with a linear scheduler which decays from the initial learning rate to 0.Echoing findings from [39] on model performance saturation with few data instances, we randomly sample 10% of each task's data to scale experiments. For all the datasets we used, data for each client is partitioned into training, validation, and testing sets in a ratio of 8:1:1. In the initial setup of the natural instruction dataset, each task is accompanied by its definition, along with a positive example and a negative example. To adapt this dataset for model fine-tuning through instruction tuning, we transform the structure so that the task definition serves as the direct instruction for each data instance. This restructuring is illustrated by comparing the original and modified data formats in the natural instruction dataset, as shown in Table 7 and Table 8 . This adjustment ensures that each instance is now explicitly aligned with its instructional context, facilitating a more straightforward and effective fine-tuning process.We explore the influence of insertion layers of LoRA on the performance of the fine-tuning. While [16] and [8] add LoRA on all the attention layers, [43] adds LoRA on all the linear layers, i.e., the attention layers and Feed-Forward Network layers. We experimented with both approaches to adding LoRA. From Table 9 , we demonstrate that adding LoRA to both the attention layers and Feed-Forward Network layers boosts generalization performance and adopt this setting in our experiments. The statistical test results presented in Table 11 provide a quantitative comparison between the performance of FlexLoRA across different FL baselines to be integrated, namely FedAvg, FedIT, and SLoRA. The p-values obtained from the statistical tests are used to determine whether the observed differences in performance are statistically significant. Overall, the results indicate that FlexLoRA is a robust approach that can yield performance improvements in federated tuning environments for LLMs, especially in scenarios characterized by a resource-heterogeneous distribution of client resources.Figure 7 illustrates the performance of global models on natural language tasks when trained with FlexLoRA in the FedIT setting. In line with results from the FedAvg setting, global models that leverage heterogeneous ranks in the FedIT setting also perform better on tasks involving logical relationships between sentences, such as overlap extraction, textual entailment, cause-effect classification, and dialogue act recognition. However, the global models underperform in word-level analysis tasks, such as word analogy and keyword tagging. We hypothesize that this underperformance in the FedIT setting is due to the local clients being trained with the Adam optimizer rather than the conventional SGD optimizer. The top decoding block of the LLaMA transformer, which processes word-level information, likely overfits specific information because of the momentum component inherent to the Adam optimizer.Similar to Figure 6 (b), Figure 8 also verifies the empirical information loss from SVD by illustrating the singular value distribution and error ratio of k proj weights. k proj weights shares similar trends with q proj weights regarding the singular value distribution and error ratio. Table 12 provides the result of standard FedAvg vs FlexLoRA-integrated FedAvg under uniform and heavy-tail-light resource distributions. The foundation model includes DataJucier(1.3B) and LLaMA 3 (8B). Table 13 below shows the single client performance under a small rank of 8 and a large rank of 200. This data shows that performance improves with higher LoRA ranks uniformly regardless of the tasks assigned to each client, motivating us to allocate the highest feasible rank given a client's resource budget. Figure 9 illustrates all the clients' performance under rank 8 and rank 200. Besides directly applying LoRA to each clients, we also investigate the efficacy of existing dynamic LoRA methods in central learning. We adapts concept of ReLoRA [24] into our setting and explore its utility when combining with FlexLoRA. ReLoRA trains LoRA for several epochs, uploads the LoRA weight to the pretrained weights, then initializes a new LoRA module and trains based on the new frozen weights, repeating this process for several LoRA modules. We create a baseline by incorporating the concept of ReLoRA, uploading aggregated LoRA weights to the pretrained models after several communication rounds.However, we observed a slower convergence speed than the regular FlexLoRA. We show in Table 14 below that incorporating the step of "uploading LoRA weights" does not lead to better performance. This suggests that observations and improvements seen in centralized dynamic LoRA methods may not translate directly to federated learning due to the unique challenges posed by heterogeneity in FL, which underlines the importance of tailored solutions like FlexLoRA for effectively managing heterogeneity in FL environments. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.https://huggingface.co/datajuicer/LLaMA-1B-dj-refine-150B 


Decision: This study introduces FlexLoRA which is a federated finetune algorithm based on LoRA. It operates by maintaining a full LoRA module on a central server while decomposing it via SVD as heterogeneous rank LoRA which are tailored to specific client tasks and their capabilities. The authors have conducted thorough experiments and provided a detailed theoretical analysis to demonstrate the proposed method's efficacy and scalability.

All reviewers are positive to the paper and agree that the paper makes significant contributions including it is well-motivated and straightforward and provides comprehensive experiments.

The main concerns raised by reviewers including privacy issues of FlexLoRA, its overhead and lack of a few baselines were properly addressed during the discussion phase.

Overall, I agree with the reviewers and recommend its acceptance

Reviews: review: paper_summary: paper_summary: This paper proposes a LoRA-based federated fine-tuning algorithm called FlexLoRA. FlexLoRA keeps full-size LoRA modules on the server, and decomposes them into heterogeneous rank LoRA modules according to the client task and capacity. The heterogeneous local LoRA modules will be transferred back to full-size and then aggregated on the server. FlexLoRA can make LoRA-based federated fine-tuning more flexible and easier to be implemented.
strength_weakness: strength_weakness: (1)The main concern is the privacy issue of FlexLoRA. Since the server needs to preset all the local LoRA ranks, the server needs to get access to some client information such as the amount of data, computation resources, ..., etc. Will this threaten the privacy of local data?
(2)The experiment shows results of language model around 1B. Why larger models (e.g., llama) and benchmarks (e.g., mmlu, gsm8k) are not shown in the body part. There are some results in the appendix but may need more detailed discussion in the experiment section.
questions: questions: please see the weaknesses
limitations: limitations: N/A
paper_summary: paper_summary: This paper proposes FlexLora to tackle FL's heterogeneous resources and data distribution problem. FlexLora allows for dynamic adjustment of local Lora ranks by employing SVD for weight distribution, improving the global model’s generalization ability.
strength_weakness: strength_weakness: 1.  This paper claims to follow a cross-device setting, but cross-device setting usually involves thousands of clients, while the client number in the experiment of this paper is much smaller. If the method can only be verified in a small-scale experiment, cross-device setting is not a very rigorous statement. Cross-silo setting might be more suitable for this paper.
2.  One of the main contributions of this paper is allowing for dynamic adjustment of local Lora ranks in FL. Some dynamic/flexible Lora methods already exist in central learning. I wonder whether these methods can also be applied to tackle the heterogeneity of downstream tasks in FL. Related discussions or experiments are suggested here.
questions: questions: Please refer to weakness.
limitations: limitations: The authors have discussed the limitations of the model and client scale. For the small client scale, claiming a cross-device setting is not so rigorous. Cross-silo setting might be more suitable.
paper_summary: paper_summary: Under the framework of federated fine-tuning large model, this paper proposes a simple and effective LoRA aggregation method with different ranks, which mainly focuses on the problem of client resource heterogeneity. Specifically, the full-size LoRA is first obtained, and then after normal aggregation, it is decomposed into a matrix BA with the corresponding rank of each client by SVD. Extensive experiments demonstrate the effectiveness of the proposed approach when client resources are inconsistent.
strength_weakness: strength_weakness: 1. Does it increase the communication cost to convert matrix BA to full-size LoRA before uploading?
2. It seems that the configurations of ranks in Table 1 are not very diverse. Can the proposed method adapt to very diverse and more flexible configurations of ranks?
3. Different ranks seem to be specific to client resource heterogeneity, but do not seem to be particularly optimized for task heterogeneity?
4. While there are performance comparisons with the very related HETLoRA which is used to solve the resource heterogeneity problem, should comparison results with the rest of the state of the art FL methods also be given.
questions: questions: As shown in the weaknesses section, we won't add anything more here.
limitations: limitations: This paper lacks a detailed discussion of the limitations and shortcomings of the proposed method.
paper_summary: paper_summary: The use of LoRA in FL is challenged by the heterogeneity of downstream tasks and available resources among clients. Traditional FL methods often use the smallest viable LoRA rank for all clients for aggregation compatibility, which makes it hard to capture the full diversity of client contributions and fully utilize ample client resources. To fully leverage heterogeneous client resources for enhancing the global model's generalization ability, this paper proposes FlexLoRA, which enables the mixture of diverse LoRA weights across individual clients to account for local client resources and task differences by allocating the highest feasible rank given a client's resource budget to ensure all clients contribute effectively regardless of resource capability, with heterogeneous aggregation and redistribution of weights through SVD. Extensive experiments and theoretical analysis verify the effectiveness and scalability of the proposed FlexLoRA.
strength_weakness: strength_weakness: 1.	Numerous typos remain in this paper, including but not limited to: 1) "${Client_j}$" should be corrected to "${Client&nbsp;j}$" in Figure 2; 2) "Effect" should be corrected to "effect" in line 249; 3) "r=200 2" should be corrected to "r=200" in Table 6.
2.	In contrast to HETLORA [8], the proposed FlexLoRA needs to perform SVD instead of redistributing it directly before redistributing the aggregated global LoRA weight to clients. This seems to be because it provides better initialization for the low-rank matrices B and A. However, the authors do not mention this point, and this practice has already been presented in FeDeRA [42], so this may not be considered one of this paper's contributions. 
3.	Some relevant work in this field is missing, including but not limited to: 
[1] Liping Yi, Han Yu, Gang Wang, Xiaoguang Liu, Xiaoxiao Li. FedLoRA: Model-Heterogeneous Personalized Federated Learning with LoRA Tuning. arXiv preprint arXiv:2310.13283, 2023.
[2] Shangchao Su, Bin Li, Xiangyang Xue. FedRA: A Random Allocation Strategy for Federated Tuning to Unleash the Power of Heterogeneous Clients. arXiv preprint arXiv:2310.13283, 2023.
questions: questions: 1.	The method described in Section 3 omits too many preliminaries and algorithmic details, which is reader-unfriendly.
limitations: limitations: 1.	The proposed approach essentially fails into the area of federated fine-tuning of LLMs with parameter-efficient fine-tuning (PEFT) techniques. However, this paper does not discuss any limitations that may be imposed by PEFT especially LoRA used in this paper, e.g., PEFT methods sacrifice performance due to the parameter update is limited to a much lower-dimensional subspace, especially when data across clients are non-IID.

Comments: title: Thank for the response
comment: Dear reviewer,

Thank you for your reply and for making the necessary changes! We noticed that the review modification timestamp in the system may be inconsistent. Please let us know if there have been any oversight changes or if you have any further questions. We will do our best to address all your comments. Thank you once again for your time and effort!
I have read the author rebuttal and made any necessary changes to my review.
Dear Reviewer,

Thank you very much for your response! We greatly appreciate your support throughout the review process. Your constructive feedback have been instrumental in helping us refine our work.
I very appreciate the authors' response. My concerns have been well addressed, and hence I am willing to raise my score.
Dear Reviewer,

Thanks for your quick reply! We are sincerely appreciate for your support in raising the score. Your constructive feedback has been invaluable in enhancing the quality of our work.
The author has addressed my concerns, and I agree with raising the score.