Thinking Forward: Memory-Efficient Federated Finetuning of Language Models

本文介绍了一种名为Spry的联邦学习算法，将大型语言模型的可训练权重分配给参与的客户端，使每个客户端使用正向模式自动微分计算梯度，从而更接近真实梯度。Spry算法在联邦学习中实现了低内存占用、高准确性和快速收敛。作者理论上证明了，对于客户端之间的同质数据分布，Spry中的全局梯度是真实全局梯度的无偏估计，而异质性会增加估计的偏差。作者还推导了Spry的收敛速度，表明梯度随着联邦学习轮数的增加而减少，说明在异质性的极限情况下也能实现收敛。实验结果表明，与反向传播相比，Spry在许多语言任务、模型和联邦学习设置下降低了1.4-7.1倍的内存占用，同时达到了可比较的准确性。在与最先进的零阶方法进行比较时，Spry将收敛时间降低了1.2-20.3倍，准确性提高了5.2-13.5\%。在使用LoRA对Llama2-7B进行微调时，与反向传播的峰值内存使用量33.9GB相比，Spry仅使用了6.2GB的内存。对于OPT13B，内存占用从76.5GB降至10.8GB。Spry使得在普通移动和边缘设备上实现以前无法实现的联邦学习部署成为可能。作者提供了Spry的源代码，可在https://github.com/Astuary/Spry获取。

本论文旨在解决在联邦学习（FL）设置下，使用反向传播进行大型语言模型（LLM）微调所需的过多中间激活内存的问题。

本文提出了Spry算法，将LLM的可训练权重分配给参与的客户端，使每个客户端使用正向自动微分来计算梯度，从而更接近真实梯度估计，从而达到低内存占用、高准确性和快速收敛的效果。

Spry算法可以在训练期间将内存占用减少1.4-7.1倍，同时在各种语言任务、模型和FL设置下达到相当的准确性。在与最先进的零阶方法的比较中，Spry可以将收敛时间减少1.2-20.3倍，并获得5.2-13.5％的更高准确性。Spry算法使得以前无法在移动和边缘设备上进行FL部署变得可行。该算法的源代码已经开源。

与此相关的研究包括：《Federated Learning with Non-IID Data: An Empirical Study》、《Federated Learning: Challenges, Methods, and Future Directions》等。