论文
大批量何时有助于扩展 LLM 强化学习?
When Do Larger Batches Help Scale LLM Reinforcement Learning?
摘要
较大的批次会减少每次更新的随机梯度的方差,因此通常有望加速训练。然而,这种统计优势是否会转化为更短的目标实现时间仍不清楚,因为每次更新都会消耗更多样本,并且可能需要更长的时间来执行。我们研究了 大语言模型 强化学习中的这种权衡。我们通过沿自然轴比较学习和执行来区分其算法和系统效果。在算法层面,我们比较相同累积样本数的配置,同时重新调整与批次相关的超参数。在有限的批量大小范围内,此过程产生一个近似批量大小不变的族,其成员遵循类似的样本索引学习轨迹。在系统级别,我们利用了采样轨迹生成和训练之间的计算不对称性:自回归生成通常在低并发下受内存带宽限制,而训练工作大约随着处理的词元数量而扩展。结合这两个视图会产生一个直接决策规则:仅当吞吐量增益超过其样本到目标损失时,较大批量配置才会缩短达到目标的时间。 GRPO 和 PPO 的实验支持了这种分解的双方。在算法层面,Adam 的平方根学习率缩放在有限的批量大小范围内产生近似批量大小不变的学习曲线。在系统层面,更大的批量可将固定硬件上的生成吞吐量提高高达 2.29 倍。在 GRPO 中,将更高的吞吐量与学习率重新调整相结合可将实现目标的时间缩短高达 29%,而在不重新调整的情况下增加批次速度会更慢,尽管吞吐量更高。