Kesteral
Back home

Fine-tuning

Fine-Tune Llama 3.2 Vision-Language Model on Custom Datasets

Mar 20, 2024 · 9 min read

Llama 3.2, a powerful multimodal large language model (LLM) from Meta AI, pushes the boundaries of AI by letting machines understand both visual and textual information. The pre-trained model is impressive out of the box, but fine-tuning it on your own dataset can significantly improve its performance for your use case.

In this guide, we walk through fine-tuning the Llama 3.2 Vision-Language Model (VLM) on a custom dataset, covering everything from setting up your environment to testing the fine-tuned model.

Key steps in fine-tuning Llama 3.2 VLM

  • Define your use case.
  • Set up the development environment.
  • Prepare the dataset.
  • Fine-tune the VLM using TRL and SFTTrainer.
  • Test the fine-tuned model.

1. Defining your multimodal use case

Before fine-tuning, be clear about the problem you are solving. For VLMs, that usually means tasks that combine visual and textual data. As an example, imagine an e-commerce platform where sellers upload product images and metadata, and you need a system that automatically writes detailed product descriptions from that input.

Why fine-tune? Pre-trained VLMs may perform well, but they may not fully capture the unique attributes of your dataset or use case. Fine-tuning adapts the model to your specific needs, such as generating SEO-optimized product descriptions or handling advanced image-based tasks.

A note on hardware: Llama 3.2 11B Vision needs at least 24 GB of GPU memory for efficient training. 32 GB or 40 GB is better, especially for large batch sizes or complex datasets.

2. Setting up the development environment

Install the necessary libraries first. The stack includes PyTorch 2.4.0, transformers 4.45.1, datasets 3.0.1, accelerate 0.34.2, evaluate 0.4.3, bitsandbytes 0.44.0, trl 0.11.1, peft 0.13.0, and qwen_vl_utils.

Next, log in to Hugging Face to access the model and dataset with the login helper: from huggingface_hub import login, then login(token="YOUR_HF_TOKEN"). When you create your API key, agree to share your contact information so you can access the Llama 3.2 vision model.

3. Preparing the dataset

A high-quality dataset is crucial for effective fine-tuning. For this tutorial we use the Amazon Product Descriptions VLM dataset created by philschmid, built for fine-tuning vision-language models on e-commerce description tasks. It pairs product images with product names, categories, and descriptions; images are JPEGs accessed via URLs.

We load and prepare the data with the Hugging Face datasets library. The code defines a prompt template for generating descriptions, sets a system message that positions the model as an expert product writer, and adds a format_data function that structures each sample into system, user, and assistant messages — with the image URL placed in the user content and the product name and category inserted into the prompt.

The format_data function handles the main preprocessing: it builds a conversation-like structure of system, user, and assistant messages, includes the image URLs in the user message, and inserts the product names and categories into the prompt template.

This dataset does not need heavy augmentation, but for other use cases you might consider random cropping or resizing of images, text augmentation such as synonym replacement or random insertion and deletion, or generating extra descriptions with other LLMs. These techniques help the model generalize, especially on smaller datasets.

4. Fine-tuning with TRL, SFTTrainer, and Unsloth

For fine-tuning we use Unsloth, a library that optimizes LLM training. Unsloth can speed up the process and cut memory usage, which makes it easier to fine-tune large models like Llama 3.2 Vision. Install it with the Unsloth package, then initialize the model.

We load the Llama 3.2 11B Vision model through Unsloth's FastLanguageModel with a 2048-token sequence length, bfloat16 precision, and 4-bit quantization. The quantization and bfloat16 reduce memory use and can speed up training.

Next we set up LoRA (Low-Rank Adaptation) with alpha 16, dropout 0.05, rank 8, targeting the q_proj and v_proj modules. LoRA, implemented through the PEFT (Parameter-Efficient Fine-Tuning) library, sharply reduces the number of trainable parameters, which makes fine-tuning large models far more efficient.

The training configuration uses an SFTConfig with 3 epochs, batch size 4, gradient accumulation over 8 steps, gradient checkpointing, the AdamW optimizer, a 2e-4 learning rate, bfloat16 precision, and TensorBoard logging. Gradient checkpointing and accumulation keep memory use in check during training.

Finally, a collate_fn applies the chat template and processes the vision data into batches. We build an SFTTrainer from TRL with the model, dataset, tokenizer, and LoRA config, apply Unsloth optimizations via FastLanguageModel.get_peft_model, and call trainer.train() to start fine-tuning. This setup makes it possible to fine-tune Llama 3.2 Vision on consumer-grade hardware.

5. Saving the fine-tuned model

After fine-tuning, save the PEFT weights. They hold only the changes made during fine-tuning, so they are much smaller than the full model and easy to store and share. Use trainer.model.save_pretrained() to export the weights to a directory.

Save the tokenizer to the same directory with tokenizer.save_pretrained() so you can tokenize inputs correctly at inference time. Optionally, push both to the Hugging Face Hub with HfApi.upload_folder(), replacing the repo path with your own username and model name.

6. Loading the fine-tuned model for inference

To run inference, load the base Llama 3.2 model from Meta, apply the saved PEFT weights, and import the matching tokenizer and processor. Put the model in evaluation mode and move it to the available GPU or CPU.

For an example, fetch a product image via URL, build a prompt with the product name and category, process the inputs, and generate up to 100 tokens with sampling before decoding the output. The steps are: load the image, prepare the prompt, process the input, generate the description, then decode and print it.

For a sample chair, the fine-tuned model produced a detailed, SEO-friendly description highlighting adjustable lumbar support, a breathable mesh backrest, and customizable armrests — showing its ability to write compelling e-commerce content from an image and basic product information.

Conclusion

We walked through fine-tuning the Llama 3.2 Vision-Language Model for product description generation: setting up the environment, preparing the Amazon Product Descriptions VLM dataset, fine-tuning with LoRA and Unsloth, saving the weights and tokenizer, and loading the model for inference. You now have a custom-tuned VLM that writes SEO-optimized descriptions from images and basic product data.

Remember that fine-tuning is iterative. Experiment with different hyperparameters, dataset sizes, or model architectures to get the best results for your use case.

Originally published on the FutureSmart AI blog.

Have an idea to build?

One engineer. Your idea live in 21 days. You own the code.

Book a call