StateSet API Docs

Introduction

Welcome to StateSet’s revolutionary Reinforcement Learning platform, powered by Group Relative Policy Optimization (GRPO). This represents a fundamental shift in how we train AI models - moving beyond simple next-word prediction to models that learn through exploration, evaluation, and optimization toward specific goals.

Key Insight: Traditional fine-tuning maximizes next-word prediction probability. GRPO maximizes reward functions - teaching models not just what to say, but how to achieve optimal outcomes.

Why Reinforcement Learning?

Goal-Oriented Learning

Models learn to maximize specific objectives rather than just mimicking training data

Exploration & Discovery

Models generate multiple solutions and learn from comparing outcomes

Continuous Improvement

Every interaction becomes a learning opportunity through reward optimization

How GRPO Trains Your Model

The Training Process

Step-by-Step Breakdown

Multiple Response Generation

For each question-answer pair, the model generates multiple possible responses (e.g., 8-16 variations)

# Example: Model generates variations for "How to handle a refund?"
responses = [
    "I'd be happy to help with your refund...",
    "Let me process that refund for you...",
    "I understand you need a refund...",
    # ... 5-13 more variations
]

Response Evaluation

Each response is evaluated using sophisticated reward functions

rewards = []
for response in responses:
    reward = evaluate_response(response, criteria={
        'helpfulness': 0.4,
        'accuracy': 0.3,
        'tone': 0.2,
        'efficiency': 0.1
    })
    rewards.append(reward)

Baseline Calculation

The average reward serves as a baseline for comparison

baseline = sum(rewards) / len(rewards)
relative_rewards = [r - baseline for r in rewards]

Weight Updates

Model weights are updated to reinforce above-average responses

# Responses better than average get reinforced
# Responses worse than average get discouraged
model.update_weights(responses, relative_rewards)

Training Scale

Small Dataset
Standard Training
Advanced Training

300 rows × 1 epoch = 300 training steps

Quick iteration
Rapid prototyping
Initial model validation

Understanding Reward Functions & Verifiers

The Distinction

Verifier

Binary Evaluation

Determines correct/incorrect
No numerical scoring
Can execute code for validation

def verify_math(question, answer):
    if question == "2+2" and answer == "4":
        return True
    return False

Reward Function

Numerical Scoring

Assigns scores (-∞ to +∞)
Considers multiple criteria
Guides optimization direction

def reward_function(question, answer, verified):
    score = 0
    if verified:
        score += 2
    if len(answer) < 100:  # Brevity bonus
        score += 0.5
    return score

Reward Function Design

The power of GRPO lies in well-designed reward functions that capture your exact objectives:

class CustomerServiceReward:
    def __init__(self):
        self.criteria = {
            'resolution': 0.35,      # Did it solve the problem?
            'empathy': 0.25,         # Was it understanding?
            'clarity': 0.20,         # Was it clear?
            'efficiency': 0.10,      # Was it concise?
            'policy_compliance': 0.10 # Did it follow rules?
        }
    
    def calculate(self, response, context):
        scores = {}
        
        # Resolution scoring
        scores['resolution'] = self.check_resolution(response, context)
        
        # Empathy detection
        scores['empathy'] = self.measure_empathy(response)
        
        # Clarity assessment
        scores['clarity'] = self.evaluate_clarity(response)
        
        # Length efficiency
        scores['efficiency'] = min(1.0, 50 / len(response.split()))
        
        # Policy adherence
        scores['policy_compliance'] = self.check_policies(response, context)
        
        # Weighted sum
        total = sum(scores[k] * self.criteria[k] for k in scores)
        
        return total, scores  # Return total and breakdown

Critical: Poorly designed reward functions can degrade model performance. Always test reward functions thoroughly before full-scale training.

Group Relative Policy Optimization (GRPO)

The Innovation

Traditional RL algorithms like PPO require training a separate “critic” model to estimate value. GRPO eliminates this overhead:

Traditional PPO
GRPO Innovation

Drawbacks:

Train two models
More memory usage
Slower convergence

How GRPO Works

Group Sampling

Generate multiple solutions for each problem

# Instead of one response per prompt
responses = model.generate(
    prompt="How should I handle an angry customer?",
    num_returns=8  # Generate 8 variations
)

Reward Assignment

Evaluate each solution’s quality

rewards = []
for response in responses:
    reward = reward_function(response)
    rewards.append(reward)
# rewards = [0.8, 1.2, 0.6, 1.5, 0.9, 1.1, 0.7, 1.3]

Baseline Calculation

Use group average as baseline

baseline = sum(rewards) / len(rewards)  # 1.0125

Policy Update

Reinforce above-average, discourage below-average

for response, reward in zip(responses, rewards):
    advantage = reward - baseline
    if advantage > 0:
        # Increase probability of this response pattern
        model.reinforce(response, advantage)
    else:
        # Decrease probability of this response pattern
        model.discourage(response, -advantage)

Configuration Deep Dive

Key Parameters

# GRPO Configuration
actor_rollout:
  ref:
    rollout:
      n: 8  # Generate 8 responses per prompt (critical for GRPO)

data:
  train_batch_size: 32  # Number of prompts per batch

actor_rollout_ref:
  actor:
    ppo_mini_batch_size: 256  # Mini-batch for weight updates
    ppo_epochs: 4  # Update epochs per trajectory set
    clip_ratio: 0.2  # GRPO clip range for stable training
    
    # GRPO-specific settings
    use_kl_loss: true  # Enable KL regularization
    kl_loss_coef: 0.001  # KL penalty strength
    kl_loss_type: "kl"  # Options: kl, abs, mse, low_var_kl, full
    
    # Loss aggregation
    loss_agg_mode: "token-mean"  # Stable for long responses

algorithm:
  adv_estimator: "grpo"  # Use GRPO instead of GAE

Configuration Profiles

Conservative
Balanced
Aggressive

# For critical applications
clip_ratio: 0.1
kl_loss_coef: 0.01
ppo_epochs: 2
rollout.n: 4

# For most use cases
clip_ratio: 0.2
kl_loss_coef: 0.001
ppo_epochs: 4
rollout.n: 8

# For rapid learning
clip_ratio: 0.3
kl_loss_coef: 0.0001
ppo_epochs: 6
rollout.n: 16

Practical Implementation

Example: Customer Service Agent

from StateSet.rl import GRPOTrainer, RewardFunction
from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import LoraConfig, get_peft_model, TaskType
from trl import GRPOTrainer, GRPOConfig
from datasets import Dataset
from sklearn.model_selection import train_test_split

# Model setup with LoRA for efficient training
def setup_model(model_name="Qwen/Qwen2.5-7B-Instruct"):
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(
        model_name, 
        trust_remote_code=True,
        padding_side="left"  # Critical for generation
    )
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
    
    # Load model with mixed precision
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        device_map='auto',
        trust_remote_code=True
    )
    
    # Add LoRA adapters for efficient training
    lora_config = LoraConfig(
        r=8,
        lora_alpha=16,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
        task_type=TaskType.CAUSAL_LM,
    )
    
    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()  # Typically <1% of total
    
    return model, tokenizer

