titleLoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches
authorsShaokun Zhang, Yifan Zhang, Jian Hu, Yueying Li, Hao Zhang, Binfeng Xu, Jan Kautz, Yi Dong
date2026-10-05
venuearXiv (preprint, not peer reviewed)
idarXiv:2610.06647
affiliationNVIDIA
upvotes45 (Hugging Face Daily Papers, 2026-11-01)
linksabs pdf hf
abstractReinforcement 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. To prevent overly large updates from disrupting learning, we complement gradient compression with predicted-KL step control, which estimates policy changes before applying each update and adjusts its magnitude accordingly. Across reasoning tasks, LoGRA reduces average training memory by up to 45.7\% without sacrificing performance. It also enables stable training of a 27B-parameter model for over 1,100 steps on a single eight-GPU node, where dense Adam runs out of memory, making previously memory-infeasible RL training practical. Code is available in the https://github.com/skzhang1/labs-molt/tree/logra/examples/scripts/logra{Molt library}.

cite

@misc{zhang2026logra,
  title={LoGRA: Scaling LLM Reinforcement Learning with Low-Rank Gradient Sketches},
  author={Shaokun Zhang and Yifan Zhang and Jian Hu and Yueying Li and Hao Zhang and Binfeng Xu and Jan Kautz and Yi Dong},
  year={2026},
  eprint={2610.06647},
  archivePrefix={arXiv},
  url={https://arxiv.org/abs/2610.06647}
}