PaperlayerView exact revision on arXiv

Source-linked AI summary

FLoRA: Federated Fine-Tuning Large Language Models with Heterogeneous Low-Rank Adaptations

Ziyao Wang, Zheyu Shen, Yexiao He, Guoheng Sun, Hongyi Wang, Lingjuan Lyu, Ang Li

arXiv:2409.05976v1cs.LGcs.DC

TL;DR

Federated LLM fine-tuning must accommodate private, distributed data and heterogeneous client resources, while existing LoRA aggregation can be mathematically inaccurate and rank-incompatible. FLORA uses stacking-based aggregation to eliminate aggregation noise and support heterogeneous LoRA ranks. It outperforms state-of-the-art methods in both homogeneous and heterogeneous settings, with only marginally higher total communication than FedIT.

  • Problem

    Existing federated LoRA fine-tuning methods introduce mathematically inaccurate aggregation noise and cannot support heterogeneous LoRA ranks across clients.

  • Method

    FLORA aggregates local LoRA modules by stacking them, producing noise-free updates while accommodating heterogeneous ranks.

  • Results

    FLORA outperforms FedIT across evaluated homogeneous settings and maintains performance in heterogeneous settings, while remaining comparable in total communication costs.

  • Takeaways & Limitations

    FLORA enables accurate, lightweight federated fine-tuning of LLMs across clients with heterogeneous data and computational resources.

  • Takeaways & Limitations

    Stacking increases update communication overhead, and stacking uploaded modules creates a privacy concern requiring random rank-1 submodule ordering or standard privacy mechanisms.

Abstract

from arXiv · show

The rapid development of Large Language Models (LLMs) has been pivotal in advancing AI, with pre-trained LLMs being adaptable to diverse downstream tasks through fine-tuning. Federated learning (FL) further enhances fine-tuning in a privacy-aware manner by utilizing clients' local data through in-situ computation, eliminating the need for data movement. However, fine-tuning LLMs, given their massive scale of parameters, poses challenges for clients with constrained and heterogeneous resources in FL. Previous methods employed low-rank adaptation (LoRA) for efficient federated fine-tuning but utilized traditional FL aggregation strategies on LoRA adapters. These approaches led to mathematically inaccurate aggregation noise, reducing fine-tuning effectiveness and failing to address heterogeneous LoRAs. In this work, we first highlight the mathematical incorrectness of LoRA aggregation in existing federated fine-tuning methods. We introduce a new approach called FLORA that enables federated fine-tuning on heterogeneous LoRA adapters across clients through a novel stacking-based aggregation method. Our approach is noise-free and seamlessly supports heterogeneous LoRA adapters. Extensive experiments demonstrate FLORA' s superior performance in both homogeneous and heterogeneous settings, surpassing state-of-the-art methods. We envision this work as a milestone for efficient, privacy-preserving, and accurate federated fine-tuning of LLMs. Our code is available at https://github.com/ATP-1010/FederatedLLM.

1 Introduction

Federated fine-tuning preserves data privacy while adapting LLMs, but existing LoRA-based aggregation introduces mathematical noise and cannot accommodate heterogeneous client ranks. FLORA addresses these issues with noise-free aggregation and supports heterogeneous LoRA configurations.

  • Motivation: Federated learning enables clients to fine-tune LLMs collaboratively without moving private local data.This addresses scenarios where downstream data is distributed across devices, such as hospitals.
  • Prior Work: FedIT reduces computation and communication costs by training and communicating only local LoRA modules.Clients fine-tune LoRA modules locally before sending them to the server for aggregation.
  • Limitations: FedIT independently averages local A and B modules, introducing mathematical errors and noise into the global LoRA update.The resulting inaccurate aggregation can hinder convergence and increase fine-tuning costs.
  • Limitations: FedIT cannot aggregate LoRA modules with heterogeneous ranks required by clients with different data distributions and hardware resources.Clients may need different ranks to avoid overfitting, preserve generalization, or fit available memory.
  • FLORA: FLORA provides noise-free LoRA aggregation and supports heterogeneous ranks across clients.Experiments on multiple LLMs and benchmarks report superior performance in both homogeneous and heterogeneous settings.

2 Preliminaries

LoRA reduces LLM fine-tuning costs by representing parameter updates with low-rank matrices, but FedIT’s independent averaging of those matrices is mathematically inaccurate. This produces noisy updates, slower convergence, higher costs, and no support for heterogeneous ranks.

  • LoRA: LoRA represents a target-module update as ∆W = BA using low-rank matrices A and B instead of directly updating W.A has dimensions r×n and B has dimensions m×r, while the resulting update matches W’s dimensions.
  • LoRA: LoRA rank r is much smaller than the target dimensions, substantially reducing trainable parameters and GPU memory usage.For LLaMA-7B attention modules with rank 16, trainable parameters are reduced to 0.78% of the pretrained model.
  • FedIT: FedIT applies weighted averaging to local LoRA matrices A and B independently after clients fine-tune them locally.This integrates FedAvg with LoRA while restricting training and communication to the adapter modules.
  • FedIT: Independent averaging is inaccurate because only BA represents each client’s model update, whereas separate averaging creates cross-client intermediate terms.These terms introduce unexpected noise into the aggregated update.
  • FedIT: The noisy aggregation grows with client count, while applying scaling factors to both A and B produces p_k^2 coefficients and further increases error.The resulting aggregation slows convergence and raises computation costs.
  • Heterogeneous LoRA: FedIT cannot support heterogeneous LoRA ranks, although data complexity and client memory constraints may require different ranks.Higher ranks can overfit simpler local data, while smaller ranks can lack generalization capacity.