# Define comprehensive reward function
class CustomerServiceReward:
    def __init__(self, expected_responses=None):
        self.empathy_keywords = ['sorry', 'understand', 'help', 'happy', 'glad', 'assist']
        self.action_keywords = ['visit', 'email', 'click', 'check', 'provide']
        self.expected_responses = expected_responses or {}
        
    def compute_reward(self, query, response, expected=None):
        response_lower = response.lower()
        
        # Similarity to expected response (if available)
        similarity_reward = 0.0
        if expected:
            import difflib
            similarity = difflib.SequenceMatcher(None, response_lower, expected.lower()).ratio()
            similarity_reward = similarity
        
        # Empathy score
        empathy_score = sum(1 for word in self.empathy_keywords if word in response_lower) / len(self.empathy_keywords)
        
        # Action-oriented score
        action_score = sum(1 for word in self.action_keywords if word in response_lower) / len(self.action_keywords)
        
        # Length penalty (concise but complete)
        word_count = len(response.split())
        if word_count < 10:
            length_penalty = -0.5
        elif word_count > 100:
            length_penalty = -0.3
        else:
            length_penalty = 0.0
        
        # Weighted combination
        total_reward = (
            0.4 * similarity_reward +
            0.3 * empathy_score +
            0.2 * action_score +
            0.1 * (1.0 + length_penalty)
        )
        
        return {
            "total": total_reward,
            "similarity": similarity_reward,
            "empathy": empathy_score,
            "action": action_score,
            "length_penalty": length_penalty
        }

# Configure GRPO training
training_config = GRPOConfig(
    output_dir="./checkpoints/grpo",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,  # Effective batch size = 8
    learning_rate=1e-5,
    num_train_epochs=1,
    max_grad_norm=0.5,
    warmup_steps=50,
    fp16=True,  # Mixed precision training
    logging_steps=10,
    save_steps=100,
    eval_steps=50,
    
    # GRPO specific parameters
    beta=0.0,  # KL penalty coefficient
    num_generations=4,  # Generate 4 responses per prompt
    num_iterations=1,
    
    # Generation parameters
    max_prompt_length=128,
    max_completion_length=128,
    temperature=0.7,
    top_p=0.9,
    
    # Reproducibility
    seed=42,
)

# Train with proper data splitting
def train_customer_service_model(data, eval_split=0.1):
    # Split data
    train_data, eval_data = train_test_split(
        data, 
        test_size=eval_split, 
        random_state=42,
        stratify=[d['task_type'] for d in data]  # Maintain task distribution
    )
    
    # Create datasets
    train_dataset = Dataset.from_list([
        {"prompt": f"Customer: {d['query']}\nAssistant:"}
        for d in train_data
    ])
    eval_dataset = Dataset.from_list([
        {"prompt": f"Customer: {d['query']}\nAssistant:"}
        for d in eval_data
    ])
    
    # Setup reward function with expected responses
    reward_model = CustomerServiceReward({
        f"Customer: {d['query']}\nAssistant:": d['expected_response']
        for d in train_data
    })
    
    def reward_fn(completions, prompts):
        rewards = []
        for completion, prompt in zip(completions, prompts):
            expected = reward_model.expected_responses.get(prompt)
            reward_dict = reward_model.compute_reward(
                prompt.split("Customer: ")[1].split("\nAssistant:")[0],
                completion,
                expected
            )
            rewards.append(reward_dict["total"])
        return rewards
    
    # Initialize trainer
    trainer = GRPOTrainer(
        model=model,
        args=training_config,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        reward_funcs=reward_fn,
    )
    
    # Train
    trainer.train()
    
    return trainer.model

Real-World Training Pipeline

# Complete training pipeline with evaluation
import logging
from pathlib import Path
import torch
import wandb

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler('grpo_training.log')
    ]
)
logger = logging.getLogger(__name__)

# Environment configuration
CONFIG = {
    'max_examples': 5000,
    'batch_size': 2,
    'gradient_accumulation_steps': 4,
    'learning_rate': 1e-5,
    'num_epochs': 1,
    'max_grad_norm': 0.5,
    'warmup_steps': 50,
    'checkpoint_dir': './checkpoints/grpo',
    'use_mixed_precision': True,
    'model_name': 'Qwen/Qwen2.5-7B-Instruct',
    'max_length': 256,
    'max_new_tokens': 128,
    'lora_r': 8,
    'lora_alpha': 16,
    'lora_dropout': 0.05,
    'num_generations': 4,
    'eval_split_size': 0.1,
    'seed': 42,
}

# Set random seeds for reproducibility
def set_seed(seed: int):
    import random
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

