[Paper Review] Federated Full-Parameter Tuning of Billion-Sized Language Models with Communication Cost under 18 Kilobytes

Zhen Qin, Daoyuan Chen+4 more|arXiv (Cornell University)|Dec 11, 2023

DNA and Biological Computing8 citations

TL;DR

The paper presents FedKSeed, a zeroth-order federated fine-tuning method that enables full-parameter tuning of billion-sized LLMs with extremely low communication per round (under 18 KB) by using a fixed small seed pool and scalar gradients, plus a non-uniform seed sampling enhancement FedKSeed-Pro.

ABSTRACT

Pre-trained large language models (LLMs) need fine-tuning to improve their responsiveness to natural language instructions. Federated learning offers a way to fine-tune LLMs using the abundant data on end devices without compromising data privacy. Most existing federated fine-tuning methods for LLMs rely on parameter-efficient fine-tuning techniques, which may not reach the performance height possible with full-parameter tuning. However, federated full-parameter tuning of LLMs is a non-trivial problem due to the immense communication cost. This work introduces FedKSeed that employs zeroth-order optimization with a finite set of random seeds. It significantly reduces transmission requirements between the server and clients to just a few random seeds and scalar gradients, amounting to only a few thousand bytes, making federated full-parameter tuning of billion-sized LLMs possible on devices. Building on it, we develop a strategy enabling probability-differentiated seed sampling, prioritizing perturbations with greater impact on model accuracy. Experiments across six scenarios with various LLMs, datasets and data partitions demonstrate that our approach outperforms existing federated LLM fine-tuning methods in both communication efficiency and new task generalization.

Motivation & Objective

• Demonstrate feasibility of full-parameter fine-tuning of billion-sized LLMs in federated settings with minimal communication.
• Reduce server-client transmission by avoiding full model parameter exchange and backpropagation-based memory usage.
• Develop seed-based perturbation strategies that maintain or improve accuracy under federated, non-IID data.
• Provide theoretical and empirical analysis of seed reuse in a zeroth-order optimization framework.

Proposed method

• Use zeroth-order optimization (ZOO) with a finite set of random seeds to generate perturbations for full-parameter tuning.
• Maintain a server-side pool of K seeds and a scalar-gradient accumulator to reconstruct the latest global model without transmitting full parameters.
• Clients perform local updates using two-point gradient estimates tied to chosen seeds and report seed-gradient histories instead of full model updates.
• Aggregate per-seed scalar gradients on the server to update the global perturbation projection and reconstruct the latest model locally via a seed-based update rule.
• Introduce non-uniform seed sampling (FedKSeed-Pro) by weighting seeds according to estimated gradient magnitude, to boost efficiency and accuracy.

Experimental results

Research questions

• RQ1 Can full-parameter tuning of billion-sized LLMs be made communication-efficient enough for practical federated deployment?
• RQ2 How does limiting the seed pool to K perturbations affect convergence and accuracy compared to unlimited-seed approaches?
• RQ3 Does non-uniform seed sampling improve accuracy and reduce required seed cardinality in FedKSeed-Pro?
• RQ4 What is the theoretical relationship between seed reuse in ZOO-based FL and convergence guarantees?
• RQ5 How do FedKSeed and FedKSeed-Pro perform across diverse LLMs, datasets, and data-partition settings compared to PEFT-based federated methods?

Key findings

• FedKSeed achieves full-parameter tuning with per-round communication under 18 kilobytes by encoding perturbations as seeds and transmitting only scalar gradients.
• Convergence analysis shows seed diversity does not affect convergence under the proposed framework, enabling seed reuse without sacrificing theoretical guarantees.
• Non-uniform seed sampling (FedKSeed-Pro) consistently improves Rouge-L scores and reduces the needed seed cardinality, achieving up to 7.26% relative improvement over baselines on average across scenarios.
• Empirical results across six scenarios with 1B–3B models show FedKSeed and FedKSeed-Pro outperform PEFT-based federated fine-tuning baselines in both accuracy and communication/memory efficiency.
• FedKSeed-Pro can achieve higher accuracy with smaller K (e.g., 1024–4096) and outperforms FedKSeed in synchronization efficiency in several settings.