FedKRSO: Communication and Memory Efficient Federated Fine-Tuning of Large Language Models

TL;DR

FedKRSO introduces a novel federated fine-tuning method for large language models that significantly reduces communication and memory costs while maintaining high performance, suitable for resource-constrained edge devices.

Contribution

The paper proposes FedKRSO, a new federated fine-tuning approach that uses random subspace optimization and accumulator updates to improve efficiency and performance.

Findings

Achieves high accuracy on GLUE benchmark in federated settings.

Reduces communication and memory overhead compared to existing methods.

Closely matches the performance of full fine-tuning in federated scenarios.

Abstract

Fine-tuning is essential to adapt general-purpose large language models (LLMs) to domain-specific tasks. As a privacy-preserving framework to leverage decentralized data for collaborative model training, Federated Learning (FL) is gaining popularity in LLM fine-tuning, but remains challenging due to the high cost of transmitting full model parameters and computing full gradients on resource-constrained clients. While Parameter-Efficient Fine-Tuning (PEFT) methods are widely used in FL to reduce communication and memory costs, they often sacrifice model performance compared to FFT. This paper proposes FedKRSO (Federated $K$-Seed Random Subspace Optimization), a novel method that enables communication and memory efficient FFT of LLMs in federated settings. In FedKRSO, clients update the model within a shared set of random low-dimension subspaces generated by the server to save memory usage. Furthermore, instead of transmitting full model parameters in each FL round, clients send only the model update accumulators along the subspaces to the server, enabling efficient global model aggregation and dissemination. By using these strategies, FedKRSO can substantially reduce communication and memory overhead while overcoming the performance limitations of PEFT, closely approximating the performance of federated FFT. The convergence properties of FedKRSO are analyzed rigorously under general FL settings. Extensive experiments on the GLUE benchmark across diverse FL scenarios demonstrate that FedKRSO achieves both superior performance and low communication and memory overhead, paving the way towards on federated LLM fine-tuning at the resource-constrained edge.