PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Zhen Qin, Daoyuan Chen, Bingchen Qian, Bolin Ding, Yaliang Li, Shuiguang Deng

arXiv:2312.06353v5cs.LGcs.DC

TL;DR

Full-parameter tuning is promising for federated LLM tuning because PEFT may not consistently match its accuracy, especially with statistically heterogeneous client data. FedKSeed uses ZOO with a finite set of K random seeds and corresponding scalar gradients, then assigns differentiated sampling probabilities to candidate seeds. Across 6 scenarios, the proposed method achieved a 7.26% average relative Rouge-L improvement over the best-performing practical baseline on held-out tasks and reduced communication costs by around a thousandfold.

  • Problem

    Full-parameter tuning is promising for federated LLM tuning because PEFT may not consistently match its accuracy, especially with statistically heterogeneous client data.

    Source: §1 · p. 1

  • Method

    FedKSeed uses ZOO with a finite set of K random seeds and corresponding scalar gradients, then assigns differentiated sampling probabilities to candidate seeds.

    Source: §1 · p. 2

  • Results

    Across 6 scenarios, the proposed method achieved a 7.26% average relative Rouge-L improvement over the best-performing practical baseline on held-out tasks and reduced communication costs by around a thousandfold.

    Source: §1 · p. 2

  • Takeaways & Limitations

    FedKSeed makes full-parameter tuning of billion-sized LLMs feasible on federated devices with communication below 18 kilobytes per round.

    Source: §1 · p. 2, §6 · p. 9

    Existing ZOO-based federated approaches had been validated only on models with no more than 10 million parameters.

    Source: §2 · p. 3

Abstract

from arXiv · show

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.

1. Introduction

Federated learning could enable privacy-protected LLM tuning on private device data, but full-parameter tuning is blocked by communication, memory, and synchronization costs. FedKSeed addresses these constraints with finite seed reuse and differentiated seed sampling, achieving stronger held-out-task accuracy with much lower communication.

  • Existing PEFT-based federated methods may not match full-parameter accuracy, especially with statistically heterogeneous client data.Source: §1 · p. 1
  • Full-parameter federated tuning is difficult because communication scales with model size, while BP can require over 20GB of memory for a 1.3-billion-parameter LLM.Source: §1 · pp. 1–2
  • Prior seed-based ZOO methods either neglect substantial downlink costs or require update replication whose computation grows with communication rounds.Source: §1 · p. 2, Figure 2 · p. 2
  • FedKSeed transmits only K random seeds and corresponding scalar gradients, enabling billion-sized full-parameter LLM tuning with less than 18 kilobytes per round.Source: §1 · p. 2
  • Probability-differentiated seed sampling prioritizes more influential perturbations, reducing the candidate-seed pool and accelerating synchronization with the latest global model.Source: §1 · p. 2
  • 7.26% average relative Rouge-L improvement over the best-performing practical baseline was achieved on held-out tasks across 6 scenarios, while communication costs fell by around a thousandfold.Source: §1 · p. 2

2. Related Work

Prior federated LLM work largely uses PEFT or communication-saving ZOO variants, but these methods remain constrained by accuracy, scale, memory, communication, or synchronization costs. FedKSeed instead targets full-parameter tuning while reducing communication and memory requirements.

  • PEFT-based federated methods include LoRA approaches, but LoRA can underperform full-parameter tuning under non-IID client data.Source: §2 · p. 2
  • Existing FL studies with ZOO were validated only on small models with no more than 10 million parameters.Source: §2 · pp. 2–3
  • Prior seed-based methods distribute the latest model, incur tremendous latest-model computation, or rely on BP with substantial memory use.Source: §2 · p. 3
  • FwdLLM optimizes only client-side uplink and requires several hundred GB of communication to tune an approximately 300-million-parameter LLM.Source: §2 · p. 3
  • FedKSeed enables full-parameter federated LLM tuning and reports higher accuracy while reducing communication and memory costs relative to approaches tailored for federated LLM tuning.Source: §2 · p. 3

3. Problem Formulation

The paper formulates federated fine-tuning as collaborative optimization of pretrained model weights over private client datasets. It contrasts BP-based gradient acquisition with ZOO, which estimates gradients through forward propagation using random perturbations and scalar gradients.

  • Federated fine-tuning begins from pretrained weights w0 and seeks to collaboratively tune model weights w across clients with private datasets.Source: §3 · p. 3
  • Clients perform multiple local optimization steps from the server model in each round, after which the server aggregates their local models.Source: §3 · p. 3
  • ZOO estimates gradients through forward propagations rather than the direct gradient calculation used by BP.Source: §3 · p. 3
  • A ZOO gradient estimate uses a random perturbation z drawn from N(0, I_d) and a perturbation scale ϵ.Source: §3 · p. 3
  • The estimated gradient is represented as a scalar gradient multiplied by the perturbation, bg = bϱ · z.Source: §3 · p. 3

