大语言模型训练数据混合的内存高效小批量核集

使用更大的小批量可提高收敛速度并获得更佳性能。然而，由于大型语言模型（LLM）的大规模GPU内存需求，使用大小批量变得不可行。为此，一种有效方法是寻找与更大小批量梯度 closely match 的小型小批量核集。然而，由于语言数据中小来源的高度不平衡、采用Adam优化器以及LLM的极大梯度维度，这种方法在LLM上变得不可行且无效。本文通过提出用于训练大语言模型的核集方法（CoLM）来解决上述挑战。首先，我们指出，由梯度匹配找到的小批量核集以高概率不包含小来源的代表性示例，因此对于获得最佳性能，必须将小来源的所有示例包含在小批量核集中。其次，我们对梯度按其历史指数进行归一化，以便使用Adam进行训练时找到小批量核集。最后，我们利用零阶方法找到最后一个V投影矩阵的平滑梯度，并稀疏化其以保留归一化梯度幅值最大的维度。我们将CoLM应用于在MathInstruct和SuperGLUE基准上对Phi-2、Phi-3、Zephyr和Llama-3模型进行LoRA微调。显著地，CoLM将微调所需内存降低50%，甚至能超过使用4倍大小批量的性能。此外，CoLM可无缝集成到诸如LoRA等现有内存高效训练方法中，进一步降低大语言模型训练的内存需求。我们的代码已公开于 https://github.com/BigML-CS-UCLA/CoLM。