Molmo 2: State-of-the-art video understanding, pointing, and tracking

This repository is for training and using Ai2's open vision language models, Molmo2 and MolmoPoint. Molmo2 is state-of-the-art among open-source models and demonstrates exceptional new capabilities in point-driven grounding in single image, multi-image, and video tasks as shown below. MolmoPoint is an extension with a new architecture for pointing. This README is mostly concerned with Molmo2, see MolmoPoint for how to train MolmoPoint.

See our blog post or our paper for more details about Molmo2. Huggingface models can be found here.

Table of Contents

• Setup
• Installation
• Docker
• Downloading Data
• Downloading Pretrained Models
• Visualizing Data
• Environment
• Training and Evaluations
• Checkpoints
• Pre-Training
• SFT Training
• Long-Context SFT Training
• Evaluation
• Context Parallel
• Transformers and vLLM
• Convert Checkpoint to Hugging Face Format
• Transformers Inference
  • Image Inference Example
  • Video Inference Example
• MolmoPoint Transformers Inference
  • Image Inference Example
  • Video Inference Example
• Fast Inference with vLLM
  • Install Vision Process Package
  • Install vLLM (>= 0.15.0)
  • Run vLLM inference (Gradio Demo)
• Code
• Data Pipeline
• Message Trees
• Packing

Setup

Installation

We recommend using python >= 3.11 First install PyTorch according to the instructions specific to your operating system.

To install dependencies, run:

git clone https://github.com/allenai/molmo2.git
cd molmo2
pip install torchcodec
pip install -e .[all]

It's recommended to install torchcodec separately since it has some complex dependencies that can break if installed in combination with the others as done using install -e .[all]

Docker

We provide a container with the dependencies (but not the code) pre-installed, pull it with: docker pull ghcr.io/allenai/molmo2:latest

Downloading Data

Molmo2 uses a mix of huggingface datasets and custom data stored in MOLMO_DATA_DIR.

For example, if you want to store the data in /data/molmo you could set

export MOLMO_DATA_DIR=/data/molmo
export HF_HOME=/data/molmo/huggingface

See here for more info on where the huggingface data is stored.

We provide a script to download most datasets:

python3 scripts/download_datasets.py all --n_proc 8

Downloading can be resumed if canceled or an error occurs mid-download.

Some datasets need to be manually downloaded, often due licensing agreements. See the relevant classes for their locations and download instructions. These include: - DocQA, InfoQA, and SceneText need to be downloaded from https://rrc.cvc.uab.es.
- LVBench needs to be downloaded from https://huggingface.co/datasets/zai-org/LVBench. - MLVU and LongVideoBench have HuggingFace user agreements that must be accepted before the download scripts will work - The nturgbd subset of MVBench needs to be manually downloaded. - Tracking datasets that require manual download: Ref-YT-VOS, YTVIS, ReVOS, LaSOT, Molmo2VideoTrack, and etc. See olmo/data/academic_video_track_datasets.py and olmo/data/molmo2_video_track_datasets.py for download instructions.

The download scripts will throw an error and provide instructions if those files are not found.

To download a specific dataset provide the dataset or class name as follows:

python3 scripts/download_datasets.py ChartQa --n-procs 12

You can also download by group:

# Download image academic benchmarks
python3 scripts/download_datasets.py image_academic

# Download multiple specific datasets
python3 scripts/download_datasets.py text_vqa doc_qa chart_qa

# Download video academic benchmarks
python3 scripts/download_datasets.py video_academic

# Download all video tracking datasets (MOT + SOT)
python3 scripts/download_datasets.py video_tracking --n-procs 8

Available groups: image_academic, video_academic, pixmo, image_pointing, video_pointing, video_tracking, demo.

Downloading Pretrained Models for Training from scratch

Pretrained models can be downloaded and prepared with scripts/prepare_pretrained_model.py

For example:

python scripts/prepare_pretrained_model.py qwen3_4b_instruct
python scripts/prepare_pretrained_model.py siglip2

This will download the checkpoint, convert it into a compatible format, and save a sharded version in the location specified by the corresponding config olmo/model_configs.py for fast loading.

Visualizing Data

Once downloaded, datasets can be visualized by using the scripts/dataset_visualize.py script:

python3 scripts/dataset_visualize.py chart_qa /path/to/viz/dir

This script will build a HTML file to show what the data looks like after pre-processing.

Environment

Generally training runs should use these flags:

HF_DATASETS_OFFLINE=1
OLMO_SHARED_FS=1
HF_ACCESS_TOKEN=YOUR_HF_KEY
OPENAI_API_KEY=YOUR_OPENAI_API_KEY
WANDB_API_KEY=YOUR_WANDB_KEY
OMP_NUM_THREADS=8

HF_DATASETS_OFFLINE stops HF from sending tons of requests to the HF dataset hub even though the data is already download.

OLMO_SHARED_FS tell the codes to assume, for multi-nodes jobs, you are saving to a shared file system.

HF_ACCESS_TOKEN might be used to download the tokenizer, OPENAI_API_KEY might be used in some evaluations, and WANDB_API_KEY is for wandb logging.

