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

Jiamu Bai; Daoyuan Chen; Bingchen Qian; Liuyi Yao; Yaliang Li

Source-linked AI summary

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

Jiamu Bai, Daoyuan Chen, Bingchen Qian, Liuyi Yao, Yaliang Li

arXiv:2402.11505v2Feb 18, 2024cs.CLcs.AI

TL;DR

Federated LLM fine-tuning must handle clients with heterogeneous tasks, data, and resources without letting the least-resourced participants constrain others. FlexLoRA aggregates heterogeneous LoRA updates and redistributes weights with SVD, achieving better downstream performance across heterogeneous settings while relying on stated analytical assumptions.

Problem
Heterogeneous client resources and tasks create a bucket effect in federated LLM fine-tuning, limiting the use of larger LoRA ranks and zero-shot generalization.
Method
FlexLoRA synthesizes a full-size LoRA weight from client contributions and uses SVD-based redistribution to accommodate dynamic local ranks and heterogeneous resources.
Results
FlexLoRA achieves 3.1% and 4% improvements in zero-shot Rouge-L and language-understanding tasks across thousands-client heterogeneous NLP settings.
Takeaways & Limitations
FlexLoRA can be applied in a plug-and-play manner to several SOTA FL baselines and leverages larger local LoRA ranks for broader global-model generalization.
Takeaways & Limitations
The analysis assumes Lipschitz conditions and bounded LoRA weights, including a bound on error from SVD rank approximation.

Abstract

from arXiv · show

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 ``bucket 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.

1 Introduction

Federated LLM fine-tuning combines privacy-preserving distributed data with parameter-efficient LoRA updates, but heterogeneous tasks and resources expose a bucket effect. FlexLoRA addresses this mismatch by aggregating diverse ranks and achieves stronger generalization across large-scale heterogeneous NLP settings.

Federated learning enables privacy-preserving LLM fine-tuning across distributed clients, while LoRA reduces computational and communication burdens.
Resource and task heterogeneity can force conventional FL to use the smallest viable LoRA rank, underutilizing capable clients and weakening generalization.
FlexLoRA mixes client LoRA weights, constructs a full-size global weight, and uses SVD to redistribute rank-adjusted weights to clients.
3.1% and 4% improvements were reported for zero-shot Rouge-L and language-understanding tasks across thousands of clients and heterogeneous resource distributions.
The method is presented as scalable, theoretically supported, and applicable to existing LoRA-based FL methods in cross-device settings.

2 Related Work

Related work applies parameter-efficient fine-tuning methods, especially LoRA, within federated learning to reduce communication costs and improve efficiency. Data and resource heterogeneity remain important challenges, with existing solutions addressing them separately or through heterogeneous-rank LoRA aggregation.

Parameter-Efficient Fine-tuning of LLMs: PEFT methods such as adapter and prefix tuning reduce the computation and storage demands of traditional LLM fine-tuning, while this work employs LoRA for its simplicity and performance.The aggregation scheme is also described as extendable to other PEFT methods by replacing LoRA weights.
PEFT in Federated Learning: Federated PEFT integrates methods such as LoRA into local model updates to minimize communication costs and maximize efficiency.Prior work includes LoRA-based local updates aggregated with FedAvg and combinations of sparse and LoRA fine-tuning.
Data and Resource Heterogeneity in FL: Data and resource heterogeneity remain significant FL challenges, while prior approaches address either data or resource heterogeneity outside the LLM context.HETLORA concurrently allows heterogeneous LoRA ranks through zero-padding and truncation during aggregation and distribution.

3 Methodology of FlexLoRA

FlexLoRA addresses heterogeneous tasks and client resources by aggregating LoRA updates at full-matrix level, then redistributing rank-constrained approximations according to each client’s capacity. Its analysis links higher ranks and more participating clients to improved generalization, while experiments motivate allocating the highest feasible local rank.

