Mini-RL

A minimal PyTorch-based training and inference framework for small language models with reinforcement learning from human feedback (RLHF). Supports DPO (Direct Preference Optimization) and GRPO (Group Relative Policy Optimization) training modes.

Features

~1M parameter transformer model runnable on CPU
Modern architecture: RMSNorm, Rotary Position Embeddings (RoPE), SwiGLU activation
DPO training: Direct preference optimization without reward modeling
GRPO training: Group relative policy optimization (DeepSeek-R1 style) without critic networks
Character-level tokenizer: No external dependencies, works out of the box
Simple data format: JSON Lines for easy data preparation
Minimal dependencies: Only requires PyTorch

Installation

# Clone or navigate to the project
cd mini-rl

# Install dependencies
pip install -r requirements.txt

Requirements:

Python 3.8+
PyTorch >= 2.0

Quick Start

Training with DPO

python train.py --mode dpo --data data/dpo_sample.jsonl --epochs 10 --output dpo_model.pt

Training with GRPO

python train.py --mode grpo --data data/grpo_sample.jsonl --epochs 10 --output grpo_model.pt

Inference

# Single prompt
python inference.py --checkpoint model.pt --prompt "Hello, how are you?"

# Interactive chat mode
python inference.py --checkpoint model.pt

# With custom generation settings
python inference.py --checkpoint model.pt --prompt "Hello" \
    --max-tokens 100 --temperature 0.8 --top-p 0.9

Data Formats

DPO Training Data

JSON Lines format with preference pairs:

{"prompt": "What is 2+2?", "chosen": "2+2 equals 4.", "rejected": "2+2 equals 5."}
{"prompt": "What color is the sky?", "chosen": "The sky is blue.", "rejected": "The sky is green."}

GRPO Training Data

JSON Lines format with grouped responses and rewards:

{"prompt": "What is 3*4?", "responses": ["12", "The answer is 12.", "11"], "rewards": [1.0, 0.9, 0.0]}
{"prompt": "Say hello:", "responses": ["Hello!", "Hi there!", "Goodbye!"], "rewards": [1.0, 0.9, 0.1]}

Model Architecture

The model follows a GPT-style decoder-only transformer with modern improvements:

Component	Description
Normalization	RMSNorm (Root Mean Square Layer Normalization)
Position Encoding	Rotary Position Embeddings (RoPE)
Activation	SwiGLU in feed-forward layers
Attention	Multi-head causal self-attention

Default Configuration (~1M parameters)

Parameter	Value
Vocabulary Size	260 (256 ASCII + 4 special tokens)
Hidden Dimension	128
Number of Layers	4
Attention Heads	4
Feed-forward Dimension	512
Max Sequence Length	512
Dropout	0.1

Training Options

Common Arguments

--mode          Training mode: dpo or grpo (required)
--data          Path to training data JSONL (required)
--eval-data     Path to evaluation data JSONL (optional)
--epochs        Number of training epochs (default: 10)
--batch-size    Batch size (default: 8)
--lr            Learning rate (default: 1e-4)
--device        Device: cpu or cuda (default: cpu)
--output        Output checkpoint path
--checkpoint    Resume from checkpoint

Model Arguments

--d-model       Hidden dimension (default: 128)
--n-layers      Number of transformer layers (default: 4)
--n-heads       Number of attention heads (default: 4)
--d-ff          Feed-forward dimension (default: 512)
--max-seq-len   Maximum sequence length (default: 512)
--dropout       Dropout rate (default: 0.1)

DPO-Specific Arguments

--beta              KL penalty coefficient (default: 0.1)
--reference-free    Use reference-free DPO variant
--label-smoothing   Label smoothing factor (default: 0.0)

GRPO-Specific Arguments

--group-size            Number of responses per prompt (default: 4)
--clip-eps              PPO clipping epsilon (default: 0.2)
--kl-coef               KL divergence coefficient (default: 0.01)
--normalize-advantages  Normalize advantages by std

