Your First Multimodal LLM Fine-Tune: QLoRA + Gemma, Step-by-Step

If you’re not a member, you can still access the article for free here.

This comprehensive tutorial will guide you through the process of fine-tuning a Large Language Model (LLM) using Quantized Low-Rank Adaptation (QLoRA). We’ll cover everything from setting up your environment to preparing your dataset, training the model, and finally, testing its inference capabilities. This guide is designed for beginners, so every concept and line of code will be explained in detail.

What is Quantized Low-Rank Adaptation (QLoRA)?

Before we dive into the practical steps, let’s understand what QLoRA is and why it’s so beneficial for fine-tuning LLMs.

Large Language Models (LLMs): Imagine a super-smart computer program that can understand and generate human-like text. That’s essentially what an LLM is. These models are trained on massive amounts of text data from the internet, allowing them to perform various language-related tasks like writing articles, answering questions, translating languages, and even generating creative content.

Fine-tuning: While LLMs are powerful, they are often general-purpose. To make them excel at a specific task (like generating product descriptions for e-commerce), we “fine-tune” them. This involves training the already-trained LLM on a smaller, more specific dataset relevant to our task. Think of it like taking a highly skilled generalist and giving them specialized training for a particular job.

The Challenge of Fine-tuning LLMs: LLMs are enormous! They have billions of parameters (think of these as the model’s internal “knobs” that it adjusts during training). Fine-tuning such a large model usually requires a lot of computational resources, like powerful graphics processing units (GPUs) and a lot of memory. This can be a barrier for many researchers and developers.

Quantized Low-Rank Adaptation (QLoRA) to the Rescue! QLoRA is a clever technique that significantly reduces the computational resources needed for fine-tuning LLMs while still achieving high performance. Here’s how it works:

Quantization (4-bit): In QLoRA, the original, pre-trained LLM is “quantized” to 4-bit precision.

• Precision: In computing, precision refers to the number of bits used to represent a number. Higher precision means more bits, which can lead to more accurate calculations but also requires more memory and computational power.
• 4-bit Quantization: This means that the numbers (weights) inside the LLM, which are typically represented with higher precision (e.g., 16-bit or 32-bit), are compressed to use only 4 bits. This drastically reduces the memory footprint of the model.
• Frozen Weights: Once the model is quantized, its original weights are “frozen.” This means they are not updated during the fine-tuning process. This is a crucial step that saves a lot of computation.

Trainable Adapter Layers (LoRA): Instead of training the entire massive LLM, QLoRA introduces small, additional layers called “adapter layers” (based on a technique called LoRA, or Low-Rank Adaptation).

• Adapter Layers: Think of these as tiny, specialized add-ons to the original LLM. They are much smaller than the full model.
• Only Adapters are Trained: During fine-tuning, only these small adapter layers are trained. They learn to adjust the output of the frozen, quantized base model to better suit the specific fine-tuning task. Because these layers are so much smaller, training them requires significantly less computational power and memory.

Merging or Keeping Separate: After training, the weights of these adapter layers can either be:

• Merged: Combined with the original, frozen base model. This creates a single, updated model that can be used for inference (making predictions).
• Kept as a Separate Adapter: The adapter layers can be kept as a separate, small file. When you want to use the fine-tuned model, you load the original base model and then “attach” these adapter layers on top.

In essence, QLoRA is like giving a massive, powerful brain a tiny, specialized “skill module” that it can quickly learn and attach, rather than retraining the entire brain from scratch. This makes fine-tuning LLMs much more accessible and efficient.

Setup Development Environment

The first crucial step is to prepare your environment by installing all the necessary libraries. These libraries provide the tools and functionalities we’ll need to fine-tune our LLM.

# Install Pytorch & other libraries
%pip install "torch>=2.4.0" tensorboard torchvision

# Install Gemma release branch from Hugging Face
%pip install "transformers>=4.51.3"

# Install Hugging Face libraries
%pip install  --upgrade \
  "datasets==3.3.2" \
  "accelerate==1.4.0" \
  "evaluate==0.4.3" \
  "bitsandbytes==0.45.3" \
  "trl==0.15.2" \
  "peft==0.14.0" \
  "pillow==11.1.0" \
  protobuf \
  sentencepiece

Explanation of each line:

• %pip install ...: This is a command commonly used in Jupyter notebooks or Google Colab to install Python packages. The % at the beginning indicates that it's a "magic command" specific to these environments, allowing you to run shell commands directly.

**%pip install "torch>=2.4.0" tensorboard torchvision**

• torch>=2.4.0: This installs the PyTorch library, which is a fundamental open-source machine learning framework. It's used for building and training neural networks. The >=2.4.0 specifies that we need version 2.4.0 or newer. PyTorch is the backbone for many operations in this fine-tuning process.
• tensorboard: This installs TensorBoard, a visualization tool provided by TensorFlow (another machine learning framework, but often used alongside PyTorch for visualization). TensorBoard helps us visualize training progress, metrics, and model graphs, which is incredibly useful for debugging and understanding how our model is learning.
• torchvision: This is a PyTorch library that provides datasets, model architectures, and image transformations specifically designed for computer vision tasks. While our primary task is text generation, this library might be used by the Gemma model or underlying components for image processing.

**%pip install "transformers>=4.51.3"**

• transformers>=4.51.3: This installs the Hugging Face Transformers library. This is a hugely popular library that provides pre-trained models, tokenizers, and training utilities for a wide range of natural language processing (NLP) tasks. It's the core library for interacting with models like Gemma. The >=4.51.3 ensures we have a compatible and recent version.

**%pip install --upgrade \ ...**

• --upgrade: This flag tells pip to upgrade the specified packages to their latest versions if they are already installed. This is good practice to ensure you have the most up-to-date features and bug fixes.
• "datasets==3.3.2": Installs the Hugging Face datasets library, specifically version 3.3.2. This library is essential for easily loading, processing, and sharing datasets for machine learning. We will use it to load our product description dataset.
• "accelerate==1.4.0": Installs the Hugging Face accelerate library, version 1.4.0. This library helps to easily run PyTorch training scripts on various hardware configurations (e.g., multiple GPUs, distributed training) without changing your code significantly. It "accelerates" your training.
• "evaluate==0.4.3": Installs the Hugging Face evaluate library, version 0.4.3. This library provides easy access to various metrics for evaluating machine learning models. While not explicitly used in the training loop provided, it's generally useful for assessing model performance.
• "bitsandbytes==0.45.3": Installs the bitsandbytes library, version 0.45.3. This library is crucial for QLoRA because it provides highly optimized routines for 8-bit and 4-bit quantization, which are at the heart of QLoRA's memory efficiency.
• "trl==0.15.2": Installs the Hugging Face trl (Transformer Reinforcement Learning) library, version 0.15.2. This library provides tools for training large language models with reinforcement learning from human feedback (RLHF) and, more relevant to us, includes the SFTTrainer for supervised fine-tuning.
• "peft==0.14.0": Installs the Hugging Face peft (Parameter-Efficient Fine-Tuning) library, version 0.14.0. This library is specifically designed to enable efficient fine-tuning techniques like LoRA and QLoRA. It handles the creation and management of the adapter layers.
• "pillow==11.1.0": Installs the Pillow library, version 11.1.0. Pillow is a powerful image processing library in Python. It's used here because our dataset includes images, and Pillow will handle loading and manipulating these image files (e.g., converting them to RGB format).
• protobuf: Installs the protobuf library. Protocol Buffers are Google's language-neutral, platform-neutral, extensible mechanism for serializing structured data. It's often used for efficient data interchange and might be a dependency for some of the other installed libraries, especially those from Google or related to data handling.
• sentencepiece: Installs the sentencepiece library. SentencePiece is an unsupervised text tokenizer and detokenizer. It's often used by modern LLMs (including Gemma) to break down text into smaller units (tokens) that the model can understand. This is essential for both input processing and output generation.

