LoGRA:用低秩梯度草图降低大模型强化学习的显存门槛
导语
强化学习后训练正在成为提升大语言模型推理能力的重要路径,但它的计算与显存成本往往高于常规监督微调。训练过程中不仅要保存模型参数,还要处理优化器状态、梯度以及策略更新所需的额外信息。对于更大规模的模型,传统的稠密 Adam 优化可能很快触及硬件上限。NVIDIA 研究团队提出的 LoGRA,试图从“如何保存梯度信号”这一环节压缩强化学习的内存开销。
核心方法
LoGRA 的关键思路是把强化学习更新过程中有用的梯度信息保留为低秩梯度草图,而不是始终保存完整的稠密梯度。低秩表示规模更小,因而可以降低训练过程中的内存占用;同时,这些紧凑表示不只用于计算模型更新,还可支持更高效的策略同步。
不过,单纯压缩梯度可能带来另一类风险:如果压缩后的更新幅度过大,当前策略可能偏离过快,进而扰乱强化学习过程。为此,LoGRA 引入预测 KL 步长控制机制。系统会在真正应用更新前估计策略变化,并据此调整更新幅度,让梯度压缩与训练稳定性控制配合工作。
论文摘要披露的实验结果包括:
- 在推理任务上,LoGRA 将平均训练内存最高降低 45.7%,且未牺牲任务表现;
- 在稠密 Adam 已无法容纳的条件下,LoGRA 支持 27B 参数模型在单个八 GPU 节点上稳定训练超过 1,100 步;
- 相关实现已在 Molt 库中开源,代码位于项目的 LoGRA 示例脚本目录。
意义与局限
LoGRA 的价值不只是节省若干显存,更在于把一部分原本受硬件限制的强化学习实验转化为可执行方案。对于研究团队而言,这可能降低大模型 RL 后训练对集群规模的依赖,也为在有限 GPU 资源上探索更大模型提供新的工程路径。其设计还表明,梯度压缩不能脱离更新稳定性单独讨论,压缩效率与策略变化控制需要同时考虑。
但现有素材主要提供了方法概述和摘要级结果,尚未给出不同低秩配置、任务类型、基线设置及额外通信成本的完整对比。因此,45.7% 的节省幅度应理解为论文报告的最高结果,而不是所有训练场景的固定收益。LoGRA 是否适用于更多 RL 算法、不同模型架构和更长训练过程,仍需要结合完整论文与代码进一步验证。
评论
正在确认登录状态……
正在加载评论……