Sparse-RL: Breaking the Memory Wall in LLM Reinforcement Learning via Stable Sparse Rollouts

TL;DR

Sparse-RL tackles the memory bottleneck of KV caches during long-horizon RL rollout for LLMs by uncovering a policy mismatch among dense, sparse, and learner policies. It introduces Sparsity-Aware Rejection Sampling to filter anomalous compressed-rollout trajectories and Importance-based Reweighting to correct off-policy bias, enabling stable training under a fixed KV budget. Empirical results across multiple models and benchmarks show Sparse-RL matches or surpasses dense baselines while achieving substantial memory savings, and its sparse-inference performance is robust due to Sparsity-Aware Training. The approach offers a universal framework for memory-efficient RL in LLMs, compatible with various compression methods and scalable to larger models.

Abstract

Reinforcement Learning (RL) has become essential for eliciting complex reasoning capabilities in Large Language Models (LLMs). However, the substantial memory overhead of storing Key-Value (KV) caches during long-horizon rollouts acts as a critical bottleneck, often prohibiting efficient training on limited hardware. While existing KV compression techniques offer a remedy for inference, directly applying them to RL training induces a severe policy mismatch, leading to catastrophic performance collapse. To address this, we introduce Sparse-RL empowers stable RL training under sparse rollouts. We show that instability arises from a fundamental policy mismatch among the dense old policy, the sparse sampler policy, and the learner policy. To mitigate this issue, Sparse-RL incorporates Sparsity-Aware Rejection Sampling and Importance-based Reweighting to correct the off-policy bias introduced by compression-induced information loss. Experimental results show that Sparse-RL reduces rollout overhead compared to dense baselines while preserving the performance. Furthermore, Sparse-RL inherently implements sparsity-aware training, significantly enhancing model robustness during sparse inference deployment.

Figures (7)

• Figure 1: Impact of Sparse-RL on training stability. We compare training dynamics between naive GRPO and Sparse-RL during Qwen2.5-7B training with R-KV compression. While naive GRPO suffers from a reward collapse and gradient spikes, Sparse-RL maintains stable reward growth with gradient norms remaining consistently low (appearing near the x-axis relative to the massive spikes of the naive baseline).
• Figure 2: Comparison of training curves between GRPO-Dense and GRPO + Sparse-RL (with R-KV) on Qwen2.5-3B over 400 training steps. Left: Average rewards per step. Middle: Average response length (tokens). Right: Policy entropy.
• Figure 3: Comparison of the Mismatch KL between GRPO-Dense and GRPO + Sparse-RL (with R-KV) on Qwen2.5-7B.
• Figure 4: Ablation study on KV cache budget size. Performance of Sparse-RL (with R-KV) on Qwen2.5-3B across different budget levels on MATH500 and Olympiad-Bench. The red dashed line represents the dense rollouts baseline (FullKV).
• Figure 5: Dynamics of rejection rate on Qwen2.5-3B trained by GRPO + Sparse-RL (with R-KV).