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

LoGRA. Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches. By Shaokun Zhang and seven others, from NVIDIA. Published on the fifth of October. Its code repository is worked out: about one hundred and twenty source files, with a readme and a licence.

Reinforcement learning has greatly advanced what large language models can do. Its memory demands, though, remain a barrier to broader adoption. In practice that comes down to whether a training run fits in the hardware at hand, and past a certain model size it does not.

LoGRA is an approach to RL post-training that cuts memory by keeping the useful learning signal in low-rank gradient sketches. The bet is that the signal that matters can be held in a far more compact form than the full gradient. Those compact representations then do two jobs. They drive the model updates, and they make policy synchronization efficient, so the same saving serves both halves of the loop.

Alongside the compression sits a guard against overly large updates, which can disrupt learning. It is called predicted-KL step control. Before each update is applied, it estimates how much the policy will change, and it adjusts the update's magnitude accordingly. The step is sized by its predicted effect on the policy, rather than taken as it comes.

Across reasoning tasks, LoGRA reduces average training memory by up to forty-five point seven percent without sacrificing performance. It also trains a twenty-seven billion parameter model stably for more than one thousand one hundred steps on a single eight-GPU node. On that same node, dense Adam runs out of memory. Code is available in the Molt library.

The second result carries the point. Saving memory matters most when it turns a run that does not fit into one that does, and the authors frame LoGRA as doing exactly that: making previously memory-infeasible RL training practical on one machine, with a step controller keeping the compressed updates stable along the way.