Mini-R1 是一个旨在复现 DeepSeek R1 “aha moment” 的强化学习教程项目,核心是使用 Group Relative Policy Optimization (GRPO) 和 Countdown Game 来训练开源模型,使其自主发展出自我验证与搜索能力。
用途 本项目主要用于教育和实验,帮助开发者理解如何通过强化学习提升大语言模型的推理能力。
主要能力
- 使用 GRPO 算法训练模型,无需价值函数模型,通过组内分数估计基线,降低内存和计算开销。
- 在 Countdown Game 任务上训练模型,该游戏要求使用随机抽取的数字和基本算术运算达到或接近目标数字。
- 通过自定义奖励函数(格式奖励和准确性奖励)引导模型学习正确的输出格式和解题逻辑。
- 支持分布式训练,结合 Deepspeed 和 vLLM 加速生成和训练过程。
- 提供训练脚本、配置文件以及训练过程中的观察结果和成功推理样本。
使用方式
- 环境设置:安装 Hugging Face 库、PyTorch、vLLM、TRL、transformers 和 datasets。使用 Hugging Face Hub 进行模型版本管理,需注册并登录。
- 生成训练样本:使用 Jiayi-Pan/Countdown-Tasks-3to4 数据集,模型采用 Qwen/Qwen2.5-3B-Instruct。
- 教育性训练:提供交互式 Jupyter Notebook,展示如何使用 TRL 的 GRPOTrainer 和 Q-LoRA 训练模型,包含格式奖励和准确性奖励函数定义。
- 分布式训练:提供 run_r1_grpo.py 脚本和 recipes/grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml 配置文件,可在多 GPU 或 SLURM 集群上运行。配置需设置 num_processes 为 GPU 数量减一,vllm_device 为最后一个 GPU 索引。
- 训练命令:使用 DeepSpeed 和 vLLM 运行训练脚本,如
accelerate launch --config_file recipes/accelerate_configs/deepspeed_zero3.yaml run_r1_grpo.py --config recipes/grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml。 - 结果观察:训练脚本保存完成样本到 completion_samples 文件夹,包括所有完成和成功完成的样本。模型检查点每 25 步保存,可在 philschmid/qwen-2.5-3b-r1-countdown 查看。
适用场景
- 学习强化学习在 LLM 推理中的应用,特别是 GRPO 算法。
- 研究如何通过规则奖励模型引导模型自主学习推理格式。
- 实验小规模模型(如 3B 参数)在特定任务上的推理能力提升。
- 需要多 GPU 环境进行分布式训练的实验。
限制
- 强化学习训练非常缓慢且计算密集。在单张 L4 GPU 上使用 Q-LoRA、批次大小 1 和每样本 2 次生成,单步训练超过 20 分钟。
- 优化后使用 4x H100 80GB,每样本 8 次生成,单步约 45-60 秒,450 步训练约 6 小时。
- 模型性能提升有限:约 50 步学会正确格式,100 步成功率约 25%,200 步约 40%,450 步约 50%。
- 训练不稳定:初始学习率 1e-6 和 beta 0.04 导致约 150 步后不稳定,调整学习率为 5e-7 和 beta 为 0.001 后改善。
- 未测试增加 num_generations 从 8 到 64 的影响。
- 模型可能从“词语推理”转向“程序化执行”,原因可能是模型规模不足、奖励函数定义不够好或任务单一。
- 未提供具体收费、价格或平台信息。