+1 (726) 227-4060

Fine-Tuning a Vision-Language Model for Document Extraction

Most teams that need "AI on images" in 2026 no longer train a classifier. They take an open vision-language model (VLM), show it a few thousand examples of their own documents, screenshots or product photos, and fine-tune it to emit exactly the structure their system needs. It is faster than building a bespoke CNN pipeline, it survives layout changes that break template-based OCR, and it runs on a single GPU.

This tutorial fine-tunes an open VLM on a document-extraction task with PyTorch, LoRA and TRL, then evaluates and serves it. The recipe transfers directly to visual QA, screenshot understanding, defect inspection and image captioning.

1. Pick a model, and check what "open" means

The current open VLM families worth starting from are Qwen3-VL (2B through 30B+, strong OCR and document understanding, Apache-2.0), Llama 3.2 Vision (11B/90B, community licence with a usage threshold), Gemma 3 (4B/12B/27B multimodal, custom licence), InternVL and Pixtral. For document and screenshot work, start with a Qwen-VL variant in the 2-8B range: it fits on one 24 GB GPU with LoRA, and the accuracy gap to the 30B tier is usually smaller than the gap between a good and a bad dataset.

Two licence questions before you spend an engineer-week: may you use outputs commercially, and may you use outputs to train other models? The answers differ across these families.

Architecturally they are all the same shape: a vision encoder produces patch embeddings, a projector maps them into the language model's embedding space, and a decoder-only LLM does the rest. That matters because it tells you what to train. Almost always: freeze the vision encoder, train LoRA adapters on the language model's attention and MLP projections, and optionally train the projector. Fine-tuning the vision encoder on a few thousand examples usually degrades it.

2. Build the dataset as chat messages

VLM fine-tuning data is conversational. Each example is an image plus a user instruction plus the exact output you want. Be ruthless about output format: if downstream code parses JSON, every target must be valid JSON with the same key order and the same behaviour for missing fields.

def to_example(record):
    return {
        "messages": [
            {"role": "user", "content": [
                {"type": "image", "image": record["image_path"]},
                {"type": "text",  "text": "Extract the invoice as JSON with keys: "
                                          "vendor, invoice_number, date, currency, total. "
                                          "Use null for missing fields."},
            ]},
            {"role": "assistant", "content": [
                {"type": "text", "text": json.dumps(record["target"], ensure_ascii=False)},
            ]},
        ]
    }

Practical dataset rules we apply on client work:

  • 500 examples is a pilot, 2,000-5,000 is a product. Below a few hundred you are measuring noise.
  • Hold out by document source, not at random. Random splits leak layouts between train and test and flatter your numbers.
  • Include the hard cases deliberately: rotated scans, multi-page, handwriting, missing fields, two invoices on one page. A model that never saw a missing field will hallucinate one.
  • Cap image resolution. VLMs tokenize images into hundreds or thousands of visual tokens; a 4000px scan can cost more tokens than the text answer. Most processors expose min_pixels / max_pixels or a resolution argument. Setting a sane ceiling is often the single biggest throughput win of the whole run.

3. Load the model and attach LoRA

import torch
from transformers import AutoProcessor, AutoModelForImageTextToText
from peft import LoraConfig

MODEL = "Qwen/Qwen2.5-VL-7B-Instruct"  # or your chosen VLM

processor = AutoProcessor.from_pretrained(MODEL, min_pixels=256*28*28, max_pixels=1280*28*28)
model = AutoModelForImageTextToText.from_pretrained(
    MODEL, dtype=torch.bfloat16, attn_implementation="flash_attention_2", device_map="cuda:0"
)

peft_config = LoraConfig(
    r=32, lora_alpha=64, lora_dropout=0.05, bias="none", task_type="CAUSAL_LM",
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                    "gate_proj", "up_proj", "down_proj"],
)

Note what is not in target_modules: anything inside the vision tower. If memory is tight, load the base in 4-bit (BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)) and you are doing QLoRA, exactly as in Fine-Tuning Llama with QLoRA — the only VLM-specific change is keeping the vision encoder and projector out of the quantization skip-list debate by leaving them in BF16.

4. Train with TRL

TRL's SFTTrainer handles multimodal chat datasets, but you must supply a collator that runs the processor over images and text together and masks the prompt tokens from the loss.

from trl import SFTConfig, SFTTrainer