# Post-training evaluation
def run_post_training_evaluation(model, tokenizer, eval_examples, reward_model, num_samples=5):
    """Evaluate model performance on held-out examples"""
    logger.info("Running post-training evaluation...")
    
    eval_results = []
    for i, example in enumerate(eval_examples[:num_samples]):
        prompt = f"Customer: {example['query']}\nAssistant:"
        inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=CONFIG['max_length'])
        inputs = {k: v.to(model.device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = model.generate(
                **inputs,
                max_new_tokens=CONFIG['max_new_tokens'],
                temperature=0.7,
                top_p=0.9,
                pad_token_id=tokenizer.pad_token_id,
                eos_token_id=tokenizer.eos_token_id,
            )
        
        generated = tokenizer.decode(outputs[0][len(inputs['input_ids'][0]):], skip_special_tokens=True)
        reward_dict = reward_model.compute_reward(example['query'], generated, example['expected_response'])
        
        eval_results.append({
            'query': example['query'],
            'generated': generated,
            'expected': example['expected_response'],
            'reward': reward_dict
        })
        
        logger.info(f"Sample {i+1}:")
        logger.info(f"  Query: {example['query']}")
        logger.info(f"  Generated: {generated[:100]}...")
        logger.info(f"  Reward: {reward_dict['total']:.4f} (similarity: {reward_dict['similarity']:.4f})")
    
    return eval_results

# Main training function
def main():
    # Initialize wandb for experiment tracking
    wandb.init(
        project="customer-service-grpo",
        config=CONFIG,
        name=f"grpo-{datetime.now().strftime('%Y%m%d-%H%M%S')}"
    )
    
    # Set seed
    set_seed(CONFIG['seed'])
    
    # Load and prepare data
    data = load_customer_service_data()  # Your data loading function
    
    # Train model
    model = train_customer_service_model(data, eval_split=CONFIG['eval_split_size'])
    
    # Save model
    model.save_pretrained(CONFIG['checkpoint_dir'])
    tokenizer.save_pretrained(CONFIG['checkpoint_dir'])
    
    # Run evaluation
    eval_results = run_post_training_evaluation(
        model, tokenizer, eval_data, reward_model, 
        num_samples=5
    )
    
    # Log results to wandb
    for i, result in enumerate(eval_results):
        wandb.log({
            f"eval/sample_{i}/reward": result['reward']['total'],
            f"eval/sample_{i}/similarity": result['reward']['similarity'],
        })
    
    logger.info("✨ Training completed successfully!")

Monitoring Training Progress

# Real-time metrics during training
trainer.on_step = lambda metrics: print(f"""
Step {metrics.step}:
  Average Reward: {metrics.avg_reward:.3f}
  Reward Std: {metrics.reward_std:.3f}
  KL Divergence: {metrics.kl_div:.4f}
  Response Quality: {metrics.quality_score:.2%}
""")

# Track detailed metrics with wandb
wandb.log({
    "train/reward_mean": metrics.avg_reward,
    "train/reward_std": metrics.reward_std,
    "train/kl_divergence": metrics.kl_div,
    "train/learning_rate": trainer.optimizer.param_groups[0]['lr'],
    "train/gradient_norm": metrics.grad_norm,
})

Best Practices

1. Data Preparation

Quality Over Quantity

Focus on high-quality, diverse examples

def validate_conversation(conv):
    """Ensure data quality before training"""
    # Check structure
    if not isinstance(conv, dict) or 'messages' not in conv:
        return False
    
    messages = conv['messages']
    if not isinstance(messages, list) or len(messages) < 2:
        return False
    
    # Validate each message
    for msg in messages:
        if not all(k in msg for k in ['role', 'content']):
            return False
        if msg['role'] not in ['user', 'assistant', 'system']:
            return False
        if not msg['content'].strip():
            return False
    
    return True

Task Classification

Classify examples by task type for balanced training

def classify_task_type(query, response):
    """Categorize customer service tasks"""
    query_lower = query.lower()
    
    task_mapping = {
        'order_tracking': ['tracking', 'order', 'shipment', 'delivery'],
        'returns': ['return', 'exchange', 'refund'],
        'product_inquiry': ['shade', 'color', 'match', 'size'],
        'technical_issue': ['damaged', 'broken', 'issue', 'problem'],
        'general': []  # Default category
    }
    
    for task_type, keywords in task_mapping.items():
        if any(word in query_lower for word in keywords):
            return task_type
    
    return 'general'

Train/Eval Split

Always maintain a held-out evaluation set

from sklearn.model_selection import train_test_split

# Stratified split to maintain task distribution
train_examples, eval_examples = train_test_split(
    examples, 
    test_size=0.1,  # 10% for evaluation
    random_state=42,
    stratify=[e.task_type for e in examples]
)

2. Reward Function Design

Start Simple

Begin with basic reward functions and gradually add complexity

# Start with this
def simple_reward(response):
    return 1.0 if is_correct(response) else -1.0

# Evolve to this
def complex_reward(response, context):
    return weighted_sum([
        correctness_score(response),
        style_score(response),
        efficiency_score(response),
        context_relevance(response, context)
    ])

Test Extensively

Validate reward functions before training

def test_reward_function(reward_fn, test_cases):
    for prompt, good_response, bad_response in test_cases:
        good_score = reward_fn(prompt, good_response)
        bad_score = reward_fn(prompt, bad_response)
        assert good_score > bad_score, f"Reward function failed on {prompt}"

Balance Criteria

Avoid over-optimizing for single metrics

# Bad: Single metric
def bad_reward(response):
    return -len(response)  # Only optimizes for brevity

# Good: Balanced metrics
def good_reward(response):
    return 0.6 * quality + 0.3 * brevity + 0.1 * style

3. Performance Optimization

# Memory-efficient training with gradient accumulation
trainer = GRPOTrainer(
    gradient_accumulation_steps=4,  # Simulate larger batch
    gradient_checkpointing=True,    # Trade compute for memory
    fp16=True,                      # Mixed precision
    dataloader_num_workers=4,       # Parallel data loading
)

# Multi-GPU training
if torch.cuda.device_count() > 1:
    # Automatically handled by trainer with proper config
    training_config.ddp_find_unused_parameters = False
    training_config.dataloader_num_workers = 4

4. Production Deployment

# Load trained model for inference
def load_trained_model(checkpoint_dir):
    from peft import PeftModel
    
    # Load base model
    base_model = AutoModelForCausalLM.from_pretrained(
        CONFIG['model_name'],
        torch_dtype=torch.float16,
        device_map='auto'
    )
    
    # Load LoRA weights
    model = PeftModel.from_pretrained(base_model, checkpoint_dir)
    model = model.merge_and_unload()  # Merge for faster inference
    
    # Load tokenizer
    tokenizer = AutoTokenizer.from_pretrained(checkpoint_dir)
    
    return model, tokenizer

# Inference function
def generate_response(model, tokenizer, query):
    prompt = f"Customer: {query}\nAssistant:"
    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
    inputs = {k: v.to(model.device) for k, v in inputs.items()}
    
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=128,
            temperature=0.7,
            top_p=0.9,
            do_sample=True,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id,
        )
    
    response = tokenizer.decode(outputs[0][len(inputs['input_ids'][0]):], skip_special_tokens=True)
    return response.strip()

Real-World Impact

Case Study: Customer Support

Before GRPO

Generic responses
65% resolution rate
3.2/5 satisfaction
High escalation rate

After GRPO

Context-aware responses
89% resolution rate
4.6/5 satisfaction
40% fewer escalations

Performance Metrics

# Measure improvement
baseline_model = load_model("base")
grpo_model = load_model("grpo_trained")

metrics = evaluate_models(baseline_model, grpo_model, test_set)

print(f"""
Performance Improvement:
- Task Success: +{metrics.success_delta:.1%}
- User Satisfaction: +{metrics.satisfaction_delta:.1%}
- Response Quality: +{metrics.quality_delta:.1%}
- Efficiency: +{metrics.efficiency_delta:.1%}
""")

Getting Started

Define Your Objective

What behavior do you want to optimize for?

objective = "Maximize customer satisfaction while resolving issues efficiently"

Design Reward Function

Translate objectives into measurable rewards

reward = CustomerSatisfactionReward(
    weights={'resolution': 0.5, 'tone': 0.3, 'efficiency': 0.2}
)

Prepare Training Data

Collect quality examples (300+ recommended)

data = prepare_training_data(
    source="support_tickets",
    min_quality_score=0.8
)

Configure & Train

Start with balanced settings

model = train_grpo_model(
    data=data,
    reward_fn=reward,
    profile="balanced"
)

Evaluate & Deploy

Test thoroughly before production

if evaluate(model).meets_criteria():
    deploy_to_production(model)

Advanced Topics

Multi-Objective Optimization

class MultiObjectiveReward:
    def __init__(self, objectives):
        self.objectives = objectives
    
    def compute(self, response, context):
        scores = {}
        for name, (weight, function) in self.objectives.items():
            scores[name] = function(response, context) * weight
        
        # Pareto optimization
        if self.is_pareto_optimal(scores):
            bonus = 0.2
        else:
            bonus = 0
        
        return sum(scores.values()) + bonus

Curriculum Learning

# Start with easy tasks, gradually increase difficulty
curriculum = [
    {"difficulty": "easy", "epochs": 1, "reward_scale": 1.0},
    {"difficulty": "medium", "epochs": 2, "reward_scale": 0.8},
    {"difficulty": "hard", "epochs": 3, "reward_scale": 0.6}
]

for stage in curriculum:
    model = trainer.train(
        data=filter_by_difficulty(data, stage["difficulty"]),
        epochs=stage["epochs"],
        reward_scale=stage["reward_scale"]
    )

Distributed Training

# Setup for multi-GPU training
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_distributed_training():
    """Initialize distributed training environment"""
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    
    if local_rank != -1:
        # Initialize process group
        torch.cuda.set_device(local_rank)
        dist.init_process_group(backend="nccl")
        logger.info(f"Initialized distributed training: rank {local_rank}/{world_size}")
        
        # Only log from main process
        if local_rank != 0:
            logging.getLogger().setLevel(logging.WARNING)
    
    return local_rank, world_size

# Launch distributed training
# torchrun --nproc_per_node=4 train_grpo.py

Error Handling & Robustness

# Robust data loading with fallback
def load_data_with_fallback(file_path: str, max_examples: int = None):
    """Load data with automatic fallback to sample data"""
    
    def get_fallback_data():
        """Return sample data for testing/development"""
        return [
            {
                "messages": [
                    {"role": "user", "content": "Where is my order?"},
                    {"role": "assistant", "content": "I'd be happy to help track your order. Could you please provide your order number?"}
                ]
            },
            # Add more examples...
        ]
    
    # Try loading actual data
    if not os.path.exists(file_path):
        logger.warning(f"Data file not found: {file_path}, using fallback data")
        return get_fallback_data()
    
    try:
        data = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    try:
                        conv = json.loads(line)
                        if validate_conversation(conv):
                            data.append(conv)
                    except json.JSONDecodeError:
                        continue
                    
                if max_examples and len(data) >= max_examples:
                    break
        
        if not data:
            logger.warning("No valid data loaded, using fallback")
            return get_fallback_data()
            
        return data
        
    except Exception as e:
        logger.error(f"Error loading data: {e}")
        return get_fallback_data()

# Model loading with automatic fallback
class ModelManager:
    def __init__(self, model_name: str, fallback_model: str = "gpt2"):
        self.model_name = model_name
        self.fallback_model = fallback_model
        
    def load_model_and_tokenizer(self):
        try:
            # Try loading requested model
            return self._load_model(self.model_name)
        except Exception as e:
            logger.error(f"Failed to load {self.model_name}: {e}")
            if self.model_name != self.fallback_model:
                logger.info(f"Falling back to {self.fallback_model}")
                return self._load_model(self.fallback_model)
            else:
                raise

Advanced GRPO Configuration

# Comprehensive GRPO configuration with all options
advanced_config = GRPOConfig(
    # Basic training parameters
    output_dir="./checkpoints/grpo",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=1e-5,
    num_train_epochs=1,
    
    # Optimization
    max_grad_norm=0.5,
    warmup_steps=50,
    weight_decay=0.01,
    adam_epsilon=1e-8,
    
    # Mixed precision & performance
    fp16=True,
    bf16=False,  # Use FP16 instead of BF16
    gradient_checkpointing=True,
    
    # Logging & checkpointing
    logging_steps=10,
    save_steps=100,
    eval_steps=50,
    save_total_limit=3,
    load_best_model_at_end=True,
    metric_for_best_model="eval_reward",
    
    # GRPO specific
    beta=0.0,  # KL penalty (0 = no penalty)
    num_generations=4,  # Responses per prompt
    num_iterations=1,  # GRPO iterations per batch
    
    # Generation config
    max_prompt_length=128,
    max_completion_length=128,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
    
    # Advanced options
    use_liger_loss=False,  # Experimental loss function
    ddp_find_unused_parameters=False,  # For multi-GPU
    dataloader_num_workers=4,
    
    # Reproducibility
    seed=42,
    data_seed=42,
)

# Dynamic configuration based on hardware
def get_optimal_config():
    """Automatically configure based on available hardware"""
    config = GRPOConfig()
    
    # Adjust batch size based on GPU memory
    if torch.cuda.is_available():
        gpu_memory = torch.cuda.get_device_properties(0).total_memory
        if gpu_memory > 40 * 1024**3:  # 40GB+ (A100)
            config.per_device_train_batch_size = 8
            config.gradient_accumulation_steps = 1
        elif gpu_memory > 20 * 1024**3:  # 20GB+ (A6000)
            config.per_device_train_batch_size = 4
            config.gradient_accumulation_steps = 2
        else:  # Smaller GPUs
            config.per_device_train_batch_size = 1
            config.gradient_accumulation_steps = 8
    
    # Enable mixed precision on capable GPUs
    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 7:
        config.fp16 = True
    
    return config

Experiment Tracking & Analysis

# Comprehensive experiment tracking
class ExperimentTracker:
    def __init__(self, project_name: str, use_wandb: bool = True):
        self.project_name = project_name
        self.use_wandb = use_wandb and self._init_wandb()
        self.metrics_history = []
        
    def _init_wandb(self):
        try:
            import wandb
            wandb.login(key=os.getenv('WANDB_API_KEY'))
            wandb.init(
                project=self.project_name,
                config=CONFIG,
                name=f"grpo-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
                tags=["grpo", "customer-service", CONFIG['model_name'].split('/')[-1]]
            )
            return True
        except Exception as e:
            logger.warning(f"Wandb init failed: {e}")
            return False
    
    def log_metrics(self, step: int, metrics: dict):
        """Log training metrics"""
        self.metrics_history.append({"step": step, **metrics})
        
        if self.use_wandb:
            import wandb
            wandb.log(metrics, step=step)
        
        # Also log to file for backup
        with open("metrics.jsonl", "a") as f:
            f.write(json.dumps({"step": step, **metrics}) + "\n")
    
    def log_generation_samples(self, samples: list):
        """Log example generations"""
        if self.use_wandb:
            import wandb
            table = wandb.Table(columns=["Query", "Generated", "Expected", "Reward"])
            for sample in samples:
                table.add_data(
                    sample['query'],
                    sample['generated'],
                    sample['expected'],
                    sample['reward']['total']
                )
            wandb.log({"generation_samples": table})
    
    def create_summary_report(self):
        """Generate training summary"""
        if not self.metrics_history:
            return
        
        import pandas as pd
        df = pd.DataFrame(self.metrics_history)
        
        summary = {
            "total_steps": len(df),
            "final_reward": df['reward'].iloc[-1],
            "max_reward": df['reward'].max(),
            "avg_reward": df['reward'].mean(),
            "reward_improvement": df['reward'].iloc[-1] - df['reward'].iloc[0],
        }
        
        logger.info("Training Summary:")
        for key, value in summary.items():
            logger.info(f"  {key}: {value:.4f}")
        
        return summary

Model Deployment Strategies

# Production deployment with optimization
class ProductionDeployment:
    def __init__(self, checkpoint_dir: str):
        self.checkpoint_dir = checkpoint_dir
        
    def prepare_for_deployment(self):
        """Optimize model for production inference"""
        from peft import PeftModel
        import torch.quantization as quantization
        
        # Load model
        base_model = AutoModelForCausalLM.from_pretrained(
            CONFIG['model_name'],
            torch_dtype=torch.float16,
            device_map='auto'
        )
        
        # Merge LoRA weights
        model = PeftModel.from_pretrained(base_model, self.checkpoint_dir)
        model = model.merge_and_unload()
        
        # Optional: Quantization for faster inference
        if CONFIG.get('quantize_for_deployment', False):
            model = quantization.quantize_dynamic(
                model, {torch.nn.Linear}, dtype=torch.qint8
            )
        
        # Save optimized model
        output_dir = f"{self.checkpoint_dir}_production"
        model.save_pretrained(output_dir)
        
        # Also save tokenizer
        tokenizer = AutoTokenizer.from_pretrained(self.checkpoint_dir)
        tokenizer.save_pretrained(output_dir)
        
        return output_dir
    
    def create_inference_api(self, model_path: str):
        """Create FastAPI endpoint for model inference"""
        from fastapi import FastAPI
        from pydantic import BaseModel
        
        app = FastAPI()
        
        # Load model once at startup
        model, tokenizer = load_trained_model(model_path)
        
        class Query(BaseModel):
            text: str
            max_length: int = 128
            temperature: float = 0.7
        
        @app.post("/generate")
        async def generate(query: Query):
            response = generate_response(
                model, tokenizer, query.text,
                max_length=query.max_length,
                temperature=query.temperature
            )
            return {"response": response}
        
        return app

Troubleshooting Guide

Out of Memory Errors

# Solutions for OOM errors

# 1. Reduce batch size
config.per_device_train_batch_size = 1
config.gradient_accumulation_steps = 8

# 2. Enable gradient checkpointing
config.gradient_checkpointing = True

# 3. Use LoRA with smaller rank
lora_config.r = 4  # Instead of 8 or 16

# 4. Reduce max sequence length
config.max_length = 128  # Instead of 256+

# 5. Use CPU offloading
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    offload_folder="offload",
    offload_state_dict=True
)