Accepting Gemma Terms of Use and Hugging Face Login

Before you can use the Gemma model, there’s an important step: accepting its terms of use and logging into your Hugging Face account.

Why this is important:

• Terms of Use: Large pre-trained models often come with specific licenses and terms of use that you must agree to before downloading and using them. This ensures you understand the conditions of their usage.
• Hugging Face Token: Hugging Face Hub is a platform where many pre-trained models and datasets are hosted. To download some models (especially larger or gated ones like Gemma), and to push your fine-tuned model back to the Hub, you need to be authenticated. A Hugging Face Token acts as your password for programmatic access.

Steps:

1. Accept Gemma License:

• Go to the model page on Hugging Face: http://huggingface.co/google/gemma-3-4b-pt (or the appropriate page for other Gemma models if you're using a different variant).
• Look for an “Agree and access repository” button and click it. This signifies your acceptance of the terms.

2. Get a Hugging Face Token:

• If you don’t have a Hugging Face account, create one.
• Once logged in, go to your profile settings on Hugging Face.
• Navigate to “Access Tokens” or “Tokens.”
• Generate a new token. Crucially, make sure this token has “write access.” This is important because later in the tutorial, we will push our fine-tuned model to the Hugging Face Hub, and write access is required for that.

Code for Hugging Face Login:

from google.colab import userdata
from huggingface_hub import login

# Login into Hugging Face Hub
hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab
login(hf_token)

Explanation of each line:

**from google.colab import userdata**

• This line imports the userdata module from google.colab. This is specifically for users running their code in Google Colab. Google Colab provides a secure way to store sensitive information like API keys or tokens using "Colab secrets." This method allows you to retrieve your Hugging Face token without hardcoding it directly into your script, which is a good security practice.

**from huggingface_hub import login**

• This line imports the login function from the huggingface_hub library. This function is used to authenticate with the Hugging Face Hub.

**hf_token = userdata.get('HF_TOKEN') # If you are running inside a Google Colab**

• This line attempts to retrieve your Hugging Face token from Colab secrets. You would have previously saved your Hugging Face token under the name HF_TOKEN in your Colab environment's secrets manager. If you are not using Google Colab (e.g., running on your local machine or another cloud environment), you would typically replace this line with: hf_token = "YOUR_HUGGING_FACE_WRITE_TOKEN_HERE" (but never share your token publicly).

**login(hf_token)**

• This line calls the login function, passing your retrieved Hugging Face token (hf_token). This action authenticates your current Python session with the Hugging Face Hub, allowing you to download gated models and upload your fine-tuned model.

Create and Prepare the Fine-tuning Dataset

The quality and structure of your dataset are paramount for effective fine-tuning. This section focuses on understanding the importance of a well-defined use case and then demonstrates how to load and format a multimodal dataset for our specific task.

Understanding Your Use Case:

Before you even think about data, you need to clearly define what you want your fine-tuned LLM to do.

• Why is this important? The dataset you create or choose must directly align with your desired outcome. If your model needs to write short, catchy headlines, your dataset should contain examples of short, catchy headlines paired with their corresponding inputs. If your dataset contains long, academic papers, your model won’t learn to write effective short headlines.
• Example Use Case (from the text): “Fine-tuning a Gemma model to generate concise, SEO-optimized product descriptions for an e-commerce platform, specifically tailored for mobile search.”
• Concise: The descriptions should be brief.
• SEO-optimized: They should include keywords that help products appear higher in search results.
• E-commerce platform: The language should be appropriate for selling products online.
• Mobile search: Implies brevity and ease of reading on small screens.

Choosing a Dataset:

The tutorial uses the philschmid/amazon-product-descriptions-vlm dataset.

• Why this dataset? It’s chosen because it directly fits the defined use case. It contains Amazon product descriptions, including product images and categories, which serve as the inputs our model uses to generate descriptions.
• Multimodal: The “VLM” in the dataset name stands for “Vision-Language Model,” indicating that it contains both visual (images) and textual (descriptions, product names, categories) data. This is crucial as Gemma is a vision-capable LLM.

Multimodal Conversation Structure for TRL

Hugging Face TRL (Transformer Reinforcement Learning library) has a specific way of representing multimodal conversations. This structure tells the processing class how to handle different types of content (text and images).

The required format is a list of “messages,” where each message has a “role” (e.g., “system”, “user”, “assistant”) and “content.” The content itself is a list, where each item can be either text or an image.

{
  "messages": [
    {
      "role": "system",
      "content": [{"type": "text", "text": "You are..."}]
    },
    {
      "role": "user",
      "content": [
        {"type": "text", "text": "..."},
        {"type": "image"}
      ]
    },
    {
      "role": "assistant",
      "content": [{"type": "text", "text": "..."}]
    }
  ]
}

Explanation of each line:

• **"messages":** The top-level key for a single conversation turn. It contains a list of message objects.
• **"role":** Indicates who is speaking.
• "system": Provides initial instructions or context to the model (e.g., "You are an expert product description writer.").
• "user": Represents the input or query from the user. In our case, this will include the product name, category, and image.
• "assistant": Represents the expected output or response from the model. In our case, this will be the desired product description.
• **"content":** A list of content elements within a message.
• {"type": "text", "text": "..."}: Represents a text part of the message.
• {"type": "image"}: Crucially, this tells the processing class that an image needs to be loaded here. The actual image data will be stored separately (as Pil.Image objects).

Loading and Formatting the Dataset

Now, let’s look at the Python code to load the dataset and transform it into the required multimodal conversation format.

from datasets import load_dataset
from PIL import Image

# System message for the assistant
system_message = "You are an expert product description writer for Amazon."
# User prompt that combines the user query and the schema
user_prompt = """Create a Short Product description based on the provided <PRODUCT> and <CATEGORY> and image.Only return description. The description should be SEO optimized and for a better mobile search experience.
<PRODUCT>{product}
</PRODUCT>
<CATEGORY>{category}
</CATEGORY>"""
# Convert dataset to OAI messages
def format_data(sample):
    return {
        "messages": [
            {
                "role": "system",
                "content": [{"type": "text", "text": system_message}],
            },
            {
                "role": "user",
                "content": [
                    {
                        "type": "text",
                        "text": user_prompt.format(
                            product=sample["Product Name"],
                            category=sample["Category"],
                        ),
                    },
                    {
                        "type": "image",
                        "image": sample["image"], # The PIL.Image object is stored here
                    },
                ],
            },
            {
                "role": "assistant",
                "content": [{"type": "text", "text": sample["description"]}],
            },
        ],
    }