4. The proposed FedKSeed

FedKSeed enables federated full-parameter tuning through a finite seed set, preserving similar convergence while reducing model-synchronization computation. FedKSeed-Pro further uses non-uniform seed sampling to reduce the seed pool and potentially improve accuracy.

4.1. Overview

  • FedKSeed stores K candidate seeds and scalar-gradient accumulators, allowing clients to reconstruct and update full models without transmitting model parameters.The server maintains no model parameters; clients exchange seed-gradient histories and reconstruct the latest global model locally.Source: §4.1 · p. 3, §4.2 · p. 4

4.2. Federated Full-Parameter Tuning by Limited Seeds

  • At most K iterations are needed to reconstruct the latest global model, reducing the computation complexity to O(d) when K is fixed.Seed reuse merges updates associated with the same seed, replacing the potentially unbounded update sequence required by unrestricted seed sampling.Source: §4.2 · p. 4

4.2.1. THEORETICAL SUPPORT FOR SEED REUSE

  • FedKSeed and FedZO show similar convergence because seed-restriction error is bounded and becomes small as K and the total number of perturbations increase.Experiments on Natural Instructions report similar test-loss convergence trends for FedKSeed and FedZO (b1 = b2 = 1).Source: §4.2.1 · pp. 5–6

4.2.2. SELECTION OF K

  • Accuracy decreases when K is too small, while increasing K beyond the intrinsic dimension yields diminishing or adverse effects; empirical values around 1024-4096 are supported.FedKSeed operates in a K-dimensional random subspace, and the paper notes that several thousand seeds work well for 1B- and 3B-parameter models.Source: §4.2.2 · p. 6

4.3. Sampling Seeds with Non-uniform Probabilities

  • FedKSeed-Pro samples seeds with non-uniform probabilities based on perturbation importance, reducing the required seed cardinality without sacrificing and sometimes improving accuracy.Importance is estimated from the average absolute scalar gradient, then normalized to determine each seed’s sampling probability.Source: §4.3 · pp. 6–7

5. Experiments

Across six federated scenarios, FedKSeed and FedKSeed-Pro achieve leading accuracy while reducing communication and memory overheads through limited seed sampling. Sensitivity studies show that performance depends on balancing candidate-seed cardinality and zeroth-order hyperparameters, while FedKSeed-Pro generally benefits from non-uniform sampling.

5.1. Experimental Setup

  • The experiments use Natural Instructions and Dolly-15K with held-out-task evaluation on DataJuicer-1.3B and LLaMA-3B under feature- and label-skewed federated partitions.Rouge-L is the evaluation metric, with 5% of clients participating per round and 40 or 60 total rounds depending on the dataset.Source: §5.1 · p. 7

5.2. Comparisons on Accuracy Performance

  • FedKSeed and FedKSeed-Pro achieve the top two performances among practical approaches across all six scenarios, with FedKSeed-Pro improving 3.06% over FedIT on Dolly-15K (α = 0.5) with LLaMA-3B.The gains are attributed to full-parameter tuning and are generally larger for LLaMA-3B than DataJuicer-1.3B.Source: §5.2 · p. 7, Table 2 · p. 8
  • FedKSeed and FedKSeed-Pro do not trail FedAvg, FedZO, or FedMeZO in accuracy while offering communication and memory advantages over these full-parameter references.The reference methods incur excessive communication overheads, whereas the proposed methods avoid noticeable accuracy loss from restricted seed sampling.Source: §5.2 · p. 7, Table 3 · p. 8
  • For models with 1B to 3B parameters, the preferred candidate-seed cardinality is K ∈ [1024, 4096], balancing expressiveness against computational waste.Too many seeds reduce the data available to determine each seed's step size, whereas too few restrict perturbation expressiveness.Source: §5.2 · pp. 7–8, Figure 7 · p. 9
  • FedKSeed-Pro outperforms FedKSeed in five of six scenarios, showing that probability-differentiated sampling of more important seeds can improve accuracy.The exception is Dolly-15K (α = 5.0) with DataJuicer-1.3B, where the methods are comparable.Source: §5.2 · p. 8

