GRPO-Zero Documentation

The project uses the CountDown task, where the model must generate a mathematical expression using a provided set of numbers that evaluates to a target number.

The model is trained to follow a specific reasoning format:

<think>Model step by step reasoning</think>
<answer>Final answer</answer>

The reward function is composed of two parts:

1. Format Reward: A reward of 0.1 is given if the model correctly uses the <think> and <answer> tags; otherwise, it receives 0.
2. Answer Reward: A reward of 1 is given if the final answer uses each provided number exactly once and correctly evaluates to the target value; otherwise, it receives 0.

GRPO (Group Relative Policy Optimization) trains LLMs by sampling multiple answers for each question and using their relative rewards to define advantages, eliminating the need for a value estimation network.

GRPO:Zero implements several specific improvements:

• Token-level policy gradient loss: Every token is equally weighted in the loss.
• Removing KL Divergence: The KL divergence is omitted from the policy gradient loss, which reduces GPU memory usage because a reference policy network is no longer required.
• Overlong episode filtering: Controlled by the skip_unfinished_episodes setting, this skips episodes that exceed context limits to stabilize training.

To use GRPO:Zero, you need to initialize the environment using uv, install git-lfs for large file handling, and download the required dataset and pretrained model.

Follow these steps:

1. Initialize the environment using uv.
2. Install git-lfs to handle large model and dataset files.
3. Download the dataset from Hugging Face.
4. Download the pretrained model (specifically Qwen2.5-3B-Instruct).

# initialize the environment
pip install uv
uv sync

# install git-lfs
apt update; apt install git-lfs -y; git lfs install

# download the dataset
git clone https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4

# download the pretrained model
git clone https://huggingface.co/Qwen/Qwen2.5-3B-Instruct

You can start training the model using the train.py script via uv run.

Depending on your hardware, choose one of the following commands:

• Standard Training: For high VRAM GPUs (e.g., an A40 with 48GB VRAM).
• Low VRAM Training: For GPUs with 24GB VRAM (e.g., an RTX 4090). This uses a specific configuration that offloads the optimizer to the CPU to save memory.

Note: You can enable overlong episode filtering by setting skip_unfinished_episodes to true in your configuration to stabilize training by skipping unfinished episodes that exceed context length limits.

# train the model
uv run train.py

# train the model with a 24GB VRAM GPU (e.g., an RTX 4090 GPU)
uv run train.py --config config_24GB.yaml