BAdam：一种面向大语言模型的内存高效全参数优化方法
机器学习
2024-11-18 v3
摘要
本工作提出了 BAdam，一种结合块坐标下降（BCD）框架与 Adam 更新规则的优化方法。BAdam 为大语言模型的全参数微调提供了一种内存高效的途径。我们在确定性情况下对 BAdam 进行了理论收敛性分析。在实验中，我们分别使用单张 RTX3090-24GB GPU 和 4 张 A100-80GB GPU 应用 BAdam 对 Llama 3-8B 和 Llama 3-70B 模型进行微调。结果证实了 BAdam 在内存使用、运行时间和优化能力方面的高效性。此外，基于 MT-bench 和数学基准的下游性能评估表明，BAdam 优于 LoRA 等现有的内存高效基线方法。结果还表明，与 Adam 相比，BAdam 能够取得相当甚至更优的性能。最后，使用 SGD 更新规则的消融研究说明了 BCD 对微调 LLM 的适用性。我们的代码可以轻松集成到任何基于 PyTorch 的代码库中，并可在 https://github.com/Ledzy/BAdam 获取。
引用
@article{arxiv.2404.02827,
title = {BAdam: A Memory Efficient Full Parameter Optimization Method for Large Language Models},
author = {Qijun Luo and Hengxu Yu and Xiao Li},
journal= {arXiv preprint arXiv:2404.02827},
year = {2024}
}
备注
Accepted for Publication in Conference on Neural Information Processing Systems, 2024