解决大模型强化学习训练中的内存瓶颈问题,实现高效稳定的大规模采样。
Infinite Sampling: Efficient and Stable Grouped RL Training for Large Language Models
- 将大样本组拆分为小批次,分轮生成降低显存占用
- 通过连续采样和调度优化,吞吐提升25%以上,峰值内存减半
- 适合资源受限环境下大规模语言模型微调的用户
基于群体的强化学习算法(如GRPO)在利用人类反馈微调大语言模型方面表现有效。然而,为每个提示生成并存储多个响应会带来显著内存开销,尤其当样本组规模增大时,限制了在硬件受限环境下的可扩展性。本文提出Infinite Sampling框架,通过解耦组大小与显存使用,实现高效稳定的GRPO训练。其核心包括:(1) 微采样组,将大组分解为内存可承受的生成轮次;(2) 连续采样,跨组交错生成以提高资源利用率;(3) 基于长度感知的调度器,结合基于令牌的序列长度预测与两阶段策略:全局分组采用FPTAS,运行时补填采用SJF。实验表明,微采样组使峰值内存使用量降低超过50%(如在Qwen3-1.7B上从21.55 GB降至10.64 GB)。在此基础上,Infinite Sampling相较朴素微采样方法提升吞吐超25%,减少解码步数,同时保持完整生成长度与低内存占用。混合调度策略确保了在真实显存约束下更大群体的高效稳定训练。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。