OMP_NUM_THREADS is for torch.

Training and Evaluations

Molmo2 training has three stages:

1. Pre-Training — Train on image captioning, NLP, and image pointing using launch_scripts/pretrain.py. Start from pretrained LLM + ViT weights.
2. SFT — Multitask supervised fine-tuning on the full mixture (QA, pointing, tracking, video, etc.) using launch_scripts/sft.py. Start from a pretrained checkpoint.
3. Long-Context SFT — Continue SFT with longer sequences (36k+ tokens, 384 frames) for improved video understanding. Uses the same launch_scripts/sft.py with increased --seq_len.

Each stage produces a checkpoint that feeds into the next. We release checkpoints at each stage (see below).

Checkpoints

We release model weights after pre-training, SFT, and long-context SFT in a format compatible with this codebase. The Long-Context SFT Checkpoint matches the hugging face repo checkpoints, but have a slightly different format. The config files are backwards-compatible with this repo but might not match exactly.

HF Model	Pretrained Checkpoint	SFT Checkpoint	Long-Context SFT Checkpoint
Molmo2-4B	Pretrain	SFT	Long-Context SFT
Molmo2-8B	Pretrain	SFT	Long-Context SFT
Molmo2-O-7B	Pretrain	SFT	Long-Context SFT

To use these checkpoints download them, untar them, and they can be evaluated or used as a starting point for fine-tuning. For example:

wget https://storage.googleapis.com/oe-training-public/Molmo2-1225/Molmo2-8B.tar
tar -xf Molmo2-8B.tar

Pre-Training

The main pretraining script is scripts/run_trainer.py. To train a model you can either construct a config file to pass to it, or call one of the higher-level scripts in launch_scripts which will construct a low-level config from some higher-level settings and then invoke the train script for you.

For pretraining, we use launch_scripts/pretrain.py. This will run the first stage of training on image captioning, NLP, and image pointing - see the paper for details.

To start a debugging run:

torchrun --nproc-per-node=1 launch_scripts/pretrain.py debug --save_folder=/path/to/save/folder

To train with the Qwen3 4B Instruct LLM and the SigLIP vision encoder:

WANDB_API_KEY=key torchrun --nproc-per-node=8 launch_scripts/pretrain.py qwen3_4b_instruct \
  --wandb.name=run_name --wandb.entity=entity --wandb.project=project \
  --save_folder=/path/to/save/folder

Molmo2-8B uses qwen3_8b, Molmo2-4B uses qwen3_4b_instruct, and Molmo2-O-7B uses olmo3_7b_instruct.

Under-the-hood, the launch_scripts/pretrain.py constructs a TrainerConfig object and then runs it. For fine-grained control, CLI args can be used to override parts of the TrainerConfig, for example, to run without wandb, use:

torchrun --nproc-per-node=8 launch_scripts/pretrain.py qwen2_7b \
  --wandb=null --save_folder=/path/to/save/folder

See TrainerConfig for a full list of configurable args.

The script defaults matches Molmo 2's default, the LLM flags determines which version of Molmo 2 is being trained.

Note: use --data.num_workers=1 if you encounter a DataLoader "Bus error" (or increase shared memory)

SFT Training

Multitask training can be done with launch_scripts/sft.py, for example:

WANDB_API_KEY=key torchrun --nproc-per-node=8 launch_scripts/sft.py /path/to/pretrained/model molmo2 \
  --wandb.name=run_name --wandb.entity=entity --wandb.project=project \
  --save_folder=/path/to/save/folder

Here /path/to/pretrained/model points to a model checkpoint to start from (typically a pretrained model) and molmo2 refers to what training mixture to use.

To launch a debug run:

torchrun --nproc-per-node=1 launch_scripts/sft.py /path/to/pretrained/model debug --debug --save_folder=dbg --save_overwrite

This will run a lightweight version of the model and a small dataset to allow easier debugging.

For MolmoPoint, we used launch_scripts/sft.py /path/to/pretrained/molmo_point/model molmo_point instead.

Long-Context SFT Training

Long-context training is done with the same script. If you have B200s you can run it like this:

torchrun --nproc-per-node=8 launch_scripts/sft.py /path/to/sft/checkpoint molmo2 \
  --max_duration=2000 --device_batch_size=1 --data.num_workers=4 --seq_len=36864 \
  --model.mm_preprocessor.video.max_frames=384 --model.llm.max_sequence_length=36864

For smaller memory GPUs or for even higher frame counts, you might require context parallelism.

Evaluation

Evaluation can be done with the launch_scripts/eval.py script.

Note that the vLLM version of Molmo will be significantly faster for inference, but most of our numbers were reported using the results of this local evaluation.

To eval on a single task:

torchrun --nproc-per-node 8 launch_scripts/eval.py Molmo2-4B --task=chart_qa --save_to_checkpoint_dir

This will save the metrics and predictions in the save directory. Future calls to the eval script will re-use cached metrics if they exist. See EvalConfig for additional config options.

The --loss flag can be used to compute the loss instead of doi