Open-Source AI Cookbook

26 / 120·notebook

Chapter 26

Fine-tuning SmolVLM with TRL on a consumer GPU

NotebookPython 367 cells

Fine-tuning SmolVLM with TRL on a consumer GPU

In this recipe, we’ll demonstrate how to fine-tune a smol 🤏 Vision Language Model (VLM) using the Hugging Face ecosystem, leveraging the powerful [Transformer Reinforcement Learning library (TRL)](https://huggingface.co/docs/trl/index). This step-by-step guide will enable you to customize VLMs for your specific tasks, even on consumer GPUs.

🌟 Model & Dataset Overview

In this notebook, we will fine-tune the SmolVLM model using the ChartQA dataset. SmolVLM is a highly performant and memory-efficient model, making it an ideal choice for this task. The ChartQA dataset contains images of various chart types paired with question-answer pairs, offering a valuable resource for enhancing the model's visual question-answering (VQA) capabilities. These skills are crucial for a range of practical applications, including data analysis, business intelligence, and educational tools.

💡 Note: The instruct model we are fine-tuning has already been trained on this dataset, so it is familiar with the data. However, this serves as a valuable educational exercise for understanding fine-tuning techniques. For a complete list of datasets used to train this model, check out this document.

📖 Additional Resources

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

Multimodal Recipes in Cookbook: Explore practical recipes for multimodal models, including RAG pipelines and fine-tuning. We already have a recipe for fine-tuning a VLM with TRL, so refer to it for more details.
TRL Community Tutorials: A treasure trove of tutorials to deepen your understanding of TRL and its applications.

With these resources, you’ll be equipped to dive deeper into the world of VLMs and push the boundaries of what they can achieve!

This notebook is tested using a L4 GPU.

Smol VLMs comparison

1. Install Dependencies

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

In [1]python · cell 5

python

!pip install -U -q git+https://github.com/huggingface/trl.git bitsandbytes peft qwen-vl-utils trackio
# Tested with trl==0.22.0.dev0, bitsandbytes==0.47.0, peft==0.17.1, qwen-vl-utils==0.0.11, trackio==0.2.8

Output

  Installing build dependencies ... [?25l[?25hdone
  Getting requirements to build wheel ... [?25l[?25hdone
  Preparing metadata (pyproject.toml) ... [?25l[?25hdone
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m61.3/61.3 MB[0m [31m40.0 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m504.9/504.9 kB[0m [31m38.0 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m844.5/844.5 kB[0m [31m56.4 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m59.6/59.6 MB[0m [31m40.9 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m324.6/324.6 kB[0m [31m28.8 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m40.0/40.0 MB[0m [31m60.5 MB/s[0m eta [36m0:00:00[0m
[?25h  Building wheel for trl (pyproject.toml) ... [?25l[?25hdone

In [2]python · cell 6

python

!pip install -q flash-attn --no-build-isolation

Output

[?25l     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m0.0/8.4 MB[0m [31m?[0m eta [36m-:--:--[0m
[2K     [91m━━━━━━[0m[91m╸[0m[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m1.4/8.4 MB[0m [31m43.6 MB/s[0m eta [36m0:00:01[0m
[2K     [91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m[91m╸[0m [32m8.4/8.4 MB[0m [31m133.9 MB/s[0m eta [36m0:00:01[0m
[2K     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m8.4/8.4 MB[0m [31m92.3 MB/s[0m eta [36m0:00:00[0m
[?25h  Preparing metadata (setup.py) ... [?25l[?25hdone
  Building wheel for flash-attn (setup.py) ... [?25l[?25hdone

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

In [3]python · cell 8

python

from huggingface_hub import notebook_login

notebook_login()

Output

VBox(children=(HTML(value='<center> <img\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…

2. Load Dataset 📁

We’ll load the HuggingFaceM4/ChartQA dataset, which provides chart images along with corresponding questions and answers—perfect for fine-tuning visual question-answering models.

We’ll create a system message to make the VLM act as a chart analysis expert, giving concise answers about chart images.

In [4]python · cell 11

python

system_message = """You are a Vision Language Model specialized in interpreting visual data from chart images.
Your task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.
The charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.
Focus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary."""

We’ll format the dataset into a chatbot structure, with the system message, image, user query, and answer for each interaction.

💡For more tips on using this model, check out the Model Card.

In [5]python · cell 13

python

def format_data(sample):
    return {
      "images": [sample["image"]],
      "messages": [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": system_message
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": sample["image"],
                },
                {
                    "type": "text",
                    "text": sample['query'],
                }
            ],
        },
        {
            "role": "assistant",
            "content": [
                {
                    "type": "text",
                    "text": sample["label"][0]
                }
            ],
        },
        ]
    }

For educational purposes, we’ll load only 10% of each split in the dataset. In a real-world scenario, you would load the entire dataset.

In [6]python · cell 15

python

from datasets import load_dataset

dataset_id = "HuggingFaceM4/ChartQA"
train_dataset, eval_dataset, test_dataset = load_dataset(dataset_id, split=['train[:10%]', 'val[:10%]', 'test[:10%]'])

Output

/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:86: UserWarning: 
Access to the secret `HF_TOKEN` has not been granted on this notebook.
You will not be requested again.
Please restart the session if you want to be prompted again.
  warnings.warn(

README.md:   0%|          | 0.00/852 [00:00<?, ?B/s]

data/train-00000-of-00003-49492f364babfa(…):   0%|          | 0.00/219M [00:00<?, ?B/s]

data/train-00001-of-00003-7302bae5e425bb(…):   0%|          | 0.00/311M [00:00<?, ?B/s]

data/train-00002-of-00003-194c9400785577(…):   0%|          | 0.00/315M [00:00<?, ?B/s]

data/val-00000-of-00001-0f11003c77497969(…):   0%|          | 0.00/50.2M [00:00<?, ?B/s]

data/test-00000-of-00001-e2cd0b7a0f9eb20(…):   0%|          | 0.00/68.9M [00:00<?, ?B/s]

Generating train split:   0%|          | 0/28299 [00:00<?, ? examples/s]

Generating val split:   0%|          | 0/1920 [00:00<?, ? examples/s]

Generating test split:   0%|          | 0/2500 [00:00<?, ? examples/s]

Let’s take a look at the dataset structure. It includes an image, a query, a label (the answer), and a fourth feature that we’ll be discarding.

In [7]python · cell 17

python

train_dataset

Output

Dataset({
    features: ['image', 'query', 'label', 'human_or_machine'],
    num_rows: 2830
})

Now, let’s format the data using the chatbot structure. This will set up the interactions for the model.

In [8]python · cell 19

python

train_dataset = [format_data(sample) for sample in train_dataset]
eval_dataset = [format_data(sample) for sample in eval_dataset]
test_dataset = [format_data(sample) for sample in test_dataset]

In [9]python · cell 20

python

train_dataset[200]

Output

{'images': [<PIL.PngImagePlugin.PngImageFile image mode=RGB size=308x369>],
 'messages': [{'role': 'system',
   'content': [{'type': 'text',
     'text': 'You are a Vision Language Model specialized in interpreting visual data from chart images.\nYour task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.\nThe charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.\nFocus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary.'}]},
  {'role': 'user',
   'content': [{'type': 'image',
     'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=308x369>},
    {'type': 'text',
     'text': 'Is the rightmost value of light brown graph 58?'}]},
  {'role': 'assistant', 'content': [{'type': 'text', 'text': 'No'}]}]}

3. Load Model and Check Performance! 🤔

Now that we’ve loaded the dataset, it’s time to load the HuggingFaceTB/SmolVLM-Instruct, a 2B parameter Vision Language Model (VLM) that offers state-of-the-art (SOTA) performance while being efficient in terms of memory usage.

For a broader comparison of state-of-the-art VLMs, explore the WildVision Arena and the OpenVLM Leaderboard, where you can find the best-performing models across various benchmarks.

In [10]python · cell 23

python

import torch
from transformers import Idefics3ForConditionalGeneration, AutoProcessor

model_id = "HuggingFaceTB/SmolVLM-Instruct"

Next, we’ll load the model and the tokenizer to prepare for inference.

In [11]python · cell 25

python

model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2",
)

processor = AutoProcessor.from_pretrained(model_id)

Output

config.json: 0.00B [00:00, ?B/s]

model.safetensors:   0%|          | 0.00/4.49G [00:00<?, ?B/s]

generation_config.json:   0%|          | 0.00/136 [00:00<?, ?B/s]

processor_config.json:   0%|          | 0.00/68.0 [00:00<?, ?B/s]

chat_template.json:   0%|          | 0.00/429 [00:00<?, ?B/s]

preprocessor_config.json:   0%|          | 0.00/486 [00:00<?, ?B/s]

tokenizer_config.json: 0.00B [00:00, ?B/s]

vocab.json: 0.00B [00:00, ?B/s]

merges.txt: 0.00B [00:00, ?B/s]

tokenizer.json: 0.00B [00:00, ?B/s]

added_tokens.json:   0%|          | 0.00/92.0 [00:00<?, ?B/s]

special_tokens_map.json: 0.00B [00:00, ?B/s]

To evaluate the model's performance, we’ll use a sample from the dataset. First, let’s inspect the internal structure of this sample to understand how the data is organized.

In [12]python · cell 27

python

train_dataset[1]

Output

{'images': [<PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>],
 'messages': [{'role': 'system',
   'content': [{'type': 'text',
     'text': 'You are a Vision Language Model specialized in interpreting visual data from chart images.\nYour task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.\nThe charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.\nFocus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary.'}]},
  {'role': 'user',
   'content': [{'type': 'image',
     'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>},
    {'type': 'text',
     'text': 'How many values are below 40 in Unfavorable graph?'}]},
  {'role': 'assistant', 'content': [{'type': 'text', 'text': '6'}]}]}

We’ll use the sample without the system message to assess the VLM's raw understanding. Here’s the input we will use:

In [13]python · cell 29

python

train_dataset[0]['messages'][1:2]

Output

[{'role': 'user',
  'content': [{'type': 'image',
    'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>},
   {'type': 'text', 'text': 'Is the value of Favorable 38 in 2015?'}]}]

Now, let’s take a look at the chart corresponding to the sample. Can you answer the query based on the visual information?

In [14]python · cell 31

python

train_dataset[0]['images'][0]

Output

<PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>

Let’s create a method that takes the model, processor, and sample as inputs to generate the model's answer. This will allow us to streamline the inference process and easily evaluate the VLM's performance.

In [15]python · cell 33

python

def generate_text_from_sample(model, processor, sample, max_new_tokens=1024, device="cuda"):
    # Prepare the text input by applying the chat template
    text_input = processor.apply_chat_template(
        sample['messages'][1:2],  # Use the sample without the system message
        add_generation_prompt=True
    )

    image_inputs = []
    image = sample['images'][0]
    #image = sample[1]['content'][0]['image']
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image_inputs.append([image])

    # Prepare the inputs for the model
    model_inputs = processor(
        #text=[text_input],
        text=text_input,
        images=image_inputs,
        return_tensors="pt",
    ).to(device)  # Move inputs to the specified device

    # Generate text with the model
    generated_ids = model.generate(**model_inputs, max_new_tokens=max_new_tokens)

    # Trim the generated ids to remove the input ids
    trimmed_generated_ids = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    # Decode the output text
    output_text = processor.batch_decode(
        trimmed_generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )

    return output_text[0]  # Return the first decoded output text

In [16]python · cell 34

python

output = generate_text_from_sample(model, processor, train_dataset[1])
output

Output

' 3.'

It seems like the model is referencing the wrong line, causing it to fail. To improve its performance, we can fine-tune the model with more relevant data to ensure it better understands the context and provides more accurate responses.

Remove Model and Clean GPU

Before we proceed with training the model in the next section, let's clear the current variables and clean the GPU to free up resources.

In [17]python · cell 37

python

import gc
import time

def clear_memory():
    # Delete variables if they exist in the current global scope
    if 'inputs' in globals(): del globals()['inputs']
    if 'model' in globals(): del globals()['model']
    if 'processor' in globals(): del globals()['processor']
    if 'trainer' in globals(): del globals()['trainer']
    if 'peft_model' in globals(): del globals()['peft_model']
    if 'bnb_config' in globals(): del globals()['bnb_config']
    time.sleep(2)

    # Garbage collection and clearing CUDA memory
    gc.collect()
    time.sleep(2)
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    time.sleep(2)
    gc.collect()
    time.sleep(2)

    print(f"GPU allocated memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
    print(f"GPU reserved memory: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

clear_memory()

Output

GPU allocated memory: 0.01 GB
GPU reserved memory: 0.02 GB

4. Fine-Tune the Model using TRL

4.1 Load the Quantized Model for Training ⚙️

Next, we’ll load the quantized model using bitsandbytes. If you want to learn more about quantization, check out this blog post or this one.

In [18]python · cell 40

python

from transformers import BitsAndBytesConfig

# BitsAndBytesConfig int-4 config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load model and tokenizer
model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config,
    _attn_implementation="flash_attention_2",
)
processor = AutoProcessor.from_pretrained(model_id)

4.2 Set Up QLoRA and SFTConfig 🚀

Next, we’ll configure QLoRA for our training setup. QLoRA allows efficient fine-tuning of large models by reducing the memory footprint. Unlike traditional LoRA, which uses low-rank approximation, QLoRA further quantizes the LoRA adapter weights, leading to even lower memory usage and faster training.

To boost efficiency, we can also leverage a paged optimizer or 8-bit optimizer during QLoRA implementation. This approach enhances memory efficiency and speeds up computations, making it ideal for optimizing our model without sacrificing performance.

In [19]python · cell 42

python

from peft import LoraConfig, get_peft_model

# Configure LoRA
peft_config = LoraConfig(
    r=8,
    lora_alpha=8,
    lora_dropout=0.1,
    target_modules=['down_proj','o_proj','k_proj','q_proj','gate_proj','up_proj','v_proj'],
    use_dora=True,
    init_lora_weights="gaussian"
)

# Apply PEFT model adaptation
peft_model = get_peft_model(model, peft_config)

# Print trainable parameters
peft_model.print_trainable_parameters()

Output

trainable params: 11,269,248 || all params: 2,257,542,128 || trainable%: 0.4992

We will use Supervised Fine-Tuning (SFT) to improve our model's performance on the specific task. To achieve this, we'll define the training arguments with the SFTConfig class from the TRL library. SFT leverages labeled data to help the model generate more accurate responses, adapting it to the task. This approach enhances the model's ability to understand and respond to visual queries more effectively.

In [27]python · cell 44

python

from trl import SFTConfig

# Configure training arguments using SFTConfig
training_args = SFTConfig(
    output_dir="smolvlm-instruct-trl-sft-ChartQA",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    warmup_steps=50,
    learning_rate=1e-4,
    weight_decay=0.01,
    logging_steps=25,
    save_strategy="steps",
    save_steps=25,
    save_total_limit=1,
    optim="adamw_torch_fused",
    bf16=True,
    push_to_hub=True,
    report_to="none",
    max_length=None
)

In [21]python · cell 45

python

import trackio

trackio.init(
    project="smolvlm-instruct-trl-sft-ChartQA",
    name="smolvlm-instruct-trl-sft-ChartQA",
    config=training_args,
    space_id=training_args.output_dir + "-trackio"
)

Output

* Trackio project initialized: smolvlm-instruct-trl-sft-ChartQA
* Trackio metrics will be synced to Hugging Face Dataset: sergiopaniego/smolvlm-instruct-trl-sft-ChartQA-trackio-dataset
* Creating new space: https://huggingface.co/spaces/sergiopaniego/smolvlm-instruct-trl-sft-ChartQA-trackio
* View dashboard by going to: https://huggingface.co/spaces/sergiopaniego/smolvlm-instruct-trl-sft-ChartQA-trackio

<trackio.run.Run at 0x7fcb74bcde50>

4.3 Training the Model 🏃

We will define the SFTTrainer, which is a wrapper around the transformers.Trainer class and inherits its attributes and methods. This class simplifies the fine-tuning process by properly initializing the PeftModel when a PeftConfig object is provided. By using SFTTrainer, we can efficiently manage the training workflow and ensure a smooth fine-tuning experience for our Vision Language Model.

In [28]python · cell 48

python

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    processing_class=processor,
)

Output

/usr/local/lib/python3.12/dist-packages/peft/mapping_func.py:73: UserWarning: You are trying to modify a model with PEFT for a second time. If you want to reload the model with a different config, make sure to call `.unload()` before.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/peft/tuners/tuners_utils.py:196: UserWarning: Already found a `peft_config` attribute in the model. This will lead to having multiple adapters in the model. Make sure to know what you are doing!
  warnings.warn(

Time to Train the Model! 🎉

In [ ]python · cell 50

python

trainer.train()

Output

Casting fp32 inputs back to torch.bfloat16 for flash-attn compatibility.
`loss_type=None` was set in the config but it is unrecognised.Using the default loss: `ForCausalLMLoss`.

<IPython.core.display.HTML object>

[ 3/177 01:23 < 4:01:02, 0.01 it/s, Epoch 0.01/1]

Step	Training Loss

Let's save the results 💾

In [ ]python · cell 52

python

trainer.save_model(training_args.output_dir)

5. Testing the Fine-Tuned Model 🔍

Now that our Vision Language Model (VLM) is fine-tuned, it's time to evaluate its performance! In this section, we'll test the model using examples from the ChartQA dataset to assess how accurately it answers questions based on chart images. Let's dive into the results and see how well it performs! 🚀

Let's clean up the GPU memory to ensure optimal performance 🧹

In [ ]python · cell 55

python

clear_memory()

Output

GPU allocated memory: 16.34 GB
GPU reserved memory: 18.69 GB

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

In [ ]python · cell 57

python

model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2",
)

processor = AutoProcessor.from_pretrained(model_id)

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.

In [ ]python · cell 59

python

adapter_path = "sergiopaniego/smolvlm-instruct-trl-sft-ChartQA"
model.load_adapter(adapter_path)

Let's evaluate the model on an unseen sample.

In [ ]python · cell 61

python

test_dataset[20]['messages'][:2]

Output

[{'role': 'system',
  'content': [{'type': 'text',
    'text': 'You are a Vision Language Model specialized in interpreting visual data from chart images.\nYour task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.\nThe charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.\nFocus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary.'}]},
 {'role': 'user',
  'content': [{'type': 'image',
    'image': <PIL.PngImagePlugin.PngImageFile image mode=RGBA size=850x600>},
   {'type': 'text', 'text': 'Which line represents data about boys?'}]}]

In [ ]python · cell 62

python

test_dataset[20]['images'][0]

Output

<PIL.PngImagePlugin.PngImageFile image mode=RGBA size=850x600>

In [ ]python · cell 63

python

output = generate_text_from_sample(model, processor, test_dataset[20])
output

Output

' Child Labor (Boys, World, 2000-2012)'

The model has successfully learned to respond to the queries as specified in the dataset. We've achieved our goal! 🎉✨

💻 I’ve developed an example application to test the model, which you can find here. You can easily compare it with another Space featuring the pre-trained model, available here.

In [ ]python · cell 66

python

from IPython.display import IFrame

IFrame(src="https://sergiopaniego-smolvlm-trl-sft-chartqa.hf.space", width=1000, height=800)

Output

<IPython.lib.display.IFrame at 0x7adbb46f9120>

6. Continuing the Learning Journey 🧑‍🎓️

To further enhance your skills with multimodal models, I recommend checking out the resources shared at the beginning of this notebook or revisiting the section with the same name 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.

← 返回课程

Open-Source AI Cookbook

26 / 120·notebook

Chapter 26

Fine-tuning SmolVLM with TRL on a consumer GPU

NotebookPython 367 cells

Fine-tuning SmolVLM with TRL on a consumer GPU

Authored by: Sergio Paniego

🌟 Model & Dataset Overview

📖 Additional Resources

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

Multimodal Recipes in Cookbook: Explore practical recipes for multimodal models, including RAG pipelines and fine-tuning. We already have a recipe for fine-tuning a VLM with TRL, so refer to it for more details.
TRL Community Tutorials: A treasure trove of tutorials to deepen your understanding of TRL and its applications.

With these resources, you’ll be equipped to dive deeper into the world of VLMs and push the boundaries of what they can achieve!

This notebook is tested using a L4 GPU.

Smol VLMs comparison

1. Install Dependencies

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

In [1]python · cell 5

python

!pip install -U -q git+https://github.com/huggingface/trl.git bitsandbytes peft qwen-vl-utils trackio
# Tested with trl==0.22.0.dev0, bitsandbytes==0.47.0, peft==0.17.1, qwen-vl-utils==0.0.11, trackio==0.2.8

Output

  Installing build dependencies ... [?25l[?25hdone
  Getting requirements to build wheel ... [?25l[?25hdone
  Preparing metadata (pyproject.toml) ... [?25l[?25hdone
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m61.3/61.3 MB[0m [31m40.0 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m504.9/504.9 kB[0m [31m38.0 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m844.5/844.5 kB[0m [31m56.4 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m59.6/59.6 MB[0m [31m40.9 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m324.6/324.6 kB[0m [31m28.8 MB/s[0m eta [36m0:00:00[0m
[2K   [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m40.0/40.0 MB[0m [31m60.5 MB/s[0m eta [36m0:00:00[0m
[?25h  Building wheel for trl (pyproject.toml) ... [?25l[?25hdone

In [2]python · cell 6

python

!pip install -q flash-attn --no-build-isolation

Output

[?25l     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m0.0/8.4 MB[0m [31m?[0m eta [36m-:--:--[0m
[2K     [91m━━━━━━[0m[91m╸[0m[90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m1.4/8.4 MB[0m [31m43.6 MB/s[0m eta [36m0:00:01[0m
[2K     [91m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m[91m╸[0m [32m8.4/8.4 MB[0m [31m133.9 MB/s[0m eta [36m0:00:01[0m
[2K     [90m━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━[0m [32m8.4/8.4 MB[0m [31m92.3 MB/s[0m eta [36m0:00:00[0m
[?25h  Preparing metadata (setup.py) ... [?25l[?25hdone
  Building wheel for flash-attn (setup.py) ... [?25l[?25hdone

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

In [3]python · cell 8

python

from huggingface_hub import notebook_login

notebook_login()

Output

VBox(children=(HTML(value='<center> <img\nsrc=https://huggingface.co/front/assets/huggingface_logo-noborder.sv…

2. Load Dataset 📁

We’ll load the HuggingFaceM4/ChartQA dataset, which provides chart images along with corresponding questions and answers—perfect for fine-tuning visual question-answering models.

We’ll create a system message to make the VLM act as a chart analysis expert, giving concise answers about chart images.

In [4]python · cell 11

python

system_message = """You are a Vision Language Model specialized in interpreting visual data from chart images.
Your task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.
The charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.
Focus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary."""

We’ll format the dataset into a chatbot structure, with the system message, image, user query, and answer for each interaction.

💡For more tips on using this model, check out the Model Card.

In [5]python · cell 13

python

def format_data(sample):
    return {
      "images": [sample["image"]],
      "messages": [
        {
            "role": "system",
            "content": [
                {
                    "type": "text",
                    "text": system_message
                }
            ],
        },
        {
            "role": "user",
            "content": [
                {
                    "type": "image",
                    "image": sample["image"],
                },
                {
                    "type": "text",
                    "text": sample['query'],
                }
            ],
        },
        {
            "role": "assistant",
            "content": [
                {
                    "type": "text",
                    "text": sample["label"][0]
                }
            ],
        },
        ]
    }

For educational purposes, we’ll load only 10% of each split in the dataset. In a real-world scenario, you would load the entire dataset.

In [6]python · cell 15

python

from datasets import load_dataset

dataset_id = "HuggingFaceM4/ChartQA"
train_dataset, eval_dataset, test_dataset = load_dataset(dataset_id, split=['train[:10%]', 'val[:10%]', 'test[:10%]'])

Output

/usr/local/lib/python3.12/dist-packages/huggingface_hub/utils/_auth.py:86: UserWarning: 
Access to the secret `HF_TOKEN` has not been granted on this notebook.
You will not be requested again.
Please restart the session if you want to be prompted again.
  warnings.warn(

README.md:   0%|          | 0.00/852 [00:00<?, ?B/s]

data/train-00000-of-00003-49492f364babfa(…):   0%|          | 0.00/219M [00:00<?, ?B/s]

data/train-00001-of-00003-7302bae5e425bb(…):   0%|          | 0.00/311M [00:00<?, ?B/s]

data/train-00002-of-00003-194c9400785577(…):   0%|          | 0.00/315M [00:00<?, ?B/s]

data/val-00000-of-00001-0f11003c77497969(…):   0%|          | 0.00/50.2M [00:00<?, ?B/s]

data/test-00000-of-00001-e2cd0b7a0f9eb20(…):   0%|          | 0.00/68.9M [00:00<?, ?B/s]

Generating train split:   0%|          | 0/28299 [00:00<?, ? examples/s]

Generating val split:   0%|          | 0/1920 [00:00<?, ? examples/s]

Generating test split:   0%|          | 0/2500 [00:00<?, ? examples/s]

Let’s take a look at the dataset structure. It includes an image, a query, a label (the answer), and a fourth feature that we’ll be discarding.

In [7]python · cell 17

python

train_dataset

Output

Dataset({
    features: ['image', 'query', 'label', 'human_or_machine'],
    num_rows: 2830
})

Now, let’s format the data using the chatbot structure. This will set up the interactions for the model.

In [8]python · cell 19

python

train_dataset = [format_data(sample) for sample in train_dataset]
eval_dataset = [format_data(sample) for sample in eval_dataset]
test_dataset = [format_data(sample) for sample in test_dataset]

In [9]python · cell 20

python

train_dataset[200]

Output

{'images': [<PIL.PngImagePlugin.PngImageFile image mode=RGB size=308x369>],
 'messages': [{'role': 'system',
   'content': [{'type': 'text',
     'text': 'You are a Vision Language Model specialized in interpreting visual data from chart images.\nYour task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.\nThe charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.\nFocus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary.'}]},
  {'role': 'user',
   'content': [{'type': 'image',
     'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=308x369>},
    {'type': 'text',
     'text': 'Is the rightmost value of light brown graph 58?'}]},
  {'role': 'assistant', 'content': [{'type': 'text', 'text': 'No'}]}]}

3. Load Model and Check Performance! 🤔

For a broader comparison of state-of-the-art VLMs, explore the WildVision Arena and the OpenVLM Leaderboard, where you can find the best-performing models across various benchmarks.

In [10]python · cell 23

python

import torch
from transformers import Idefics3ForConditionalGeneration, AutoProcessor

model_id = "HuggingFaceTB/SmolVLM-Instruct"

Next, we’ll load the model and the tokenizer to prepare for inference.

In [11]python · cell 25

python

model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2",
)

processor = AutoProcessor.from_pretrained(model_id)

Output

config.json: 0.00B [00:00, ?B/s]

model.safetensors:   0%|          | 0.00/4.49G [00:00<?, ?B/s]

generation_config.json:   0%|          | 0.00/136 [00:00<?, ?B/s]

processor_config.json:   0%|          | 0.00/68.0 [00:00<?, ?B/s]

chat_template.json:   0%|          | 0.00/429 [00:00<?, ?B/s]

preprocessor_config.json:   0%|          | 0.00/486 [00:00<?, ?B/s]

tokenizer_config.json: 0.00B [00:00, ?B/s]

vocab.json: 0.00B [00:00, ?B/s]

merges.txt: 0.00B [00:00, ?B/s]

tokenizer.json: 0.00B [00:00, ?B/s]

added_tokens.json:   0%|          | 0.00/92.0 [00:00<?, ?B/s]

special_tokens_map.json: 0.00B [00:00, ?B/s]

To evaluate the model's performance, we’ll use a sample from the dataset. First, let’s inspect the internal structure of this sample to understand how the data is organized.

In [12]python · cell 27

python

train_dataset[1]

Output

{'images': [<PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>],
 'messages': [{'role': 'system',
   'content': [{'type': 'text',
     'text': 'You are a Vision Language Model specialized in interpreting visual data from chart images.\nYour task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.\nThe charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.\nFocus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary.'}]},
  {'role': 'user',
   'content': [{'type': 'image',
     'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>},
    {'type': 'text',
     'text': 'How many values are below 40 in Unfavorable graph?'}]},
  {'role': 'assistant', 'content': [{'type': 'text', 'text': '6'}]}]}

We’ll use the sample without the system message to assess the VLM's raw understanding. Here’s the input we will use:

In [13]python · cell 29

python

train_dataset[0]['messages'][1:2]

Output

[{'role': 'user',
  'content': [{'type': 'image',
    'image': <PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>},
   {'type': 'text', 'text': 'Is the value of Favorable 38 in 2015?'}]}]

Now, let’s take a look at the chart corresponding to the sample. Can you answer the query based on the visual information?

In [14]python · cell 31

python

train_dataset[0]['images'][0]

Output

<PIL.PngImagePlugin.PngImageFile image mode=RGB size=422x359>

In [15]python · cell 33

python

def generate_text_from_sample(model, processor, sample, max_new_tokens=1024, device="cuda"):
    # Prepare the text input by applying the chat template
    text_input = processor.apply_chat_template(
        sample['messages'][1:2],  # Use the sample without the system message
        add_generation_prompt=True
    )

    image_inputs = []
    image = sample['images'][0]
    #image = sample[1]['content'][0]['image']
    if image.mode != 'RGB':
        image = image.convert('RGB')
    image_inputs.append([image])

    # Prepare the inputs for the model
    model_inputs = processor(
        #text=[text_input],
        text=text_input,
        images=image_inputs,
        return_tensors="pt",
    ).to(device)  # Move inputs to the specified device

    # Generate text with the model
    generated_ids = model.generate(**model_inputs, max_new_tokens=max_new_tokens)

    # Trim the generated ids to remove the input ids
    trimmed_generated_ids = [
        out_ids[len(in_ids):] for in_ids, out_ids in zip(model_inputs.input_ids, generated_ids)
    ]

    # Decode the output text
    output_text = processor.batch_decode(
        trimmed_generated_ids,
        skip_special_tokens=True,
        clean_up_tokenization_spaces=False
    )

    return output_text[0]  # Return the first decoded output text

In [16]python · cell 34

python

output = generate_text_from_sample(model, processor, train_dataset[1])
output

Output

' 3.'

Remove Model and Clean GPU

Before we proceed with training the model in the next section, let's clear the current variables and clean the GPU to free up resources.

In [17]python · cell 37

python

import gc
import time

def clear_memory():
    # Delete variables if they exist in the current global scope
    if 'inputs' in globals(): del globals()['inputs']
    if 'model' in globals(): del globals()['model']
    if 'processor' in globals(): del globals()['processor']
    if 'trainer' in globals(): del globals()['trainer']
    if 'peft_model' in globals(): del globals()['peft_model']
    if 'bnb_config' in globals(): del globals()['bnb_config']
    time.sleep(2)

    # Garbage collection and clearing CUDA memory
    gc.collect()
    time.sleep(2)
    torch.cuda.empty_cache()
    torch.cuda.synchronize()
    time.sleep(2)
    gc.collect()
    time.sleep(2)

    print(f"GPU allocated memory: {torch.cuda.memory_allocated() / 1024**3:.2f} GB")
    print(f"GPU reserved memory: {torch.cuda.memory_reserved() / 1024**3:.2f} GB")

clear_memory()

Output

GPU allocated memory: 0.01 GB
GPU reserved memory: 0.02 GB

4. Fine-Tune the Model using TRL

4.1 Load the Quantized Model for Training ⚙️

Next, we’ll load the quantized model using bitsandbytes. If you want to learn more about quantization, check out this blog post or this one.

In [18]python · cell 40

python

from transformers import BitsAndBytesConfig

# BitsAndBytesConfig int-4 config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# Load model and tokenizer
model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    quantization_config=bnb_config,
    _attn_implementation="flash_attention_2",
)
processor = AutoProcessor.from_pretrained(model_id)

4.2 Set Up QLoRA and SFTConfig 🚀

In [19]python · cell 42

python

from peft import LoraConfig, get_peft_model

# Configure LoRA
peft_config = LoraConfig(
    r=8,
    lora_alpha=8,
    lora_dropout=0.1,
    target_modules=['down_proj','o_proj','k_proj','q_proj','gate_proj','up_proj','v_proj'],
    use_dora=True,
    init_lora_weights="gaussian"
)

# Apply PEFT model adaptation
peft_model = get_peft_model(model, peft_config)

# Print trainable parameters
peft_model.print_trainable_parameters()

Output

trainable params: 11,269,248 || all params: 2,257,542,128 || trainable%: 0.4992

In [27]python · cell 44

python

from trl import SFTConfig

# Configure training arguments using SFTConfig
training_args = SFTConfig(
    output_dir="smolvlm-instruct-trl-sft-ChartQA",
    num_train_epochs=1,
    per_device_train_batch_size=4,
    gradient_accumulation_steps=4,
    warmup_steps=50,
    learning_rate=1e-4,
    weight_decay=0.01,
    logging_steps=25,
    save_strategy="steps",
    save_steps=25,
    save_total_limit=1,
    optim="adamw_torch_fused",
    bf16=True,
    push_to_hub=True,
    report_to="none",
    max_length=None
)

In [21]python · cell 45

python

import trackio

trackio.init(
    project="smolvlm-instruct-trl-sft-ChartQA",
    name="smolvlm-instruct-trl-sft-ChartQA",
    config=training_args,
    space_id=training_args.output_dir + "-trackio"
)

Output

* Trackio project initialized: smolvlm-instruct-trl-sft-ChartQA
* Trackio metrics will be synced to Hugging Face Dataset: sergiopaniego/smolvlm-instruct-trl-sft-ChartQA-trackio-dataset
* Creating new space: https://huggingface.co/spaces/sergiopaniego/smolvlm-instruct-trl-sft-ChartQA-trackio
* View dashboard by going to: https://huggingface.co/spaces/sergiopaniego/smolvlm-instruct-trl-sft-ChartQA-trackio

<trackio.run.Run at 0x7fcb74bcde50>

4.3 Training the Model 🏃

In [28]python · cell 48

python

from trl import SFTTrainer

trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=peft_config,
    processing_class=processor,
)

Output

/usr/local/lib/python3.12/dist-packages/peft/mapping_func.py:73: UserWarning: You are trying to modify a model with PEFT for a second time. If you want to reload the model with a different config, make sure to call `.unload()` before.
  warnings.warn(
/usr/local/lib/python3.12/dist-packages/peft/tuners/tuners_utils.py:196: UserWarning: Already found a `peft_config` attribute in the model. This will lead to having multiple adapters in the model. Make sure to know what you are doing!
  warnings.warn(

Time to Train the Model! 🎉

In [ ]python · cell 50

python

trainer.train()

Output

Casting fp32 inputs back to torch.bfloat16 for flash-attn compatibility.
`loss_type=None` was set in the config but it is unrecognised.Using the default loss: `ForCausalLMLoss`.

<IPython.core.display.HTML object>

[ 3/177 01:23 < 4:01:02, 0.01 it/s, Epoch 0.01/1]

Step	Training Loss

Let's save the results 💾

In [ ]python · cell 52

python

trainer.save_model(training_args.output_dir)

5. Testing the Fine-Tuned Model 🔍

Let's clean up the GPU memory to ensure optimal performance 🧹

In [ ]python · cell 55

python

clear_memory()

Output

GPU allocated memory: 16.34 GB
GPU reserved memory: 18.69 GB

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

In [ ]python · cell 57

python

model = Idefics3ForConditionalGeneration.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    _attn_implementation="flash_attention_2",
)

processor = AutoProcessor.from_pretrained(model_id)

In [ ]python · cell 59

python

adapter_path = "sergiopaniego/smolvlm-instruct-trl-sft-ChartQA"
model.load_adapter(adapter_path)

Let's evaluate the model on an unseen sample.

In [ ]python · cell 61

python

test_dataset[20]['messages'][:2]

Output

[{'role': 'system',
  'content': [{'type': 'text',
    'text': 'You are a Vision Language Model specialized in interpreting visual data from chart images.\nYour task is to analyze the provided chart image and respond to queries with concise answers, usually a single word, number, or short phrase.\nThe charts include a variety of types (e.g., line charts, bar charts) and contain colors, labels, and text.\nFocus on delivering accurate, succinct answers based on the visual information. Avoid additional explanation unless absolutely necessary.'}]},
 {'role': 'user',
  'content': [{'type': 'image',
    'image': <PIL.PngImagePlugin.PngImageFile image mode=RGBA size=850x600>},
   {'type': 'text', 'text': 'Which line represents data about boys?'}]}]

In [ ]python · cell 62

python

test_dataset[20]['images'][0]

Output

<PIL.PngImagePlugin.PngImageFile image mode=RGBA size=850x600>

In [ ]python · cell 63

python

output = generate_text_from_sample(model, processor, test_dataset[20])
output

Output

' Child Labor (Boys, World, 2000-2012)'

The model has successfully learned to respond to the queries as specified in the dataset. We've achieved our goal! 🎉✨

💻 I’ve developed an example application to test the model, which you can find here. You can easily compare it with another Space featuring the pre-trained model, available here.

In [ ]python · cell 66

python

from IPython.display import IFrame

IFrame(src="https://sergiopaniego-smolvlm-trl-sft-chartqa.hf.space", width=1000, height=800)

Output

<IPython.lib.display.IFrame at 0x7adbb46f9120>

6. Continuing the Learning Journey 🧑‍🎓️

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