沙丘社区
论文

LoGRA：使用低秩梯度草图扩展LLM强化学习

LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches

Shaokun Zhang; Yifan Zhang; Jian Hu; Yueying Li; Hao Zhang; Binfeng Xu; Jan Kautz; Yi Dong
AI 基础设施分布式训练基础设施

摘要

强化学习 (RL) 极大地提升了大语言模型 (LLM) 的功能，但其记忆需求仍然是更广泛采用的障碍。我们引入了 LoGRA，这是一种 RL 后训练方法，它通过在低秩梯度草图中保留有用的学习信号来减少记忆。这些紧凑的表示支持模型更新和高效的策略同步。为了防止过大的更新扰乱学习，我们用预测 KL 步长控制来补充梯度压缩，该控制在应用每次更新之前估计策略的变化并相应地调整其幅度。在推理任务中，LoGRA 在不牺牲性能的情况下将平均训练记忆降低了高达 45.7%。它还可以在单个 8 个 GPU 节点上实现 27B-参数模型的稳定训练超过 1,100 个步骤，其中密集的 Adam 耗尽了记忆，使得之前记忆不可行的 RL 训练变得实用。代码可在\href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{Molt 库}中找到。

LoGRA：使用低秩梯度草图扩展LLM强化学习的原论文方法或结果图
图 5：梯度压缩的组件。 (a) 训练性能作为投影等级的函数； (b) 梯度-固定和更新的 Rademacher 基的重建余弦。 (c) 对刷新的 Rademacher 基和高斯基进行相同的诊断。热图行表示 Transformer 层，列表示等级，具有共享的 $[0,0.4]$ 色标；较高的值表示与密集的梯度更紧密地对齐。