grpo-rl-training | SkillBundle

GRPO/RL Training with TRL

Expert-level guidance for implementing Group Relative Policy Optimization (GRPO) using the Transformer Reinforcement Learning (TRL) library. This skill provides battle-tested patterns, critical insights, and production-ready workflows for fine-tuning language models with custom reward functions.

When to Use This Skill

Use GRPO training when you need to:

Enforce specific output formats (e.g., XML tags, JSON, structured reasoning)
Teach verifiable tasks with objective correctness metrics (math, coding, fact-checking)
Improve reasoning capabilities by rewarding chain-of-thought patterns
Align models to domain-specific behaviors without labeled preference data
Optimize for multiple objectives simultaneously (format + correctness + style)

Do NOT use GRPO for:

Simple supervised fine-tuning tasks (use SFT instead)
Tasks without clear reward signals
When you already have high-quality preference pairs (use DPO/PPO instead)

Core Concepts

1. GRPO Algorithm Fundamentals

Key Mechanism:

Generates multiple completions for each prompt (group size: 4-16)
Compares completions within each group using reward functions
Updates policy to favor higher-rewarded responses relative to the group

Critical Difference from PPO:

No separate reward model needed
More sample-efficient (learns from within-group comparisons)
Simpler to implement and debug

Mathematical Intuition:

For each prompt p:
  1. Generate N completions: {c₁, c₂, ..., cₙ}
  2. Compute rewards: {r₁, r₂, ..., rₙ}
  3. Learn to increase probability of high-reward completions
     relative to low-reward ones in the same group

2. Reward Function Design Philosophy

Golden Rules:

Compose multiple reward functions - Each handles one aspect (format, correctness, style)
Scale rewards appropriately - Higher weight = stronger signal
Use incremental rewards - Partial credit for partial compliance
Test rewards independently - Debug each reward function in isolation

Reward Function Types:

Type	Use Case	Example Weight
Correctness	Verifiable tasks (math, code)	2.0 (highest)
Format	Strict structure enforcement	0.5-1.0
Length	Encourage verbosity/conciseness	0.1-0.5
Style	Penalize unwanted patterns	-0.5 to 0.5

Implementation Workflow

Step 1: Dataset Preparation

Critical Requirements:

Prompts in chat format (list of dicts with 'role' and 'content')
Include system prompts to set expectations
For verifiable tasks, include ground truth answers as additional columns

Example Structure:

from datasets import load_dataset, Dataset
 
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
[Your step-by-step thinking]

Pro Tips:

Use one-shot or few-shot examples in system prompt for complex formats
Keep prompts concise (max_prompt_length: 256-512 tokens)
Validate data quality before training (garbage in = garbage out)

Step 2: Reward Function Implementation

Template Structure:

Example 1: Correctness Reward (Math/Coding)

Example 2: Format Reward (Structured Output)

Example 3: Incremental Format Reward (Partial Credit)

Critical Insight: Combine 3-5 reward functions for robust training. Order matters less than diversity of signals.

Step 3: Training Configuration

Memory-Optimized Config (Small GPU)

High-Performance Config (Large GPU)

Critical Hyperparameters:

Parameter	Impact	Tuning Advice
`num_generations`	Group size for comparison	Start with 8, increase to 16 if GPU allows
`learning_rate`	Convergence speed/stability	5e-6 (safe), 1e-5 (faster, riskier)
`max_completion_length`	Output verbosity	Match your task (512 for reasoning, 256 for short answers)
`gradient_accumulation_steps`	Effective batch size	Increase if GPU memory limited

Step 4: Model Setup and Training

Standard Setup (Transformers)

Unsloth Setup (2-3x Faster)

Critical Training Insights

1. Loss Behavior (EXPECTED PATTERN)

Loss starts near 0 and INCREASES during training
This is CORRECT - loss measures KL divergence from initial policy
Model is learning (diverging from original behavior to optimize rewards)
Monitor reward metrics instead of loss for progress

2. Reward Tracking

Key metrics to watch:

reward: Average across all completions
reward_std: Diversity within groups (should remain > 0)
kl: KL divergence from reference (should grow moderately)

Healthy Training Pattern:

Warning Signs:

Reward std → 0 (model collapsing to single response)
KL exploding (> 0.5) (diverging too much, reduce LR)
Reward stuck (reward functions too harsh or model capacity issue)

3. Common Pitfalls and Solutions

Problem	Symptom	Solution
Mode collapse	All completions identical	Increase `num_generations`, add diversity penalty
No learning	Flat rewards	Check reward function logic, increase LR
OOM errors	GPU memory exceeded	Reduce `num_generations`, enable gradient checkpointing
Slow training	< 1 it/s	Enable `use_vllm=True`, use Unsloth, reduce seq length
Format ignored	Model doesn't follow structure	Increase format reward weight, add incremental rewards