5.3. Comparisons on Overheads

  • FedKSeed and FedKSeed-Pro incur the least communication and memory costs by transmitting seeds instead of trainable parameters and using in-place zeroth-order optimization instead of backpropagation.FedKSeed-Pro can use fewer seeds for the same accuracy and accelerates synchronization because latest-model computation grows with the number of candidate seeds.Source: §5.3 · p. 8, Figure 8 · p. 9

5.4. Hyper-parameter Sensitivity

  • In the Dolly-15K (α = 0.5) LLaMA-3B case, both the learning rate η and perturbation scale ϵ should avoid excessively large values.Smaller ϵ improves gradient approximation until half-precision underflow becomes possible, while overly large η can cause aggressive updates or divergence.Source: §5.4 · p. 8, Figure 9 · p. 9

5.5. Comparisons in Various Federated Scenarios

  • Accuracy generally increases with more available data and stabilizes as more clients participate per round, while the proposed methods consistently outperform baselines across federated settings.The comparison uses the various-scenario evaluations summarized in Figure 10.Source: §5.5 · pp. 8–9

6. Conclusion

The paper concludes that FedKSeed enables federated full-parameter tuning of billion-sized LLMs with communication below 18 kilobytes per round. FedKSeed-Pro further improves accuracy and synchronization by assigning differentiated sampling probabilities to seeds.

6. Conclusion

  • FedKSeed enables full-parameter tuning of billion-sized LLMs in federated learning with per-round communication costs below 18 kilobytes.It transmits only a limited set of candidate seeds and corresponding scalar gradients rather than full model parameters.Source: §6 · p. 9
  • FedKSeed-Pro assigns differentiated sampling probabilities according to seed importance, reducing the number of required seeds while improving accuracy over FedKSeed.The strategy also speeds obtaining the latest model through fewer candidate seeds.Source: §6 · p. 9
  • Experiments show that the approaches surpass federated LLM-tuning baselines in unseen-task accuracy, communication cost, and memory footprint.The conclusion reports these gains jointly across real-world datasets.Source: §6 · p. 9
  • The paper identifies decentralized federated fine-tuning as a potential research direction because communication costs are especially critical in that setting.This direction is presented as future work rather than an evaluated contribution.Source: §6 · p. 9

Appendix

  • The appendices provide technical distinctions, notation, algorithms, theorem assumptions and proofs, implementation details, communication calculations, and supplementary experiments.Supplementary experiments cover seed efficiency, convergence, seed probabilities, and local-training time.Source: Appendix · p. 13

A. Detailed Technical Comparisons

FedKSeed addresses federated full-parameter tuning by combining finite seed reuse with zeroth-order updates, reducing communication and memory burdens while retaining practical synchronization. FedKSeed-Pro further samples seeds non-uniformly according to their importance.

A. Detailed Technical Comparisons

  • Existing seed-based methods either retain costly model synchronization or incur unsustainable computation, memory, or evaluation limitations for billion-sized LLMs.Unrestricted candidate-seed methods still distribute the latest model, while alternatives require extensive repeated updates or backpropagation.Source: A. Detailed Technical Comparisons · pp. 13–14
  • Less than 18 kilobytes per round enables federated full-parameter tuning of billion-sized LLMs while reducing memory consumption to inference-level requirements.FedKSeed transmits K seeds and scalar gradients instead of model parameters.Source: A. Detailed Technical Comparisons · p. 14

B. Summarization of Notations

C. The Algorithm of the Proposed FedKSeed and FedKSeed-Pro

  • The algorithm initializes candidate seeds and probabilities, lets clients reconstruct the latest model and sample perturbations, then returns seed-gradient information for server aggregation.The underlined algorithm components are detachable and apply only to FedKSeed-Pro.Source: C. The Algorithm of the Proposed… · p. 15

D. Detailed Assumptions of Theorem 1

The convergence analysis relies on bounded loss, smooth objectives, bounded stochastic-gradient second moments, and bounded local-global gradient dissimilarity. These assumptions combine standard stochastic-optimization conditions with a condition capturing federated statistical heterogeneity.

