⏶4
前缀分组器:通过共享前缀前向传播实现高效GRPO训练
发表
由
Zikang Liu 提交
作者: Zikang Liu, Tongtian Yue, Yepeng Tang, Longteng Guo, Junxian Cai, Qingbin Liu, Xi Chen, Jing Liu
摘要
群体相对策略优化(GRPO)通过计算共享相同输入前缀的候选输出之间的相对比较梯度来增强策略学习。尽管GRPO有效,但在处理长共享前缀时会引入大量的计算开销,因为每个组成员都必须冗余地编码该前缀。这种低效率在长上下文学习场景中成为主要的扩展性瓶颈。我们提出了Prefix Grouper,这是一种高效的GRPO训练算法,它通过“共享前缀前向”策略消除了冗余的前缀计算。具体而言,通过将自注意力重构为两部分,我们的方法使共享前缀只需编码一次,同时保留了完全可微分性并与端到端训练兼容。我们提供了理论和经验证据,证明Prefix Grouper与标准GRPO在训练上是等效的:它产生相同的前向输出和后向梯度,确保优化动态和最终策略性能保持不变。经验上,我们的实验证实Prefix Grouper在显著降低训练计算成本的同时取得了一致的结果,尤其是在长前缀场景中。所提出的方法是完全即插即用的:它与现有的基于GRPO的架构兼容,可以作为直接替代品无缝集成到当前的训练流程中,无需结构修改,只需对输入构建和注意力计算进行最少改动。Prefix Grouper使得在相同计算预算下可以使用更大的组大小,从而提高了GRPO在更复杂任务和更大模型上的可扩展性。代码现已在 https://github.com/johncaged/PrefixGrouper 提供。
我们提出了“PrefixGrouper”,一个基于共享前缀前向传播的高效GRPO训练工具。它即插即用,只需对现有代码库进行最小程度的修改,即可实现计算量减少、设备内存消耗降低和训练加速。我们提供了详细教程,欢迎使用!