LLM2D
用于大型语言模型内存高效训练的小批量核心集
Mini-batch Coresets for Memory-efficient Training of Large Language Models
作者: Dang Nguyen, Wenhan Yang, Rathul Anand, Yu Yang, Baharan Mirzasoleiman
发布日期: 10/11/2024
arXiv ID: oai:arXiv.org:2407.19580v2

摘要

使用更大的小批量训练可以提高收敛速度,并能带来更好的性能。然而,由于大型语言模型 (LLM) 对 GPU 内存要求很高,使用大批量训练对它们来说变得不可行。为了解决这个问题,一种有效的方法是找到与更大批量梯度密切匹配的小批量核心集。然而,由于语言数据中来源的高度不平衡性、Adam 优化器的使用以及 LLM 非常大的梯度维度,这种方法对于 LLM 来说变得不可行且无效。在这项工作中,我们通过提出用于训练 LLM 的核心集 (CoLM) 来解决上述挑战。首先,我们证明了通过梯度匹配找到的小批量核心集 w.h.p. 不包含小来源的代表性示例,因此将小来源的所有示例包含在小批量核心集中对于获得最佳性能至关重要。其次,我们通过其历史指数对梯度进行归一化,以找到用于 Adam 训练的小批量核心集。最后,我们利用零阶方法找到最后一个 V 投影矩阵的平滑梯度,并将其稀疏化以保留具有最大归一化梯度幅度的维度。我们将 CoLM 应用于使用 LoRA 在 MathInstruct 和 SuperGLUE 基准测试中微调 Phi-2、Phi-3 和 Zephyr。值得注意的是,CoLM 将微调的内存需求降低了 2 倍,甚至超过了使用 4 倍更大批量进行训练的效果。值得注意的是,CoLM 可以轻松地与现有的内存高效训练方法(例如 LoRA)堆叠。