nanoVLM: A Minimalist PyTorch Toolkit for Training Vision Language Models

Hugging Face has introduced nanoVLM, a minimalist toolkit written in pure PyTorch designed to demystify the process of training Vision Language Models (VLMs). Inspired by Andrej Karpathy's nanoGPT, nanoVLM provides a readable, lightweight codebase that allows users to train a VLM on a free-tier Google Colab notebook.

Vision Language Model (VLM) Fundamentals

A Vision Language Model is a multi-modal architecture capable of processing both image and text inputs to generate text outputs. While VLMs can be used for diverse tasks such as image captioning, object detection, and semantic segmentation, nanoVLM specifically focuses on Visual Question Answering (VQA) as its primary training objective.

Technical Architecture

nanoVLM utilizes a modular architecture that aligns two pre-trained backbones through a projection layer:

  • Vision Backbone: Uses Google's SigLIP (google/siglip-base-patch16-224) vision encoder.
  • Language Backbone: Follows the Llama 3 architecture, specifically utilizing HuggingFaceTB/SmolLM2-135M.
  • Modality Projection Module: This module aligns the vision and text modalities. It transforms image embeddings from the vision backbone into embeddings compatible with the language model's embedding layer. This process involves a pixel shuffle operation—which reduces the number of image tokens to lower computational costs and increase training speed—followed by a linear layer.

Training Workflow and Implementation

The training process is managed via train.py, which handles dataset loading, model initialization, and optimization.

Data Pipeline

The get_dataloaders function leverages the Hugging Face load_dataset API to load, shuffle, and split datasets. It utilizes custom datasets (VQADataset, MMStarDataset) and collators (VQACollator, MMStarCollator) to prepare the data.

Optimization Strategy

To balance the training of pre-trained backbones with a newly initialized projector, nanoVLM employs a dual learning rate (LR) strategy:

  • Higher LR: Applied to the Modality Projector (MP) to facilitate rapid learning.
  • Lower LR: Applied to the encoder/decoder stack to preserve the pre-existing knowledge in the backbones.

Training Loop and Monitoring

Training incorporates torch.autocast for mixed precision and a cosine learning rate schedule with linear warmup. Performance is monitored via token throughput (tokens/sec) and, if enabled, Weights & Biases (wandb) for tracking batch loss, validation loss, and accuracy.

Inference and Pre-trained Model

Hugging Face has provided a pre-trained nanoVLM model (nanoVLM-222M) published to the Hub. This model was trained for approximately 6 hours on a single H100 GPU using 1.7 million samples from the cauldron dataset.

Users can run inference using the generate.py script, which follows this logical flow:

  1. Initialization: Load the model, tokenizer, and image processor.
  2. Processing: Tokenize the text prompt and process the image into a tensor.
  3. Generation: Execute model.generate to produce text output.
  4. Decoding: Convert the generated tokens back into human-readable text using batch_decode.

Getting Started

To begin training, users can clone the repository and run the training script:

# Clone the repo
git clone https://github.com/huggingface/nanoVLM.git

# Execute the training script
python train.py

Sources