FedKRSO: Communication and Memory Efficient Federated Fine-Tuning of Large Language Models
Guohao Yang, Tongle Wu, Yuanxiong Guo, Ying Sun, Yanmin Gong
TL;DR
Federated fine-tuning of large language models is hindered by the memory and communication demands of full-parameter fine-tuning. FedKRSO introduces a seed-based random subspace optimization framework that compresses local gradients into a low-dimensional subspace and aggregates updates via seed-specific accumulators, achieving memory and communication efficiency close to PEFT methods while approaching federated FFT performance. The authors provide convergence guarantees for the non-convex FL setting and validate the approach on GLUE with RoBERTa under IID and non-IID data, showing superior performance to LoRA-based baselines and substantial efficiency gains over FFT. This work offers a practical path to deploying collaborative LLM fine-tuning on edge devices with limited resources, enabling privacy-preserving, scalable adaptation across domains.
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.