D. Detailed Assumptions of Theorem 1

  • The analysis assumes the global loss is lower bounded, while local, client-expected, and global objectives are L-smooth.These are stated as the loss-boundary and objective-smoothness assumptions.Source: D. Detailed Assumptions of Theorem 1 · p. 15
  • The stochastic gradient’s second-order moment is bounded, providing a variance-related condition for the optimization analysis.This is Assumption 4.Source: D. Detailed Assumptions of Theorem 1 · p. 16
  • Local-global gradient dissimilarity is bounded to represent the statistical heterogeneity of client data in federated learning.The bound uses a positive constant ch.Source: D. Detailed Assumptions of Theorem 1 · p. 16
  • Assumptions 2–4 are standard in stochastic-optimization analysis, whereas Assumption 5 addresses non-IID federated settings and Theorem 1 follows Fang et al. (2022) with lower-order terms omitted.The expected client loss is defined as the expectation of local loss over each client’s dataset.Source: D. Detailed Assumptions of Theorem 1 · pp. 15–16

E.1. Proof of Lemma 1

The analytical sections explain how finite seed sampling approximates unrestricted perturbations and why the candidate-seed count affects optimization. The implementation details specify reproducible prompt and evaluation choices.

E.1. Proof of Lemma 1

  • Chebyshev’s and Hoeffding’s inequalities bound the behavior of perturbations sampled first from a finite Gaussian-derived set and then reused uniformly.The proof uses two sampling stages: generating K variables and drawing Γ perturbations from them.Source: §E.1 · p. 16

E.2. Proof of Theorem 2

  • Theorem 2 compares the two-point gradient estimator under finite seed restrictions with the estimator using unrestricted seeds, using objective smoothness to control the difference.The restricted estimator is expressed through losses evaluated at positively and negatively perturbed models.Source: §E.2 · pp. 16–17

F.1. Additional Analytical Support to Principle 1

  • Larger K generally supplies more independent perturbation directions and makes the optimization constraints more informative, whereas K = 1 can leave the problem fundamentally unoptimizable.The argument treats the perturbation matrix as a constraint system whose rank is typically K when d is much larger than K.Source: §F.1 · p. 17

F.2. Additional Analytical Support to Principle 2

  • For each perturbation direction, FedKSeed estimates a step size from local data; fewer candidate seeds reuse each direction more often and probabilistically enlarge its effective batch size.Increasing the seed pool expands the optimization space and can make random search more difficult.Source: §F.2 · p. 17

G. Implementation Details

  • The implementation section documents prompt templates, dataset processing, experimental settings, and reproducibility choices, including Alpaca-based prompting and fixed maximum sequence length.The supplied implementation passages also include the prompt-template and implementation-detail organization.Source: Table 5 · p. 18, Table 6 · p. 18, G. Implementation Details · p. 18

G.1. Datasets & Evaluation Metrics

Experiments use held-out-task evaluations on Natural Instructions and Dolly-15K under distinct non-IID client partitions. Reproducibility choices standardize prompting, hardware, optimization, and candidate-seed initialization.

G.1. Datasets & Evaluation Metrics

  • Natural Instructions uses 738 training-task clients and 119 held-out test tasks, while Dolly-15K reserves its last task for evaluation across 200 clients.The two datasets provide task-held-out evaluation settings.Source: §G.1 · p. 18
  • Dolly-15K varies label-distribution skew with Dirichlet α = 0.5 and α = 5.0, where lower α indicates stronger skew.This creates non-IID scenarios with different degrees of label heterogeneity.Source: §G.1 · p. 19

G.2. Prompt Template

  • Inputs use the Alpaca prompt template, a 1,024-token maximum, and greedy decoding during evaluation.Instances exceeding the maximum length are ignored.Source: §G.1 · p. 19, §G.2 · p. 19

G.3. Experimental Platforms

  • Performance evaluations use Rouge-L with DataJuicer-1.3B and LLaMA-3B on RTX 3090 and A100 platforms using 16-bit pretrained models.The setup is selected for consistency across approaches, although FedKSeed methods require less memory than the available platforms.Source: §G.3 · p. 19

G.4. Implementations

  • FedKSeed and FedKSeed-Pro use batch size 1, specified learning-rate and perturbation-scale searches, and randomly initialized candidate seeds with client-size-proportional aggregation.An efficient alternative implementation was not integrated for reproducing the reported Rouge-L scores.Source: §G.4 · p. 19, §G.5 · p. 19

H.1. Experimental Support for Principle 2

Across the three studies, seed counts show diminishing Rouge-L gains beyond K=4096, while FedKSeed variants converge within 40 rounds and FedKSeed-Pro concentrates sampling on more important perturbations.

H.1. Experimental Support for Principle 2

  • When K exceeds 4096, additional seeds yield negligible marginal Rouge-L gains, supporting Principle 2 and choosing a smaller K to limit synchronization time.Below K=1024, each additional seed gives consistently positive gains; between 1024 and 4096, gains are negligible in some scenarios.Source: §H.1 · pp. 19–21
  • FedKSeed-Pro can achieve comparable accuracy with fewer candidate seeds than FedKSeed, so Table 2 uses fewer candidates when peak accuracy is not the sole objective.Figure 11 reports marginal improvement rather than absolute accuracy, while Figure 7 supports the fewer-seed comparison.Source: §H.1 · pp. 20–21, Figure 11 · p. 20

H.2. Convergence Study

  • FedKSeed, FedKSeed-Pro, and the baselines except FedPrompt converge within 40 rounds on LLaMA-3B, with the proposed methods converging faster than expected for full-parameter ZOO tuning.FedPrompt converges more slowly and underperforms other baselines; the discussion attributes ZOO convergence to local effective rank rather than full model size in adequately pretrained models.Source: §H.2 · p. 21, Figure 12 · p. 21

H.3. Illustrations of Seed Probabilities in FedKSeed-Pro

  • A relatively small subset of candidate seeds receives higher sampling probabilities under FedKSeed-Pro, and average scalar-gradient amplitude correlates positively with perturbation importance for model accuracy.The comparison uses DataJuicer-1.3B and LLaMA-3B across Natural Instructions and Dolly-15K settings.Source: §H.3 · p. 21, Figure 13 · p. 22
  • FedKSeed-Pro uses scalar-gradient amplitude rather than cosine or Euclidean similarity to quantify seed importance because high-dimensional Gaussian perturbations are typically nearly orthogonal.Near-zero similarity makes distance-based methods poorly suited to distinguishing the importance of random perturbations.Source: §H.3 · p. 22

H.4. Training Efficiency

FedKSeed and FedKSeed-Pro incur only minimal additional per-step training time compared with baselines, despite enabling substantially more trainable parameters.

H.4. Training Efficiency

  • FedKSeed and FedKSeed-Pro incur minimal additional per-step training time overhead compared with the baselines.Training times are compared across approaches and LLMs on the same NVIDIA V100 GPU and Intel Xeon CPU platform.Source: §H.4 · p. 22, Figure 14 · p. 23
  • The extra computation supports expanding the number of trainable parameters by several orders of magnitude, improving global-model accuracy while reducing memory and communication overhead.The paper argues that time constraints are less strict than memory and communication constraints, making this trade-off worthwhile.Source: §H.4 · pp. 22–23

I. Detailed Calculation of Communication Overhead

The communication analysis gives exact per-round costs for FedKSeed variants and connects their low-dimensional updates to reduced server burden, decentralized fine-tuning, and checkpoint storage.

I. Detailed Calculation of Communication Overhead

  • For baselines, the accounting includes per-round uploads and downloads of trainable parameters, whereas FedKSeed variants transmit seeds, scalar-gradient accumulators, and gradient histories.The comparison uses 32-bit floating-point values and 32-bit integer seeds; half precision would halve all listed costs uniformly.Source: §I · p. 23
  • FedKSeed requires 17,988 bytes per client per round, combining 16,388 bytes of downlink data with 1,600 bytes of uploaded gradient history.The downlink encodes candidate seeds and 4,096 accumulated scalar gradients, while the uplink contains 200 seed-gradient pairs.Source: §I · p. 23
  • FedKSeed-Pro requires 9,796 bytes per round with DataJuicer-1.3B and 17,988 bytes with LLaMA-3B under the stated K settings.The respective candidate-seed counts are K=1,024 and K=2,048.Source: §I · p. 24

J.1. Alleviating the Burden of Aggregation

  • FedKSeed and FedKSeed-Pro reduce server computation and communication complexity from O(md) to O(mK), where d is model size and m is the number of active clients.This reduction can allow an FL server to operate with substantially fewer computational and communication resources.Source: §J.1 · p. 24

J.2. Enabling Possibility to Decentralized Federated Fine-Tuning of LLMs

  • Their reduced communication costs make decentralized federated fine-tuning of LLMs possible when clients would otherwise transmit model parameters to multiple recipients.The motivation is to avoid extra communication caused by decentralized aggregation over unstable or long-distance connections.Source: §J.2 · p. 24

J.3. Alleviating the Burden of Saving Checkpoints

  • Model snapshots can be represented by accumulated scalar-gradient arrays containing 2K scalars, substantially reducing checkpoint storage consumption.This supports preserving rollback snapshots during unfavorable optimization trajectories or potential attacks.Source: §J.3 · p. 25

Related papers

Loading 2312.06353v5…