Train Grpo Example

1#!/usr/bin/env python3
2# /// script
3# requires-python = ">=3.10"
4# dependencies = [
5#     "trl>=0.12.0",
6#     "transformers>=4.36.0",
7#     "accelerate>=0.24.0",
8#     "trackio",
9# ]
10# ///
11
12"""
13Production-ready GRPO training example for online RL.
14
15GRPO (Group Relative Policy Optimization) is an online RL method that
16optimizes relative to group performance. Best for tasks with automatic
17reward signals like code execution or math verification.
18
19Usage with hf_jobs MCP tool:
20    hf_jobs("uv", {
21        "script": '''<paste this entire file>''',
22        "flavor": "a10g-large",
23        "timeout": "4h",
24        "secrets": {"HF_TOKEN": "$HF_TOKEN"},
25    })
26
27Or submit the script content directly inline without saving to a file.
28
29Note: For most GRPO use cases, the TRL maintained script is recommended:
30    https://raw.githubusercontent.com/huggingface/trl/main/examples/scripts/grpo.py
31"""
32
33import trackio
34from datasets import load_dataset
35from trl import GRPOTrainer, GRPOConfig
36
37
38# Load dataset (GRPO uses prompt-only format)
39dataset = load_dataset("trl-lib/math_shepherd", split="train")
40print(f"✅ Dataset loaded: {len(dataset)} prompts")
41
42# Training configuration
43config = GRPOConfig(
44    # CRITICAL: Hub settings
45    output_dir="qwen-grpo-math",
46    push_to_hub=True,
47    hub_model_id="username/qwen-grpo-math",
48    hub_strategy="every_save",
49
50    # Training parameters
51    num_train_epochs=1,
52    per_device_train_batch_size=4,
53    gradient_accumulation_steps=4,
54    learning_rate=1e-6,
55
56    # Logging & checkpointing
57    logging_steps=10,
58    save_strategy="steps",
59    save_steps=100,
60    save_total_limit=2,
61
62    # Optimization
63    warmup_ratio=0.1,
64    lr_scheduler_type="cosine",
65
66    # Monitoring
67    report_to="trackio",  # Integrate with Trackio
68    project="meaningful_project_name", # project name for the training name (trackio)
69    run_name="baseline-run", #Descriptive name for this training run
70
71)
72
73# Initialize and train
74# Note: GRPO requires an instruct-tuned model as the base
75trainer = GRPOTrainer(
76    model="Qwen/Qwen2.5-0.5B-Instruct",
77    train_dataset=dataset,
78    args=config,
79)
80
81print("🚀 Starting GRPO training...")
82trainer.train()
83
84print("💾 Pushing to Hub...")
85trainer.push_to_hub()
86
87
88print("✅ Complete! Model at: https://huggingface.co/username/qwen-grpo-math")
89print("📊 View metrics at: https://huggingface.co/spaces/username/trackio")