Unstable Training

# Solutions for training instability

# 1. Reduce learning rate
config.learning_rate = 5e-6  # Instead of 1e-5

# 2. Increase warmup steps
config.warmup_steps = 100  # Instead of 50

# 3. Add KL penalty
config.beta = 0.01  # Instead of 0.0

# 4. Clip gradients more aggressively
config.max_grad_norm = 0.3  # Instead of 0.5

# 5. Use smaller generation count
config.num_generations = 2  # Instead of 4+

Poor Generation Quality

# Improve generation quality

# 1. Adjust generation parameters
generation_config = {
    "temperature": 0.8,  # Slightly higher for diversity
    "top_p": 0.95,       # Broader sampling
    "top_k": 50,         # Limit vocabulary
    "repetition_penalty": 1.1,  # Reduce repetition
}

# 2. Improve reward function
def enhanced_reward(response):
    # Add more nuanced scoring
    scores = {
        "relevance": check_relevance(response),
        "completeness": check_completeness(response),
        "tone": check_tone(response),
        "grammar": check_grammar(response),
    }
    return weighted_average(scores)

# 3. Filter training data
high_quality_data = [
    ex for ex in data 
    if len(ex['response'].split()) > 20  # Min length
    and len(ex['response'].split()) < 150  # Max length
    and validate_quality(ex)
]

