Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward
Group Relative Policy Optimization (GRPO) enhances policy learning by
computing gradients from relative comparisons among candidate outputs that
share a common input prefix. Despite its effectiveness, GRPO introduces
substantial computational overhead when processing long shared prefixes, which
must be redundantly encoded for each group member. This inefficiency becomes a
major scalability bottleneck in long-context learning scenarios. We propose
Prefix Grouper, an efficient GRPO training algorithm that eliminates redundant
prefix computation via a Shared-Prefix Forward strategy. In particular, by
restructuring self-attention into two parts, our method enables the shared
prefix to be encoded only once, while preserving full differentiability and
compatibility with end-to-end training. We provide both theoretical and
empirical evidence that Prefix Grouper is training-equivalent to standard GRPO:
it yields identical forward outputs and backward gradients, ensuring that the
optimization dynamics and final policy performance remain unchanged.
Empirically, our experiments confirm that Prefix Grouper achieves consistent
results while significantly reducing the computational cost of training,
particularly in long-prefix scenarios. The proposed method is fully
plug-and-play: it is compatible with existing GRPO-based architectures and can
be seamlessly integrated into current training pipelines as a drop-in
replacement, requiring no structural modifications and only minimal changes to
input construction and attention computation. Prefix Grouper enables the use of
larger group sizes under the same computational budget, thereby improving the
scalability of GRPO to more complex tasks and larger models. Code is now
available at https://github.com/johncaged/PrefixGrouper
68469df13ec10bdd8ab4db92
Zikang Liu
68469df13ec10bdd8ab4db93
Tongtian Yue
68469df13ec10bdd8ab4db94
Yepeng Tang
68469df13ec10bdd8ab4db95
Longteng Guo
68469df13ec10bdd8ab4db96
Junxian Cai
68469df13ec10bdd8ab4db97
Qingbin Liu
68469df13ec10bdd8ab4db98
Xi Chen
68469df13ec10bdd8ab4db99
Jing Liu
Prefix Grouper: Efficient GRPO Training through Shared-Prefix Forward
Group Relative Policy Optimization (GRPO) enhances policy learning by
computing gradients from relative comparisons among candidate outputs that
share a common input prefix. Despite its effectiveness, GRPO introduces
substantial computational overhead when processing long shared prefixes, which
must be redundantly encoded for each group member. This inefficiency becomes a
major scalability bottleneck in long-context learning scenarios. We propose
Prefix Grouper, an efficient GRPO training algorithm that eliminates redundant
prefix computation via a Shared-Prefix Forward strategy. In particular, by
restructuring self-attention into two parts, our method enables the shared
prefix to be encoded only once, while preserving full differentiability and
compatibility with end-to-end training. We provide both theoretical and
empirical evidence that Prefix Grouper is training-equivalent to standard GRPO:
it yields identical forward outputs and backward gradients, ensuring that the
optimization dynamics and final policy performance remain unchanged.
Empirically, our experiments confirm that Prefix Grouper achieves consistent
results while significantly reducing the computational cost of training,
particularly in long-prefix scenarios. The proposed method is fully
plug-and-play: it is compatible with existing GRPO-based architectures and can
be seamlessly integrated into current training pipelines as a drop-in
replacement, requiring no structural modifications and only minimal changes to
input construction and attention computation. Prefix Grouper enables the use of
larger group sizes under the same computational budget, thereby improving the
scalability of GRPO to more complex tasks and larger models. Code is now
available at https://github.com/johncaged/PrefixGrouper
Prefix Grouper reduces computational overhead in GRPO by encoding shared prefixes only once, improving scalability in long-context scenarios without altering training dynamics or policy performance.
Group Relative Policy Optimization (GRPO) self-attention Shared-Prefix Forward strategy computational overhead long-context learning scenarios differentiability end-to-end training training-equivalent