52CV 2025-06-10 21:32 江苏
代码开源，方便好用
关注公众号，发现CV技术之美
本文分享一篇最新的技术报告《Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward》 ，聚焦于强化学习优化算法GRPO的改进，旨在消除相同前缀情况下的冗余计算，可使计算量和内存占用均大幅降低，而且即插即用。技术报告、代码均已开源。
背景：GRPO训练的效率瓶颈
Group Relative Policy Optimization(GRPO)是一种强大的强化学习优化算法，其核心思想是：针对同一个输入前缀（Prompt/上下文）生成多个候选响应（Rollout），然后通过比较这些响应的相对优劣（如Reward或人工偏好）来计算策略梯度，进而更新模型。这种方法避免了显式值函数估计，已被广泛应用于指令遵循、思维链推理、多模态任务等场景。
然而，标准GRPO训练存在一个显著的效率瓶颈：在模型的前向计算（Forward Pass）中，每个候选响应都需要独立且完整地编码整个共享的输入前缀。当输入前缀较长（例如包含复杂指令、多模态信息或长文档上下文），并且组大小（Group Size ，即候选响应数量）较大时，这种对相同前缀的冗余计算会消耗海量的计算资源（FLOPs）和GPU内存，成为GRPO方法在长上下文任务中扩展（如使用更大或更长前缀）的主要障碍。
PrefixGrouper——高效GRPO训练算法
为解决这一核心痛点，团队创新性地提出了PrefixGrouper算法。该算法基于“共享前缀前向计算”（Shared-Prefix Forward）的核心策略，实现了GRPO训练效率的飞跃。
特点：共享前缀，一次编码
传统方法(Repeated-Prefix Forward)：将每个候选响应与共享前缀拼接成独立序列输入模型，导致前缀被重复编码次。
PrefixGrouper(Shared-Prefix Forward)：将共享前缀仅输入一次，然后将所有候选响应顺序拼接在其后，形成单个长序列输入模型。
算法的关键在于如何处理这种拼接序列中的自注意力计算。PrefixGrouper将自注意力计算巧妙地拆分为两部分：
前缀自注意力(Prefix Self-Attn)：仅在前缀的tokens之间计算自注意力，更新其表示。
后缀拼接注意力(Suffix Concat-Attn)：每个响应后缀中的tokens，其Query向量仅来自，但其Key和Value向量则来自。
这种分组注意力(Grouped Attention)机制，确保了注意力计算的等价性和高效性。
理论证明：训练等效性(Training Equivalence)
团队证明了PrefixGrouper在前向输出（Forward Outputs）和反向梯度（Backward Gradients）上等效于标准的Repeated-Prefix Forward GRPO训练，这意味着采用PrefixGrouper不会改变GRPO的优化过程和最终学到的策略性能。其带来的效率提升是“免费”的，没有性能妥协的代价。
实验结果：效率显著提升
固定前缀长度下（4096、8192、16384）不同前后缀长度比例（前缀长度/后缀长度）的FLOPs结果。
固定前缀长度下（4096、8192、16384）不同前后缀长度比例（前缀长度/后缀长度）的显存占用结果。
计算量(FLOPs)锐减：理论分析与实验验证均表明，在共享前缀长度远大于单个响应长度的长上下文典型场景下，PrefixGrouper能将自注意力部分的计算量降低到传统方法的。点式操作（MLP，投影层）的计算量也大幅降低。
内存占用大幅降低：由于避免了冗余前缀状态的存储，GPU内存消耗显著下降，尤其在较大时优势更明显。
无缝集成，即插即用(Plug-and-Play)
PrefixGrouper是一个实现层面的优化。它兼容现有的基于Transformer和GRPO的训练框架，只需调整输入数据的组织方式和修改注意力计算的实现（通常是少量代码改动），即可无缝集成到现有训练流程中。另外，该方法对计算设备友好，支持各种训练设备（GPU、NPU等）和attention算法。
代码不仅开源，而且已经有了详细的使用教程，欢迎尝试！
END
欢迎加入「强化学习」交流群👇备注：RL
