LLM2VLM: Finetuning LLMs into Vision-Language Models

A complete system for turning a large language model into a vision-language model (VLM) by adding a frozen vision encoder, a trainable MLP projector, and LoRA adapters. Supports 1-8 GPU distributed training via DDP. Built and tested on AMD MI308X GPUs with ROCm.

The default configuration finetunes GPT-OSS-20B (a 20B-parameter Mixture-of-Experts LLM) with SigLIP2-SO400M as the vision encoder, trained on LLaVA-Instruct-150K plus 7 OCR-focused datasets from LLaVA-OneVision-Data. The system includes MXFP4-to-bf16 manual dequantization to work around ROCm incompatibility with MXFP4 quantized weights (native MXFP4 Triton kernels lack autograd backward support required for training).

Architecture

Image (variable resolution, up to 5 tiles of 384x384)
  -> SigLIP2-SO400M (frozen, 27 layers, hidden=1152)
  -> PseudoDeepStack: extract layers {8, 17, 26}, concat -> [B, N, 3456]
  -> MLP Projector: Linear(3456,3456) -> GELU -> Linear(3456,2880), trainable
  -> GPT-OSS-20B (frozen MoE: 24 layers, hidden=2880, 32 experts top-4, LoRA on q/k/v/o_proj)
  -> Text output

• Multi-crop vision: High-resolution images are split into up to 4 tiles + 1 overview (5 x 729 = 3645 image tokens), enabling text reading and fine-grained detail recognition
• 729 image tokens per tile (27x27 patches from 384/14) replace <image> placeholder token embeddings with projected vision features
• LoRA applied to attention projections with separate learning rates (2e-5 for LoRA vs 2e-4 for projector)
• MXFP4 dequantization handles GPT-OSS-20B's Microscaling FP4 quantized expert weights, which require bf16 dequantization for training (native MXFP4 Triton kernels lack autograd backward support)
• Gradient checkpointing reduces activation memory ~4x, enabling training within GPU memory limits

Files

File	Purpose
config.py	Configuration dataclass + OptimalConfigurator for auto-tuning batch size, LoRA rank, and gradient accumulation based on GPU memory
model.py	Core VLM architecture: MXFP4 dequantization, multi-crop PseudoDeepStack feature extraction, MLP projector, LLM2VLM model class with dual stop token generation
data.py	Mixed dataset loader: LLaVA-Instruct-150K + 7 OCR datasets with multi-crop image processing, chat template formatting, label masking, and image token expansion
train.py	DDP training loop (1-8 GPUs) with gradient checkpointing, cosine schedule with warmup, gradient accumulation, bf16 autocast
inference.py	Inference from a saved training checkpoint
export.py	Export merged model to HuggingFace-ready format (sharded safetensors + custom modeling code for trust_remote_code)
test.py	Test inference from exported HF-ready model; single image or batch test on all images in a folder
launch_train.sh	Convenience launcher for multi-GPU training via torchrun

Requirements

Hardware

• AMD Instinct MI308X GPU (192 GB HBM3) or equivalent with ROCm support
• Tested on 1-8 GPUs with DDP (DistributedDataParallel)

Software

Package	Version
ROCm	7.1+
PyTorch	2.10+ (ROCm build)
Transformers	5.2+
PEFT	0.18+
bitsandbytes	0.49+
accelerate	1.13+
unsloth	2026.2+ (optional)
safetensors	latest
Pillow	latest

pip install torch transformers peft bitsandbytes accelerate safetensors pillow

Models (from HuggingFace)

By default, the following model files are expected in the models/ directory:

Path	HuggingFace Source
models/google/siglip2-so400m-patch14-384	google/siglip2-so400m-patch14-384
models/openai/gpt-oss-20b	openai/gpt-oss-20b

Download them with the HuggingFace CLI:

mkdir -p models/google models/openai
huggingface-cli download google/siglip2-so400m-patch14-384 --local-dir models/google/siglip2-so400m-patch14-384
huggingface-cli download openai/gpt-oss-20b --local-dir models/openai/gpt-oss-20b

Datasets (from HuggingFace)

By default, the following datasets are expected in the datasets/ directory:

Path	HuggingFace Source	Description
datasets/coco2014/train2014/	COCO 2014 Train Images	Training images (~83K JPEG files)
datasets/liuhaotian/LLaVA-Instruct-150K/	liuhaotian/LLaVA-Instruct-150K	150K instruction-following conversations referencing COCO images
datasets/lmms-lab/LLaVA-OneVision-Data/	lmms-lab/LLaVA-OneVision-Data	OCR-focused datasets (7 subsets, see below)

OCR datasets (enabled by default with --use_ocr_data):

Subset	Description
textocr(gpt4v)	Text detection and recognition
llavar_gpt4_20k	Text-rich image understanding
st_vqa(cauldron,llava_format)	Scene text VQA
textcaps	Text-aware image captioning
chartqa(cauldron,llava_format)	Chart understanding
infographic_vqa_llava_format	Infographic VQA
sroie	Receipt OCR

mkdir -p datasets
huggingface-cli download liuhaotian/LLaVA-Instruct-150K --local-dir datasets/liuhaotian/LLaVA-Instruct-150K
huggingface-cli download lmms-lab/LLaVA-OneVision-Data --local-dir datasets/lmms-lab/LLaVA-OneVision-Data
# For COCO 2014, download and extract train2014.zip into datasets/coco2014/

Model and dataset paths can be overridden in config.py.

Usage

Training

# Single GPU
python train.py

# Multi-GPU (auto-detect all GPUs)
bash launch_train.sh

# Multi-GPU (specific count)
bash launch_train.sh 4

# Multi-GPU with specific GPUs
CUDA_VISIBLE_DEVICES=0,1,2,3 bash launch_train.sh 4

# With auto-tuned hyperparameters
bash launch_train.sh 8 --optimal

# Without OCR data
bash launch_train.sh --no_ocr_data

# Without multi-crop (fixed 384x384)
bash launch_train.sh --no_multi_crop

# With QLoRA (lower memory)
python train.py --use_qlora

# Disable gradient checkpointing (uses more memory)
python train.py --no_gradient_checkpointing

Gradient accumulation is automatically adjusted for multi-GPU to maintain the same effective batch size (default 32):

• 1 GPU: batch=4, accum=8, effective=32
• 2 GPUs: batch=4, accum=4, effective=32
• 4 GPUs: batch=4, accum=2, effective=32
• 8 GPUs: batch=4, accum=1, effective=32

Inference from Checkpoint

CUDA_VISIBLE_DEVICES=0 python inference.py \
  --checkpoint ./outputs/final_model \
  --image photo.jpg \
  --prompt "Describe this image in detail."

Export to HuggingFace Format

Exports the trained model (merged LoRA + dequantized weights) to a self-contained HuggingFace-compatible directory with sharded safetensors and custom modeling code:

CUDA_VISIBLE_DEVICES=0 python export.py \
  --checkpoint ./outputs/final_model \
  --output_dir ./outputs/hf_final_model

The exported model is fully self-contained — it includes embedded vision and LLM configs, so it loads without the original model directories.

Test Exported Model

# Single image
CUDA_VISIBLE_DEVICES=0 python test.py \
  --model_dir ./outputs/hf_final_model \
  --image photo.jpg

# Batch test — all images in a folder (loads model once)
CUDA_VISIBLE_DEVICES=0 python test.py \
  --model_dir ./outputs/hf_final_model \
  --batch_test /path/to/images

# Batch test — default COCO dir, limit to 10 images
CUDA_VISIBLE_DEVICES=0 python test.py \
  --model_dir ./outputs/hf_final_model \
  --batch_test --num_images 10

# Batch test with custom prompt
CUDA_VISIBLE_DEVICES=0 python test.py \
  --model_dir ./outputs/hf_final_model \
  --batch_test /path/to/images \
  --prompt "Read all text in this image."

Supported image formats: jpg, jpeg, png, bmp, tiff, webp, gif.

Key Technical Details

MXFP4 Dequantization for ROCm

