Knapsack RL: Compute-Efficient Reinforcement Learning via Heterogeneous Rollout Allocation
Abstract
Reinforcement learning (RL) fine-tuning of Large Language Models (LLMs) is compute-intensive because each prompt requires generating multiple responses, or rollouts. To make the best use of GPU resources, the key question is how to allocate rollout-generation jobs across prompts. Existing methods typically use uniform allocation, assigning every prompt the same rollout budget. This is inefficient and ineffective: easy prompts are oversampled after they are already solved, while hard prompts receive too little exploration. In policy optimization methods such as Group Relative Policy Optimization (GRPO), both cases often yield near-zero gradients and limited learning progress. We address this problem by formulating rollout allocation as a computeconstrained resource allocation task, where each prompt-budget choice has an expected learning value and rollout cost. Based on this view, we propose Knapsack RL, a practical framework that uses knapsack optimization to assign heterogeneous, task-specific rollout budgets. It directs exploration toward prompts that benefit most from additional sampling. Applied to GRPO, Knapsack RL increases the effective-gradient ratio by up to 40%, enables larger budgets for challenging prompts, and improves mathematical reasoning by 2–4 points on average, with peak gains of up to 9 points. Achieving similar performance with uniform allocation requires about 2× more compute, demonstrating a simple and practical path toward scaling RL fine-tuning for LLMs.
Lay Summary
Training large language models with reinforcement learning is expensive because the model must generate many candidate answers for each training question. Today, many systems give every question the same number of attempts, even though some questions are already easy while others need more exploration. This paper studies how to use the same computing budget more wisely. We propose Knapsack RL, a method that gives different numbers of attempts to different questions. It gives fewer attempts to questions that are already solved or unlikely to help, and more attempts to questions where extra answers are likely to improve learning. Across several mathematical reasoning tests, this strategy improves performance compared with giving every question the same budget. It also shows that similar performance with uniform allocation may require about twice as much generation compute.