向前思考：内存高效的联邦语言模型微调
机器学习
2024-10-23 v2
摘要
在联邦学习（FL）中对大型语言模型（LLM）进行微调日益重要，因为它允许资源受限的设备使用私有数据进行微调。然而，使用反向传播进行微调需要大量内存（尤其是中间激活），对于资源受限的设备构成了挑战。虽然前向模式自动微分（AD）可以显著减少激活的内存占用，但我们观察到直接将其应用于 LLM 微调会导致收敛缓慢和精度差。本文引入 Spry，一种 FL 算法，将 LLM 的可训练权重在参与客户端之间划分，使得每个客户端使用前向模式 AD 计算的梯度是对真实梯度的更接近的估计。Spry 实现了低内存占用、高精度和快速收敛。我们形式化地证明，在客户端之间数据分布均匀的情况下，Spry 中的全局梯度是真实全局梯度的无偏估计；而数据异构性会增加估计的偏差。我们还推导出 Spry 的收敛速率，表明梯度随 FL 轮数的反比例减小，表明在异构性限制下达到收敛。实验上，Spry 在训练期间的内存占用相较于反向传播降低 1.4-7.1 倍，在广泛的语言任务、模型和 FL 设置下实现了可比的精度。Spry 将收敛时间缩短 1.2-20.3 倍，准确率比零阶方法高出 5.2-13.5%。在对 Llama2-7B 使用 LoRA 进行微调时，与反向传播的峰值内存占用 33.9GB 相比，Spry 仅占用 6.2GB 峰值内存。对于 OPT13B，内存占用从 76.5GB 降至 10.8GB。Spry 使得以前在普通边缘设备上实现的 FL 部署成为可能。我们的源代码可在 https://github.com/Astuary/Spry 获取。
引用
@article{arxiv.2405.15551,
title = {Thinking Forward: Memory-Efficient Federated Finetuning of Language Models},
author = {Kunjal Panchal and Nisarg Parikh and Sunav Choudhary and Lijun Zhang and Yuriy Brun and Hui Guan},
journal= {arXiv preprint arXiv:2405.15551},
year = {2024}
}
备注
Accepted to NeurIPS 2024