def process_vision_info(messages: list[dict]) -> list[Image.Image]:
    image_inputs = []
    # Iterate through each conversation
    for msg in messages:
        # Get content (ensure it's a list)
        content = msg.get("content", [])
        if not isinstance(content, list):
            content = [content]
        # Check each content element for images
        for element in content:
            if isinstance(element, dict) and (
                "image" in element or element.get("type") == "image"
            ):
                # Get the image and convert to RGB
                if "image" in element:
                    image = element["image"]
                else:
                    image = element
                image_inputs.append(image.convert("RGB"))
    return image_inputs
# Load dataset from the hub
dataset = load_dataset("philschmid/amazon-product-descriptions-vlm", split="train")
# Convert dataset to OAI messages
# need to use list comprehension to keep Pil.Image type, .mape convert image to bytes
dataset = [format_data(sample) for sample in dataset]
print(dataset[345]["messages"])

Explanation of each line:

**from datasets import load_dataset**

• Imports the load_dataset function from the Hugging Face datasets library. This is the primary function to download and load datasets from the Hugging Face Hub.

**from PIL import Image**

• Imports the Image module from the Pillow (PIL) library. This is used to work with image objects, as the dataset provides images as PIL.Image objects.

**system_message = "You are an expert product description writer for Amazon."**

• Defines a string variable system_message. This will be included in the "system" role of our conversational prompt. It sets the persona and task for the LLM, guiding its behavior.

**user_prompt = """..."""**

• Defines a multi-line string variable user_prompt. This is a template for the user's input to the model.
• It contains clear instructions for the model (“Create a Short Product description…”).
• It uses placeholders {product} and {category} which will be filled with actual product names and categories from our dataset.
• The <PRODUCT> and <CATEGORY> tags are XML-like structures. These are often used to clearly delineate different pieces of information within a prompt, making it easier for the LLM to parse and understand.
• “Only return description.” is a crucial instruction to ensure the model doesn’t generate extra conversational filler.
• “SEO optimized and for a better mobile search experience.” reinforces the specific requirements of our use case.

**def format_data(sample):**

• Defines a Python function named format_data that takes a sample (a single entry/row from our dataset) as input. This function will transform each sample into the TRL multimodal conversation format.

**return { "messages": [ ... ] }**

• This is the core of the format_data function. It constructs the dictionary representing a single conversation following the TRL multimodal format.
• System Message Part:

{
     "role": "system",
     "content": [{"type": "text", "text": system_message}], },

• This creates the system message. It sets the role to "system" and its content is a list containing a single dictionary with type: "text" and the system_message string.
• User Message Part:

{
     "role": "user",
     "content": [
         { 
             "type": "text",
             "text": user_prompt.format(
                 product=sample["Product Name"], 
                category=sample["Category"],
             ),
         },
         {
             "type": "image",
             "image": sample["image"],
         },
     ],
 },

• This creates the user message.
• Its role is "user".
• Its content is a list with two elements:
• The first element is a dictionary for the text part: type: "text" and the text field is populated by calling user_prompt.format(). This replaces {product} and {category} placeholders with the actual "Product Name" and "Category" from the current sample.
• The second element is a dictionary for the image part: type: "image" and, crucially, image: sample["image"]. This is where the actual PIL.Image object from the dataset sample is placed.
• Assistant Message Part:

{
     "role": "assistant",
     "content": [{"type": "text", "text": sample["description"]}], },

• This creates the assistant message, which represents the desired output.
• Its role is "assistant".
• Its content is a list with a single dictionary for the text part, where the text field is the "description" from the current sample. This is what our model will learn to generate.

**def process_vision_info(messages: list[dict]) -> list[Image.Image]:**

• Defines a utility function process_vision_info. This function takes a list of messages (in the TRL conversation format) and extracts all the PIL.Image objects from them.
• messages: list[dict]: Type hint indicating that messages is expected to be a list of dictionaries.
• -> list[Image.Image]: Type hint indicating that the function is expected to return a list of PIL.Image.Image objects.
• image_inputs = []: Initializes an empty list to store the extracted images.
• for msg in messages:: Iterates through each message in the conversation.
• content = msg.get("content", []): Safely retrieves the "content" list from the message. If "content" is missing, it defaults to an empty list.
• if not isinstance(content, list): content = [content]: Ensures content is always a list, even if it somehow was a single item.
• for element in content:: Iterates through each element (text or image) within the message's content.
• if isinstance(element, dict) and ("image" in element or element.get("type") == "image"):: Checks if the current element is a dictionary and if it either directly contains an "image" key (where the PIL Image object is) or if its "type" key is "image". This identifies image elements.
• if "image" in element: image = element["image"] else: image = element: Retrieves the image. If the image is directly stored under the "image" key, use that. Otherwise, assume the element itself is the image (this handles cases where the structure might be slightly different but still represents an image).
• image_inputs.append(image.convert("RGB")): Appends the extracted PIL.Image object to image_inputs list. image.convert("RGB") is crucial: it converts the image to the RGB color mode, which is a standard format expected by many models.

**dataset = load_dataset("philschmid/amazon-product-descriptions-vlm", split="train")**

• Loads the specified dataset from the Hugging Face Hub.
• split="train": Specifies that we want to load the "train" split of the dataset, which is the data used for training our model.

**dataset = [format_data(sample) for sample in dataset]**

• This is a Python list comprehension. It’s a concise way to create a new list by applying the format_data function to each sample (row) in the original dataset.
• Crucial Note: The comment # need to use list comprehension to keep Pil.Image type, .mape convert image to bytes is very important. The dataset.map() function (which is commonly used with Hugging Face datasets) can sometimes alter the type of data, specifically converting PIL.Image objects to bytes. For multimodal models that directly expect PIL.Image objects, using a list comprehension ensures that the image objects retain their original PIL.Image type, preventing issues later in the pipeline when the model's processor expects them.

**print(dataset[345]["messages"])**

• Prints the messages content of the 346th sample (index 345) in our newly formatted dataset. This is a good way to visually inspect if the data has been transformed correctly into the desired multimodal conversation structure.

Fine-tune Gemma using TRL and the SFTTrainer

This is the core section where we configure and initiate the fine-tuning process using the SFTTrainer from the Hugging Face TRL library.

Introduction to SFTTrainer

The SFTTrainer (Supervised Fine-Tuning Trainer) is a powerful tool provided by the Hugging Face trl library specifically designed for fine-tuning LLMs. It extends the Trainer class from the transformers library, meaning it inherits all the robust features of the standard Trainer (like logging, evaluation, checkpointing, etc.) and adds specialized functionalities for supervised fine-tuning.

