Training a small VLM

TST, HongKong
Locally Hosted Annotation Model
Annotation Pipeline
Python Code

import os
import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoProcessor, AutoModelForVision2Seq, TrainingArguments
from trl import SFTTrainer

# ==========================================
# CONFIGURATION
# ==========================================
MODEL_ID = "HuggingFaceTB/SmolVLM-256M-Instruct"
JSONL_DATASET_PATH = "vlm_training_dataset.jsonl"
OUTPUT_DIR = "./smolvlm_finetuned"
BATCH_SIZE = 2             # Keep low for consumer hardware; scale up if VRAM allows
GRAD_ACCUM_STEPS = 4       # Simulates a larger batch size (2 * 4 = 8 effective batch size)
NUM_EPOCHS = 3

# ==========================================
# 1. LOAD AND PREPARE DATASET
# ==========================================
# Hugging Face's datasets handles parsing the local JSONL file automatically
dataset = load_dataset("json", data_files=JSONL_DATASET_PATH, split="train")

# We must map our customized JSONL format to the exact dict structure Hugging Face expects
def format_conversations(example):
    formatted_messages = []
    # Translate our ShareGPT style ("from": "human"/"gpt", "value": "...") 
    # to Hugging Face standard conversation dictionaries
    for message in example["conversations"]:
        role = "user" if message["from"] == "human" else "assistant"
        formatted_messages.append({
            "role": role,
            "content": [{"type": "text", "text": message["value"]}]
        })
    return {"messages": formatted_messages, "images": [example["image"]]}

# Process dataset
formatted_dataset = dataset.map(format_conversations)

# Split into train and validation sets
split_dataset = formatted_dataset.train_test_split(test_size=0.1, seed=42)
train_dataset = split_dataset["train"]
eval_dataset = split_dataset["test"]

# ==========================================
# 2. INITIALIZE PROCESSOR & MODEL
# ==========================================
processor = AutoProcessor.from_pretrained(MODEL_ID)

# Configure LoRA to freeze base weights and train only small adapter layers (saves massive VRAM)
lora_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=["q_proj", "v_proj", "k_proj", "o_proj"], # Targets the visual and text attention layers
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM"
)

# Load base model in 8-bit precision (or use device_map="auto" for raw float16/bfloat16)
model = AutoModelForVision2Seq.from_pretrained(
    MODEL_ID,
    torch_dtype=torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16,
    device_map="auto"
)

# ==========================================
# 3. CONFIGURE TRAINING RUN
# ==========================================
training_args = TrainingArguments(
    output_dir=OUTPUT_DIR,
    num_train_epochs=NUM_EPOCHS,
    per_device_train_batch_size=BATCH_SIZE,
    gradient_accumulation_steps=GRAD_ACCUM_STEPS,
    learning_rate=2e-4,
    weight_decay=0.01,
    logging_steps=10,
    evaluation_strategy="steps",
    eval_steps=50,
    save_strategy="steps",
    save_steps=100,
    save_total_limit=2,
    fp16=not torch.cuda.is_bf16_supported(),
    bf16=torch.cuda.is_bf16_supported(),
    remove_unused_columns=False, # Crucial: Let TRL handle custom multimodal elements
    report_to="none"            # Change to "wandb" if you want live visualization
)

# ==========================================
# 4. EXECUTE TRAINING
# ==========================================
trainer = SFTTrainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,
    peft_config=lora_config,
    processor=processor, # SFTTrainer natively knows how to coordinate image/text tokens
)

print("Starting SmolVLM Fine-tuning...")
trainer.train()

# Save final adapter weights and processing configs
trainer.save_model(OUTPUT_DIR)
processor.save_pretrained(OUTPUT_DIR)
print(f"Training Complete! Finetuned weights successfully saved to: {OUTPUT_DIR}")
Python Code​

Python Code