Intrinsic Dimension and Generalization: LoRA rank affects the task-specific solution subspace and the model’s ability to generalize across tasks.LoRA represents updates with compact matrices A and B, where rank r captures intrinsic dimension.
Motivation: Uniform ranks can underfit global heterogeneity, while the bucket effect makes large uniform ranks infeasible for least-resourced clients.Small ranks capture only subsets of task-specific features, whereas conventional aggregation is constrained by the lowest participating rank.
Heterogeneous-Rank Aggregation: FlexLoRA forms client LoRA matrices, computes a weighted average of heterogeneous-rank weights, and applies SVD to the resulting global weight.This avoids directly averaging decomposed matrices whose ranks differ.
Heterogeneous-Rank Aggregation: SVD redistributes each global weight using the client’s rank, selecting the corresponding top singular vectors and values for low-rank reconstruction.The reconstructed client weight is incorporated into local LoRA weights before subsequent local training rounds.
Resource-Aware Rank Allocation: FlexLoRA allocates the highest feasible rank under each client’s resource budget because larger ranks generally improve generalization.The method adds lightweight SVD overhead while retaining compatibility with standard federated training procedures.
Generalization Analysis: Under its Lipschitz and bounded-approximation assumptions, increasing client rank reduces SVD error and the required sample size for effective generalization.The analysis also reports that increasing the number of clients improves generalization with an O(1/|C|) effect, stronger than the logarithmic rank-related error effect.

4 Experiments

Experiments simulate heterogeneous client resources and NLP tasks, evaluating FlexLoRA against homogeneous- and heterogeneous-rank baselines. Across unseen-client generalization, downstream tasks, scalability, and system costs, FlexLoRA generally improves performance while leveraging resource heterogeneity.

Experimental setup: The study assigns clients heterogeneous LoRA configurations and unique tasks from a dataset spanning over 1600 tasks and 76 NLP task types.Four LoRA configuration types model resource heterogeneity, while task partitioning creates heterogeneous client data distributions.
Experimental setup: FlexLoRA is compared with FedAvg, FedIT, SLoRA, and HETLORA across four resource-heterogeneity scenarios and unseen-client zero-shot evaluation.The evaluation measures global-model generalization to newly sampled clients with unseen data distributions.
Unseen-client generalization: FlexLoRA consistently outperforms HETLORA across resource distributions, while heterogeneous-rank methods generally outperform homogeneous-rank baselines on unseen-client Rouge-L.The comparison attributes the advantage to allowing higher-rank clients to exploit their capacity and leveraging information from heterogeneous configurations.
Resource-distribution effects: 2.14% for FedAvg, 0.86% for FedIT, and 1.94% for SLoRA are the average gains from integrating FlexLoRA over homogeneous-rank implementations.The largest improvements occur under heavy-tail-strong and normal resource distributions.
Cross-task generalization: At most 4% improvement occurs for overlap extraction, while textual entailment, cause-effect classification, and dialogue act recognition improve by around 2.5%.These gains are reported for tasks requiring logical relationships between sentences.

5 Conclusion

FlexLoRA addresses resource and data heterogeneity in federated LLM fine-tuning by scaling local LoRA ranks and aggregating client contributions. The paper supports its effectiveness through theory and experiments, while identifying large-scale LLaMA-3 evaluation as future work.

Conclusion: FlexLoRA uses larger local LoRA ranks so clients with different resources can contribute meaningfully to a better-generalized global model.The paper supports this claim with theoretical analysis and extensive experiments.
Limitations: Because of resource limitations, the study does not test LLaMA-3 in thousands-client scenarios and leaves that evaluation for future work.
Additional experiments: Appendix F reports improvement percentages and significance tests comparing FlexLoRA with baselines across resource-distribution types.
Additional experiments: Appendices G–J provide task-specific improvements, SVD effects, mixed-task results, and higher-rank performance across NLP tasks.

A Implementation of Figure 1

The Figure 1 experiments compare naive FedIT across LoRA ranks and test FlexLoRA under an extreme heavy-tail resource distribution. The appendix also points to pseudocode describing FlexLoRA’s federated procedure.

Figure 1 setup: Figure 1 plots zero-shot test loss for naive FedIT with LoRA ranks 1, 8, and 200.
Figure 1 setup: The FlexLoRA experiment assigns rank 200 to every client except one client assigned rank 8, forming an extreme heavy-tail scenario.
Figure 1 setup: The Figure 1 experiments use the same hyperparameters and federated-learning settings as Table 2.
Implementation: The appendix summarizes FlexLoRA’s pseudocode in Algorithms 1 and 2.