def collate(examples):
    texts, images = [], []
    for ex in examples:
        texts.append(processor.apply_chat_template(ex["messages"], tokenize=False))
        images.append(load_images(ex["messages"]))     # list of PIL images per example
    batch = processor(text=texts, images=images, return_tensors="pt", padding=True)
    labels = batch["input_ids"].clone()
    labels[labels == processor.tokenizer.pad_token_id] = -100
    for tid in processor.tokenizer.convert_tokens_to_ids(IMAGE_TOKENS):
        labels[labels == tid] = -100                   # never predict image placeholders
    batch["labels"] = labels
    return batch

args = SFTConfig(
    output_dir="out/vlm-invoices",
    num_train_epochs=2,
    per_device_train_batch_size=1,
    gradient_accumulation_steps=8,
    learning_rate=1e-4,
    lr_scheduler_type="cosine",
    warmup_ratio=0.03,
    bf16=True,
    gradient_checkpointing=True,
    gradient_checkpointing_kwargs={"use_reentrant": False},
    logging_steps=10,
    save_strategy="epoch",
    remove_unused_columns=False,          # required: images live in extra columns
    dataset_kwargs={"skip_prepare_dataset": True},
)

trainer = SFTTrainer(model=model, args=args, train_dataset=train_ds,
                     data_collator=collate, peft_config=peft_config)
trainer.train()

remove_unused_columns=False and skip_prepare_dataset=True are the two flags people miss; without them the trainer strips the image column and you train a text-only model on captions.

Smoke-test on 20 examples for 10 steps before launching the real run. A VLM fine-tune that dies at hour three on an OOM caused by one oversized scan is the most common way to waste a day.

5. Evaluate on the thing you actually care about

Loss is not the metric. For structured extraction, evaluate three levels:

  1. Parse rate — what fraction of outputs are valid JSON with the expected keys. Should be ~100% after fine-tuning; if it is not, your targets were inconsistent.
  2. Field accuracy — exact match per field, normalized (case, whitespace, date format, currency symbols). Report per field: total and date usually lag vendor.
  3. Document accuracy — the fraction of documents where every field is right. This is the number the business feels, and it is always lower than field accuracy.
def evaluate(model, processor, ds):
    ok_parse = ok_doc = 0
    for ex in ds:
        pred = generate(model, processor, ex)          # greedy, max_new_tokens=512
        try:
            obj = json.loads(pred)
        except json.JSONDecodeError:
            continue
        ok_parse += 1
        ok_doc += int(all(normalize(obj.get(k)) == normalize(v)
                          for k, v in ex["target"].items()))
    return ok_parse / len(ds), ok_doc / len(ds)

Always run the same harness against the base model before fine-tuning, and against a commercial API if one is a candidate. Without those two baselines you cannot say whether the fine-tune earned its keep. Decode greedily (do_sample=False) for structured tasks; sampling buys you nothing but invalid JSON.

6. Serve it

Merge the adapters and serve with vLLM, which supports the major VLM families through the same OpenAI-compatible API used for text:

vllm serve ./merged-vlm-invoices --limit-mm-per-prompt image=2 --max-model-len 8192

Clients send image_url content parts with base64 data URLs. Two operational notes specific to vision: visual tokens dominate your context budget, so set --limit-mm-per-prompt and resize server-side rather than trusting callers; and throughput is far more sensitive to image resolution than to text length, so measure with production-sized images. The rest — continuous batching, metrics, autoscaling — is as covered in Serving an LLM with vLLM. Quantizing the merged model to FP8 or INT4 is worthwhile at volume; see BF16, FP8, INT8 and validate accuracy with the harness from step 5, not by eye.

Failure modes we see

  • Repetition or truncation on long documents. Usually a max-token or resolution ceiling, not a training problem.
  • The model ignores the image. Check that image tokens survived collation and that labels masked the prompt; if the loss falls normally but predictions are layout-blind, the images were dropped.
  • Great validation, poor production. Nearly always a random split that leaked layouts, or production images captured at a different resolution or compression than training data.
  • Hallucinated fields. Add explicit negative examples with null targets; a VLM that has never been rewarded for saying "missing" will invent a value.
  • Silent drift after a supplier changes their template. Sample and score production outputs weekly. Nobody owns this by default, which is one of the recurring problems in why enterprises still get PyTorch wrong.

When not to fine-tune a VLM

If your documents are genuinely fixed-layout and clean, a classical OCR-plus-template pipeline is cheaper and more auditable. If you need a bounding box for every extracted value, a detection model still beats a generative one. And if you have fewer than a few hundred labelled examples, spend the week on prompt engineering and few-shot examples against a strong base model first — then fine-tune once you have production data worth learning from.

Multimodal extraction and visual understanding sit across our deep learning model development, LLM fine-tuning and inference optimization practices. If you have a document, screenshot or image workflow you would like scoped, contact us — a two-week pilot on your own data usually settles the build-or-buy question for good.