Skip to content

README.md

Latest commit

 

History

History
255 lines (186 loc) · 7.46 KB

README.md

File metadata and controls

255 lines (186 loc) · 7.46 KB
Raw

Math Mind: GRPO Training for Mathematical Reasoning

This repository implements Group Relative Policy Optimization (GRPO) for training language models on mathematical reasoning tasks, specifically using the GSM8K dataset.

Overview

This project fine-tunes language models (e.g., Qwen2.5-3B-Instruct) using GRPO to improve their mathematical reasoning capabilities. The training uses reward functions that evaluate:

  1. Format correctness: Checks if the model follows the structured reasoning format
  2. Answer correctness: Validates if the final numerical answer is correct
  3. Stepwise reasoning: Evaluates the correctness and progression of intermediate steps

Project Structure

math-mind/
├── src/
│   ├── config.py              # Configuration constants and hyperparameters
│   ├── data/
│   │   ├── __init__.py
│   │   └── processing.py      # Dataset loading and preprocessing
│   ├── rewards/
│   │   ├── __init__.py
│   │   ├── format_rewards.py  # Format checking rewards
│   │   ├── correctness_rewards.py  # Answer correctness rewards
│   │   ├── stepwise_rewards.py     # Stepwise reasoning rewards
│   │   ├── equation_parser.py      # Equation parsing utilities
│   │   └── combined_rewards.py     # Combined reward function
│   ├── training/
│   │   ├── __init__.py
│   │   └── trainer.py         # Model setup and training logic
│   └── inference/
│       ├── __init__.py
│       └── inference.py       # Inference and evaluation utilities
├── train.py                   # Main training script
├── inference.py               # Main inference script
├── requirements.txt           # Python dependencies
└── README.md                  # This file

Installation

1. Clone the repository

git clone <your-repo-url>
cd math-mind

2. Create a virtual environment

python -m venv venv
source venv/bin/activate  # On Windows: venv\Scripts\activate

3. Install dependencies

pip install -r requirements.txt

Note: The requirements include PyTorch with CUDA support. Adjust the PyTorch installation in requirements.txt based on your CUDA version. For CPU-only systems, use:

pip install torch --index-url https://download.pytorch.org/whl/cpu

Optional dependencies:

  • flash-attn: For faster training (requires compatible CUDA setup)
  • vllm: For faster generation in distributed training

4. Setup HuggingFace authentication

You'll need a HuggingFace token to download models and datasets:

huggingface-cli login

Or set the HF_TOKEN environment variable:

export HF_TOKEN=your_token_here

5. Setup Wandb (optional, for experiment tracking)

wandb login

Usage

Training

Train a model using GRPO on the GSM8K dataset:

python train.py \
    --model-name "Qwen/Qwen2.5-3B-Instruct" \
    --output-dir "qwen25-3b-grpo-onlyfinalreward-scaling1" \
    --hub-model-id "your-username/qwen25-3b-grpo-onlyfinalreward-scaling1" \
    --train-limit 500 \
    --eval-limit 50 \
    --wandb-entity "your-wandb-entity"

Arguments:

  • --model-name: HuggingFace model identifier (default: Qwen/Qwen2.5-3B-Instruct)
  • --output-dir: Output directory for checkpoints and final model (required)
  • --hub-model-id: Optional HuggingFace Hub model ID for pushing checkpoints
  • --dataset-name: Dataset name (default: openai/gsm8k)
  • --train-limit: Limit number of training examples (optional, for quick testing)
  • --eval-limit: Limit number of evaluation examples (default: 50)
  • --wandb-entity: Wandb entity/username (optional)
  • --run-name: Run name for wandb (optional)

Example with minimal settings:

python train.py --output-dir "my-grpo-model" --eval-limit 20

Inference

Run inference on a trained model:

python inference.py \
    --model-id "your-username/qwen25-3b-grpo-onlyfinalreward-scaling1" \
    --output-file "outputs/test_results.json" \
    --limit 100 \
    --max-new-tokens 1024

Arguments:

  • --model-id: HuggingFace model ID or local path (required)
  • --output-file: Path to save inference results JSON file (required)
  • --dataset-name: Dataset name (default: openai/gsm8k)
  • --split: Dataset split to use (default: test)
  • --limit: Limit number of examples to process (optional)
  • --max-new-tokens: Maximum number of new tokens to generate (default: 1024)

Example:

python inference.py \
    --model-id "thegrey07/qwen25-3b-grpo-onlyfinalreward-scaling1" \
    --output-file "outputs/my_results.json" \
    --limit 50

Configuration

Training hyperparameters and other settings can be modified in src/config.py:

  • Format tokens: REASONING_START, REASONING_END, ANSWER_SEPARATOR, EQUATION_SEPARATOR
  • System prompt: SYSTEM_PROMPT
  • Training config: TRAINING_CONFIG (learning rate, batch size, max steps, etc.)
  • LoRA config: LORA_CONFIG (rank, alpha, dropout, target modules)
  • Reward scaling: SCALING_FACTOR

Reward Functions

The repository includes several reward functions in src/rewards/:

  1. check_format_rewards: Validates structured reasoning format
  2. check_answer_correctness: Checks final numerical answer correctness
  3. reward_stepwise: Evaluates stepwise reasoning quality
  4. combined_reward_fn: Combines all reward components

You can modify train.py to use different reward functions:

from src.rewards import combined_reward_fn, check_answer_correctness

# Use only final answer reward (current default)
reward_funcs=[check_answer_correctness]

# Or use combined rewards
reward_funcs=[combined_reward_fn]

Output Format

The model is trained to produce structured reasoning outputs:

<begin_steps>
<STEP 1>
Explanation...
@@@@ 10 * 7 = 70 @@@@
</STEP>

<STEP 2>
Explanation...
@@@@ 70 / 2 = 35 @@@@
</STEP>
</end_steps>

#### 35

Hardware Requirements

  • GPU: Recommended (CUDA-capable GPU with at least 16GB VRAM for 3B models)
  • Memory: 32GB+ RAM recommended
  • Storage: ~10GB for models and checkpoints

Training on CPU is possible but very slow and not recommended.

Troubleshooting

Out of Memory Errors

  • Reduce per_device_train_batch_size in src/config.py
  • Reduce max_completion_length or max_prompt_length
  • Use smaller models (e.g., Qwen2.5-0.5B-Instruct)
  • Enable gradient checkpointing (already enabled by default)

Installation Issues

  • Ensure your CUDA version matches PyTorch requirements
  • For flash-attn, you may need to install from source if pre-built wheels aren't available
  • Some dependencies may have platform-specific requirements

Training Instability

  • Adjust beta (KL coefficient) in TRAINING_CONFIG
  • Reduce learning rate
  • Adjust reward scaling factor

Contributing

Contributions are welcome! Please:

  1. Fork the repository
  2. Create a feature branch
  3. Make your changes
  4. Add tests if applicable
  5. Submit a pull request

Acknowledgments

  • Based on GRPO (Group Relative Policy Optimization) from the DeepSeekMath paper
  • Uses TRL (Transformers Reinforcement Learning) library from HuggingFace
  • Trained on GSM8K dataset from OpenAI

References