Mini-R1 is a reinforcement learning tutorial project aimed at reproducing the DeepSeek R1 "aha moment", with the core being the use of Group Relative Policy Optimization (GRPO) and the Countdown Game to train open-source models so that they autonomously develop self-verification and search capabilities.
Purpose This project is mainly for education and experimentation, helping developers understand how to improve the reasoning capabilities of large language models through reinforcement learning.
Main Capabilities
- Use the GRPO algorithm to train models without a value function model, estimating the baseline from within-group scores, reducing memory and computational overhead.
- Train models on the Countdown Game task, which requires using randomly drawn numbers and basic arithmetic operations to reach or get close to a target number.
- Guide the model to learn the correct output format and problem-solving logic through custom reward functions (format reward and accuracy reward).
- Support distributed training, combining Deepspeed and vLLM to accelerate the generation and training process.
- Provide training scripts, configuration files, as well as observations during training and successful reasoning samples.
How to Use
- Environment Setup: Install the Hugging Face libraries, PyTorch, vLLM, TRL, transformers, and datasets. Use Hugging Face Hub for model version management, which requires registration and login.
- Generate Training Samples: Use the Jiayi-Pan/Countdown-Tasks-3to4 dataset, and the model uses Qwen/Qwen2.5-3B-Instruct.
- Educational Training: Provide an interactive Jupyter Notebook showing how to train models using TRL's GRPOTrainer and Q-LoRA, including definitions of format reward and accuracy reward functions.
- Distributed Training: Provide the run_r1_grpo.py script and the recipes/grpo-qwen-2.5-3b-deepseek-r1-countdown.yaml configuration file, which can be run on multiple GPUs or a SLURM cluster. The configuration requires setting num_processes to the number of GPUs minus one, and vllm_device to the index of the last GPU.
- Training Commands: Run the training script using DeepSpeed and vLLM, such as
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. - Result observation: The training script saves completed samples to the completion_samples folder, including all completed and successfully completed samples. Model checkpoints are saved every 25 steps and can be viewed at philschmid/qwen-2.5-3b-r1-countdown.
Applicable scenarios
- Learning the application of reinforcement learning in LLM reasoning, especially the GRPO algorithm.
- Studying how to guide a model to autonomously learn a reasoning format through a rule-based reward model.
- Experiments on improving the reasoning ability of small-scale models (such as 3B parameters) on specific tasks.
- Experiments that require a multi-GPU environment for distributed training.
Limitations
- Reinforcement learning training is very slow and computationally intensive. On a single L4 GPU using Q-LoRA, batch size 1, and 2 generations per sample, a single training step takes more than 20 minutes.
- After optimization using 4x H100 80GB and 8 generations per sample, a single step takes about 45-60 seconds, and 450 steps of training take about 6 hours.
- The model performance improvement is limited: it learns the correct format in about 50 steps, with a success rate of about 25% at 100 steps, about 40% at 200 steps, and about 50% at 450 steps.
- Training is unstable: an initial learning rate of 1e-6 and beta 0.04 lead to instability after about 150 steps, and adjusting the learning rate to 5e-7 and beta to 0.001 improves it.
- The effect of increasing num_generations from 8 to 64 was not tested.
- The model may shift from "word reasoning" to "programmatic execution," possibly because the model scale is insufficient, the reward function is not defined well enough, or the task is too单一.
- No specific fees, prices, or platform information is provided.