Key features of SFTTrainer:

• Dataset formatting: It simplifies handling conversational and instruction-based datasets, making it easier to prepare your data for training.
• Training on completions only: It can intelligently train the model to generate only the “assistant” (completion) part of the conversation, ignoring the “user” prompts during loss calculation. This prevents the model from just copying the input.
• Packing datasets: For more efficient training, it can “pack” multiple short sequences into a single, longer sequence. This reduces padding and maximizes GPU utilization.
• Parameter-efficient fine-tuning (PEFT) support: It has built-in support for techniques like QLoRA, making it seamless to integrate these memory-efficient methods.
• Preparing model and tokenizer: It can automatically handle adding special tokens (like [BOS], [EOS], [IMG]) to the tokenizer and model, which are crucial for conversational models and vision-language models.

Loading Model and Tokenizer with Quantization Configuration

Now, let’s load our Gemma model and its corresponding processor (which handles both tokenization and image processing) and set up the 4-bit quantization.

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig
# Hugging Face model id
model_id = "google/gemma-3-4b-pt" # or `google/gemma-3-12b-pt`, `google/gemma-3-27-pt`
# Check if GPU benefits from bfloat16
if torch.cuda.get_device_capability()[0] < 8:
    raise ValueError("GPU does not support bfloat16, please use a GPU that supports bfloat16.")
# Define model init arguments
model_kwargs = dict(
    attn_implementation="eager", # Use "flash_attention_2" when running on Ampere or newer GPU
    torch_dtype=torch.bfloat16, # What torch dtype to use, defaults to auto
    device_map="auto", # Let torch decide how to load the model
)
# BitsAndBytesConfig int-4 config
model_kwargs["quantization_config"] = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=model_kwargs["torch_dtype"],
    bnb_4bit_quant_storage=model_kwargs["torch_dtype"],
)
# Load model and processor
model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)
processor = AutoProcessor.from_pretrained("google/gemma-3-4b-it")

Explanation of each line:

**import torch**

• Imports the PyTorch library, which is fundamental for tensor operations and deep learning in general.

**from transformers import AutoProcessor, AutoModelForImageTextToText, BitsAndBytesConfig**

• AutoProcessor: This class from transformers automatically loads the appropriate preprocessor (tokenizer + image processor) for a given model ID. It simplifies handling multimodal inputs.
• AutoModelForImageTextToText: This class automatically loads a model suitable for image-to-text and text-to-image tasks (vision-language models). Gemma is one such model.
• BitsAndBytesConfig: This class is from the bitsandbytes library and is used to define the configuration for 4-bit (or 8-bit) quantization.

**model_id = "google/gemma-3-4b-pt"**

• Defines the model_id string. This is the identifier for the specific Gemma model we want to load from the Hugging Face Hub. gemma-3-4b-pt refers to a Gemma model with 3.4 billion parameters, pre-trained. Other variants like 12b or 27b could be used depending on your GPU memory.

**if torch.cuda.get_device_capability()[0] < 8:**

• This is a check to ensure your GPU is compatible with bfloat16 (Brain Floating Point 16-bit) precision.
• torch.cuda.get_device_capability()[0] returns the major version of the CUDA compute capability for the current GPU.
• GPUs with compute capability 8.0 or higher (e.g., NVIDIA Ampere architecture and newer, like A100, 30-series GPUs) natively support bfloat16, which offers a good balance between precision and memory efficiency.
• raise ValueError(...): If your GPU doesn't support bfloat16, this line will stop the execution and print an error message, as bfloat16 is used in the quantization configuration.
• **model_kwargs = dict(...)**
• Creates a dictionary model_kwargs to store arguments that will be passed when loading the model.
• attn_implementation="eager": Specifies the attention mechanism to use. "eager" is a standard implementation. For better performance on compatible GPUs, flash_attention_2 can be used. Flash Attention is an optimized attention algorithm that significantly speeds up and reduces memory usage for attention computations.
• torch_dtype=torch.bfloat16: Sets the data type for the model's computations to bfloat16. This is a mixed-precision training technique that saves memory and speeds up training without a significant loss in performance on compatible hardware.
• device_map="auto": This tells Hugging Face transformers to automatically determine how to load the model across available devices (e.g., multiple GPUs or CPU if no GPU is available). It tries to optimize memory usage.

**model_kwargs["quantization_config"] = BitsAndBytesConfig(...)**

• This line adds a quantization_config key to our model_kwargs dictionary, setting up the 4-bit quantization.
• load_in_4bit=True: This is the core setting for QLoRA. It tells the bitsandbytes library to load the pre-trained model weights in 4-bit precision.
• bnb_4bit_use_double_quant=True: Enables "double quantization." This further quantizes the 4-bit quantization constants, resulting in even greater memory savings (typically 0.4 bits per parameter) without compromising performance.
• bnb_4bit_quant_type="nf4": Specifies the quantization data type as "NF4" (NormalFloat4). This is a custom 4-bit data type introduced by the QLoRA paper, which is empirically shown to be very effective for LLMs.
• bnb_4bit_compute_dtype=model_kwargs["torch_dtype"]: Sets the data type that will be used for computation during the forward and backward passes. Here, it's set to bfloat16, meaning while the weights are stored in 4-bit, the actual calculations will happen in bfloat16 for better numerical stability and performance.
• bnb_4bit_quant_storage=model_kwargs["torch_dtype"]: Sets the data type for the storage of the quantized weights. This is also set to bfloat16 for consistency with computation.

**model = AutoModelForImageTextToText.from_pretrained(model_id, **model_kwargs)**

• This line loads the pre-trained Gemma model.
• AutoModelForImageTextToText.from_pretrained(model_id, ...): Downloads the model specified by model_id from the Hugging Face Hub.
• **model_kwargs: This unpacks all the key-value pairs from our model_kwargs dictionary and passes them as individual arguments to the from_pretrained function. This includes our attn_implementation, torch_dtype, device_map, and most importantly, the quantization_config.

**processor = AutoProcessor.from_pretrained("google/gemma-3-4b-it")**

• Loads the preprocessor (tokenizer and image processor) for the Gemma model.
• "google/gemma-3-4b-it": This model ID is specifically for the "instruction-tuned" version of Gemma's processor, which is often optimized for chat-like interactions and multimodal inputs.

Configuring LoRA Adapters

With the base model loaded and quantized, the next step is to define the LoraConfig. This configuration tells the peft library how to create and attach the small, trainable adapter layers.

from peft import LoraConfig
peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.05,
    r=16,
    bias="none",
    target_modules="all-linear",
    task_type="CAUSAL_LM",
    modules_to_save=[
        "lm_head",
        "embed_tokens",
    ],
)

Explanation of each line:

**from peft import LoraConfig**

• Imports the LoraConfig class from the peft (Parameter-Efficient Fine-Tuning) library. This class is used to define the parameters for the LoRA adapter layers.

**peft_config = LoraConfig(...)**

