Fine-tuning SmolVLM using direct preference optimization (DPO) with TRL on a consumer GPU

Authored by: Sergio Paniego

In this recipe, we’ll guide you through fine-tuning a smol 🤏 Vision Language Model (VLM) with Direct Preference Optimization (DPO) using the Transformer Reinforcement Learning (TRL) library to demonstrate how you can tailor VLMs to suit your specific needs, even when working with consumer-grade GPUs.

We’ll fine-tune SmolVLM using a preference dataset to help the model align with desired outputs. SmolVLM is a highly performant and memory-efficient model, making it an ideal choice for this task. If you’re new to Preference Optimization for language or vision-language models, check out this blog for an in-depth introduction.

The dataset we’ll use is HuggingFaceH4/rlaif-v_formatted, which contains pairs of prompt + image along with a chosen and rejected answer for each pair. The goal of this fine-tuning process is to make the model consistently prefer the chosen answers from the dataset, reducing hallucinations.

This notebook has been tested using an NVIDIA L4 GPU.

1. Install Dependencies

Let’s start by installing the essential libraries we’ll need for fine-tuning! 🚀

Authenticate with your Hugging Face account to save and share your model directly from this notebook 🗝️.

2. Load Dataset 📁

We’ll work with the HuggingFaceH4/rlaif-v_formatted dataset, which provides pairs of prompt + image along with a chosen and rejected answers for each pair. This structured format is ideal for training models with Direct Preference Optimization (DPO).

The dataset is already preformatted for this task. If you’re working with a custom dataset, you’ll need to preprocess it into the same format.

In this example, we'll use a subset of the dataset to demonstrate the process. However, in a real-world scenario, you should utilize the full dataset for better performance.

We will ensure all the images are RGB formatted:

Let’s explore an example from the dataset to better understand its structure and the type of data we’re working with.

3. Fine-Tune the Model using TRL

3.1 Load the Quantized Model for Training ⚙️

Let's first load a quantized version of the SmolVLM-Instruct model using bitsandbytes, and let's also load the processor. We'll use SmolVLM-Instruct.

3.2 Set Up QLoRA and DPOConfig 🚀

In this step, we’ll configure QLoRA for our training setup. QLoRA is a powerful fine-tuning technique designed to reduce the memory footprint, making it possible to fine-tune large models efficiently, even on limited hardware.

QLoRA builds upon traditional LoRA (Low-Rank Adaptation) by introducing quantization for the adapter weights. This enhancement leads to significantly lower memory usage and faster training, making it an ideal choice for resource-constrained environments.

Next, we will configure the training options using DPOConfig.

We will define the training arguments for Direct Preference Optimization (DPO) with the DPOTrainer class from the TRL library.

DPO uses labeled preference data to guide the model toward generating responses that align with preferences. TRL's DPOTrainer will tokenize the dataset before training and save it to disk. This process can consume significant disk space, depending on the amount of data used for training. Plan accordingly to avoid running out of storage.

This step may take a while, so feel free to relax and enjoy the process! 😄

Time to train the model! 🎉

Let's save the results 💾

4. Testing the Fine-Tuned Model 🔍

With our Vision Language Model (VLM) fine-tuned, it’s time to evaluate its performance! In this section, we’ll test the model using examples from the HuggingFaceH4/rlaif-v_formatted dataset. Let’s dive into the results and assess how well the model aligns with the preferred responses! 🚀

Before we begin, let’s clean up the GPU memory to ensure smooth and optimal performance. 🧹

We will reload the base model using the same pipeline as before.

We will attach the trained adapter to the pretrained model. This adapter contains the fine-tuning adjustments made during training, enabling the base model to leverage the new knowledge while keeping its core parameters intact. By integrating the adapter, we enhance the model's capabilities without altering its original structure.

Let's evaluate the model on an unseen sample.

Let’s create a common function that we can call with different samples to streamline the testing process. This function will allow us to evaluate the model’s performance on multiple examples efficiently without needing to rewrite code for each one. By using this reusable function, we can quickly assess how well the model performs across a variety of inputs.

Now, we’re ready to call the function and evaluate the model! 🚀

The model is now able to generate responses based on the provided image and prompt. For tasks like this, it’s useful to compare your model's performance against a benchmark to see how much it has improved and how it stacks up against other options. For more information and details on this comparison, check out this post.

💻 I’ve developed an example application to test the model, which you can find here.

Since here we only run an example training with a subset of the dataset, for the Space I've used the official Hugging Face DPO fine tuned model. You can easily compare it with another Space featuring the pre-trained model, available here.

5. Continuing the Learning Journey 🧑‍🎓️

Expand your knowledge of Vision Language Models and related tools with these resources:

• Multimodal Recipes in the Cookbook: Discover practical recipes for multimodal models, including Retrieval-Augmented Generation (RAG) pipelines and fine-tuning. We’ve already published a recipe for fine-tuning a smol VLM with TRL using SFT, which complements this guide perfectly—check it out for additional details.

• TRL Community Tutorials: Explore a rich collection of tutorials that dive into the intricacies of TRL and its real-world applications.

You can also revisit the Continuing the Learning Journey section in Fine-Tuning a Vision Language Model (Qwen2-VL-7B) with the Hugging Face Ecosystem (TRL).

These resources will help deepen your knowledge and expertise in multimodal learning.