Infinite Sampling: Efficient and Stable Grouped RL Training for Large Language Models

📄 arXiv: 2506.22950v1 📥 PDF

作者: Liangyu Wang, Huanyi Xie, Xinhai Wang, Tianjin Huang, Mengdi Li, Di Wang

分类: cs.LG

发布日期: 2025-06-28

💡 一句话要点

提出Infinite Sampling框架，解决LLM分组强化学习训练中内存瓶颈问题。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大型语言模型 强化学习 分组训练 内存优化 吞吐量提升

📋 核心要点

  1. 基于分组的强化学习算法在微调LLM时有效，但生成和存储多组回复导致内存开销巨大，限制了硬件受限情况下的扩展性。
  2. Infinite Sampling框架将组大小与GPU内存解耦，通过微采样分组、连续采样和长度感知调度，实现高效稳定的GRPO训练。
  3. 实验表明，该方法显著降低了内存使用，提高了吞吐量，并在GPU内存约束下实现了高效稳定的GRPO训练。

📝 摘要（中文）

本文提出Infinite Sampling框架，旨在解决使用基于分组的强化学习算法（如GRPO）微调大型语言模型（LLM）时，因每个prompt生成和存储多个response而导致的大量内存开销问题。Infinite Sampling通过将组大小与GPU内存使用解耦，实现高效且稳定的GRPO训练。该框架包含：(1) 将大组分解为内存可行的轮次进行微采样分组；(2) 交叉进行跨组生成以提高利用率的连续采样；(3) 结合token条件序列长度预测和两阶段计划（通过FPTAS进行全局分组，通过SJF进行运行时重新填充）的长度感知调度器。实验表明，微采样分组可将峰值内存使用量降低50%以上，而Infinite Sampling相比朴素的微采样分组方法，吞吐量提高了25%以上，同时保持了完整长度的生成和内存使用率。

🔬 方法详解

问题定义：现有基于分组的强化学习算法，如GRPO，在微调大型语言模型时，需要为每个prompt生成多个response，这导致了巨大的内存开销，尤其是在组大小增加时。这种内存瓶颈限制了算法在硬件资源有限的环境中的可扩展性。现有方法的痛点在于无法在保证训练效果的同时，有效控制内存占用。

核心思路：Infinite Sampling的核心思路是将大组采样分解为多个小的、内存可承受的微采样轮次。通过连续采样，交错进行不同组的生成，提高GPU的利用率。同时，利用长度感知调度器，根据token条件序列长度预测，动态调整采样策略，在保证生成质量的前提下，优化内存使用和吞吐量。

技术框架：Infinite Sampling框架主要包含三个模块：微采样分组（Micro Sampling Groups）、连续采样（Continuous Sampling）和长度感知调度器（Length-aware Scheduler）。微采样分组将大的采样组分解为多个小的、内存可承受的轮次。连续采样交错进行不同组的生成，提高GPU利用率。长度感知调度器则结合token条件序列长度预测和两阶段计划（全局分组和运行时重新填充），动态调整采样策略。

关键创新：该方法最重要的创新点在于将组大小与GPU内存使用解耦。通过微采样分组，可以在不增加内存开销的情况下，实现更大的有效组大小。连续采样和长度感知调度器进一步优化了GPU利用率和吞吐量。与现有方法相比，Infinite Sampling能够在内存受限的环境中，更高效地进行LLM的强化学习微调。

关键设计：长度感知调度器是关键设计之一，它包含token条件序列长度预测模型，用于预测生成序列的长度。两阶段计划包括：(1) 全局分组，使用完全多项式时间近似方案（FPTAS）进行初始分组；(2) 运行时重新填充，使用最短作业优先（SJF）算法动态调整分组。这些设计旨在平衡内存使用、吞吐量和生成质量。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，Micro Sampling Groups可以将峰值内存使用量降低超过50%，例如在Qwen3-1.7B模型上，从21.55 GB降低到10.64 GB。此外，Infinite Sampling相比于朴素的微采样分组方法，吞吐量提高了超过25%，在保持完整长度生成和内存使用率的同时，减少了解码步骤。

🎯 应用场景

该研究成果可应用于各种需要使用强化学习微调大型语言模型的场景，尤其是在计算资源受限的环境中。例如，在边缘设备上部署LLM，或者在预算有限的实验室中进行LLM研究。该方法可以降低训练成本，提高训练效率，加速LLM的开发和应用。

📄 摘要（原文）

Group-based reinforcement learning algorithms such as Group Reward Policy Optimization (GRPO) have proven effective for fine-tuning large language models (LLMs) with human feedback. However, generating and storing multiple responses per prompt incurs substantial memory overhead, especially as the sample group size increases, limiting scalability under constrained hardware. We propose Infinite Sampling, a framework that enables efficient and stable GRPO training by decoupling group size from GPU memory usage. It consists of: (1) micro sampling groups that decompose large groups into memory-feasible rounds; (2) continuous sampling that interleaves generation across groups to improve utilization; and (3) a length-aware scheduler combining token-conditioned sequence length prediction with a two-stage plan: global grouping via FPTAS and runtime refill via SJF. Experiments show that our Micro Sampling Groups reduce peak memory usage by over 50% compared to full-group decoding (e.g., from 21.55 GB to 10.64 GB on Qwen3-1.7B). Building on this, Infinite Sampling improves throughput by over 25% compared to the naive micro sampling group method, reducing decoding steps while maintaining full-length completions and memory usage. Our hybrid scheduling ensures efficient and stable GRPO training with larger groups under realistic GPU memory constraints.