LoGRA cuts RL post-training memory with low-rank gradient sketches

Reinforcement learning methods such as Group Relative Policy Optimization (GRPO) can improve mathematical reasoning in language models, yet their training loops retain gradients, activations, and optimizer state alongside model weights. The LoGRA paper introduces low-rank gradient sketches for RL post-training. Across the evaluated reasoning tasks, the authors report up to 45.7% lower average training memory and a 27B-parameter model trained for more than 1,100 steps on one eight-GPU node.

Lower memory use can determine whether a training run fits on existing hardware. In the paper’s 27B experiment, dense Adam exhausted available memory under the same tested configuration.

Why Adam fills the GPU

Model weights account for one part of an RL training job’s memory budget. For a 7B-parameter model, the paper’s example breaks down the core allocations as follows:

• BF16 model weights: approximately 14 GB.
• Adam moment buffers: approximately 56 GB for two FP32 states.
• Additional allocations: gradients, activations, rollout data, and framework buffers.

Together, the BF16 weights and Adam moments occupy about 70 GB before the training loop allocates the remaining tensors. The moment buffers alone consume four times as much memory as the BF16 weights.

Reward-based training often derives an update from a small scalar signal attached to each generated response. A binary verifier, for example, may record whether an answer is correct even though backpropagation ordinarily produces a value for every trainable parameter. LoGRA tests the hypothesis that the useful update directions occupy a much smaller space than the dense gradient tensor.