Recommended NVIDIA paper on memory-efficient post-training RL.
A very interesting technique for scaling RL for LLMs.
The main takeaway is that you can update all of a model's weights with RL and still use less memory than LoRA. At 1.5B, LoGRA averages 7.2 GiB of training memory, while rank-256 LoRA needs 13.2 GiB. Pass@4 is about the same (81.8% and 81.4%), and Pass@1 is about 2 points lower.
If GPU memory is what stops you from running RL on a bigger model, this method is worth a look.
LoGRA never builds a full gradient for the weight matrices it compresses.
As microbatches accumulate, each weight gradient is compressed with a rank-256 random projection that is redrawn every step, and the update is rebuilt from two small matrices. Before each step, it predicts how much the policy will change and scales the step to stay within a KL budget.
Average training memory per GPU drops 45.7% at 7B (31.8 to 17.3 GiB), with peak accuracy close to dense Adam (72.3% and 72.5% Pass@1).
It also trains a 27B model for more than 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory.
Paper: https://arxiv.org/abs/2610.06647
Chat with Paper: https://academy.dair.ai/papers/logra-scaling-llm-reinforcement-learning-with-low-rank-gradient-sketches-2610.06647