LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
Shaokun Zhang, Jan Kautz, Yi Dong and colleagues at NVIDIA propose LoGRA, which cuts RL post-training memory by accumulating gradients as low-rank random sketches instead of full gradient buffers, and adds predicted-KL step control to keep updates stable.
Ask this paper
Gradient sketches. A rank-256 random-sign projection, redrawn each step and held fixed across that step's microbatches, compresses each selected weight gradient as microbatches accumulate, so no full gradient buffer is stored for those matrices. The update is reconstructed as the product of two small matrices, which also makes policy synchronization to rollout workers cheaper.
Step control. Before applying an update, LoGRA predicts the resulting KL change in the policy and scales the step to stay within an annealed budget. This replaces learning-rate annealing as the stability mechanism.
Memory. Average training memory falls 21.8% at 1.5B (9.18 to 7.18 GiB) and 45.7% at 7B (31.82 to 17.29 GiB) with peak accuracy close to or above dense Adam (7B Pass@1 72.33% vs 72.48%; 1.5B Pass@1 67.87% vs 63.77%). At 1.5B it also uses less memory than rank-256 LoRA (7.18 vs 13.21 GiB) at similar Pass@4.
27B on one node. LoGRA trains a 27B model for more than 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory.
Abstract
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}.