3 Proposed Method: FLoRA

FLORA aggregates client LoRA updates by stacking corresponding A and B modules, producing an accurate global update while supporting heterogeneous ranks. The workflow initializes and fine-tunes local modules, stacks them at the server, and redistributes the resulting global modules to clients.

  • 3.1 Stacking-based Noise-free Aggregation: FLORA computes the sum of local LoRA products as the product of stacked B and A matrices, enabling noise-free aggregation.The server stacks all B_k modules horizontally and A_k modules vertically so their product represents the summed local updates.
  • 3.1 Stacking-based Noise-free Aggregation: The stacking operator ⊕ vertically stacks successive A modules and horizontally stacks successive B modules.This orientation preserves matrix multiplication while combining client-specific LoRA factors.
  • 3.2 FLORA: Stacking-based Federated Fine-tuning for Heterogeneous LoRA: FLORA supports heterogeneous LoRA ranks without constraining clients beyond sharing the pretrained model dimensions m and n.This allows local modules to use different ranks while remaining compatible with global aggregation.
  • 3.2 FLORA: Stacking-based Federated Fine-tuning for Heterogeneous LoRA: Each FLORA round disseminates the pretrained model, initializes client-specific LoRA modules, locally fine-tunes them, and sends them to the server for stacking.Clients initialize modules according to local data complexity and available resources; rank adaptation itself is outside this paper's scope.
  • 3.2 FLORA: Stacking-based Federated Fine-tuning for Heterogeneous LoRA: The scaling factor p_k is applied to only one of A_k or B_k, preventing its unintended squaring in the final update BA.This design preserves the intended weighting of each local update and maintains noise-free aggregation.
  • 3.2 FLORA: Stacking-based Federated Fine-tuning for Heterogeneous LoRA: FLORA's stacked global modules may increase per-round communication overhead, but the paper reports this overhead is negligible because LoRA modules are small and few rounds are typically needed.The pretrained model remains the dominant initialization communication cost.

4 Experiments

FLORA is evaluated across Llama-based models, QA and chat tasks, and homogeneous or heterogeneous LoRA settings. It consistently outperforms FedIT, supports heterogeneous ranks without degradation, and reveals setting-dependent effects of scaling factors and rank choices.

  • Experiment Setup: Experiments use TinyLlama, Llama, and Llama2 models on QA and chat tasks evaluated with MMLU and MT-bench.The study samples 10 clients under the non-IID setting and applies LoRA to self-attention layers.
  • Homogeneous LoRA: FLORA consistently outperforms FedIT across evaluated models and tasks with homogeneous LoRA rank 16.For example, Llama fine-tuned on Wizard reaches an MT-bench score of 4.21 versus FedIT’s 3.07.
  • Heterogeneous LoRA: FLORA accommodates heterogeneous ranks [64, 32, 16, 16, 8, 8, 4, 4, 4, 4] without performance degradation, whereas zero-padding sharply degrades FedIT.Llama fine-tuned on Alpaca reaches 7.97% MMLU accuracy with zero-padding.
  • The Impact of Scaling Factor: No single scaling factor is optimal across settings because performance depends on the dataset, task, and model.The experiments compare scaling factors 0.01, 0.05, 0.1, and 0.2 using MMLU accuracy and MT-bench score.
  • The Impact of Heterogeneous LoRA Ranks: Figure 4 compares global-model performance with local-model performance across LoRA ranks, showing the global model usually outperforms local models.The exception is TinyLlama fine-tuned on Wizard, where the rank-32 client slightly exceeds the global model.
  • The Impact of Heterogeneous LoRA Ranks: LoRA rank 8 delivers strong performance across models and datasets, while extreme-rank performance varies by model and dataset.Rank 64 underperforms smaller ranks for TinyLlama on Wizard but excels for Llama.

5 Discussion

FLORA adds modest communication relative to FedIT while retaining substantial savings versus full fine-tuning, and addresses privacy concerns through randomized rank-1 stacking.

  • Communication Overhead: FLORA transmits slightly more parameters than FedIT but significantly fewer than full fine-tuning across three communication rounds.The initial full-model transmission dominates communication, while later rounds mainly transmit smaller LoRA updates.
  • Communication Overhead: The overall communication impact of FLORA remains marginal and comparable to FedIT despite its additional update transmission.The paper attributes this to the relatively small size of subsequent LoRA updates compared with the initial model transmission.
  • Trade-off: FLORA's communication increase over FedIT is presented alongside improved fine-tuning effectiveness and support for heterogeneous LoRA ranks.These properties motivate FLORA as a preferable federated fine-tuning solution within the paper's evaluated scope.
  • Privacy Preservation: FLORA splits client LoRA modules into rank-1 sub-modules and stacks them in random order to prevent recovery of individual client modules.The approach is also compatible with encryption and differential privacy.

