Adaptive Federated LoRA in Heterogeneous Wireless Networks with Independent Sampling

TL;DR

This paper introduces an adaptive federated LoRA method with independent client sampling that optimizes training efficiency in heterogeneous wireless networks, significantly reducing wall-clock convergence time for federated fine-tuning of large language models.

Contribution

It provides a new convergence bound for federated LoRA with independent sampling, and develops an adaptive bandwidth allocation and joint optimization scheme to minimize training time.

Findings

Reduces wall-clock training time compared to existing methods.

Effectively handles heterogeneity in client resources and system bandwidth.

Demonstrates significant efficiency improvements across various models and datasets.

Abstract

Federated LoRA has emerged as a promising technique for efficiently fine-tuning large language models (LLMs) on distributed devices by reducing the number of trainable parameters. However, existing approaches often inadequately overlook the theoretical and practical implications of system and data heterogeneity, thereby failing to optimize the overall training efficiency, particularly in terms of wall-clock time. In this paper, we propose an adaptive federated LoRA strategy with independent client sampling to minimize the convergence wall-clock time of federated fine-tuning under both computation and communication heterogeneity. We first derive a new convergence bound for federated LoRA with arbitrary and independent client sampling, notably without requiring the stringent bounded gradient assumption. Then, we introduce an adaptive bandwidth allocation scheme that accounts for heterogeneous client resources and system bandwidth constraints. Based on the derived theory, we formulate and solve a non-convex optimization problem to jointly determine the LoRA sketching ratios and sampling probabilities, aiming to minimize wall-clock convergence time. An efficient and low-complexity algorithm is developed to approximate the solution. Finally, extensive experiments demonstrate that our approach significantly reduces wall-clock training time compared to state-of-the-art methods across various models and datasets.