• Creates an instance of LoraConfig and assigns it to the peft_config variable. All the arguments inside define how the LoRA adapters will be structured and integrated.

**lora_alpha=16**

• lora_alpha is a scaling factor for the LoRA update. It's used in conjunction with r (LoRA rank) to scale the weight updates. A higher alpha can mean stronger adaptation. The QLoRA paper suggests specific values, and 16 is a common choice.

**lora_dropout=0.05**

• lora_dropout is the dropout probability applied to the LoRA layers. Dropout is a regularization technique that randomly "turns off" a fraction of neurons during training to prevent overfitting. A value of 0.05 means 5% of the LoRA layer connections will be randomly dropped during each training step.

**r=16**

• r is the "rank" of the LoRA update matrices. This is a critical hyperparameter in LoRA. A lower rank r means fewer parameters are added, leading to smaller adapter layers and more memory savings. A higher r means more expressive adapters but also more parameters. 16 is a commonly used value that balances performance and efficiency.

**bias="none"**

• This specifies whether bias parameters in the LoRA layers should be trained. "none" means no bias terms will be trained, which is common in LoRA as it often yields good results without adding extra parameters. Other options might include "all" (train all bias terms) or "lora_only" (train bias terms only in the LoRA layers).

**target_modules="all-linear"**

• This is a crucial parameter that determines which layers in the pre-trained model the LoRA adapters will be attached to.
• "all-linear" means LoRA adapters will be added to all linear (dense) layers found in the model. Linear layers are very common in transformer architectures, and applying LoRA to them is generally effective. You could also specify a list of specific layer names if you wanted more fine-grained control.

**task_type="CAUSAL_LM"**

• Specifies the type of task the model is performing. CAUSAL_LM stands for Causal Language Modeling, which is the task of predicting the next token in a sequence, given the preceding tokens. This is the underlying task for most text generation models, including Gemma.

**modules_to_save=["lm_head", "embed_tokens"]**

• This is an important parameter for models where certain parts of the original model (even if frozen) are critical for correct behavior and should be saved alongside the LoRA adapters.
• lm_head: This is typically the final linear layer of a language model that projects the model's hidden states to the vocabulary size to make predictions for the next token. While the base model's weights are frozen, the lm_head might need to be "active" or considered during fine-tuning.
• embed_tokens: These are the embedding layers that convert input tokens (words or subword units) into numerical representations (vectors) that the model can process. Keeping these active or savable ensures proper input processing.
• By including these, we ensure that when we save our fine-tuned model, these critical components are also saved correctly, allowing for seamless inference later.

Defining Hyperparameters and Data Collator

Before training, we need to set various hyperparameters that control the training process and define a custom collate_fn (data collator) to properly prepare our multimodal data for the model.

from trl import SFTConfig
args = SFTConfig(
    output_dir="gemma-product-description",     # directory to save and repository id
    num_train_epochs=1,                         # number of training epochs
    per_device_train_batch_size=1,              # batch size per device during training
    gradient_accumulation_steps=4,              # number of steps before performing a backward/update pass
    gradient_checkpointing=True,                # use gradient checkpointing to save memory
    optim="adamw_torch_fused",                  # use fused adamw optimizer
    logging_steps=5,                            # log every 5 steps
    save_strategy="epoch",                      # save checkpoint every epoch
    learning_rate=2e-4,                         # learning rate, based on QLoRA paper
    bf16=True,                                  # use bfloat16 precision
    max_grad_norm=0.3,                          # max gradient norm based on QLoRA paper
    warmup_ratio=0.03,                          # warmup ratio based on QLoRA paper
    lr_scheduler_type="constant",               # use constant learning rate scheduler
    push_to_hub=True,                           # push model to hub
    report_to="tensorboard",                    # report metrics to tensorboard
    gradient_checkpointing_kwargs={
        "use_reentrant": False
    },  # use reentrant checkpointing
    dataset_text_field="",                      # need a dummy field for collator
    dataset_kwargs={"skip_prepare_dataset": True},  # important for collator
)
args.remove_unused_columns = False # important for collator
# Create a data collator to encode text and image pairs
def collate_fn(examples):
    texts = []
    images = []
    for example in examples:
        image_inputs = process_vision_info(example["messages"])
        text = processor.apply_chat_template(
            example["messages"], add_generation_prompt=False, tokenize=False
        )
        texts.append(text.strip())
        images.append(image_inputs)
    # Tokenize the texts and process the images
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
    # The labels are the input_ids, and we mask the padding tokens and image tokens in the loss computation
    labels = batch["input_ids"].clone()
    # Mask image tokens
    image_token_id = [
        processor.tokenizer.convert_tokens_to_ids(
            processor.tokenizer.special_tokens_map["boi_token"]
        )
    ]
    # Mask tokens for not being used in the loss computation
    labels[labels == processor.tokenizer.pad_token_id] = -100
    labels[labels == image_token_id] = -100
    labels[labels == 262144] = -100 # This specific ID is likely another special image-related token for Gemma.
    batch["labels"] = labels
    return batch

Explanation of SFTConfig (Hyperparameters):

**from trl import SFTConfig**

• Imports the SFTConfig class, which is used to define the training arguments for the SFTTrainer.

**args = SFTConfig(...)**

• Creates an instance of SFTConfig and assigns it to the args variable.

**output_dir="gemma-product-description"**

• Specifies the directory where the model checkpoints (saved states of the model during training) and the final trained model will be saved. This name will also be used as the repository ID when pushing to the Hugging Face Hub.

**num_train_epochs=1**

• Sets the number of training epochs. An “epoch” means one complete pass through the entire training dataset. For fine-tuning, often 1 to 3 epochs are sufficient.

**per_device_train_batch_size=1**

• Defines the batch size per GPU (or CPU) during training. A batch is a group of samples processed together. A smaller batch size requires less memory. Here, it’s set to 1, indicating a very small batch size, likely due to memory constraints with large models.

**gradient_accumulation_steps=4**

• This is a technique to simulate a larger batch size when per_device_train_batch_size is small due to memory limits. The gradients are accumulated over 4 steps (meaning 4 batches of size 1 are processed), and then a single optimization step is performed. This effectively creates a "virtual" batch size of 1 * 4 = 4.

**gradient_checkpointing=True**

• Enables gradient checkpointing. This is a memory-saving technique where instead of storing all intermediate activations during the forward pass (which takes a lot of memory), only a subset are stored. The remaining activations are recomputed during the backward pass. This trades computation time for memory reduction.

**optim="adamw_torch_fused"**

• Specifies the optimizer to use. adamw_torch_fused refers to a fused (highly optimized) implementation of the AdamW optimizer from PyTorch. AdamW is a popular and effective optimizer for training neural networks.

**logging_steps=5**

• Determines how often training metrics (like loss) are logged. Here, metrics will be logged every 5 steps (batches).

**save_strategy="epoch"**

• Defines when checkpoints are saved. "epoch" means the model will be saved at the end of each training epoch.

**learning_rate=2e-4**

