Post training a VLM for reasoning with GRPO using TRL

Authored by: Sergio Paniego

🚨 WARNING: This notebook is resource-intensive and requires substantial computational power. If you're running this in Colab, it will utilize an A100 GPU.

In this recipe, we'll demonstrate how to post-train a Vision Language Model (VLM) using GRPO for adding reasoning capabilities to a VLM using the Hugging Face ecosystem, specifically with the Transformer Reinforcement Learning library (trl).

We'll be fine-tuning Qwen2.5-VL-3B-Instruct using a subset of the lmms-lab/multimodal-open-r1-8k-verified dataset. This dataset includes images with problem descriptions along with their solution and thinking trace to reach that solution. We'll leverage this data format, along with the GRPO reward functions, to teach the model how to reason to reach the solution.

1. Install Dependencies

Let's start by installing the essential libraries we'll need for fine-tuning. We'll install trl from source, as the VLM GRPO trainer hasn't been included in an official release at the time of writing.

Authenticate using your Hugging Face 🤗 account to save and share the trained model.

2. Load Dataset 📁

We leverage lmms-lab/multimodal-open-r1-8k-verified for this recipe. This dataset contains 8k multimodal RL training examples focused on math reasoning. This data was created using GPT4o and includes image, problem, solution, original question and original answer for each sample. It was created in this project.

For our particular case where we want the model to learn to reason using images, we use image and problem as input and solution as output.

For this educational resource, we'll only use 5% of the dataset and divide it into train and test sets to make it faster to train. In a real training, we'd use the full dataset.

We'll load the dataset and divide it.

Let's check the structure of the dataset.

Let's check one sample:

In addition to the problem and image columns, we also include a custom system prompt to tell the model how we'd like the generation.

The system prompt is extracted from DeepSeek R1. Refer to this previous recipe for more details.

We convert the dataset samples into conversation samples, including the system prompt and one image and problem description per sample, since this is how the GRPO trainer expects them.

We also set padding_side="left" to ensure that generated completions during training are concatenated directly after the prompt, which is essential for GRPO to correctly compare token-level probabilities between preferred and rejected responses.

Let's take a look at a converted example:

We'll remove the the columns that we don't need for training.

We can check that the columns are now gone.

3. Post-Training the VLM Using GRPO

The diagram below highlights the main differences between PPO (Proximal Policy Optimization) and GRPO (Group Relative Policy Optimization), specifically the removal of the value model in GRPO. For more detailed information on the key differences, you can refer to this further explanation.

To implement the training pipeline, we leverage trl, Hugging Face's library for reinforcement learning, which provides a streamlined interface and built-in support for key training algorithms. In our case, we use the GRPOConfig and GRPOTrainer classes. A crucial step in this process is defining custom reward functions that guide the model's behavior and help it align with our specific objectives.

But first, let's load the model. In this case, we use Qwen/Qwen2.5-VL-3B-Instruct, a powerful VLM developed by Qwen. For better results, it would be important to consider models with a larger number of parameters.

Others examples of VLM projects that include reasoning capabilities are:

• GLM-4.1V-9B-Thinking

• VLM-R1 models

• R1-V

3.1 Loading the Baseline Model

Let's load the baseline model first. As previously introduced, Qwen/Qwen2.5-VL-3B-Instruct.

3.2 Configuring LoRA

We'll leverage LoRA for training the model, so let's configure it.

3.3 Loading Reward Functions

For the reward component of the system, we can use either pretrained reward models or reward functions defined directly in code. For training, the DeepSeek-R1 authors used an accuracy-based reward model that evaluates whether the response is correct, alongside a format-based reward that ensures the model places its reasoning process between <think> </think> tags. You can find more details here. We can simply define and implement these reward functions as generic Python functions.

In this case, we will utilize the following reward functions, directly extracted from the Open R1 implementation:

1. Format Enforcement: Ensures that the generation follows a specific format using <think> </think> <answer> </answer> tags for reasoning.

2. Solution Accuracy: Verifies whether the solution to the problem is correct, comparing it to the solution column in the dataset.

3.4 Configuring GRPO Training Parameters

Next, let's configure the training parameters for GRPO. We recommend experimenting with the max_completion_length, num_generations, and max_prompt_length parameters.

It'd be interesting to play with the max_completion_length, num_generations, and max_prompt_length params in order to find the best training combination.

The parameter selection has been adjusted to fit within the hardware limitations of a Google Colab session. To observe the full potential of reward improvements, especially in the second objective function, and to further improve the model's reasoning capabilities in a real-world scenario, a more ambitious setup would be required. This would involve larger models, an increased number of generations, and a high-quality, diverse dataset.

3.5 Training the Model 🏃

Now, let's configure the trainer and start training the model!

In this case, we pass the two reward functions we previously defined to the trainer, in addition with the model, trainings arguments and dataset.

Below, you'll find a diagram of the training procedure we'll be reproducing, which is extracted from the Open-R1 project.

Time to train the model!

We can review the training metrics directly in TensorBoard on the [model page]((https://huggingface.co/sergiopaniego/Qwen2.5-VL-3B-Instruct-Thinking/tensorboard). While the loss curve might look a bit off, the reward results tell a clearer story: the model steadily improves, increasing the amount of reward it receives over time.

Now, let's save the results in our account 💾

4. Check the Model Performance

Now that we've our model trained, we can check it's performance to evaluate it qualitatively.

We recommend restarting your session in order to free the resources used for training.

For that, we will be using the test subset of our dataset. Let's first load our trained model and it's processor.

We'll generate an auxiliary function for generating our responses. This will make it easier for us to just send a problem and image and retrieve the model response, which should include the reasoning trace and final answer.

Let's check it!

The answer seems to follow the constraints that we've added during traing using the reward functions. We can sse that the model generates something like this: <think>reasoning</think><answer>solution</answer>. Let's check the actual solution, to understand if the model is correct.

It seems like the model has already including some reasoning capabilities to their functionality! Let's also check the inference time and generated tokens, for further check on the model capabilities.

5. Continuing Your Learning Journey 🧑‍🎓

The learning journey does not stop here!

If you're eager on discovering more about GRPO, reasoning or VLMs, we can recommend some materials:

• Post training an LLM for reasoning with GRPO in TRL recipe and the linked resources that it includes
• GLM-4.1V-9B-Thinking paper
• TRL GRPO for VLMs example