Finetuning large language models (LLMs) in federated learning (FL) settings has become increasingly important as it allows resource-constrained devices to finetune a model using private data. However, finetuning LLMs using backpropagation requires excessive memory (especially from intermediate activations) for resource-constrained devices. While Forward-mode Auto-Differentiation (AD) can significantly reduce memory footprint from activations, we observe that directly applying it to LLM finetuning results in slow convergence and poor accuracy. In this paper, we introduce Spry, an FL algorithm that splits trainable weights of an LLM among participating clients, such that each client computes gradients using Forward-mode AD that are closer estimations of the true gradients. Spry achieves a low memory footprint, high accuracy, and fast convergence. We formally prove that the global gradients in Spry are unbiased estimators of true global gradients for homogeneous data distributions across clients, while heterogeneity increases bias of the estimates. We also derive Spry's convergence rate, showing that the gradients decrease inversely proportional to the number of FL rounds, indicating the convergence up to the limits of heterogeneity. Empirically, Spry reduces the memory footprint during training by 1.4-7.1X in contrast to backpropagation, while reaching comparable accuracy, across a wide range of language tasks, models, and FL settings. Spry reduces the convergence time by 1.2-20.3X and achieves 5.2-13.5% higher accuracy against state-of-the-art zero-order methods. When finetuning Llama2-7B with LoRA, compared to the peak memory consumption of 33.9GB of backpropagation, Spry only consumes 6.2GB of peak memory. For OPT13B, the reduction is from 76.5GB to 10.8GB. Spry makes feasible previously impossible FL deployments on commodity mobile and edge devices. Our source code is available for replication at https://github.com/Astuary/Spry/.
@inproceedings{Panchal24neurips,
author = {Kunjal Panchal and Nisarg Parikh and Sunav Choudhary and Lijun Zhang and Yuriy Brun and Hui Guan},
title =
{\href{http://people.cs.umass.edu/brun/pubs/pubs/Panchal24neurips.pdf}{Thinking Forward: {Memory}-Efficient Federated Finetuning of Language Models}},
booktitle = {Proceedings of the 38th Annual Conference on Neural
Information Processing Systems (NeurIPS), Advances in Neural Information
Processing Systems 38},
venue = {NeurIPS},
address = {Vancouver, BC, Canada},
month = {December},
date = {9--15},
year = {2024},
abstract = {<p>Finetuning large language models (LLMs) in federated learning (FL) settings
has become increasingly important as it allows resource-constrained devices
to finetune a model using private data. However, finetuning LLMs using
backpropagation requires excessive memory (especially from intermediate
activations) for resource-constrained devices. While Forward-mode
Auto-Differentiation (AD) can significantly reduce memory footprint from
activations, we observe that directly applying it to LLM finetuning results
in slow convergence and poor accuracy. In this paper, we introduce Spry, an
FL algorithm that splits trainable weights of an LLM among participating
clients, such that each client computes gradients using Forward-mode AD that
are closer estimations of the true gradients. Spry achieves a low memory
footprint, high accuracy, and fast convergence. We formally prove that the
global gradients in Spry are unbiased estimators of true global gradients for
homogeneous data distributions across clients, while heterogeneity increases
bias of the estimates. We also derive Spry's convergence rate, showing that
the gradients decrease inversely proportional to the number of FL rounds,
indicating the convergence up to the limits of heterogeneity. Empirically,
Spry reduces the memory footprint during training by 1.4-7.1X in contrast to
backpropagation, while reaching comparable accuracy, across a wide range of
language tasks, models, and FL settings. Spry reduces the convergence time by
1.2-20.3X and achieves 5.2-13.5% higher accuracy against state-of-the-art
zero-order methods. When finetuning Llama2-7B with LoRA, compared to the peak
memory consumption of 33.9GB of backpropagation, Spry only consumes 6.2GB of
peak memory. For OPT13B, the reduction is from 76.5GB to 10.8GB. Spry makes
feasible previously impossible FL deployments on commodity mobile and edge
devices. Our source code is available for replication at
https://github.com/Astuary/Spry/.</p>},
accept = {$\frac{4,036}{15,671} \approx 26\%$},
fundedBy = {NSF CCF-2210243, NSF DMS-2220211, NSF CNS-2224054, NSF CNS-2312396, NSF CNS-2338512, Adobe},
}