C Proof of Theorem 1

The appendix describes FlexLoRA’s federated algorithm, heterogeneous-rank server updates, theoretical proof setup, and efficiency measurements across LoRA configurations. It also documents experimental settings and supporting tables.

Proof: The proof develops covering arguments and inequalities under stated assumptions to analyze bounds on the generalization error.
Algorithm: Algorithm 1 samples clients, updates their local LoRA modules, and performs a server update across communication rounds.
Algorithm: Algorithm 2 takes client LoRA weights, ranks, dataset-size averaging constants, and global settings as inputs for the server update.
Heterogeneous aggregation: FlexLoRA computes client LoRA weights according to resource limitations and aggregates heterogeneous-rank contributions through the server procedure.
Efficiency: LoRA remains more efficient than full-parameter fine-tuning, motivating larger ranks on clients with diverse computation resources.
Experimental details: The appendix includes tables on trainable parameters, memory cost, communication-round speedup, datasets, prompts, and experimental settings.

E.2 Cross-Task Splitter Details

The experiments subsample task data and partition each dataset by client into training, validation, and testing sets. They also compare LoRA placement across attention and feed-forward layers using zero-shot Rouge-L.

Cross-task data split: Experiments randomly sample 10% of each task’s data, with each client’s data split 8:1:1 into training, validation, and testing sets.
LoRA placement: Adding LoRA to both attention and feed-forward layers improves generalization performance relative to applying it only to attention layers.
LoRA placement: Table 9 reports zero-shot Rouge-L for different LoRA insertion-layer choices under a uniform FlexLoRA resource distribution.

F Relative Improvement for Table 2 and Significant Test

FlexLoRA improves integrated FL baselines across heterogeneous resource distributions, with statistical tests indicating that several gains are significant. The strongest reported evidence appears in FedIT under Heavy Tail Strong and normal distributions.

Table 10 summarizes percentage improvements when FlexLoRA is integrated with FedAvg, FedIT, and SLoRA relative to their corresponding baselines.
FlexLoRA consistently outperforms homogeneous baselines across several resource-distribution types, with p-values below 0.05 in multiple cases.
p=0.018 under Heavy Tail Light indicates a statistically significant FedIT improvement for FlexLoRA.
Overall, the results characterize FlexLoRA as robust for performance improvement in LLM federated tuning, particularly with resource-heterogeneous clients.

G Natural Language Task Performance

In the FedIT setting, heterogeneous LoRA ranks improve global-model performance on several tasks involving logical relationships between sentences, while word-level analysis tasks show weaker performance.

Heterogeneous-rank global models perform better on overlap extraction, textual entailment, cause-effect classification, and dialogue act recognition.
The same global models underperform on word-level analysis tasks.

H More Results for Effect of SVD

The SVD analysis examines singular-value distributions and approximation errors, while task-specific comparisons assess FlexLoRA against homogeneous-rank FedIT across resource settings.

Figure 8 verifies empirical information loss from SVD using singular-value distributions and approximation-error ratios for kproj weights.
The kproj weights exhibit trends in singular-value distribution and error ratio similar to qproj weights.
Figure 7 compares task-specific improvements from FlexLoRA with homogeneous-rank FedIT across different resource-distribution settings.

I Experiment Result for Mixed Task Heterogeneity Scenario

Additional experiments compare homogeneous LoRA configurations with FlexLoRA under FedAvg and examine client performance at different ranks. Higher ranks improve single-client performance across assigned tasks, motivating resource-adaptive rank allocation.

Table 12 compares homogeneous LoRA configurations with FlexLoRA under FedAvg on DataJucier 1.3B and LLaMA-3 8B using Dolly-15K.
Higher LoRA ranks improve single-client performance uniformly across the evaluated tasks.
Table 13 reports eight single-client examples comparing homogeneous rank 8 with homogeneous rank 200 under FedIT in one round.
Figure 9 visualizes performance for all clients under rank 8 and rank 200.
The experiments use a frozen approximately 1.35-billion-parameter DataJuicer-1.3B foundation model, with client tunable size determined by each LoRA module.