6 Related Work

Related work combines parameter-efficient adaptation with federated fine-tuning to reduce resource demands and preserve privacy, but heterogeneous LoRA support remains limited.

  • Parameter-efficient Fine-tuning of LLMs: LoRA represents full-fine-tuning gradients with the product of two low-rank matrices, reducing memory use during adaptation.Other PEFT methods fine-tune biases or add pretrained adapter layers between transformer blocks.
  • Federated Fine-tuning of LLMs: FedIT combines federated learning with LoRA so clients fine-tune locally and communicate only LoRA modules, reducing computation and communication costs.Its aggregation uses the federated fine-tuning framework while retaining LoRA as the local strategy.
  • Federated Fine-tuning of LLMs: FedIT's limited support for heterogeneous LoRA is a recognized concern, while zero-padding alternatives introduce additional computing overhead.The limitation arises in federated settings involving differing local LoRA ranks.

7 Conclusion

The paper introduces FLORA to accurately aggregate heterogeneous LoRA modules through stack-based aggregation and reports stronger performance than the state-of-the-art method in both evaluated rank settings.

  • Conclusion: FLORA uses a stack-based LoRA aggregation mechanism to accurately aggregate heterogeneous LoRA modules in federated fine-tuning.The method targets practical settings where clients have heterogeneous LoRA ranks.
  • Conclusion: FLORA outperforms the state-of-the-art method in both homogeneous and heterogeneous LoRA settings.The conclusion frames the results as supporting lightweight and accurate federated fine-tuning of large language models.

A.1 Environments, Datasets and Metric

The evaluation uses multiple instruction and conversation datasets, MMLU and MT-bench metrics, specified computational settings, and selected communication-round configurations.

  • Environments: Experiments run on a 256GB AMD EPYC 7763 64-Core Processor with four NVIDIA RTX A6000 GPUs under Linux v4.18.0.These are the reported computing resources for the experiments.
  • Datasets: The evaluated training data include Dolly with 15k samples, Alpaca with 52K instruction-following examples, Wizard with 70k instruction-output pairs, and approximately 52,000 ShareGPT conversations.Wizard contains more complex instructions, while ShareGPT conversations are split into question-answering pairs.
  • Metrics: MMLU provides 14,024 questions across 57 subjects, with 1444 selected samples used to evaluate logical reasoning capabilities.The paper also evaluates chat assistants on MT-bench using GPT-4 API scores for generated conversations.
  • Hyperparameter Details: Fine-tuning uses learning rate 0.0003, batch size 128, and micro batch size 16, with fewer rounds chosen because of resource demands and MMLU overfitting concerns.The selected communication rounds and local epochs are summarized in Table 2.
  • Evaluation Settings: Table 3 reports FLORA with AdaLoRA, while Table 4 compares FLORA with baselines in Llama2.The supplied captions identify AdaLoRA as reducing rank while preserving fine-tuning effectiveness.

A.3 Supplementary Experiment Results

Experiments with Llama2 show that FLORA outperforms FedIT and Zero-padding across homogeneous and heterogeneous settings, while larger scaling factors improve fine-tuning performance.

  • Selecting an appropriate LoRA rank remains important for task performance despite the influence of data heterogeneity.AdaLoRA can adaptively adjust the rank to optimize performance and save computational resources.
  • FLORA and its baselines are evaluated on Llama2 using the Wizard and ShareGPT datasets because Llama2's QA improvement is not significant.The experiments examine both homogeneous and heterogeneous settings.
  • FLORA achieves higher MT-bench scores than FedIT and Zero-padding on both Wizard and ShareGPT in homogeneous and heterogeneous settings.
  • Higher scaling factors yield better Llama2 fine-tuning performance, following a trend similar to the Llama-7b model.

B Convergence Analysis

The convergence analysis establishes that FLORA's fine-tuning error converges to zero as communication rounds increase, whereas FedIT's noisy aggregation prevents the standard O(1/T) convergence rate.

  • FLORA's convergence analysis follows the standard FedAvg framework for non-IID data under full client participation.FedAvg converges to the global optimum at rate O(1/T) under the stated setting.
  • The analysis adds an unbiased LoRA gradient assumption linking client LoRA updates to gradients computed directly on the base model.
  • As T approaches infinity, FLORA's expected fine-tuning error δ(T) converges to zero.
  • FedIT introduces non-gradient noise by averaging LoRA modules independently, deviating from the FedAvg model-updating rule.
  • FedIT therefore fails to achieve the O(1/T) convergence rate and has impaired convergence rate and overall effectiveness.

C Limitation

The paper identifies increased server-to-client communication as a limitation of FLORA and restricts experiments to Llama models because of computational and time constraints.

  • FLORA increases communication costs because the server sends stacked LoRA modules to clients.
  • The experiments use only Llama models due to computational-resource and time constraints.The authors plan to study other types of LLM federated fine-tuning in future work.
Loading 2409.05976v1…