Reasoning Model Trainer

A comprehensive framework for training reasoning models using Supervised Fine-Tuning (SFT) and Reinforcement Learning with Human Feedback (RLHF), optimized for consumer GPUs.

Overview

This project provides a complete pipeline for building robust reasoning models that can be trained on consumer-grade hardware. It incorporates insights from leading LLM research and implements Parameter-Efficient Fine-Tuning (PEFT) techniques to enable training on limited resources.

Features

• Multi-stage Training Pipeline: SFT → Reward Model → RLHF (PPO)
• Consumer GPU Optimized: Uses LoRA, QLoRA, and gradient accumulation
• Cross-platform Support: Works on CUDA, Apple MPS, and CPU
• Synthetic Data Generation: Built-in tools for generating reasoning datasets
• Flexible Configuration: YAML-based configuration with command-line overrides
• Comprehensive Logging: Detailed training metrics and progress tracking

Quick Start

1. Installation

# Clone the repository
git clone <repository-url>
cd reasoning_model_trainer

# Install dependencies
pip install -r requirements.txt

2. Generate Synthetic Data

python generate_data.py --output_dir ./data

3. Train Models

# Step 1: Supervised Fine-Tuning (SFT)
python train.py --mode sft --config configs/sft_config.yaml

# Step 2: Train Reward Model
python train.py --mode reward --config configs/reward_config.yaml

# Step 3: RLHF Training
python train.py --mode rlhf --config configs/rlhf_config.yaml

4. Run Inference

python inference.py --sft_model_path ./sft_output/final_sft_model --prompt "What is 25 + 17?"

Architecture

The framework consists of several key components:

Core Components

• Device Manager: Automatic device detection and management
• Configuration System: Flexible YAML-based configuration
• Logging System: Comprehensive training and metrics logging

Data Handling

• Dataset Loaders: Support for SFT, reward, and RLHF data formats
• Data Collators: Efficient batching for different training modes
• Synthetic Data Generator: Built-in generation of reasoning examples
• Data Validator: Quality checks and data cleaning

Model Architecture

• Base Model Wrapper: Unified interface for causal LMs and reward models
• PEFT Integration: LoRA and QLoRA support for memory efficiency
• Policy-Value Model: Combined architecture for PPO training

Training Pipelines

• SFT Trainer: Supervised fine-tuning with PEFT
• Reward Trainer: Preference-based reward model training
• PPO Trainer: Reinforcement learning with human feedback

Configuration

SFT Configuration Example

# Model Configuration
model_name: "mistralai/Mistral-7B-v0.1"
output_dir: "./sft_output"

# Training Parameters
num_epochs: 3
per_device_train_batch_size: 1
gradient_accumulation_steps: 4
learning_rate: 1e-4

# LoRA Configuration
use_lora: true
lora_r: 16
lora_alpha: 32
lora_dropout: 0.05

# Dataset
dataset_path: "./data/synthetic_sft_dataset.jsonl"

Memory Optimization

For consumer GPUs with limited VRAM:

1. Use smaller models: Start with 7B parameter models
2. Enable LoRA: Set use_lora: true
3. Use 4-bit quantization: Set use_4bit_quantization: true
4. Reduce batch size: Use per_device_train_batch_size: 1
5. Increase gradient accumulation: Use gradient_accumulation_steps: 4-8

Data Formats

SFT Data Format (JSONL)

{"prompt": "What is 15 + 27?", "completion": " Let me solve this step by step.\n\n15 + 27 = 42\n\nTherefore, the answer is 42."}

Reward Model Data Format (JSONL)

{"prompt": "What is 15 + 27?", "chosen": " 15 + 27 = 42", "rejected": " I think it's 43."}

RLHF Prompts Format (JSONL)

{"prompt": "Solve this math problem: 15 + 27 = ?"}

Hardware Requirements

Minimum Requirements

• GPU: 8GB VRAM (with 4-bit quantization)
• RAM: 16GB system RAM
• Storage: 50GB free space

Recommended Requirements

• GPU: 16GB+ VRAM (RTX 4080, RTX 4090, or equivalent)
• RAM: 32GB+ system RAM
• Storage: 100GB+ free space

Tested Platforms

• NVIDIA GPUs with CUDA support
• Apple Silicon Macs with MPS support
• CPU-only training (slower but functional)

Advanced Usage

Custom Data

Replace the synthetic datasets with your own:

# Validate your data
python -c "from src.data.data_validator import DataValidator; validator = DataValidator(); validator.validate_dataset('your_data.jsonl', 'sft')"

# Update config files to point to your data
# Then run training as usual

Model Selection

The framework supports various base models:

• Mistral 7B: Excellent performance-to-size ratio
• Llama 2 7B/13B: Well-tested and widely supported
• Qwen 7B/14B: Strong reasoning capabilities
• TinyLlama 1.1B: For resource-constrained environments

Hyperparameter Tuning

Key hyperparameters to adjust:

• Learning Rate: Start with 1e-4 for SFT, 1e-5 for reward/RLHF
• LoRA Rank: Higher values (16-64) for more expressiveness
• Batch Size: Balance between memory and training stability
• Sequence Length: Longer sequences for complex reasoning

Troubleshooting

Common Issues

1. Out of Memory (OOM)

   • Reduce batch size to 1
   • Enable 4-bit quantization
   • Reduce sequence length
   • Use smaller LoRA rank

2. Slow Training

   • Increase batch size if memory allows
   • Use gradient accumulation
   • Enable mixed precision training

3. Poor Model Performance

   • Increase training epochs
   • Use higher quality data
   • Adjust learning rate
   • Try different base models

Performance Tips

1. Use appropriate model size for your hardware
2. Monitor GPU utilization and adjust batch size accordingly
3. Use gradient checkpointing for memory efficiency
4. Enable compilation with torch.compile() for speed

Contributing

Contributions are welcome! Please see our contributing guidelines for details.

License

This project is licensed under the MIT License - see the LICENSE file for details.

Acknowledgments

This framework incorporates insights from:

• Stanford Alpaca and LLaMA-Adapter
• Vicuna and FastChat
• DeepSeek reasoning models
• TinyZero and Journey Learning approaches
• Lightning-AI LitGPT framework

Citation

If you use this framework in your research, please cite:

@software{reasoning_model_trainer,
  title={Reasoning Model Trainer: A Framework for Consumer GPU Training},
  author={Muhammad Abiodun SULAIMAN},
  year={2024},
  url={https://github.com/Behordeun/reasoning-model-trainer}
}