具有动态稀疏性的内存高效 LLM 训练：从稳定性到实际扩展

动态稀疏训练（DST）为提高深度神经网络的训练和推理效率提供了一个有前途的范例；然而，我们发现在 大语言模型 训练中，DST 可能会遇到优化不稳定的问题，表现为拓扑更新后的损失尖峰。在这项工作中，我们表明，天真地使用基于 Adam 的标准优化器会导致新生成的参数出现冷启动问题，从而导致更新过大并破坏训练动态。为了解决这个问题，我们提出了稀疏内存高效训练（SMET），它通过优化器预热来稳定 DST，并通过密度感知学习率缩放来提高训练进度。 SMET 通过仅存储活动参数的梯度和优化器状态来进一步减少内存消耗。我们对 SMET 下的更新行为进行了理论分析，显示了优化稳定性的提高。大量实验表明，SMET 能够实现稳定、可扩展且内存高效的稀疏 预训练 或 LLM，为稀疏训练作为密集训练的实用替代方案铺平了道路。我们的代码公开于：https://github.com/QiaoXiao7282/SMET。