• Sets the learning rate, which controls the size of the steps taken by the optimizer during training. 2e-4 (0.0002) is a common learning rate for QLoRA, often derived from research papers.

**bf16=True**

• Enables bfloat16 (Brain Floating Point 16-bit) precision for training. This typically requires compatible GPU hardware (compute capability 8.0 or higher) and can significantly reduce memory usage and speed up training compared to complete 32-bit precision, while generally maintaining performance.

**max_grad_norm=0.3**

• Sets the maximum gradient norm for gradient clipping. Gradient clipping is a technique used to prevent “exploding gradients” (where gradients become extremely large, leading to unstable training). If the norm of the gradients exceeds 0.3, they are scaled down. This value is also often recommended in QLoRA papers.

**warmup_ratio=0.03**

• Sets the proportion of training steps during which the learning rate will gradually increase from 0 to its full value (learning_rate). This "warmup" period helps stabilize training at the beginning. 0.03 means 3% of the total training steps will be used for warming up.

**lr_scheduler_type="constant"**

• Specifies the learning rate scheduler. A “constant” scheduler means the learning rate will remain constant after the warmup period. Other schedulers might reduce the learning rate over time.

**push_to_hub=True**

• If True, the trained model (or its adapters) will be automatically pushed to the Hugging Face Hub under the output_dir name. This requires a Hugging Face token with write access.

**report_to="tensorboard"**

• Specifies the integration for logging training metrics. Here, metrics will be reported to TensorBoard, allowing for visual tracking of the training progress.

**gradient_checkpointing_kwargs={"use_reentrant": False}**

• Additional keyword arguments for gradient checkpointing. use_reentrant: False is often recommended for better compatibility and memory usage with some PyTorch versions and model architectures when using gradient checkpointing.

**dataset_text_field=""**

• A dummy field is needed by the SFTTrainer when using a custom data_collator that handles complex input structures (like our multimodal messages). Since our collate_fn handles everything, this field is set to an empty string.

**dataset_kwargs={"skip_prepare_dataset": True}**

• Another argument to tell the SFTTrainer that we are providing an already prepared dataset and a custom collate_fn, so it should skip its own default dataset preparation steps. This is important when you have specific data formatting needs as we do with multimodal input.**args.remove_unused_columns = False**
• By default, Trainer might try to remove columns from the dataset that are not explicitly used by the model's forward pass. Setting this to False is crucial when using a custom data_collator that relies on specific column names (like "messages" in our case) that the Trainer might not automatically recognize as "used" by the model directly, but are indeed processed by the collator.

Custom Data Collator (collate_fn)

The collate_fn is a function that takes a list of dataset samples (a batch) and prepares them into a format that the model can directly consume. For multimodal data, this is especially important as it handles both text tokenization and image processing.

**def collate_fn(examples):**

• Defines the custom collation function. examples will be a list of dictionaries, where each dictionary is one of our formatted multimodal messages from the dataset.

**texts = []**

• Initializes an empty list to store the extracted text strings from the batch.

**images = []**

• Initializes an empty list to store the extracted PIL.Image objects from the batch.

**for example in examples:**

• Iterates through each conversation (sample) in the current batch.

**image_inputs = process_vision_info(example["messages"])**

• Calls our previously defined process_vision_info utility function. This extracts all PIL.Image objects from the messages list of the current example.

**text = processor.apply_chat_template( ... )**

• This is a crucial step for preparing the text input.
• processor.apply_chat_template(example["messages"], add_generation_prompt=False, tokenize=False):
• example["messages"]: Takes the structured TRL multimodal conversation format.
• add_generation_prompt=False: We don't want to add an extra prompt for generation, as our user_prompt already guides the generation.
• tokenize=False: Very important! We do not want to tokenize the text here yet. We want to get the raw string representation of the entire conversation (system, user, and assistant parts) so that the processor can tokenize both text and images together in the next step.

**texts.append(text.strip())**

• Adds the processed text string (with leading/trailing whitespace removed using .strip()) to the texts list.

**images.append(image_inputs)**

• Adds the list of extracted PIL.Image objects for the current sample to the images list.

**batch = processor(text=texts, images=images, return_tensors="pt", padding=True)**

• This is where the actual tokenization of text and processing of images happens together.
• processor(...): The AutoProcessor instance (for Gemma) is called.
• text=texts: Provides the list of text strings (the formatted conversations).
• images=images: Provides the list of lists of PIL.Image objects. The processor handles associating the images with their corresponding textual positions.
• return_tensors="pt": Tells the processor to return PyTorch tensors.
• padding=True: Ensures that all sequences in the batch are padded to the same length so they can be processed efficiently by the model.

**labels = batch["input_ids"].clone()**

• For language modeling tasks, the labels (the target output for the model to predict) are often the same as the input_ids (the input sequence), but with specific tokens masked out so the model only predicts the "completion" part.
• .clone(): Creates a copy of the input_ids tensor. This is important because we will modify labels by masking, and we don't want to alter the original input_ids.

Masking Tokens for Loss Computation:

• The goal is to calculate the loss only on the assistant’s (desired output) part of the conversation, and also to ignore any special tokens like padding or image tokens. This prevents the model from being penalized for reproducing the prompt or padding, and ensures it focuses on generating the correct response.

image_token_id = [processor.tokenizer.convert_tokens_to_ids(processor.tokenizer.special_tokens_map["boi_token"])]

• Retrieves the token ID for the “beginning of image” (boi_token) special token. Gemma uses special tokens to denote the presence and boundaries of images within the text sequence.

labels[labels == processor.tokenizer.pad_token_id] = -100

• Sets the labels corresponding to padding tokens (pad_token_id) to -100. In Hugging Face transformers, an ignore_index of -100 for the loss function means that these positions will be ignored when calculating the loss.

labels[labels == image_token_id] = -100

• Sets the labels corresponding to the boi_token (and implicitly other image-related tokens if they share this ID or are handled similarly by the processor) to -100, so image tokens are also ignored in the loss calculation.

labels[labels == 262144] = -100

• This specific numerical ID (262144) is likely another special token used by Gemma for images or other structural elements. It's explicitly masked to ensure it's not part of the loss calculation.

**batch["labels"] = labels**

• Adds the masked labels tensor to the batch dictionary under the key "labels". This batch dictionary is what will be passed to the model for training.

Initializing and Starting the SFTTrainer

Finally, with all the pieces in place, we can create the SFTTrainer and begin the fine-tuning process.

from trl import SFTTrainer
trainer = SFTTrainer(
    model=model,
    args=args,
    train_dataset=dataset,
    peft_config=peft_config,
    processing_class=processor,
    data_collator=collate_fn,
)
# Start training, the model will be automatically saved to the Hub and the output directory
trainer.train()
# Save the final model again to the Hugging Face Hub
trainer.save_model()

Explanation of each line:

**from trl import SFTTrainer**

• Imports the SFTTrainer class.

**trainer = SFTTrainer(...)**

