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

LLM RLHF ML RL PGPO

259

2026年10月05日

强化学习已极大提升了大语言模型的能力，但其高昂的内存开销仍是阻碍其更广泛应用的主要瓶颈。本文提出 LoGRA——一种面向强化学习的后训练方法，通过在低秩梯度草图（low-rank gradient sketches）中保留关键的学习信号，显著降低内存占用。这些紧凑的表示形式既支持模型参数更新，又可实现高效的策略同步。为避免更新幅度过大而破坏学习稳定性，我们在梯度压缩的基础上进一步引入“预测 KL 步长控制”（predicted-KL step control）机制：该机制在每次参数更新前预先估计策略分布的变化程度，并据此动态调整更新步长的大小。综合运用上述全部技术后，LoGRA 在各类推理任务上平均降低了高达 45.7% 的训练内存消耗，且未对模型性能造成任何损害。此外，LoGRA 还使得一个 270 亿参数的大模型能够在单台配备八块 GPU 的服务器节点上稳定训练超过 1100 步；相比之下，采用全量参数优化器 Adam 的标准训练方式在此配置下会因显存不足而失败。由此，原本因内存限制而无法开展的强化学习训练任务变得切实可行。相关代码已开源，详见 Molt 库中的 \href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{LoGRA 示例目录}。

Reinforcement learning has greatly advanced the capabilities of large language models, but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations support both model updates and efficient policy synchronization. To prevent overly large updates from disrupting learning, we complement gradient compression with predicted-KL step control, which estimates policy changes before applying each update and adjusts its magnitude accordingly. With all techniques combined, LoGRA reduces average training memory usage by up to 45.7% across reasoning tasks without compromising performance. It also enables stable training of a 27B-parameter model for over 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory, making previously memory-infeasible RL training practical. Code is available in the \href{https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra}{Molt library}.

许愿