The memory wall in LLM reinforcement learning

Most organisations can already run a large language model. Far fewer can improve one. Reinforcement learning — the stage that turns a fluent model into one that holds up on hard, checkable work — is where the hardware bill stops being about weights and starts being about everything stored beside them. A paper from NVIDIA researchers, LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches, is useful less as a leaderboard result than as a precise account of that bill, and of one way to shrink it.

The bottleneck moved

For two years the public conversation about generative AI has been about inference: context length, tokens per second, and whether a model can be served on hardware a company already owns. Post-training sat in a different room. Supervised fine-tuning was expensive but familiar. Reinforcement learning on the model’s own answers was treated as a frontier-lab privilege.

That separation is now the wrong operating assumption. Reasoning quality — mathematics, code, tool use, multi-step procedures with a verifier — is increasingly produced by reinforcement learning with verifiable rewards, not by another pass of imitation. The organisations that will own that quality are the ones that can run the loop: generate, score, update, synchronise the policy that generates the next batch. Memory, not a missing algorithm, is what throws most teams out of that loop.

LoGRA’s authors state the asymmetry cleanly. The BF16 weights of a 7-billion-parameter Qwen model occupy about 14 GB. Adam’s two FP32 moment buffers add about 56 GB before gradients and activations are counted. A machine that can serve the model can still be unable to teach it. Hardware that is “enough to run it” is not the same specification as hardware that is “enough to improve it.”

What the method actually compresses

LoGRA is not another adapter. LoRA freezes the pretrained weights and learns a small pair of matrices beside them. LoGRA keeps the full weight matrices trainable. What it refuses to store is the full gradient.

For a weight matrix of shape d by k, a dense gradient is the same size as the matrix. LoGRA multiplies that gradient, during backpropagation, by a thin random projection of rank r — far smaller than either dimension — and accumulates only the sketch. The sketch is then multiplied back out to form a low-rank update, which is written directly into the weights. There is no persistent adapter branch. The projection is not learned. By default its entries are random signs, and it can be redrawn each step so successive updates are not trapped in one subspace.

The compression is severe, and it is the point. For a 4,096 by 4,096 matrix and rank 64, an FP32 gradient buffer falls from 64 MiB to 1 MiB. The method is applied to attention and MLP projection matrices; other parameters stay fixed. Because the sketch is accumulated from microbatches without ever materialising the full gradient, the saving shows up in the backward pass itself, not only in an optimiser state that was computed and then thrown away.

The same compact object does a second job. RL training is a two-engine system: one set of GPUs updates the policy, another generates the next responses. Those engines must agree on the weights. LoGRA sends the thin update factor and a seed, and the rollout side rebuilds the projection. The payload is proportional to d times r, not to the full matrix. Compression therefore buys training memory and policy synchronisation at once. That second saving is easy to miss in a benchmark table and expensive to miss in a real cluster, where generation and training already fight over the same nodes.

Why a smaller gradient is not automatically a safer step

A compact update can still wreck a policy. Parameter distance and behavioural distance are different quantities. A step that looks modest in weight space can move next-token probabilities enough to erase what the model had already learned. This is the old trust-region lesson from TRPO, restated for a setting where the loss signal is often a single correctness bit on a long generated answer.

LoGRA’s companion is predicted-KL step control. Before an update is applied, the method estimates how far the proposed step would move the model’s next-token distribution, using a local approximation: predicted KL grows with the square of the step scale. A budget δ is chosen. If the proposal would exceed it, the step is shrunk. If the predicted change is tiny, the step may grow, but only up to a cap. A constant budget holds that tolerance fixed. An annealed budget tightens it as training proceeds, on the view that late updates should disturb a working policy less than early ones.

The estimate is local. Realised KL can differ from the prediction. The paper is careful about that, and leaders should be too. The control is a governor, not a proof that training cannot diverge. In the authors’ longer ablations, a large constant budget degraded accuracy, while a small budget and an annealed schedule stayed stable and finished slightly ahead of an uncontrolled baseline. Conservative step size was the useful idea. Annealing was a refinement, not a magic setting.

What the experiments support — and what they do not

The evaluation is deliberately plain, which makes it easier to read. Models are Qwen variants from 1.5B to 27B parameters. The task family is mathematical reasoning, chiefly DAPO-Math with MATH-500 as the scoreboard, plus a longer run of a 27B model on Reasoning-Gym Hard. The algorithm is PPO with verifiable rewards, one response per prompt, and a simple running reward baseline. The authors chose that simplicity so the memory method would not be confused with a fancier RL recipe. Training ran on a single node of eight H100 80 GB GPUs, four for sharded training and four for generation.

On memory, the result is the one that should change a capacity plan:

• At 1.5B, average training memory fell from 9.18 to 7.18 GiB per GPU, about 22 percent, with peak accuracy on MATH-500 a few points higher than dense Adam rather than lower.
• At 7B, average memory fell from 31.82 to 17.29 GiB, a reduction of 45.7 percent. Peak Pass@1 stayed essentially tied with the dense baseline, near 72 percent.
• At 27B, dense Adam ran out of memory on the first optimiser-state allocation. LoGRA trained, using about 51.5 GiB on average, and on Reasoning-Gym Hard held a stable run past 1,100 steps. Held-out verifier score rose from roughly 40 percent to a peak near 66 percent, while answers grew shorter and less often truncated.