Logging Arguments

--log-dir             Directory for training logs (default: logs)
--experiment-name     Name for this experiment (auto-generated if not provided)
--no-logging          Disable training logging
--log-gradients       Log gradient histograms (slower, for debugging)
--log-weights         Log weight histograms (slower, for debugging)

Inference Options

--checkpoint          Path to model checkpoint (required)
--prompt              Input prompt (if omitted, enters chat mode)
--max-tokens          Maximum new tokens to generate (default: 100)
--temperature         Sampling temperature (default: 0.8)
--top-k               Top-k sampling parameter (default: 50)
--top-p               Top-p nucleus sampling (default: 0.9)
--repetition-penalty  Repetition penalty (default: 1.1)
--stream              Stream output token by token

Training Visualization

Mini-RL includes a TensorBoard-style visualization dashboard for monitoring training progress.

Launching the Dashboard

# Start the visualization server (default port: 6006)
python visualize.py --log-dir logs

# Custom port and host
python visualize.py --log-dir logs --port 8080 --host 0.0.0.0

# Don't auto-open browser
python visualize.py --log-dir logs --no-browser

Features

Scalar Charts: Interactive line charts for loss, accuracy, rewards, learning rate, and gradient norms
Smoothing: Adjustable exponential moving average smoothing for cleaner visualizations
Histograms: Parameter distribution histograms (when --log-gradients or --log-weights is enabled)
Experiment Comparison: Select multiple experiments to compare side-by-side
Configuration View: See hyperparameters for each experiment
Auto-refresh: Automatically updates every 5 seconds during training

Logged Metrics

DPO Training:

train/loss: DPO loss
train/accuracy: Preference accuracy (chosen > rejected)
train/chosen_rewards: Implicit reward for chosen responses
train/rejected_rewards: Implicit reward for rejected responses
train/reward_margin: Difference between chosen and rejected rewards
train/lr: Learning rate
train/gradient_norm: Total gradient norm
eval/*: Evaluation metrics (if eval data provided)

GRPO Training:

train/loss: GRPO loss
train/policy_loss: Policy gradient loss
train/kl_div: KL divergence from reference
train/reward_mean: Average reward in group
train/reward_std: Reward standard deviation
train/advantage_mean: Average advantage
train/ratio_mean: Policy ratio mean
train/lr: Learning rate
train/gradient_norm: Total gradient norm

Python API for Logging

from logger import TrainingLogger, LogReader

# Create logger
logger = TrainingLogger(log_dir="logs", experiment_name="my_experiment")

# Log configuration
logger.log_config(config)

# Log scalars
logger.log_scalar("train/loss", 0.5, step=100)
logger.log_scalars("train", {"loss": 0.5, "accuracy": 0.8}, step=100)

# Log histograms
logger.log_histogram("weights/layer1", tensor, step=100)

# Log model parameters
logger.log_model_params(model, step, log_gradients=True, log_weights=True)

# Close logger
logger.close()

# Read logs
reader = LogReader("logs")
experiments = reader.list_experiments()
scalars = reader.get_scalars("my_experiment")
config = reader.get_config("my_experiment")

Python API

Creating a Model

from config import ModelConfig
from model import create_model

# Default ~1M parameter model
model = create_model()

# Custom configuration
config = ModelConfig(
    d_model=256,
    n_layers=6,
    n_heads=8,
    d_ff=1024
)
model = create_model(config)

print(f"Parameters: {model.count_parameters():,}")

DPO Training

from config import ModelConfig, DPOConfig
from model import create_model
from tokenizer import CharTokenizer
from dpo_trainer import DPOTrainer
from data import create_preference_dataloader

# Setup
model = create_model()
tokenizer = CharTokenizer()
config = DPOConfig(lr=1e-4, epochs=10, device="cpu")

# Create trainer
trainer = DPOTrainer(model, tokenizer, config)

# Load data and train
dataloader = create_preference_dataloader("data.jsonl", tokenizer)
history = trainer.train(dataloader)

# Save model
trainer.save_checkpoint("model.pt")

GRPO Training

from config import ModelConfig, GRPOConfig
from model import create_model
from tokenizer import CharTokenizer
from grpo_trainer import GRPOTrainer
from data import create_grpo_dataloader

# Setup
model = create_model()
tokenizer = CharTokenizer()
config = GRPOConfig(lr=1e-4, epochs=10, group_size=4)

# Create trainer
trainer = GRPOTrainer(model, tokenizer, config)

# Load data and train
dataloader = create_grpo_dataloader("data.jsonl", tokenizer)
history = trainer.train_offline(dataloader)

Online GRPO with Custom Reward Function

from grpo_trainer import GRPOTrainer
from data import create_prompt_dataloader

def my_reward_fn(prompt: str, response: str) -> float:
    # Custom reward logic
    return 1.0 if "correct" in response.lower() else 0.0

trainer = GRPOTrainer(model, tokenizer, config, reward_fn=my_reward_fn)
prompt_loader = create_prompt_dataloader("prompts.jsonl", tokenizer)
trainer.train_online(prompt_loader)

Inference

from inference import InferenceEngine

# Load from checkpoint
engine = InferenceEngine.from_checkpoint("model.pt")

# Generate text
output = engine.generate(
    prompt="Hello",
    max_new_tokens=50,
    temperature=0.8,
    top_p=0.9
)
print(output)

File Structure

mini-rl/
├── config.py          # Configuration dataclasses
├── tokenizer.py       # Character-level tokenizer
├── model.py           # Transformer architecture
├── data.py            # Dataset and DataLoader utilities
├── dpo_trainer.py     # DPO training implementation
├── grpo_trainer.py    # GRPO training implementation
├── inference.py       # Text generation and CLI
├── train.py           # Main training entry point
├── logger.py          # Training logging utilities
├── visualize.py       # TensorBoard-style dashboard
├── test_framework.py  # Test suite
├── requirements.txt   # Dependencies
├── logs/              # Training logs (created automatically)
└── data/
    ├── dpo_sample.jsonl   # Sample DPO training data
    └── grpo_sample.jsonl  # Sample GRPO training data

Testing

Run the test suite to verify all components:

python test_framework.py

Notes

Pre-training required: DPO and GRPO are alignment methods that fine-tune existing language models. For meaningful text generation, pre-train the model on language modeling data first.
CPU-friendly: The default 1M parameter model is designed to run efficiently on CPU for experimentation and learning.
Scaling up: Increase d_model, n_layers, and d_ff for larger models. Use --device cuda for GPU acceleration.

References

License

MIT License

Name		Name	Last commit message	Last commit date
Latest commit History 1 Commit
.gitignore		.gitignore
README.md		README.md
config.py		config.py
data.py		data.py
dpo_trainer.py		dpo_trainer.py
grpo_trainer.py		grpo_trainer.py
inference.py		inference.py
logger.py		logger.py
model.py		model.py
requirements.txt		requirements.txt
test_framework.py		test_framework.py
tokenizer.py		tokenizer.py
train.py		train.py
visualize.py		visualize.py

Folders and files

Latest commit

History

Repository files navigation

Mini-RL

Features

Installation

Quick Start

Training with DPO

Training with GRPO

Inference

Data Formats

DPO Training Data

GRPO Training Data

Model Architecture

Default Configuration (~1M parameters)

Training Options

Common Arguments

Model Arguments

DPO-Specific Arguments

GRPO-Specific Arguments

Logging Arguments

Inference Options

Training Visualization

Launching the Dashboard

Features

Logged Metrics

Python API for Logging

Python API

Creating a Model

DPO Training

GRPO Training

Online GRPO with Custom Reward Function

Inference

File Structure

Testing

Notes

References

License

About

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages