Thinking Forward: Memory-Efficient Federated Finetuning of Language Models

Summary

This paper introduces SPRY, a federated learning (FL) algorithm designed to finetune large language models (LLMs) on resource-constrained devices by addressing the excessive memory requirements of traditional backpropagation methods. SPRY tackles the challenge of high memory usage from intermediate activations by utilizing Forward-mode Auto-Differentiation (AD) and splitting trainable weights among participating clients, allowing each client to compute accurate gradients with reduced memory. Theoretical analysis shows SPRY's global gradients are unbiased for homogeneous data and provides a convergence rate dependent on FL rounds and data heterogeneity. Empirical results demonstrate SPRY's efficiency, reducing memory usage by 1.4-7.1× and achieving faster convergence and higher accuracy compared to existing methods. This makes feasible the finetuning of LLMs on mobile and edge devices, significantly impacting FL deployments.

Strengths

* **S.1.** The proposed SPRY algorithm tackles a difficult task and is backed up with both empirical and theoretical results. * **S.2.** The paper is well written and the illustrations are helpful. * **S.3.** The empirical experiments are conducted on multiple datasets, models, and hardware configuration, while showing that SPRY outperforms previous works. * **S.4.** The paper provides an anonymous code repository of SPRY.

Weaknesses

* **W.1.** While SPRY shows promising results with a relatively small number of clients, the scalability to a larger number of clients, which is typical in FL scenarios, is not thoroughly investigated. The communication and computational overheads associated with increasing the number of participating clients, especially in terms of synchronization and gradient aggregation, are not discussed. This raises concerns about the practicality and efficiency of SPRY in large-scale deployments. * **W.2.** The empirical evaluation, although thorough in certain aspects, is limited in terms of the variety of datasets and models used. The evaluation focuses primarily on specific language tasks and a narrow range of LLMs. A broader evaluation including more diverse datasets and model architectures would strengthen the generalizability of the findings. Additionally, the impact of SPRY on tasks beyond natural language processing, such as computer vision or other modalities, is not explored. * **W.3.** While SPRY does outperform existing zero-order methods, it converges slower (time-wise) compared to traditional gradient based methods such as FedAvg.

Questions

* **Q.1.** Will the SPRY algorithm work in a large scale setting of 1k+ GPUs?

Rating

5

Confidence

2

Soundness

2

Presentation

3

Contribution

3

Limitations

n/a

Summary

This manuscript is focused on the memory-efficient federated finetuning of LLMs. The author first uses Forward-mode Auto-Differentiation to reduce memory. Then, the author observes that merely substituting backpropagation with Forward-mode AD in FL scenarios often results in poor accuracy and computational inefficiency. To address this challenge, the author recognizes that Forward-mode AD operates more efficiently and yields better gradient estimations when the trainable weight count is minimized. Therefore, SPRY only assigns each client a responsibility to compute gradients for only a subset of trainable weights. The experimental results show that the memory overhead can be reduced significantly. The topic is of interest and the presented numerical results seem, indeed, promising. However, there are still some questions/comments/suggestions for the current version of the paper, please refer to my comments under Questions.

Strengths

* The paper performs a number of experimental verifications as well as theoretical proofs. * The experimental results show that the method is enabled to significantly reduce GPU memory consumption while ensuring a competitive performance. * Additionally, the appendix contains valuable results.

Weaknesses

* This manuscript does not analyze the computational load of Forward-mode. * There are still some issues with the experimental setup, such as the need to consider the impact of the number of clients on performance.

Questions

* What is the computational overhead situation in each step (e.g., the training time for each gradient step)? A comparison of computational overheads is suggested in the experimental section. * Is each client's trained layer the same from start to finish? Or are the layers reassigned each communication round? * It seems that the layers for each client are fixed at the beginning. If we consider the dynamic distribution of layers, it might alleviate the issues caused by heterogeneous devices. * Communication overhead is an important metric in federal systems, and it is suggested that the authors give the communication data size in each communication round.

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

The authors give a limitation analysis in the checklist section.

Summary

This paper introduces a forward-mode AD federated learning algorithm (SPRY). They use SPRY to finetune LLMs and demonstrate a low memory footprint compared to backpropagation-based federated learning algorithms. The authors also derive SPRY’s convergence rate and provide theory behind why SPRY’s global gradients are unbiased for homogeneous data distributions across clients. The authors empirically evaluate SPRY on 8 language tasks as well as perform ablation studies.

Strengths

* The empirical evaluation on eight language tasks while testing multiple different language models is a key strength of the paper. * The paper is organized clearly with nice figures and a clear structure to the sections. The empirical evaluation section is particularly well structured. * The ablation studies further strengthen the results of the paper. * The section of peak memory consumption is clear and highlights the superior memory performance of both zero-order and first-order forward mode AD compared to backpropagation. * Splitting forward-mode across multiple layers seems a novel way to perform federated learning.

Weaknesses

* The baseline FWDLLM is described as a zero-order-based approach but upon reading this paper, it seems like it uses a first-order forward mode AD update rule in the federated setting (https://arxiv.org/pdf/2308.13894). Why does the paper categorize this approach as a zero-order-based method? Also, Figure 10 in that paper shows a better performance of all the baselines on the different data sets. Why is there such a significant discrepancy? E.g. RoBERTa-large gets a performance of much greater than 80 % for FwdLLM in the original paper for AG News and is over 60 % for Yahoo. * Equation (4) is challenging to interpret. Why is the summation over classes? What do the subscripts in square brackets mean? Does the paper define data distribution only according to how class labels are distributed among the clients? This did not appear obvious in the paper. * The implications of theorem 4.2 are not clear. Is this bound related to the variance of the gradient estimator? Generally the variance is related to $E[x^2]$.

Questions

* Are the authors aware of AsyncFGD (https://openreview.net/pdf?id=45RBLZBJid)? They also implement FGD in parallel by splitting weights across workers. They do not reach the model and experiment size of this paper, but there does appear to be some similarity in the approach. Would it be possible to identify these differences? * A general limitation of FGD as a learning method is the variance of the gradient estimator with increasing dimension of the parameter space. How does SPRY overcome this challenge? Is there a theory that updating individual layers on different clients reduces this variance? * For per-iteration communication, is there any advantage to having clients, since the server needs to generate all the tangent vectors and then update all the weights of the model. Since there is no matrix multiplication, perhaps there is no need for a GPU on the server-side? Also, generally what would the communication and compute cost be for the server? * What is the current state-of-the-art in federated learning?

Rating

6

Confidence

3

Soundness

2

Presentation

3

Contribution

3

Limitations

Authors have done a good job at describing limitations and memory usage of their approach.

Thanks for the Rebuttal

The rebuttal helped remove some of the weaknesses so I will raise my score.