• Instantiates the SFTTrainer object.
• model=model: Passes our loaded and quantized Gemma model to the trainer.
• args=args: Provides all the training hyperparameters defined in our SFTConfig object.
• train_dataset=dataset: Provides our prepared and formatted training dataset.
• peft_config=peft_config: Passes the LoRA configuration. The SFTTrainer will use this to set up the PEFT (LoRA) layers on top of the base model.
• processing_class=processor: Provides the AutoProcessor instance. The SFTTrainer uses this to handle tokenization and image processing within its internal mechanisms, complementing our custom data_collator.
• data_collator=collate_fn: Tells the SFTTrainer to use our custom collate_fn to prepare batches of data.

**trainer.train()**

• This is the magical line that starts the entire fine-tuning process! The trainer will iterate through the dataset for the specified number of epochs, perform forward and backward passes, update the LoRA adapter weights, log metrics, and save checkpoints according to the configurations we set in SFTConfig.

**trainer.save_model()**

• After the training is complete, this line explicitly saves the final state of the fine-tuned model (specifically, the LoRA adapter weights) to the output_dir and, if push_to_hub=True was set in SFTConfig, it will also push these adapters to the Hugging Face Hub.

Freeing Memory

It’s good practice to free up GPU memory after training, especially when working with large models.

# free the memory again
del model
del trainer
torch.cuda.empty_cache()

Explanation:

• del model: Deletes the Python object representing the model, releasing its memory.
• del trainer: Deletes the trainer object, releasing its memory.
• torch.cuda.empty_cache(): This explicitly clears PyTorch's CUDA memory cache. Even after deleting objects, PyTorch might hold onto allocated GPU memory in a cache for faster future allocations. This command forces it to release that cached memory back to the GPU.

Merging LoRA Adapters into the Base Model (Optional but Recommended for Inference)

When using QLoRA, you only train the small adapter layers. For convenient deployment and inference (especially with serving frameworks like vLLM or TGI), it’s often desirable to have a single, self-contained model file where the adapter weights are merged into the original base model.

Important Note: Merging adapters requires loading the original full base model into CPU memory, which can be substantial (more than 30GB as stated). If you don’t have enough CPU RAM, you can skip this step and directly proceed to “Test Model Inference” by loading the adapters on top of the base model.

from peft import PeftModel
# Load Model base model
model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)
# Merge LoRA and base model and save
peft_model = PeftModel.from_pretrained(model, args.output_dir)
merged_model = peft_model.merge_and_unload()
merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")
processor = AutoProcessor.from_pretrained(args.output_dir)
processor.save_pretrained("merged_model")

Explanation of each line:

**from peft import PeftModel**

• Imports the PeftModel class from the peft library. This class is used to load a base model and apply PEFT (LoRA) adapters on top of it.

**model = AutoModelForImageTextToText.from_pretrained(model_id, low_cpu_mem_usage=True)**

• Loads the original, unquantized base Gemma model again.
• model_id: The same ID used previously for the base model.
• low_cpu_mem_usage=True: This argument can help by trying to load the model in a more memory-efficient way (e.g., directly to GPU if possible, or by sharding).

**peft_model = PeftModel.from_pretrained(model, args.output_dir)**

• Creates a PeftModel instance. This essentially loads our trained LoRA adapter weights (which were saved in args.output_dir) and applies them on top of the newly loaded model (the original base model). At this point, peft_model is a representation of the base model with the adapters active.

**merged_model = peft_model.merge_and_unload()**

• This is the key step for merging.
• merge_and_unload(): This method takes the weights from the LoRA adapter layers and mathematically merges them directly into the corresponding layers of the base model. After this operation, the adapter layers themselves are no longer separate; their effects are integrated into the base model's weights. The method also "unloads" the adapter-specific logic, resulting in a standard transformers model. The result is a single, "merged" model that behaves like the fine-tuned model without needing separate adapter files.

**merged_model.save_pretrained("merged_model", safe_serialization=True, max_shard_size="2GB")**

• Saves the merged_model to a new directory named "merged_model".
• safe_serialization=True: This uses a more robust serialization format (often safetensors) which is generally safer, faster, and more secure than traditional PyTorch pickle files.
• max_shard_size="2GB": If the model is very large, this argument tells save_pretrained to shard (split) the model into multiple files, each no larger than 2GB. This makes it easier to download and load large models, especially on systems with limited memory or specific file size limits.

**processor = AutoProcessor.from_pretrained(args.output_dir)**

• Reloads the processor. While the processor itself isn’t “merged,” it’s good practice to load it from the fine-tuned model’s output directory, as it might contain any special tokens or configurations added during training.

**processor.save_pretrained("merged_model")**

• Saves the reloaded processor to the same “merged_model” directory. This ensures that when you load the merged_model for inference later, you can also load its associated processor from the same location, guaranteeing compatibility.

Test Model Inference and Generate Product Descriptions

After all the hard work of fine-tuning, it’s time to test our model! This section shows you how to load the fine-tuned model and use it to generate product descriptions for new inputs.

Evaluating Generative AI Models: It’s important to note that evaluating generative AI models is complex. There isn’t always one “correct” answer. For instance, multiple product descriptions could be considered good. This tutorial focuses on a “manual evaluation” or “vibe check,” which means you’ll visually inspect the generated outputs to see if they meet your expectations. For more rigorous evaluation, you would typically use human annotators or specific metrics like ROUGE, BLEU, or more advanced NLG metrics (For more information of evaluation methods this blog post).

Loading the Fine-tuned Model for Inference

We’ll load the model with the PEFT adapters attached.

import torch
# Load Model with PEFT adapter
model = AutoModelForImageTextToText.from_pretrained(
  args.output_dir, # Load from the directory where adapters were saved
  device_map="auto",
  torch_dtype=torch.bfloat16,
  attn_implementation="eager",
)
processor = AutoProcessor.from_pretrained(args.output_dir)

Explanation of each line:

**import torch**

• Ensures PyTorch is available.

**model = AutoModelForImageTextToText.from_pretrained(...)**

• Loads the pre-trained base model again, but this time, it also loads our fine-tuned LoRA adapters directly on top of it.
• args.output_dir: This is crucial! It tells the from_pretrained method where to find the saved LoRA adapter weights. The transformers library, when used with peft, is smart enough to detect these adapters in the specified directory and automatically load and apply them to the base model.
• device_map="auto": Again, automatically maps the model components to available devices for optimal memory usage.
• torch_dtype=torch.bfloat16: Loads the model in bfloat16 precision for efficient inference, assuming your GPU supports it.
• attn_implementation="eager": Specifies the attention implementation. (Again, flash_attention_2 could be considered for compatible GPUs).

**processor = AutoProcessor.from_pretrained(args.output_dir)**

• Loads the processor from the same args.output_dir. This ensures you are using the processor that was configured and saved alongside your fine-tuned adapters, maintaining compatibility, especially if any special tokens or configurations were modified during training.

Preparing a Test Sample

To test the model, we need a sample input that mimics the structure of our training data (product name, category, and an image).