Conclusion

GRPO represents a paradigm shift in AI training - from passive learning to active optimization. By defining clear objectives through reward functions and allowing models to explore multiple solutions, we create AI systems that don’t just mimic but genuinely optimize for desired outcomes.

Key Implementation Insights

Based on real-world GRPO training experience, here are the critical success factors:

Use LoRA for Efficiency

Training with LoRA adapters reduces memory usage by 90%+ while maintaining performance. Target the attention layers (q_proj, k_proj, v_proj, o_proj) for best results.

Always Split Train/Eval

Use stratified splitting to maintain task distribution. A 90/10 split provides enough evaluation data while maximizing training examples.

Design Multi-Component Rewards

Combine similarity, empathy, action-orientation, and length penalties. Weight them based on your specific use case (e.g., 40% similarity, 30% empathy).

Start Conservative

Begin with batch_size=2, gradient_accumulation=4, learning_rate=1e-5. These settings work well across different model sizes and GPUs.

Production Checklist

Data Quality

Validate all conversations have proper structure
Ensure minimum response length (>10 words)
Remove duplicates and low-quality examples
Classify by task type for balanced training

Model Configuration

Use FP16 mixed precision (not BF16)
Enable gradient checkpointing for large models
Set padding_side=“left” for proper generation
Configure LoRA with r=8, alpha=16 as starting point

Training Setup

Implement robust error handling with fallbacks
Use wandb or similar for experiment tracking
Save checkpoints frequently (every 100 steps)
Monitor reward variance for stability

Evaluation

Run post-training evaluation on held-out data
Track multiple metrics (similarity, quality, length)
Generate sample outputs for manual review
Compare against baseline model performance

Deployment

Merge LoRA weights for faster inference
Consider quantization for edge deployment
Implement proper error handling in API
Monitor inference latency and quality

Quick Start Template

Prerequisites

Before starting, ensure you have:

Python 3.8+
NVIDIA GPU with CUDA 11.0+ (for accelerated training)
Git installed
Optional: Weights & Biases account for experiment tracking

Install core dependencies:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers peft datasets trl scikit-learn wandb fastapi uvicorn

Step-by-Step Setup

# Clone the training template
git clone https://github.com/StateSet/grpo-agent
cd grpo-agent

# Create and activate virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # On Linux/Mac
# venv\Scripts\activate  # On Windows

# Install dependencies
pip install -r requirements.txt

# Prepare your data
# Data should be in JSONL format with conversations:
# {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
python scripts/prepare_data.py --input your_raw_data.json --output data/training_data.jsonl --max-examples 5000 --validate

# Configure training
# Edit config.yaml or set environment variables
export MODEL_NAME="Qwen/Qwen2.5-7B-Instruct"
export BATCH_SIZE=2
export LEARNING_RATE=1e-5
export MAX_EXAMPLES=5000
export EVAL_SPLIT=0.1
export WANDB_API_KEY="your_key_here"  # Optional

# Run training
# Use torchrun for multi-GPU: torchrun --nproc_per_node=2 train_grpo.py
python train_grpo.py --data data/training_data.jsonl --config config.yaml

# Monitor training
# View logs: tail -f grpo_training.log
# If using wandb: open the link provided in console

# Evaluate results
python evaluate.py --checkpoint checkpoints/grpo/final_model --num-samples 10 --output eval_results.json

# Deploy as API
python deploy.py --model checkpoints/grpo/final_model --port 8080 --host 0.0.0.0
# Test: curl -X POST https://yourapp.com:8080/generate -H "Content-Type: application/json" -d '{"text": "Where is my order?"}'

Customization Tips

Small Datasets: Set NUM_EPOCHS=3 and NUM_GENERATIONS=8
Large Models: Enable gradient_checkpointing in config.yaml
Debug Mode: Add —debug to train_grpo.py for verbose logging
Resume Training: Use —resume_from_checkpoint checkpoints/grpo/checkpoint-100

Common Pitfalls & Solutions

Data Loading Issues

Symptom: “No valid data loaded” or JSON decode errorsSolutions:

Ensure JSONL format (one JSON object per line)
Run validation: python scripts/validate_data.py data.jsonl
Use fallback data for testing: —use-fallback
Check encoding: All files should be UTF-8

GPU Memory Problems

Symptom: CUDA OOM errors during trainingSolutions:

Reduce per_device_train_batch_size to 1
Increase gradient_accumulation_steps to 8+
Use smaller model (e.g., “Qwen/Qwen2.5-3B-Instruct”)
Enable fp16 and gradient_checkpointing
Monitor with: nvidia-smi -l 1

Reward Function Problems

Symptom: Low/negative rewards or unstable trainingSolutions:

Normalize rewards to [-1, 1] range
Test independently: python test_reward.py —samples 10
Add epsilon to divisions: score = sum / (len + 1e-5)
Balance weights: Start with equal weights and adjust
Monitor reward distribution in wandb

Generation Quality Issues

Symptom: Repetitive or off-topic responsesSolutions:

Adjust temperature (0.7-0.9) and top_p (0.9-0.95)
Add repetition_penalty=1.2 in generation config
Increase num_generations to 8 for more exploration
Fine-tune prompt format: Add system instructions
Evaluate diversity: Compute unique n-grams in outputs

Deployment Errors

Symptom: Inference fails or slow performanceSolutions:

Merge LoRA weights before deployment
Use torch.compile(model) for PyTorch 2.0+
Set device_map=‘auto’ for multi-GPU inference
Implement batching for multiple requests
Profile with: torch.profiler

These pitfalls are based on real-world GRPO implementations - addressing them early will save significant time and resources.

Getting Started

Define Your Objective

What behavior do you want to optimize for?

objective = "Maximize customer satisfaction while resolving issues efficiently"

Design Reward Function

Translate objectives into measurable rewards