Advanced Patterns

1. Multi-Stage Training

For complex tasks, train in stages:

2. Adaptive Reward Scaling

3. Custom Dataset Integration

Deployment and Inference

Save and Merge LoRA

Inference Example

Best Practices Checklist

Before Training:

Validate dataset format (prompts as List[Dict])
Test reward functions on sample data
Calculate expected max_prompt_length from data
Choose appropriate num_generations based on GPU memory
Set up logging (wandb recommended)

During Training:

Monitor reward progression (should increase)
Check reward_std (should stay > 0.1)
Watch for OOM errors (reduce batch size if needed)
Sample generations every 50-100 steps
Validate format compliance on holdout set

After Training:

Merge LoRA weights if using PEFT
Test on diverse prompts
Compare to baseline model
Document reward weights and hyperparameters
Save reproducibility config

Troubleshooting Guide

Debugging Workflow

Isolate reward functions - Test each independently
Check data distribution - Ensure diversity in prompts
Reduce complexity - Start with single reward, add gradually
Monitor generations - Print samples every N steps
Validate extraction logic - Ensure answer parsing works

Quick Fixes

References and Resources

Official Documentation:

TRL GRPO Trainer: https://huggingface.co/docs/trl/grpo_trainer
DeepSeek R1 Paper: https://arxiv.org/abs/2501.12948
Unsloth Docs: https://docs.unsloth.ai/

Example Repositories:

Open R1 Implementation: https://github.com/huggingface/open-r1
TRL Examples: https://github.com/huggingface/trl/tree/main/examples

Usage Instructions for Agents

When this skill is loaded:

