GRPO:Zero

GRPO 训练，依赖极低（且 GPU 显存占用少！）。我们几乎从头开始实现所有内容，仅依赖 tokenizers 进行分词，以及 PyTorch 进行训练。

• 无需 transformers 和 vLLM 依赖！
• 默认配置设置为在单张 A40 GPU（48GB 显存）上运行几个小时，即可获得良好效果。（如果从 RunPod 租用 A40，每小时费用约为 $0.44。）
• 我们还支持使用 24GB 显存的 GPU（例如 RTX 4090）进行训练，方法是将优化器卸载到 CPU 上。幸运的是，这只会给训练带来很小的额外开销，因为在整个训练过程中，我们只更新策略网络几百次。
• 我们对来自 DAPO 项目 的原始 GRPO 算法进行了多项改进，包括：
  • 基于 token 的策略梯度损失：策略梯度损失中每个 token 的权重相等。
  • 移除 KL 散度：策略梯度损失中不再使用 KL 散度。这减少了 GPU 显存占用，因为我们不再需要参考策略网络。
  • 过长 episode 过滤：跳过超出上下文长度限制的未完成 episode，从而稳定训练。不过，我们默认将其关闭，以便观察模型在有限上下文长度下的学习情况。可将 skip_unfinished_episodes 设置为 true 来启用。

算法

组相对策略优化（GRPO）是由 Deepseek 提出的一种用于通过强化学习训练大型语言模型的算法。其思路很简单：对于每个问题，我们随机采样多个答案。然后将某个答案的优势定义为归一化的奖励。这样就无需价值估计网络了。具体来说，我们实现了以下算法：

1. 对于每个训练步骤，随机采样 $N$ 个问题 $q_1, q_2, \cdots, q_N$。
2. 对于每个问题 $q_i$, 采样 $M$ 个答案 $a_{i,1}, a_{i,2}, \cdots, a_{i,M}$。
3. 计算每个答案 $a_{i,j}$ 的奖励 $r_{i,j}$。
4. 计算每个问题 $q_i$ 的奖励均值和标准差。

$$ \begin{aligned} \mu_i &\leftarrow \text{mean}(r_{i,1}, r_{i,2}, \cdots, r_{i,M}) \ \sigma_i &\leftarrow \text{std}(r_{i,1}, r_{i,2}, \cdots, r_{i,M}) \end{aligned} $$

5. 对于答案 $a_{i,j}$ 中的每个 token $t$, 计算优势：

$$A_{i,j}[t] \leftarrow \frac{r_{i,j} - \mu_i}{\sigma_i}$$

6. 使用 PPO 替代目标函数计算策略梯度。为简化起见，我们每次迭代只进行一次策略更新，此时 PPO 目标的梯度等同于逐 token 的常规策略梯度估计：

$$ \nabla_\theta \log \pi_\theta(a_{i,j}[t]) \cdot A_{i,j}[t] $$

7. 使用该梯度更新策略网络 $\pi(\theta)$。然后返回步骤 1。

CountDown 任务

我们将使用 CountDown 任务 来训练 Qwen2.5 模型。给定 3 或 4 个数字以及一个目标数字，模型需要生成一个仅包含简单算术运算（+、-、*、/）的数学表达式，使其结果等于目标数字。例如：

问题：给定 1 2 3 4 和目标数字 11。请写出一个结果为 11 的表达式。
答案：1 + (2 * 3) + 4

奖励函数

为了解决 CountDown 任务，我们将使用 GRPO 算法来训练模型，在生成最终表达式之前先输出思维链推理过程。具体而言，模型被训练按照以下格式生成内容：

<think>模型逐步推理过程</think>
<answer>最终答案</answer>

奖励由两部分组成：

1. 格式奖励：当模型正确遵循包含 thinking 和 answer 标签的指定格式时，获得奖励 0.1；否则得 0。
2. 答案奖励：如果模型的最终答案恰好使用了每个给定数字一次，并且正确计算出目标值，则获得奖励 1；否则得 0。

训练

我们使用 Qwen2.5-3B-Instruct 模型进行训练。要训练模型，请运行以下命令：

# 初始化环境
pip install uv
uv sync

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

# 下载数据集
git clone https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4

# 下载预训练模型
git clone https://huggingface.co/Qwen/Qwen2.5-3B-Instruct
# 训练模型
uv run train.py
# 使用 24GB 显存的 GPU（例如 RTX 4090）训练模型
uv run train.py --config config_24GB.yaml

致谢

本项目建立在多个杰出项目的工作基础上：

• DeepSeekMath，率先提出了 GRPO 算法。
• DAPO，对原始 GRPO 算法进行了多项改进。
• TinyZero，实现了 GRPO 并创建了 CountDown-Tasks-3to4 数据集。
• nano-aha-moment，提供了清晰的 GRPO 算法实现及教程。
• Qwen2.5，开发了本项目中使用的高质量预训练模型。

GRPO-Zero 快速上手指南

GRPO-Zero 是一个极简依赖的强化学习训练框架，专为低显存环境设计。它从零实现了 GRPO 算法，仅需 tokenizers 和 pytorch，无需 transformers 或 vLLM。默认配置可在单张 A40 (48GB) 上运行数小时获得良好效果，同时也支持通过优化器卸载技术在 24GB 显存显卡（如 RTX 4090）上运行。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• GPU:
  • 推荐：NVIDIA A40 (48GB VRAM)
  • 支持：NVIDIA RTX 4090 (24GB VRAM，需开启 CPU 卸载模式)
• 软件依赖:
  • Python 3.8+
  • CUDA 兼容的 PyTorch 版本
  • git-lfs (用于下载大模型和数据集)

前置依赖安装

请确保系统已安装 git-lfs：

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

安装步骤

本项目推荐使用 uv 进行依赖管理，以获得更快的安装速度和更好的环境隔离。

1. 克隆项目代码 (假设已获取项目源码)
2. 初始化环境并安装依赖

pip install uv
uv sync

提示：国内开发者若遇到 uv sync 下载慢的问题，可尝试配置国内源（如清华源或阿里源）或在 uv 命令中指定索引 URL。

3. 下载数据集 本项目使用 CountDown 数学任务数据集：

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

加速建议：如果 HuggingFace 连接缓慢，建议使用国内镜像站（如 hf-mirror.com）或代理加速下载。

4. 下载预训练模型 本项目基于 Qwen2.5-3B-Instruct 模型：

git clone https://huggingface.co/Qwen/Qwen2.5-3B-Instruct

基本使用

场景一：标准训练 (推荐 48GB+ 显存)

适用于 A40 (48GB) 或更高显存显卡。直接运行默认配置即可开始训练。

uv run train.py

场景二：低显存训练 (24GB 显存)

适用于 RTX 4090 (24GB) 等显卡。通过加载特定配置文件，将优化器状态卸载到 CPU，以牺牲少量速度换取低显存占用。

uv run train.py --config config_24GB.yaml

训练任务说明

启动后，模型将在 CountDown 任务 上进行强化学习微调：

• 输入：给定 3-4 个数字和一个目标值。
• 输出：模型需生成包含 <think> 推理过程和 <answer> 最终表达式的回答。
• 奖励机制：
  1. 格式奖励 (0.1)：正确输出 think/answer 标签。
  2. 答案奖励 (1.0)：表达式计算结果等于目标值且数字使用正确。

训练过程中，算法会自动对每个问题采样多个答案，计算相对优势并更新策略网络，无需价值网络估计。