import requests
from PIL import Image
# Test sample with Product Name, Category and Image
sample = {
  "product_name": "Hasbro Marvel Avengers-Serie Marvel Assemble Titan-Held, Iron Man, 30,5 cm Actionfigur",
  "category": "Toys & Games | Toy Figures & Playsets | Action Figures",
  "image": Image.open(requests.get("https://m.media-amazon.com/images/I/81+7Up7IWyL._AC_SY300_SX300_.jpg", stream=True).raw).convert("RGB")
}

Explanation of each line:

**import requests**

• Imports the requests library, used for making HTTP requests (in this case, to download an image from a URL).

**from PIL import Image**

• Imports the Image module from Pillow for image handling.

**sample = { ... }**

• Creates a dictionary named sample to hold our test data.
• "product_name": The name of the product for which we want to generate a description.
• "category": The category of the product.
• "image": This is where the image for the product is loaded.
• requests.get("https://m.media-amazon.com/images/I/81+7Up7IWyL._AC_SY300_SX300_.jpg", stream=True).raw: Makes an HTTP GET request to the provided URL to fetch the image data. stream=True allows efficient streaming of binary data, and .raw gets the raw bytes.
• Image.open(...): Opens the raw image data using Pillow.
• .convert("RGB"): Converts the image to the RGB color mode, ensuring it's in a standard format expected by the model.

Generating the Description

This function encapsulates the entire inference process, from preparing the input to generating and decoding the output.

def generate_description(sample, model, processor):
    # Convert sample into messages and then apply the chat template
    messages = [
        {"role": "system", "content": [{"type": "text", "text": system_message}]},
        {"role": "user", "content": [
            {"type": "image","image": sample["image"]},
            {"type": "text", "text": user_prompt.format(product=sample["product_name"], category=sample["category"])},
        ]},
    ]
    text = processor.apply_chat_template(
        messages, tokenize=False, add_generation_prompt=True
    )
    # Process the image and text
    image_inputs = process_vision_info(messages)
    # Tokenize the text and process the images
    inputs = processor(
        text=[text],
        images=image_inputs,
        padding=True,
        return_tensors="pt",
    )
    # Move the inputs to the device
    inputs = inputs.to(model.device)
    # Generate the output
    stop_token_ids = [processor.tokenizer.eos_token_id, processor.tokenizer.convert_tokens_to_ids("<end_of_turn>")]
    generated_ids = model.generate(**inputs, max_new_tokens=256, top_p=1.0, do_sample=True, temperature=0.8, eos_token_id=stop_token_ids, disable_compile=True)
    # Trim the generation and decode the output to text
    generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]
    output_text = processor.batch_decode(
        generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return output_text[0]
# generate the description
description = generate_description(sample, model, processor)
print(description)

Explanation of generate_description function:

**def generate_description(sample, model, processor):**

• Defines a function that takes our sample data, the loaded model, and the processor as input.

**messages = [ ... ]**

• Reconstructs the TRL multimodal messages format for the input, similar to how we formatted the training data.
• system_message: The same system instruction.
• user_prompt.format(...): Fills the user prompt template with the product_name and category from the new sample.
• "image": sample["image"]: Includes the image from the new sample.

**text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)**

• This converts the messages list into a single string that includes the special tokens and formatting expected by the Gemma model for chat-based inputs.
• add_generation_prompt=True: For inference, it's often beneficial to explicitly add a generation prompt (e.g., <start_of_turn>assistant) at the end of the input sequence. This encourages the model to start generating its response (the "assistant" turn).
• tokenize=False: We still want the raw string here before full processing.

**image_inputs = process_vision_info(messages)**

• Extracts the PIL.Image objects from the messages list, just like in the collate_fn.

**inputs = processor(text=[text], images=image_inputs, padding=True, return_tensors="pt")**

• This is the same processor call used in the collate_fn. It takes the raw text string and image objects, tokenizes the text, processes the images, pads sequences, and returns PyTorch tensors ready for the model.

**inputs = inputs.to(model.device)**

• Moves the input tensors to the same device (e.g., GPU) where the model is loaded. This is essential for the model to perform computations.

**generated_ids = model.generate(...)**

• This is the core generation call!
• **inputs: Unpacks the dictionary of input tensors (like input_ids, attention_mask, pixel_values) and passes them as arguments to the generate method.
• max_new_tokens=256: Sets the maximum number of new tokens the model should generate. This prevents excessively long outputs.
• top_p=1.0: (Nucleus sampling) top_p is a parameter for text generation that controls the diversity of the generated text. A top_p of 1.0 means that all tokens with a cumulative probability up to 1.0 are considered. This generally leads to more diverse and less repetitive text.
• do_sample=True: Enables sampling-based generation. If False, the model would use deterministic decoding (like greedy search). Sampling introduces randomness, making the generated text more creative and less repetitive.
• temperature=0.8: Controls the randomness of the sampling process. A higher temperature (e.g., 1.0) leads to more random and diverse output, while a lower temperature (e.g., 0.1) makes the output more deterministic and focused. 0.8 is a common value for a balance.
• eos_token_id=stop_token_ids: Specifies a list of token IDs that, if generated, will cause the generation process to stop.
• processor.tokenizer.eos_token_id: The "end-of-sequence" token ID.
• processor.tokenizer.convert_tokens_to_ids("<end_of_turn>"): The token ID for Gemma's special "end of turn" token, which signifies the end of a conversational turn.
• disable_compile=True: Disables Torch's torch.compile optimization during generation. While torch.compile can speed up inference, it can sometimes introduce complexities or incompatibilities, especially during initial setup or with specific model configurations. Disabling it ensures smooth operation.

**generated_ids_trimmed = [out_ids[len(in_ids) :] for in_ids, out_ids in zip(inputs.input_ids, generated_ids)]**

• This line removes the input prompt from the generated sequence. The generate method returns the entire sequence (input + generated). We only want the generated part.
• It zips the original input_ids with the generated_ids. For each pair, it slices the generated_ids starting from the length of the input_ids, effectively keeping only the newly generated tokens.

**output_text = processor.batch_decode(generated_ids_trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False)**

• Decodes the generated token IDs back into human-readable text.
• skip_special_tokens=True: Ignores any special tokens (like padding, BOS, EOS) during decoding, so they don't appear in the final output string.
• clean_up_tokenization_spaces=False: Prevents the decoder from aggressively cleaning up spaces, which can sometimes remove desired spacing, especially for specific model outputs.

**return output_text[0]**

• Returns the first (and likely only) generated description from the batch.

**description = generate_description(sample, model, processor)**

• Calls our function with the defined sample, model, and processor to get the generated description.

**print(description)**

• Prints the generated product description to the console.

Final Word:

This step-by-step guide is designed to help beginners confidently fine-tune a multimodal LLM like Gemma using QLoRA — from setting up the environment to preparing data, training the model, and running inference. The code and methodology are based on the official tutorial provided by Google at ai.google.dev/gemma.

Disclaimer:

Opinions and views expressed here are solely my own and do not reflect those of my current or past employers or any affiliations.