Qwen2.5-7B-Instruct model on mathematical reasoning tasks using the math12k dataset with GRPO algorithm to improve reasoning capabilities.
If you haven’t already connected FlexAI to GitHub, run
flexai code-registry connect to set up a code registry connection. This allows FlexAI to pull repositories directly using the repository URL in training commands.Quick Start
Run GRPO training on Qwen2.5-7B with this single command:<WANDB_API_KEY_SECRET_NAME> and <HF_AUTH_TOKEN_SECRET_NAME> with your actual values.
What is EasyR1?
EasyR1 is a reinforcement learning framework specifically designed for training language models with enhanced reasoning capabilities. It implements several RL algorithms optimized for LLMs:- GRPO (Group Relative Policy Optimization): Efficient policy optimization using group-based advantage estimation
- DAPO (Data-Augmented Policy Optimization): Enhanced training with data augmentation strategies
- REINFORCE: Classic policy gradient method for LLM fine-tuning
Directory Structure
Thecode/easyR1/ directory contains:
config.yaml- Main GRPO training configurationformat_prompt/- Jinja templates for prompt formattingreward_function/- Custom reward scoring functions
1
2
Reference Baseline Examples
For pre-configured training scripts and baseline examples, refer to the EasyR1 repository. The repository provides multiple baseline configurations for different models and tasks:
Available Baselines (in EasyR1 repo)
- Mathematical Reasoning:
qwen2_5_7b_math_grpo.sh,qwen3_4b_math_grpo.sh - Geometric Reasoning (Vision-Language):
qwen2_5_vl_7b_geo3k_grpo.sh,qwen2_5_vl_7b_geo3k_dapo.sh,qwen2_5_vl_7b_geo3k_reinforce.sh - Multi-Image Tasks:
qwen2_5_vl_7b_multi_image.sh
3
Customize Your Configuration
For your specific use case, you may want to create a custom configuration. Here’s how to customize the Then update the config to reference your custom reward function:Update the config:
config.yaml:Custom Dataset
Replace the dataset configuration:Custom Reward Function
Create your own reward function incode/easyR1/reward_function/custom.py:Custom Prompt Format
Create a custom Jinja template incode/easyR1/format_prompt/custom.jinja:Create Secrets
To access HuggingFace models and datasets, you need a HuggingFace token. Use theflexai secret create command to store your HuggingFace Token as a secret:
[Optional] Pre-fetch the Model
To speed up training and avoid downloading large models at runtime, you can pre-fetch your HuggingFace model to FlexAI storage:-
Create a HuggingFace storage provider:
-
Push the model checkpoint to your storage:
Training
For RL training with EasyR1, we recommend using 1 node (8 × H100 GPUs) for 7B models to handle the actor, reference model, and rollout workers efficiently.The commands below use this repository which contains all necessary configuration files in the
code/easyR1/ directory.Standard Training: Mathematical Reasoning with GRPO
Training with Model Prefetch
Training with Custom Configuration
To use a modified configuration or different dataset, override config values:Monitoring Training Progress
You can check the status and lifecycle events of your Training Job:Training Observability with Weights & Biases
EasyR1 supports Weights & Biases (wandb) integration for detailed training metrics visualization. The configuration already includes wandb logging:Getting Training Checkpoints
Once the Training Job completes successfully, you can list all produced checkpoints:INFERENCE READY = true - these are ready for serving.
Serving the Trained Model
Deploy your RL-trained model directly from the checkpoint using FlexAI inference. Replace<CHECKPOINT_ID> with the ID from an inference-ready checkpoint:
Testing Your RL-Trained Model
Once the endpoint is running, you can test it with reasoning tasks. For our mathematical reasoning example, the model should demonstrate improved step-by-step reasoning and accurate problem-solving.Before and After Training Comparison
To illustrate the improvement from RL fine-tuning, here’s a comparison using a math problem: Problem: “If a train travels 120 miles in 2 hours, what is its average speed in miles per hour?” Base Model Response (Qwen2.5-7B-Instruct before RL training):Example API Call
Expected Results
After RL fine-tuning with EasyR1, your model should achieve:- Enhanced Reasoning: Step-by-step problem-solving with clear explanations
- Improved Accuracy: Higher success rate on reasoning tasks
- Better Generalization: Ability to apply learned reasoning patterns to new problems
- Structured Outputs: More organized and educational responses
- Explicit Step-by-Step Solutions: Clear breakdown of problem-solving process
- Higher Success Rate: Improved accuracy on math benchmarks
- Better Error Detection: Ability to identify and correct mistakes
Technical Details
Training Configuration Breakdown
Reinforcement Learning Components:- Actor Model: The model being trained (policy network)
- Reference Model: Frozen copy for KL divergence computation
- Rollout Workers: Generate multiple responses for each prompt (n=5)
- Reward Function: Evaluates response quality (custom per task)
- FSDP (Fully Sharded Data Parallel): Efficient memory usage for large models
- vLLM Integration: Fast inference during rollout generation
- Tensor Parallelism: For rollout workers (size=2)
- GRPO Algorithm: Group-based advantage estimation for stable training
- KL Penalty: Prevents model from deviating too far from base model
- Gradient Checkpointing: Reduces memory usage during backpropagation
Resource Requirements
Recommended Configuration for Qwen2.5-7B:- Nodes: 1 node (sufficient for RL training with actor + reference + rollout)
- Accelerators: 8 × H100 GPUs per node
- Memory: ~400GB+ GPU memory total (actor, reference, and rollout workers)
- Training Time: ~8-12 hours for 15 epochs
- Storage: ~50GB for checkpoints
FORCE_TORCHRUN=1: Ensures proper distributed training setup--runtime pytorch-28-vllm-0110-nvidia: PyTorch 2.8 with vLLM 0.11.0 optimized for EasyR1--repository-url: Points to the FlexAI blueprints repositoryconfig=code/easyR1/config.yaml: Main configuration file path relative to repository root
Key Configuration Parameters
Data Settings:rollout_batch_size: 512: Number of prompts per training iterationmax_prompt_length: 2048: Maximum input lengthmax_response_length: 2048: Maximum output length
adv_estimator: grpo: Choice of RL algorithmkl_coef: 1.0e-2: Strength of KL penaltyuse_kl_loss: true: Enable KL divergence loss
total_epochs: 15: Number of training epochsn_gpus_per_node: 8: GPUs per nodeval_freq: 5: Validation every 5 epochssave_freq: 5: Save checkpoint every 5 epochs
Scaling Options
- For faster training: Increase to 2 nodes (16 × H100)
- For larger models: Increase
tensor_parallel_sizefor rollout - For better exploration: Increase
rollout.n(more samples per prompt) - For memory efficiency: Enable CPU offloading (
enable_cpu_offload: true) - For different tasks: Modify reward function and prompt templates
Advanced Examples
Vision-Language Model with Geometric Reasoning
Using DAPO Algorithm
Troubleshooting
Training Job Fails to Start:- Reduce
rollout_batch_sizefrom 512 to 256 - Reduce
rollout.nfrom 5 to 3 (fewer samples per prompt) - Enable CPU offloading:
enable_cpu_offload: truein FSDP config - Reduce
tensor_parallel_sizefor rollout workers
- Verify reward function path is correct in config
- Test reward function locally before training
- Ensure reward function returns float scores for all inputs
- Check for NaN or infinite reward values
- Wait for training to complete fully
- Check
save_model_only: falsein config to include all necessary files - Verify training completed without errors
- Verify checkpoint shows
INFERENCE READY = truestatus - Check FlexAI cluster availability
- Review detailed logs with
flexai inference logs <endpoint-name>
- Verify dataset path format:
username/dataset@split - Ensure HuggingFace token has access to datasets
- Check prompt_key and answer_key match your dataset schema
- Adjust
gpu_memory_utilization(default 0.6) - Reduce
tensor_parallel_sizeif GPUs are insufficient - Enable
enforce_eager: truefor debugging
References
- EasyR1 GitHub: https://github.com/hiyouga/EasyR1
- VERL Framework: https://github.com/volcengine/verl
- FlexAI Documentation: https://docs.flex.ai
- HybridFlow Paper: https://arxiv.org/abs/2409.19256
- GRPO Algorithm: Introduced in DeepSeekMath paper - https://arxiv.org/abs/2402.03300
- GRPO Documentation: https://huggingface.co/docs/trl/grpo_trainer
Code
requirements.txt
config.yaml
format_prompt/math.jinja
reward_function/math.py
🚀 Run this on FlexAI
Managed checkpoints mean you never lose a run to preemption. Jobs launch in under 60 seconds — no infra setup, built-in observability.
Get started →Talk to us