LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
LoGRA cuts LLM reinforcement-learning memory by up to 45.7% using low-rank gradient sketches.
LoGRA is an RL post-training method that keeps useful learning signals in low-rank gradient sketches, reducing memory while still supporting model updates and policy synchronization. Predicted-KL step control estimates the policy change before each update and adjusts its magnitude. Across reasoning tasks, average training memory falls by up to 45.7% with no reported performance loss. It stably trains a 27B-parameter model for more than 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory. Code is available in the Molt library.
- Low-rank gradient sketches cut average RL training memory by up to 45.7%.
- Predicted-KL step control scales each update before it is applied.
- A 27B model trains stably for more than 1,100 steps on eight GPUs.
- Dense Adam runs out of memory in the same setup.
- Code is released in the Molt library.
From the original article
1 min read at huggingface.co
Reinforcement learning (RL) has greatly advanced the capabilities of large language models (LLMs), but its memory demands remain a barrier to broader adoption. We introduce LoGRA, an approach to RL post-training that reduces memory by retaining useful learning signals in low-rank gradient sketches. These compact representations support both model updates and efficient policy synchronization.
By Shaokun Zhang, Yifan Zhang, Jian Hu · huggingface.co
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.