Skip to main content
This experiment demonstrates how to use FlexAI to fine-tune language models using reinforcement learning (RL) techniques with EasyR1, a framework for training reasoning-capable models using GRPO (Group Relative Policy Optimization), DAPO, and REINFORCE algorithms. For illustration purposes, we’ll fine-tune the 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:
Replace <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
The framework is built on top of VERL (Versatile Efficient Reinforcement Learning), providing distributed training capabilities with FSDP and vLLM integration.

Directory Structure

The code/easyR1/ directory contains:
  • config.yaml - Main GRPO training configuration
  • format_prompt/ - Jinja templates for prompt formatting
  • reward_function/ - Custom reward scoring functions
For baseline training scripts and additional examples, refer to the EasyR1 GitHub repository.
1

Understand the Configuration

EasyR1 uses a comprehensive YAML configuration file that controls all aspects of RL training. The main configuration file is located at code/easyR1/config.yaml in this repository.

Key Configuration Sections

Data Configuration

Algorithm Settings

Worker Configuration

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
You can adapt these examples to work with FlexAI by following the training commands in this blueprint.
3

Customize Your Configuration

For your specific use case, you may want to create a custom configuration. Here’s how to customize the config.yaml:

Custom Dataset

Replace the dataset configuration:

Custom Reward Function

Create your own reward function in code/easyR1/reward_function/custom.py:
Then update the config to reference your custom reward function:

Custom Prompt Format

Create a custom Jinja template in code/easyR1/format_prompt/custom.jinja:
Update the config:

Create Secrets

To access HuggingFace models and datasets, you need a HuggingFace token. Use the flexai secret create command to store your HuggingFace Token as a secret:
Then paste your HuggingFace Token API key value. Use the same command to store your Weights & Biases (wandb) API key as a secret:
Then paste your Weights & Biases API key value.

[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:
  1. Create a HuggingFace storage provider:
  2. 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:
View the logs 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:
Look for checkpoints marked as 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:
Monitor your inference endpoint status:

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):
Issues: Correct answer but no reasoning steps shown RL Fine-tuned Model Response (after GRPO training on math12k):
Improvements: Clear reasoning steps, structured approach, educational value This demonstrates how RL training encourages the model to show its reasoning process, making it more reliable and transparent.

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
For mathematical reasoning tasks:
  • 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)
Distributed Training:
  • 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)
Optimization:
  • 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
Command Line Parameters Explained:
  • 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 repository
  • config=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 iteration
  • max_prompt_length: 2048: Maximum input length
  • max_response_length: 2048: Maximum output length
Algorithm Settings:
  • adv_estimator: grpo: Choice of RL algorithm
  • kl_coef: 1.0e-2: Strength of KL penalty
  • use_kl_loss: true: Enable KL divergence loss
Training Settings:
  • total_epochs: 15: Number of training epochs
  • n_gpus_per_node: 8: GPUs per node
  • val_freq: 5: Validation every 5 epochs
  • save_freq: 5: Save checkpoint every 5 epochs

Scaling Options

  • For faster training: Increase to 2 nodes (16 × H100)
  • For larger models: Increase tensor_parallel_size for 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:
Out of Memory Errors:
  • Reduce rollout_batch_size from 512 to 256
  • Reduce rollout.n from 5 to 3 (fewer samples per prompt)
  • Enable CPU offloading: enable_cpu_offload: true in FSDP config
  • Reduce tensor_parallel_size for rollout workers
Reward Function Errors:
  • 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
Checkpoint Not Inference Ready:
  • Wait for training to complete fully
  • Check save_model_only: false in config to include all necessary files
  • Verify training completed without errors
Endpoint Deployment Issues:
  • Verify checkpoint shows INFERENCE READY = true status
  • Check FlexAI cluster availability
  • Review detailed logs with flexai inference logs <endpoint-name>
Dataset Loading Issues:
  • Verify dataset path format: username/dataset@split
  • Ensure HuggingFace token has access to datasets
  • Check prompt_key and answer_key match your dataset schema
vLLM Rollout Errors:
  • Adjust gpu_memory_utilization (default 0.6)
  • Reduce tensor_parallel_size if GPUs are insufficient
  • Enable enforce_eager: true for debugging

References

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