LoGRA: How Low-Rank Gradient Sketches Make LLM Reinforcement Learning Fit on Real Hardware

Reinforcement learning post-training has become one of the most effective ways to improve large language model reasoning. Models like DeepSeek-R1 demonstrated that RL-based fine-tuning can unlock capabilities that supervised learning alone cannot reach. The catch: RL post-training is extraordinarily memory-hungry. Gradients and optimizer states can consume several times more GPU memory than the model weights themselves, putting serious RL runs out of reach for most practitioners.

A new paper from NVIDIA and collaborators, LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches, proposes a practical solution. By replacing full gradient buffers with compact low-rank representations — and pairing that compression with a stability mechanism called predicted-KL step control — LoGRA cuts average training memory by up to 45.7% without sacrificing task performance. The headline result: a 27B-parameter model trained stably for over 1,100 steps on a single eight-GPU node, a configuration where standard dense Adam simply runs out of memory.

Why RL Post-Training Is So Expensive

To understand what LoGRA solves, it helps to understand where the memory goes during RL training. When you fine-tune a model with a method like GRPO (Group Relative Policy Optimization), the training loop involves:

1. Rollout generation — the model samples responses, which are scored by a reward signal.
2. Gradient computation — backpropagation through the policy model to compute parameter gradients.
3. Optimizer state storage — Adam-style optimizers maintain first and second moment estimates for every parameter.

For a 7B-parameter model trained with dense Adam, the gradient and optimizer state buffers alone can push peak memory to over 30 GiB per GPU — before accounting for activations or the rollout buffer. Scaling to 27B makes this completely infeasible on a single node.

Existing workarounds like Fully Sharded Data Parallel (FSDP) and CPU offloading help by distributing or relocating state, but they don't reduce the fundamental size of the gradient tensors. LoGRA attacks the problem at the source.

The Core Idea: Gradient Sketches

LoGRA's central contribution is replacing full gradient matrices with low-rank sketches during backpropagation. Here's the mechanics:

For a weight matrix W of shape d × k, the true gradient G has the same shape. Instead of storing all d × k values, LoGRA samples a random projection matrix A of shape r × k, where r ≪ min(d, k), and accumulates a sketch:

S = G · Aᵀ   (shape: d × r)

When it's time to update the weights, the sketch is expanded back:

Ĝ = S · A   (rank-r approximation of G)

This reduces gradient storage from d × k to d × r — a factor of k/r reduction. The same compact sketch is also used for policy synchronization across nodes in distributed training, extending the memory savings to communication costs.

The approach is applied selectively to attention and MLP projection matrices, which account for the bulk of gradient memory. The authors pair this with a custom optimizer called RowAdam, which maintains running estimates of squared row magnitudes without first-moment momentum — a design that fits naturally with the low-rank update structure.

Keeping Training Stable: Predicted-KL Step Control

Gradient compression introduces a risk: an approximate gradient might produce an update that is too large, causing the policy to shift dramatically and destabilize training. This is a real concern in RL settings, where the reward landscape is non-stationary and large policy jumps can be catastrophic.

LoGRA addresses this with predicted-KL step control. Before applying any weight update, the method estimates how much the model's next-token probability distribution would change — measured by KL divergence — if the sketched update were applied at full scale. It then computes a scaling factor α that keeps the actual KL shift within a prescribed budget δ:

W ← W - α · η · Ŝ · A

where η is the nominal learning rate and α is adjusted so the predicted KL divergence stays below δ. This acts as a trust-region constraint, but one that is computed cheaply from the sketch itself rather than requiring a separate forward pass.

The result is that LoGRA can apply aggressive gradient compression without the training instability that would normally follow. The predicted-KL controller absorbs the approximation error by being conservative about step size when the sketch suggests a large policy shift.

What the Numbers Show

The authors evaluated LoGRA on Qwen models ranging from 1.5B to 27B parameters, using mathematical reasoning benchmarks. The memory results are concrete:

Model	Dense Adam	LoGRA	Reduction
1.5B	9.18 GiB	7.18 GiB	21.8%
7B	31.82 GiB	17.29 GiB	45.7%

LoGRA also compares favorably against LoRA. For the 1.5B model, LoGRA uses 7.18 GiB average versus LoRA's 13.21 GiB — LoRA adapts only a subset of parameters but still requires full gradients through frozen layers during backpropagation.

The 27B experiment is the most striking. Dense Adam fails with an out-of-memory error on a single eight-GPU H100 node. LoGRA trains stably for over 1,100 steps, with the held-out macro score improving from 39.69% to 62.94%.

How LoGRA Relates to GaLore

Readers familiar with GaLore (Gradient Low-Rank Projection) will notice the conceptual overlap. GaLore also uses low-rank gradient projections, but it targets pre-training and supervised fine-tuning, periodically recomputing the projection subspace via SVD.

LoGRA is specifically engineered for RL post-training, where the gradient landscape shifts continuously as the policy evolves. The predicted-KL step control is the key addition that makes low-rank compression viable in this noisier setting. The random projection matrix A also requires no periodic recomputation — simpler to maintain than GaLore's SVD-based subspace.

Practical Implications

LoGRA is implemented in the Molt library, making it accessible for practitioners who want to run RL post-training on models that would otherwise require multi-node clusters. The method is complementary to system-level techniques like FSDP and CPU offloading — you can combine them for even larger models.

The broader significance is that RL post-training, which has driven some of the most notable reasoning improvements in recent LLMs, becomes more accessible. A team with a single eight-GPU node can now run meaningful RL experiments on a 27B model. That's a meaningful shift in who can participate in this kind of research and development.

Conclusion

LoGRA addresses a real bottleneck in LLM development: the memory cost of reinforcement learning post-training. By sketching gradients into low-rank representations and using predicted-KL step control to keep training stable, it achieves up to 45.7% memory reduction while maintaining task performance. The ability to train a 27B model on a single node where dense Adam fails is a concrete demonstration of what this unlocks.

For practitioners working on reasoning-focused LLMs, LoGRA is worth evaluating as a drop-in replacement for standard optimizers in RL post-training pipelines. The paper and Molt library provide the implementation details needed to get started.