GPT-OSS-20B ships with MXFP4 (Microscaling FP4) quantized expert weights. The native MXFP4 Triton kernels work for inference on ROCm (after triton 3.6.x + num_sms() patching), but do not implement autograd backward pass — the expert MLP output has requires_grad=False, breaking gradient flow during training. This project implements manual MXFP4-to-bf16 dequantization:

1. Loads raw _blocks (uint8 packed FP4) and _scales (E8M0) tensors
2. Unpacks uint8 into two 4-bit nibbles, decodes E2M1 format
3. Applies E8M0 scale factors
4. Transposes expert weight layout ([experts, out, in] -> [experts, in, out])

The dequantized model is ~42 GB in bf16 (vs ~13 GB in MXFP4), but this is required for proper gradient flow through the MoE expert layers.

Optimized Model Loading

The exported model uses meta-device instantiation and shard-by-shard GPU loading:

• Model shells created on torch.device("meta") — zero memory, ~0.1s
• Safetensors loaded one shard at a time (~5 GB each), transferred directly to GPU
• Peak CPU memory: ~5 GB (vs ~84 GB with naive loading)
• Total load time: ~19 seconds

PseudoDeepStack

Extracts intermediate features from SigLIP2 encoder layers {8, 17, 26}, concatenates along the hidden dimension (3 x 1152 = 3456), then projects to the LLM's hidden dimension (2880) via a 2-layer MLP with GELU activation. This captures multi-scale visual features compared to using only the final layer output.

Multi-Crop High-Resolution Vision

For high-resolution images, the system:

1. Splits the image into a grid of up to 4 tiles (384x384 each)
2. Creates a resized overview of the full image (384x384)
3. Processes each tile + overview through SigLIP2 independently
4. Concatenates all tile features: up to 5 x 729 = 3645 image tokens

This enables text reading, chart understanding, and fine-grained detail recognition that fixed 384x384 cannot achieve. Controlled by --use_multi_crop (on by default) and --max_crop_tiles N.

OCR Training Data

Seven OCR-focused datasets from LLaVA-OneVision-Data are mixed with LLaVA-Instruct-150K during training. These datasets provide text detection, scene text VQA, chart QA, receipt OCR, and text-aware captioning. Datasets are loaded from local parquet files. Controlled by --use_ocr_data (on by default).

Gradient Checkpointing

Enabled by default (--gradient_checkpointing). Trades ~30% compute overhead for ~4x reduction in activation memory. Uses use_reentrant=False mode with a clone fix for the embedding injection to avoid in-place operation errors.

Multi-GPU DDP Training

Supports 1-8 GPUs via PyTorch DistributedDataParallel with NCCL backend (RCCL on ROCm). Each GPU holds a full model copy and processes different data. Gradient accumulation is automatically scaled down to maintain the same effective batch size. Launch via bash launch_train.sh [N] or torchrun --nproc_per_node=N train.py.

GPT-OSS Dual Stop Tokens

GPT-OSS uses two distinct termination tokens: <|end|> (ID 200007, end-of-turn delimiter) and <|return|> (ID 200002, true end-of-generation). During training, intermediate assistant turns end with <|end|> while the final turn ends with <|return|>. Both tokens are set as stop IDs during generation to ensure the model stops cleanly at the first turn boundary, preventing multi-turn hallucination.

Example Output

Image: COCO 000000039769 (two cats on a couch)

Prompt: "Describe this image in detail."

Response: "The image features two cats, a brown tabby and a black and white one, peacefully sleeping on a red blanket placed on a couch. The cats are laying side by side, one near the center and the other closer to the left side of the blanket. In addition to the cats and the blanket, there is a remote control lying on the couch near the cats, likely belonging to a TV or another electronic device in the room."

Known Limitations

• MXFP4 dequantization increases memory — 20B model uses ~42 GB in bf16 vs ~13 GB in MXFP4; required for training because native MXFP4 Triton kernels lack autograd backward support
• No automated evaluation — only qualitative testing; benchmarks (VQAv2, TextVQA, POPE) not yet integrated
• No training resumption — must restart from scratch if interrupted
• LoRA limited to attention layers — expert MLP weights use batched 3D tensors in a custom module, not nn.Linear, so PEFT cannot wrap them

License

This project is provided as-is for research and educational purposes.