R1 - Lightweight Enhanced GRPO Training Framework

A lightweight enhanced Group-based Preference Optimization (GRPO) training framework with more verifiable rewards and real-time evaluators.

English | 中文

English

Quick Start

# Setup environment
conda create -n r1 python=3.10
conda activate r1
pip install -e .

# Run training demo
python train_enhanced_demo.py

Features

• Enhanced GRPO Trainer: Extended from trl.GRPOTrainer with real-time evaluation
• Modular Reward System: Multiple reward functions (accuracy, format, reasoning, etc.)
• YAML Configuration: Flexible training parameter management
• Built-in Utilities: Logging, formatting, and data processing tools

Enhanced Features

🚀 What makes this "Enhanced"?

Compared to standard GRPO implementations, this framework provides:

Key Enhancements:

1. Advanced Reward System: 6 reward functions with intelligent bonuses

   • accuracy: Math verification + confidence detection
   • format: Structure validation + completeness scoring
   • reasoning_steps: Logic quality assessment
   • tag_count: Smart tag validation with penalty system
   • length: Optimal response length control
   • mathematical_rigor: Mathematical reasoning evaluation

2. Real-time Training Insights:

   • Live evaluation during training
   • Progress estimation with ETA
   • Best metric tracking
   • Comprehensive training summary

3. Modular Architecture: Easy to extend with new reward functions and evaluation metrics

Configuration

Basic example:

model_name_or_path: Qwen/Qwen2.5-1.5B-Instruct
dataset_name: open-r1/Mixture-of-Thoughts
learning_rate: 3.0e-6
num_train_epochs: 2
# All available reward functions
reward_funcs: [accuracy, format, reasoning_steps, tag_count, length, mathematical_rigor]
reward_weights: [4.0, 1.0, 2.0, 1.0, 0.8, 1.5]  # Customizable weights for each function

Project Structure

src/r1/
├── enhanced_grpo_trainer.py    # Main trainer implementation
├── rewards.py                  # Reward function library
├── configs.py                  # Configuration classes
├── grpo.py                     # Training script
└── utils/                      # Utility modules
    ├── logging.py
    ├── formatting.py
    └── data.py

Reward Functions

Training Scripts

• train_enhanced_demo.py - Demo with enhanced features
• src/r1/grpo.py - Standard GRPO training
• src/r1/sft.py - Supervised fine-tuning

中文

快速开始

# 环境配置
conda create -n r1 python=3.10
conda activate r1
pip install -e .

# 运行训练演示
python train_enhanced_demo.py

主要功能

• 增强GRPO训练器: 基于trl.GRPOTrainer扩展，支持实时评估
• 模块化奖励系统: 多种奖励函数（准确性、格式、推理等）
• YAML配置管理: 灵活的训练参数配置
• 内置工具库: 日志记录、格式化、数据处理工具

增强功能

🚀 "增强"体现在哪里？

相比标准GRPO实现，本框架提供：

核心增强功能：

1. 高级奖励系统: 6个奖励函数，含智能奖励机制

   • accuracy: 数学验证 + 置信度检测
   • format: 结构验证 + 完整性评分
   • reasoning_steps: 逻辑质量评估
   • tag_count: 智能标签验证，含惩罚机制
   • length: 最优响应长度控制
   • mathematical_rigor: 数学推理评估

2. 实时训练洞察:

   • 训练中实时评估
   • 进度预估和ETA
   • 最佳指标追踪
   • 全面训练总结

3. 模块化架构: 易于扩展新奖励函数和评估指标

配置示例

基础配置:

model_name_or_path: Qwen/Qwen2.5-1.5B-Instruct
dataset_name: open-r1/Mixture-of-Thoughts
learning_rate: 3.0e-6
num_train_epochs: 2
# 所有可用的奖励函数
reward_funcs: [accuracy, format, reasoning_steps, tag_count, length, mathematical_rigor]
reward_weights: [4.0, 1.0, 2.0, 1.0, 0.8, 1.5]  # 每个函数可自定义权重

项目结构

src/r1/
├── enhanced_grpo_trainer.py    # 主要训练器实现
├── rewards.py                  # 奖励函数库
├── configs.py                  # 配置类
├── grpo.py                     # 训练脚本
└── utils/                      # 工具模块
    ├── logging.py
    ├── formatting.py
    └── data.py

奖励函数

训练脚本

• train_enhanced_demo.py - 增强功能演示
• src/r1/grpo.py - 标准GRPO训练
• src/r1/sft.py - 监督微调

License

Apache-2.0