reward = CustomerSatisfactionReward(
    weights={'resolution': 0.5, 'tone': 0.3, 'efficiency': 0.2}
)

Prepare Training Data

Collect quality examples (300+ recommended)

data = prepare_training_data(
    source="support_tickets",
    min_quality_score=0.8
)

Configure & Train

Start with balanced settings

model = train_grpo_model(
    data=data,
    reward_fn=reward,
    profile="balanced"
)

Evaluate & Deploy

Test thoroughly before production

if evaluate(model).meets_criteria():
    deploy_to_production(model)

Advanced Topics

Multi-Objective Optimization

class MultiObjectiveReward:
    def __init__(self, objectives):
        self.objectives = objectives
    
    def compute(self, response, context):
        scores = {}
        for name, (weight, function) in self.objectives.items():
            scores[name] = function(response, context) * weight
        
        # Pareto optimization
        if self.is_pareto_optimal(scores):
            bonus = 0.2
        else:
            bonus = 0
        
        return sum(scores.values()) + bonus

Curriculum Learning

# Start with easy tasks, gradually increase difficulty
curriculum = [
    {"difficulty": "easy", "epochs": 1, "reward_scale": 1.0},
    {"difficulty": "medium", "epochs": 2, "reward_scale": 0.8},
    {"difficulty": "hard", "epochs": 3, "reward_scale": 0.6}
]

for stage in curriculum:
    model = trainer.train(
        data=filter_by_difficulty(data, stage["difficulty"]),
        epochs=stage["epochs"],
        reward_scale=stage["reward_scale"]
    )

Distributed Training

# Setup for multi-GPU training
import torch.distributed as dist
from torch.nn.parallel import DistributedDataParallel as DDP

def setup_distributed_training():
    """Initialize distributed training environment"""
    local_rank = int(os.environ.get("LOCAL_RANK", -1))
    world_size = int(os.environ.get("WORLD_SIZE", 1))
    
    if local_rank != -1:
        # Initialize process group
        torch.cuda.set_device(local_rank)
        dist.init_process_group(backend="nccl")
        logger.info(f"Initialized distributed training: rank {local_rank}/{world_size}")
        
        # Only log from main process
        if local_rank != 0:
            logging.getLogger().setLevel(logging.WARNING)
    
    return local_rank, world_size

# Launch distributed training
# torchrun --nproc_per_node=4 train_grpo.py

Error Handling & Robustness

# Robust data loading with fallback
def load_data_with_fallback(file_path: str, max_examples: int = None):
    """Load data with automatic fallback to sample data"""
    
    def get_fallback_data():
        """Return sample data for testing/development"""
        return [
            {
                "messages": [
                    {"role": "user", "content": "Where is my order?"},
                    {"role": "assistant", "content": "I'd be happy to help track your order. Could you please provide your order number?"}
                ]
            },
            # Add more examples...
        ]
    
    # Try loading actual data
    if not os.path.exists(file_path):
        logger.warning(f"Data file not found: {file_path}, using fallback data")
        return get_fallback_data()
    
    try:
        data = []
        with open(file_path, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    try:
                        conv = json.loads(line)
                        if validate_conversation(conv):
                            data.append(conv)
                    except json.JSONDecodeError:
                        continue
                    
                if max_examples and len(data) >= max_examples:
                    break
        
        if not data:
            logger.warning("No valid data loaded, using fallback")
            return get_fallback_data()
            
        return data
        
    except Exception as e:
        logger.error(f"Error loading data: {e}")
        return get_fallback_data()

# Model loading with automatic fallback
class ModelManager:
    def __init__(self, model_name: str, fallback_model: str = "gpt2"):
        self.model_name = model_name
        self.fallback_model = fallback_model
        
    def load_model_and_tokenizer(self):
        try:
            # Try loading requested model
            return self._load_model(self.model_name)
        except Exception as e:
            logger.error(f"Failed to load {self.model_name}: {e}")
            if self.model_name != self.fallback_model:
                logger.info(f"Falling back to {self.fallback_model}")
                return self._load_model(self.fallback_model)
            else:
                raise

Advanced GRPO Configuration

# Comprehensive GRPO configuration with all options
advanced_config = GRPOConfig(
    # Basic training parameters
    output_dir="./checkpoints/grpo",
    per_device_train_batch_size=2,
    gradient_accumulation_steps=4,
    learning_rate=1e-5,
    num_train_epochs=1,
    
    # Optimization
    max_grad_norm=0.5,
    warmup_steps=50,
    weight_decay=0.01,
    adam_epsilon=1e-8,
    
    # Mixed precision & performance
    fp16=True,
    bf16=False,  # Use FP16 instead of BF16
    gradient_checkpointing=True,
    
    # Logging & checkpointing
    logging_steps=10,
    save_steps=100,
    eval_steps=50,
    save_total_limit=3,
    load_best_model_at_end=True,
    metric_for_best_model="eval_reward",
    
    # GRPO specific
    beta=0.0,  # KL penalty (0 = no penalty)
    num_generations=4,  # Responses per prompt
    num_iterations=1,  # GRPO iterations per batch
    
    # Generation config
    max_prompt_length=128,
    max_completion_length=128,
    temperature=0.7,
    top_p=0.9,
    do_sample=True,
    
    # Advanced options
    use_liger_loss=False,  # Experimental loss function
    ddp_find_unused_parameters=False,  # For multi-GPU
    dataloader_num_workers=4,
    
    # Reproducibility
    seed=42,
    data_seed=42,
)

# Dynamic configuration based on hardware
def get_optimal_config():
    """Automatically configure based on available hardware"""
    config = GRPOConfig()
    
    # Adjust batch size based on GPU memory
    if torch.cuda.is_available():
        gpu_memory = torch.cuda.get_device_properties(0).total_memory
        if gpu_memory > 40 * 1024**3:  # 40GB+ (A100)
            config.per_device_train_batch_size = 8
            config.gradient_accumulation_steps = 1
        elif gpu_memory > 20 * 1024**3:  # 20GB+ (A6000)
            config.per_device_train_batch_size = 4
            config.gradient_accumulation_steps = 2
        else:  # Smaller GPUs
            config.per_device_train_batch_size = 1
            config.gradient_accumulation_steps = 8
    
    # Enable mixed precision on capable GPUs
    if torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 7:
        config.fp16 = True
    
    return config

Experiment Tracking & Analysis

# Comprehensive experiment tracking
class ExperimentTracker:
    def __init__(self, project_name: str, use_wandb: bool = True):
        self.project_name = project_name
        self.use_wandb = use_wandb and self._init_wandb()
        self.metrics_history = []
        
    def _init_wandb(self):
        try:
            import wandb
            wandb.login(key=os.getenv('WANDB_API_KEY'))
            wandb.init(
                project=self.project_name,
                config=CONFIG,
                name=f"grpo-{datetime.now().strftime('%Y%m%d-%H%M%S')}",
                tags=["grpo", "customer-service", CONFIG['model_name'].split('/')[-1]]
            )
            return True
        except Exception as e:
            logger.warning(f"Wandb init failed: {e}")
            return False
    
    def log_metrics(self, step: int, metrics: dict):
        """Log training metrics"""
        self.metrics_history.append({"step": step, **metrics})
        
        if self.use_wandb:
            import wandb
            wandb.log(metrics, step=step)
        
        # Also log to file for backup
        with open("metrics.jsonl", "a") as f:
            f.write(json.dumps({"step": step, **metrics}) + "\n")
    
    def log_generation_samples(self, samples: list):
        """Log example generations"""
        if self.use_wandb:
            import wandb
            table = wandb.Table(columns=["Query", "Generated", "Expected", "Reward"])
            for sample in samples:
                table.add_data(
                    sample['query'],
                    sample['generated'],
                    sample['expected'],
                    sample['reward']['total']
                )
            wandb.log({"generation_samples": table})
    
    def create_summary_report(self):
        """Generate training summary"""
        if not self.metrics_history:
            return
        
        import pandas as pd
        df = pd.DataFrame(self.metrics_history)
        
        summary = {
            "total_steps": len(df),
            "final_reward": df['reward'].iloc[-1],
            "max_reward": df['reward'].max(),
            "avg_reward": df['reward'].mean(),
            "reward_improvement": df['reward'].iloc[-1] - df['reward'].iloc[0],
        }
        
        logger.info("Training Summary:")
        for key, value in summary.items():
            logger.info(f"  {key}: {value:.4f}")
        
        return summary

Model Deployment Strategies

# Production deployment with optimization
class ProductionDeployment:
    def __init__(self, checkpoint_dir: str):
        self.checkpoint_dir = checkpoint_dir
        
    def prepare_for_deployment(self):
        """Optimize model for production inference"""
        from peft import PeftModel
        import torch.quantization as quantization
        
        # Load model
        base_model = AutoModelForCausalLM.from_pretrained(
            CONFIG['model_name'],
            torch_dtype=torch.float16,
            device_map='auto'
        )
        
        # Merge LoRA weights
        model = PeftModel.from_pretrained(base_model, self.checkpoint_dir)
        model = model.merge_and_unload()
        
        # Optional: Quantization for faster inference
        if CONFIG.get('quantize_for_deployment', False):
            model = quantization.quantize_dynamic(
                model, {torch.nn.Linear}, dtype=torch.qint8
            )
        
        # Save optimized model
        output_dir = f"{self.checkpoint_dir}_production"
        model.save_pretrained(output_dir)
        
        # Also save tokenizer
        tokenizer = AutoTokenizer.from_pretrained(self.checkpoint_dir)
        tokenizer.save_pretrained(output_dir)
        
        return output_dir
    
    def create_inference_api(self, model_path: str):
        """Create FastAPI endpoint for model inference"""
        from fastapi import FastAPI
        from pydantic import BaseModel
        
        app = FastAPI()
        
        # Load model once at startup
        model, tokenizer = load_trained_model(model_path)
        
        class Query(BaseModel):
            text: str
            max_length: int = 128
            temperature: float = 0.7
        
        @app.post("/generate")
        async def generate(query: Query):
            response = generate_response(
                model, tokenizer, query.text,
                max_length=query.max_length,
                temperature=query.temperature
            )
            return {"response": response}
        
        return app

Troubleshooting Guide

Out of Memory Errors

# Solutions for OOM errors

# 1. Reduce batch size
config.per_device_train_batch_size = 1
config.gradient_accumulation_steps = 8

# 2. Enable gradient checkpointing
config.gradient_checkpointing = True

# 3. Use LoRA with smaller rank
lora_config.r = 4  # Instead of 8 or 16

# 4. Reduce max sequence length
config.max_length = 128  # Instead of 256+

# 5. Use CPU offloading
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    device_map="auto",
    offload_folder="offload",
    offload_state_dict=True
)

