CoDistill-GRPO: A Co-Distillation Recipe for Efficient Group Relative Policy Optimization

TL;DR

CoDistill-GRPO introduces a co-distillation method that trains large and small language models simultaneously, significantly enhancing small model performance and reducing training costs on mathematical reasoning benchmarks.

Contribution

This work presents a novel co-distillation algorithm that enables efficient training of small models by leveraging large models without requiring an oracle, improving performance and speed.

Findings

Small models trained with CoDistill-GRPO outperform standard GRPO on math benchmarks.

Large models with CoDistill-GRPO nearly match standard GRPO performance while using smaller rollouts.

CoDistill-GRPO achieves approximately 18% training speedup, reducing computational costs.

Abstract

Group Relative Policy Optimization (GRPO) has emerged as a powerful algorithm for improving the reasoning capabilities of language models, but often fails to improve small models due to sparse rewards on difficult tasks. Existing works mitigate this issue by leveraging a larger model, either to provide hints for rollouts or to provide dense reward signals through knowledge distillation (KD). However, this assumes the existence of such an oracle, and training one can significantly increase total training time. In this work, we propose CoDistill-GRPO, a co-distillation algorithm that simultaneously trains a large and a small model by maximizing carefully designed GRPO objectives. The two models learn from each other: the small model uses an on-policy KD reward to learn from the large model's distribution, while the large model is updated using rollouts generated by the small model with importance reweighting, reducing the computational overhead of rollout generation. We show that CoDistill-GRPO substantially improves small model performance over standard GRPO on mathematical benchmarks across both Qwen and Llama models. Specifically, with Qwen2.5-Math-1.5B, we observe an accuracy increase of over 11.6 percentage points over the base model and an additional 6.0 percentage points over GRPO on the Minerva dataset. Interestingly, the larger model (Qwen2.5-Math-7B) trained with CoDistill-GRPO nearly matches standard GRPO performance despite training on small-model rollouts. This highlights CoDistill-GRPO as a cost-effective alternative to GRPO for larger models, yielding an approximate 18% speedup, which may be of independent interest.