Throughput, where both methods could finish, barely moved. The gain is feasibility and headroom, not a faster step. That distinction matters in a business case. LoGRA does not promise more updates per hour on a model you can already train. It promises that a model you could only serve can now be post-trained on the node you already have — and that the 7B class, which many enterprises treat as the practical ceiling, stops consuming most of an 80 GB GPU before the first useful experiment.

Rank mattered more than the flavour of the random projection. Raising rank from 4 to 256 improved short-horizon accuracy by a few points. Refreshing the projection each step, versus holding it fixed, did not show a consistent gain over the horizons they measured. The main recipe uses rank 256. Treat rank as a capacity dial: higher rank, closer to the dense gradient, more memory. It is not a hyperparameter to leave at the demo default and forget.

LoRA is a different bet

On a 1.5B comparison limited to the first 300 steps, LoRA used more training memory than LoGRA — about 13.2 GiB against 7.2 GiB — and posted a higher mean Pass@1, while LoGRA was slightly ahead on Pass@4. Update time favoured LoRA by a small margin. The paper does not claim an accuracy or speed victory. It claims a memory alternative.

For a programme owner, the architectural difference is the decision, not the table. LoRA is the right instrument when you want a reversible, shippable delta that leaves the base model untouched: a tenant-specific behaviour, a narrow task, an adapter you can load and unload. LoGRA is the right instrument when the goal is to change the model itself — to fold a reasoning policy into the weights — and the thing you cannot afford is Adam’s full state on every projection matrix. Using one because a vendor deck said “low rank” is how teams buy the wrong constraint.

The same clarity applies to the systems tricks that already exist. FSDP shards state across GPUs. Gradient checkpointing trades compute for activation memory. CPU offload moves state to the host and pays in bandwidth. LoGRA does not replace those. It compresses the gradient and the synchronisation payload so those other techniques have less to carry. A serious training platform will stack them. A slide that presents any one of them as the whole answer is selling a component as an architecture.

What a leadership team should do with this

The paper is a research result on public models and public math benchmarks, from Shaokun Zhang, Yifan Zhang, Jian Hu, Yueying Li, Hao Zhang, Binfeng Xu, Jan Kautz, and Yi Dong, with code in the Molt library. It is not a procurement specification. It is, however, specific enough to change the questions a CIO or head of AI should ask before the next reasoning programme is funded.

Separate “runs” from “learns” in the hardware plan. Ask for two numbers: memory to serve the model, and memory to run RL post-training with the optimiser you actually intend to use. If the second number is missing, the plan is an inference plan wearing a training label.

Treat the reward as the product, and the optimiser as the constraint. Verifiable rewards — a unit test, a proof checker, a schema, a policy engine — are what make this style of training honest. LoGRA lowers the cost of applying that signal. It does not invent the signal. A programme with cheap gradients and a vague reward will scale the wrong behaviour more efficiently.

Budget for the two-engine loop. Generation, scoring, update, and weight sync are one system. A memory win that ignores how the rollout policy stays current will reappear as idle GPUs and stale samples. Ask whether compressed updates are part of synchronisation, not only part of the backward pass.

Put a trust region in the runbook. Predicted KL is a practical form of a principle every RL system needs: do not let one step move the policy further than the data can justify. Whether a team uses this estimator or another, an unbounded learning rate on a reward model is how a promising curve becomes a long, confident, wrong model.

Do not generalise past the evidence. These experiments are reasoning tasks with automatic checkers, on Qwen models, on one eight-GPU node. They do not speak to preference tuning on open-ended dialogue, to multimodal policies, or to a 70B-class run. The mechanism — sketch the gradient, control the KL, write the update into the full weights — is general. The measured savings are not a universal percentage to paste into a business case.

What remains expensive after the gradient shrinks

Memory was the barrier this paper moved. It was never the only one. Evaluation sets that match production, not a public math exam, still have to be built. Reward hacking still appears the moment the checker is narrower than the behaviour you care about. Data governance, model release, and the decision of whose corrections are allowed to move the weights are organisational problems. Inference cost does not fall because training used a sketch. A shorter, better answer can reduce it; LoGRA’s 27B run did produce shorter responses, but that is an observation on one trajectory, not a guarantee.

There is also a governance consequence the authors note and that operators should take seriously. Lowering the hardware bar broadens who can post-train a capable model. That is the desirable outcome for a firm that wants reasoning quality without renting a frontier cluster. It is also why evaluation, access control, and a clear owner for the reward function have to be designed before the first long run, not after a model has been cheaply improved in a direction nobody reviewed.

Our view

The useful shift in this work is conceptual. Reinforcement learning stopped being a research luxury at the moment reasoning quality became a product requirement. It remained a luxury in practice because the optimiser state was larger than the model. LoGRA argues, with measurements, that a reward signal sparse enough to be a single correctness bit does not need a dense gradient to be useful — provided the step is kept inside a behavioural budget.

Halcer’s counsel to technology leaders is to read results like this as architecture, not as a library pin. Decide whether the programme needs an adapter or a changed model. Price the training loop separately from inference. Insist on a verifiable reward and a bound on how far any update may move the policy. Then choose the compression, the sharding, and the hardware against that design. The firms that will look current in a year are not the ones that cite the newest training paper. They are the ones that can improve a model on Monday without pretending the cluster is infinite.

This brief interprets LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches (arXiv:2610.06647). The method, experiments, and figures are the authors’ work. The operating conclusions are Halcer’s.

Talk to us about reasoning in production

← Back to Insights