HiLoRe improves training speed and memory use for group policies
HiLoRe: What to Store, Compress, or Recompute for Efficient GRPO Training
Machine Learning
Summary
Training group-relative policies for AI models often runs into memory problems because it must store many details. The authors found that existing methods waste GPU memory or spend too much time recomputing data. They developed HiLoRe, a way to decide what information to keep exactly, what to compress less precisely, and what to recalculate, all based on how important the data is for the policy update. This approach speeds up training by up to 13.5% and uses memory more efficiently without losing accuracy.
What this means in practice
- •For machine learning engineers: Improve training efficiency and memory use for group-relative policy models on GPUs by balancing precision, compression, and recomputation based on policy update sensitivity.
- •For cloud gpu operators: Optimize GPU resource utilization and reduce training costs for distributed reinforcement learning workloads using HiLoRe's adaptive memory and computation management.
Authors
Xinrui Chen, Mengyang Li, Ou Wu, Ji Zhang
Abstract
Group-relative policy optimization (GRPO) makes learner-side activations a major memory-computation bottleneck: gradient checkpointing reduces activation memory through recomputation, but fixed schedules can leave roughly 18 GB unused on a 48-GB GPU despite substantial recomputation overhead. Existing activation-management methods set state fidelity from execution cost, tensor properties, or generic compression sensitivity, without explicitly incorporating GRPO's analytic update structure into state-fidelity allocation. We formalize this dependence as policy-update exposure, linking the current GRPO loss coefficients to state-level approximation sensitivity. These coefficients are available before backward without an additional backward pass. We introduce HiLoRe, which allocates graph-attributed recovery units among high-precision storage, low-precision compression, and deterministic recomputation using measured recovery utility and update-conditioned approximation risk. It combines high-precision storage and deterministic recomputation with low-precision recovery under a calibrated risk budget. Across five model-task settings with 2K responses and memory < 1.10 times GC's per-GPU actor-update peak, HiLoRe's actor-update throughput gains reach 13.5% over GC and 7.9% over the fastest evaluated baseline, with paired mean downstream-score differences below 0.6 percentage points.