simple_GRPO is a simple open-source GRPO implementation for reproducing the reasoning capabilities of R1-like large language models. Its core loss calculation formula is referenced from Hugging Face's trl. The project's goal is to provide a minimal codebase to support: saving GPU memory to enable feasible and efficient training; quickly understanding reinforcement learning processes such as GRPO from a teaching perspective; and quickly trying many improvements, such as improved multi-answer generation, regrouping, KL penalty, and parameter tuning. An "Aha moment" can be observed in the early stages of model training.
Main capabilities and features:
- The project code is simple, only about 200 lines, spread across 2 files. It depends only on standard libraries such as deepspeed and torch, with no need for dependencies such as ray. The design allows for more complex interventions.
- The reference model is partially decoupled, allowing it to run on different GPUs (or even 4090s on different machines). This avoids running the reference model and the training model on the same GPU, prevents torch multiprocessing from creating multiple copies, and enables training a 7B model on an 80G A800.
- Training on 1 A800 GPU can be completed within 1 hour. Both Qwen2.5-7B and Qwen2.5-3B exhibit an "Aha moment" within the first 30 optimization steps.
- 2025/02/19: Added a triton implementation of the loss, with a small speedup, and it can optionally be disabled. See the simple_grpo_v1 folder.
- 2025/02/19: Added a regrouping version, implementing sampling of generated data on ref_server. See the regroup_ver folder.
- 2025/02/27: Added a vllm package to accelerate inference.
- 2025/03/24: Added the reinforce++ algorithm, with the same usage as before.
- 2025/07/23: Recommend checking out LSRL, a next-generation RL framework with a custom optimizer and streamlined architecture, for better performance.
Usage: The runtime environment is in requirements.txt. At least two GPUs are required. Warning: running the simple example script using GSM8K may require installing datasets and adding an HTTP proxy (possibly). This package is only used for simply loading GSM8K. In real use, users should prepare their own data, so datasets is not strictly required. In addition, if datasets is not installed and is not loaded in the fileDataset, no data during generation process, the training process (as well as the reference process) will keep waiting for batches. If there are three or more GPUs, there will be a better option. Run the corresponding command; this uses only one GPU to collect and run the reference model. In grpo_vllm_one.py, set the generation device index relative to the next visible device. Then open another bash.
Experimental environment and results: Hardware setup: 2×A800 (80GB) GPUs. Configuration: training uses 1 GPU with Zero-Stage 2 optimization; inference uses a dedicated 1 GPU (compatible with 3090/4090). Training performance: For Qwen2.5-3B, accuracy stabilizes above 60% after 5 optimization steps, peaking at about 70%; format compliance reaches about 100% after 30 steps. For Qwen2.5-7B, accuracy remains >90% throughout the training process; format compliance reaches 100% within 30 steps. Both metrics converge quickly. The example outputs shown (Qwen2.5-3B step 20 and Qwen2.5-7B step 20) demonstrate the reasoning process within the think and answer tags.
Known limitations:
- Answer generation may be ineffective because a group may contain all wrong or all correct answers. Group restructuring and better answer generation are needed.
- If a long chain of thought is generated, GPU memory is still tight. Groups must be split to reduce the batch size.
- These features have been implemented and are under testing, and will be available soon.
Project members: This project is led by Dr. Liang Jiaqing and Professor Xiao Yanghua of the KnowledgeWorks Lab at Fudan University. The core development team includes PhD student Han Jinyi, master's student Wang Xinyi, and other contributors. If you find the code useful, please cite it in the specified format.