Unstable Training

# Solutions for training instability

# 1. Reduce learning rate
config.learning_rate = 5e-6  # Instead of 1e-5

# 2. Increase warmup steps
config.warmup_steps = 100  # Instead of 50

# 3. Add KL penalty
config.beta = 0.01  # Instead of 0.0

# 4. Clip gradients more aggressively
config.max_grad_norm = 0.3  # Instead of 0.5

# 5. Use smaller generation count
config.num_generations = 2  # Instead of 4+

Poor Generation Quality

# Improve generation quality

# 1. Adjust generation parameters
generation_config = {
    "temperature": 0.8,  # Slightly higher for diversity
    "top_p": 0.95,       # Broader sampling
    "top_k": 50,         # Limit vocabulary
    "repetition_penalty": 1.1,  # Reduce repetition
}

# 2. Improve reward function
def enhanced_reward(response):
    # Add more nuanced scoring
    scores = {
        "relevance": check_relevance(response),
        "completeness": check_completeness(response),
        "tone": check_tone(response),
        "grammar": check_grammar(response),
    }
    return weighted_average(scores)

# 3. Filter training data
high_quality_data = [
    ex for ex in data 
    if len(ex['response'].split()) > 20  # Min length
    and len(ex['response'].split()) < 150  # Max length
    and validate_quality(ex)
]

Conclusion

Key Implementation Insights

Based on real-world GRPO training experience, here are the critical success factors:

Use LoRA for Efficiency

Training with LoRA adapters reduces memory usage by 90%+ while maintaining performance. Target the attention layers (q_proj, k_proj, v_proj, o_proj) for best results.

Always Split Train/Eval

Use stratified splitting to maintain task distribution. A 90/10 split provides enough evaluation data while maximizing training examples.

Design Multi-Component Rewards

Combine similarity, empathy, action-orientation, and length penalties. Weight them based on your specific use case (e.g., 40% similarity, 30% empathy).

Start Conservative

Begin with batch_size=2, gradient_accumulation=4, learning_rate=1e-5. These settings work well across different model sizes and GPUs.

Production Checklist

Data Quality

Validate all conversations have proper structure
Ensure minimum response length (>10 words)
Remove duplicates and low-quality examples
Classify by task type for balanced training

Model Configuration

Use FP16 mixed precision (not BF16)
Enable gradient checkpointing for large models
Set padding_side=“left” for proper generation
Configure LoRA with r=8, alpha=16 as starting point

Training Setup

Implement robust error handling with fallbacks
Use wandb or similar for experiment tracking
Save checkpoints frequently (every 100 steps)
Monitor reward variance for stability

Evaluation

Run post-training evaluation on held-out data
Track multiple metrics (similarity, quality, length)
Generate sample outputs for manual review
Compare against baseline model performance

Deployment

Merge LoRA weights for faster inference
Consider quantization for edge deployment
Implement proper error handling in API
Monitor inference latency and quality

Quick Start Template

Prerequisites

Before starting, ensure you have:

Python 3.8+
NVIDIA GPU with CUDA 11.0+ (for accelerated training)
Git installed
Optional: Weights & Biases account for experiment tracking

Install core dependencies:

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121
pip install transformers peft datasets trl scikit-learn wandb fastapi uvicorn

Step-by-Step Setup

# Clone the training template
git clone https://github.com/StateSet/grpo-agent
cd grpo-agent

# Create and activate virtual environment (recommended)
python -m venv venv
source venv/bin/activate  # On Linux/Mac
# venv\Scripts\activate  # On Windows

# Install dependencies
pip install -r requirements.txt

# Prepare your data
# Data should be in JSONL format with conversations:
# {"messages": [{"role": "user", "content": "..."}, {"role": "assistant", "content": "..."}]}
python scripts/prepare_data.py --input your_raw_data.json --output data/training_data.jsonl --max-examples 5000 --validate

# Configure training
# Edit config.yaml or set environment variables
export MODEL_NAME="Qwen/Qwen2.5-7B-Instruct"
export BATCH_SIZE=2
export LEARNING_RATE=1e-5
export MAX_EXAMPLES=5000
export EVAL_SPLIT=0.1
export WANDB_API_KEY="your_key_here"  # Optional