Read this entire file before implementing GRPO training
Start with the simplest reward function (e.g., length-based) to validate setup
Use the templates in templates/ directory as starting points
Reference examples in examples/ for task-specific implementations
Follow the workflow sequentially (don't skip steps)
Debug incrementally - add one reward function at a time

Critical Reminders:

Always use multiple reward functions (3-5 is optimal)
Monitor reward metrics, not loss
Test reward functions before training
Start small (num_generations=4), scale up gradually
Save checkpoints frequently (every 100 steps)

This skill is designed for expert-level implementation. Beginners should start with supervised fine-tuning before attempting GRPO.

GRPO/RL Training with TRL

When to Use This Skill

Use GRPO training when you need to:

Enforce specific output formats (e.g., XML tags, JSON, structured reasoning)
Teach verifiable tasks with objective correctness metrics (math, coding, fact-checking)
Improve reasoning capabilities by rewarding chain-of-thought patterns
Align models to domain-specific behaviors without labeled preference data
Optimize for multiple objectives simultaneously (format + correctness + style)

Do NOT use GRPO for:

Simple supervised fine-tuning tasks (use SFT instead)
Tasks without clear reward signals
When you already have high-quality preference pairs (use DPO/PPO instead)

Core Concepts

1. GRPO Algorithm Fundamentals

Key Mechanism:

Generates multiple completions for each prompt (group size: 4-16)
Compares completions within each group using reward functions
Updates policy to favor higher-rewarded responses relative to the group

Critical Difference from PPO:

No separate reward model needed
More sample-efficient (learns from within-group comparisons)
Simpler to implement and debug

Mathematical Intuition:

For each prompt p:
  1. Generate N completions: {c₁, c₂, ..., cₙ}
  2. Compute rewards: {r₁, r₂, ..., rₙ}
  3. Learn to increase probability of high-reward completions
     relative to low-reward ones in the same group

2. Reward Function Design Philosophy

Golden Rules:

Compose multiple reward functions - Each handles one aspect (format, correctness, style)
Scale rewards appropriately - Higher weight = stronger signal
Use incremental rewards - Partial credit for partial compliance
Test rewards independently - Debug each reward function in isolation

Reward Function Types:

Type	Use Case	Example Weight
Correctness	Verifiable tasks (math, code)	2.0 (highest)
Format	Strict structure enforcement	0.5-1.0
Length	Encourage verbosity/conciseness	0.1-0.5
Style	Penalize unwanted patterns	-0.5 to 0.5

Implementation Workflow

Step 1: Dataset Preparation

Critical Requirements:

Prompts in chat format (list of dicts with 'role' and 'content')
Include system prompts to set expectations
For verifiable tasks, include ground truth answers as additional columns

Example Structure:

from datasets import load_dataset, Dataset
 
SYSTEM_PROMPT = """
Respond in the following format:
<reasoning>
[Your step-by-step thinking]

Pro Tips:

Use one-shot or few-shot examples in system prompt for complex formats
Keep prompts concise (max_prompt_length: 256-512 tokens)
Validate data quality before training (garbage in = garbage out)

Step 2: Reward Function Implementation

Template Structure:

Example 1: Correctness Reward (Math/Coding)

Example 2: Format Reward (Structured Output)

Example 3: Incremental Format Reward (Partial Credit)

Critical Insight: Combine 3-5 reward functions for robust training. Order matters less than diversity of signals.

Step 3: Training Configuration

Memory-Optimized Config (Small GPU)

High-Performance Config (Large GPU)

Critical Hyperparameters:

Parameter	Impact	Tuning Advice
`num_generations`	Group size for comparison	Start with 8, increase to 16 if GPU allows
`learning_rate`	Convergence speed/stability	5e-6 (safe), 1e-5 (faster, riskier)
`max_completion_length`	Output verbosity	Match your task (512 for reasoning, 256 for short answers)
`gradient_accumulation_steps`	Effective batch size	Increase if GPU memory limited

Step 4: Model Setup and Training

Standard Setup (Transformers)

Unsloth Setup (2-3x Faster)

Critical Training Insights

1. Loss Behavior (EXPECTED PATTERN)

Loss starts near 0 and INCREASES during training
This is CORRECT - loss measures KL divergence from initial policy
Model is learning (diverging from original behavior to optimize rewards)
Monitor reward metrics instead of loss for progress

2. Reward Tracking

Key metrics to watch:

reward: Average across all completions
reward_std: Diversity within groups (should remain > 0)
kl: KL divergence from reference (should grow moderately)

Healthy Training Pattern:

Warning Signs:

Reward std → 0 (model collapsing to single response)
KL exploding (> 0.5) (diverging too much, reduce LR)
Reward stuck (reward functions too harsh or model capacity issue)

3. Common Pitfalls and Solutions

Problem	Symptom	Solution
Mode collapse	All completions identical	Increase `num_generations`, add diversity penalty
No learning	Flat rewards	Check reward function logic, increase LR
OOM errors	GPU memory exceeded	Reduce `num_generations`, enable gradient checkpointing
Slow training	< 1 it/s	Enable `use_vllm=True`, use Unsloth, reduce seq length
Format ignored	Model doesn't follow structure	Increase format reward weight, add incremental rewards

Advanced Patterns

1. Multi-Stage Training

For complex tasks, train in stages:

2. Adaptive Reward Scaling

3. Custom Dataset Integration

Deployment and Inference

Save and Merge LoRA

Inference Example

Best Practices Checklist

Before Training:

Validate dataset format (prompts as List[Dict])
Test reward functions on sample data
Calculate expected max_prompt_length from data
Choose appropriate num_generations based on GPU memory
Set up logging (wandb recommended)

During Training:

Monitor reward progression (should increase)
Check reward_std (should stay > 0.1)
Watch for OOM errors (reduce batch size if needed)
Sample generations every 50-100 steps
Validate format compliance on holdout set

After Training:

Merge LoRA weights if using PEFT
Test on diverse prompts
Compare to baseline model
Document reward weights and hyperparameters
Save reproducibility config

Troubleshooting Guide

Debugging Workflow

Isolate reward functions - Test each independently
Check data distribution - Ensure diversity in prompts
Reduce complexity - Start with single reward, add gradually
Monitor generations - Print samples every N steps
Validate extraction logic - Ensure answer parsing works

Quick Fixes

References and Resources

Official Documentation:

TRL GRPO Trainer: https://huggingface.co/docs/trl/grpo_trainer
DeepSeek R1 Paper: https://arxiv.org/abs/2501.12948
Unsloth Docs: https://docs.unsloth.ai/

Example Repositories:

Open R1 Implementation: https://github.com/huggingface/open-r1
TRL Examples: https://github.com/huggingface/trl/tree/main/examples

Usage Instructions for Agents

When this skill is loaded:

Read this entire file before implementing GRPO training
Start with the simplest reward function (e.g., length-based) to validate setup
Use the templates in templates/ directory as starting points
Reference examples in examples/ for task-specific implementations
Follow the workflow sequentially (don't skip steps)
Debug incrementally - add one reward function at a time

Critical Reminders:

Always use multiple reward functions (3-5 is optimal)
Monitor reward metrics, not loss
Test reward functions before training
Start small (num_generations=4), scale up gradually
Save checkpoints frequently (every 100 steps)

This skill is designed for expert-level implementation. Beginners should start with supervised fine-tuning before attempting GRPO.