Fed-pilot: Optimizing LoRA Allocation for Efficient Federated Fine-Tuning with Heterogeneous Clients

TL;DR

Fed-pilot introduces a memory-efficient federated fine-tuning framework that optimizes LoRA module allocation for heterogeneous clients, improving scalability and performance in federated learning of foundation models.

Contribution

It presents a novel knapsack-based optimization for LoRA module selection and a dynamic aggregation rule to handle heterogeneity and Non-IID data in federated fine-tuning.

Findings

Outperforms state-of-the-art methods on five datasets.

Effectively handles heterogeneous memory constraints.

Enhances federated fine-tuning scalability and efficiency.

Abstract

Federated Learning enables the fine-tuning of foundation models (FMs) across distributed clients for specific tasks; however, its scalability is limited by the heterogeneity of client memory capacities. In this work, we propose Fed-pilot, a memory-efficient federated fine-tuning framework. It enables memory-constrained clients to participate in Low-Rank Adaptation (LoRA)-based fine-tuning by training only a subset of LoRA modules locally. Fed-pilot identifies the optimal selection of trainable LoRA modules as a knapsack optimization problem, maximizing model performance under memory constraints for each client. To mitigate inconsistencies arising from heterogeneous module allocations and Non-IID data, Fed-pilot employs a novel aggregation rule that dynamically compensates for under-updated layers. Extensive experiments on five diverse datasets across various heterogeneous data settings demonstrate Fed-pilot's effectiveness and efficiency compared to state-of-the-art methods. To the best of our knowledge, this is the first study on federated fine-tuning of FMs that integrates memory-constrained optimization. The code will be publicly available.