# Run training
# Use torchrun for multi-GPU: torchrun --nproc_per_node=2 train_grpo.py
python train_grpo.py --data data/training_data.jsonl --config config.yaml

# Monitor training
# View logs: tail -f grpo_training.log
# If using wandb: open the link provided in console

# Evaluate results
python evaluate.py --checkpoint checkpoints/grpo/final_model --num-samples 10 --output eval_results.json

# Deploy as API
python deploy.py --model checkpoints/grpo/final_model --port 8080 --host 0.0.0.0
# Test: curl -X POST https://yourapp.com:8080/generate -H "Content-Type: application/json" -d '{"text": "Where is my order?"}'

Customization Tips

Small Datasets: Set NUM_EPOCHS=3 and NUM_GENERATIONS=8
Large Models: Enable gradient_checkpointing in config.yaml
Debug Mode: Add —debug to train_grpo.py for verbose logging
Resume Training: Use —resume_from_checkpoint checkpoints/grpo/checkpoint-100

Common Pitfalls & Solutions

Data Loading Issues

Symptom: “No valid data loaded” or JSON decode errorsSolutions:

Ensure JSONL format (one JSON object per line)
Run validation: python scripts/validate_data.py data.jsonl
Use fallback data for testing: —use-fallback
Check encoding: All files should be UTF-8

GPU Memory Problems

Symptom: CUDA OOM errors during trainingSolutions:

Reduce per_device_train_batch_size to 1
Increase gradient_accumulation_steps to 8+
Use smaller model (e.g., “Qwen/Qwen2.5-3B-Instruct”)
Enable fp16 and gradient_checkpointing
Monitor with: nvidia-smi -l 1

Reward Function Problems

Symptom: Low/negative rewards or unstable trainingSolutions:

Normalize rewards to [-1, 1] range
Test independently: python test_reward.py —samples 10
Add epsilon to divisions: score = sum / (len + 1e-5)
Balance weights: Start with equal weights and adjust
Monitor reward distribution in wandb

Generation Quality Issues

Symptom: Repetitive or off-topic responsesSolutions:

Adjust temperature (0.7-0.9) and top_p (0.9-0.95)
Add repetition_penalty=1.2 in generation config
Increase num_generations to 8 for more exploration
Fine-tune prompt format: Add system instructions
Evaluate diversity: Compute unique n-grams in outputs

Deployment Errors

Symptom: Inference fails or slow performanceSolutions:

Merge LoRA weights before deployment
Use torch.compile(model) for PyTorch 2.0+
Set device_map=‘auto’ for multi-GPU inference
Implement batching for multiple requests
Profile with: torch.profiler

These pitfalls are based on real-world GRPO implementations - addressing them early will save significant time and resources.

Quick Start Template

# Clone the training template
git clone https://github.com/StateSet/grpo-agent
cd grpo-agent

# Install dependencies
pip install transformers peft datasets trl scikit-learn wandb

# Prepare your data (JSONL format with conversations)
python prepare_data.py --input your_data.json --output training_data.jsonl

# Configure environment
export MODEL_NAME="Qwen/Qwen2.5-7B-Instruct"
export BATCH_SIZE=2
export LEARNING_RATE=1e-5
export MAX_EXAMPLES=5000
export WANDB_API_KEY="your_key_here"

# Run training
python train_grpo.py --data training_data.jsonl

# Evaluate results
python evaluate.py --checkpoint ./checkpoints/grpo/final_model

# Deploy model
python deploy.py --model ./checkpoints/grpo/final_model --port 8080

Documentation

Deep dive into GRPO implementation

Examples

Ready-to-run training scripts

Support

Get help with your implementation

Next Step: Ready to implement GRPO? Check out our GRPO Agent Framework for a complete implementation guide.

Transform your AI models from pattern matchers to goal achievers with StateSet’s Reinforcement Learning platform powered by GRPO.

​Introduction

​Why Reinforcement Learning?

Goal-Oriented Learning

Exploration & Discovery

Continuous Improvement

​How GRPO Trains Your Model

​The Training Process

​Step-by-Step Breakdown

​Training Scale

​Understanding Reward Functions & Verifiers

​The Distinction

Verifier

Reward Function

​Reward Function Design

​Group Relative Policy Optimization (GRPO)

​The Innovation

​How GRPO Works

​Configuration Deep Dive

​Key Parameters

​Configuration Profiles

​Practical Implementation

​Example: Customer Service Agent

​Real-World Training Pipeline

​Monitoring Training Progress

​Best Practices

​1. Data Preparation

​2. Reward Function Design

​3. Performance Optimization

​4. Production Deployment

​Real-World Impact

​Case Study: Customer Support

Before GRPO

After GRPO

​Performance Metrics

​Getting Started

​Advanced Topics

​Multi-Objective Optimization

​Curriculum Learning

​Distributed Training

​Error Handling & Robustness

​Advanced GRPO Configuration

​Experiment Tracking & Analysis

​Model Deployment Strategies

​Troubleshooting Guide

​Conclusion

​Key Implementation Insights

Use LoRA for Efficiency

Always Split Train/Eval

Design Multi-Component Rewards

Start Conservative

​Production Checklist

​Quick Start Template

​Prerequisites

​Step-by-Step Setup

​Customization Tips

​Common Pitfalls & Solutions

​Getting Started

​Advanced Topics

​Multi-Objective Optimization

​Curriculum Learning

​Distributed Training

​Error Handling & Robustness

​Advanced GRPO Configuration

​Experiment Tracking & Analysis

​Model Deployment Strategies

​Troubleshooting Guide

​Conclusion

​Key Implementation Insights

Use LoRA for Efficiency

Always Split Train/Eval

Design Multi-Component Rewards

Start Conservative

​Production Checklist

​Quick Start Template

​Prerequisites

​Step-by-Step Setup

​Customization Tips

​Common Pitfalls & Solutions

​Quick Start Template

Documentation

Introduction

Why Reinforcement Learning?

How GRPO Trains Your Model

The Training Process

Step-by-Step Breakdown

Training Scale

Understanding Reward Functions & Verifiers

The Distinction

Reward Function Design

Group Relative Policy Optimization (GRPO)

The Innovation

How GRPO Works

Configuration Deep Dive

Key Parameters

Configuration Profiles

Practical Implementation

Example: Customer Service Agent

Real-World Training Pipeline

Monitoring Training Progress

Best Practices

1. Data Preparation

2. Reward Function Design

3. Performance Optimization

4. Production Deployment

Real-World Impact

Case Study: Customer Support

Performance Metrics

Getting Started

Advanced Topics

Multi-Objective Optimization

Curriculum Learning

Distributed Training

Error Handling & Robustness

Advanced GRPO Configuration

Experiment Tracking & Analysis

Model Deployment Strategies

Troubleshooting Guide

Conclusion

Key Implementation Insights

Production Checklist

Quick Start Template

Prerequisites

Step-by-Step Setup

Customization Tips

Common Pitfalls & Solutions

Getting Started

Advanced Topics

Multi-Objective Optimization

Curriculum Learning

Distributed Training

Error Handling & Robustness

Advanced GRPO Configuration

Experiment Tracking & Analysis

Model Deployment Strategies

Troubleshooting Guide

Conclusion

Key Implementation Insights

Production Checklist

Quick Start Template

Prerequisites

Step-by-Step Setup

Customization Tips

Common Pitfalls & Solutions

Quick Start Template