MCPcopy Index your code
hub / github.com/allenai/molmo

github.com/allenai/molmo @main

Chat with this repo
repository ↗ · DeepWiki ↗ · + Follow
1,091 symbols 3,795 edges 66 files 169 documented · 15% updated 19mo ago★ 91831 open issues
What it actually does AI analysis from the code graph — generated when you open this
loading…
README

Molmo Logo

Molmo: Multimodal Open Language Model

GitHub License Blog Post Paper URL Model Checkpoints PixMo (Datasets)

Molmo is a repository for training and using Ai2's state-of-the-art multimodal open language models.

Here is a video demo of Molmo's capabilities. Try Molmo using our public demo showcasing the Molmo-7B-D model.

This codebase is based on the OLMo codebase with the addition of vision encoding and integrating generative evaluations.

Release Notes

  • [2024/12/05] 🔥 Molmo: code for modeling, training and evaluation has been released. You can find detailed technical report here.

  • [2024/11/27] 🔥 PixMo, our new collection of datasets for pre-training and fine-tuning VLMs, has been released. PixMo consists of:

  • PixMo-Cap (pre-training, fine-tuning): highly detailed dense caption dataset (roughly 200 words on average)

  • PixMo-AskModelAnything (fine-tuning): instruction-tuning data containing human-authored image-question-answer triplets
  • PixMo-CapQA (fine-tuning): synthetic instruction-tuning data, using a LLM to build QA pairs from dense captions of images
  • PixMo-Points (fine-tuning): images paired with referring expressions and annotated points, supporting grounding and counting
  • PixMo-Point-Explanations (fine-tuning): instruction-tuning data with explanations containing in-line points referring to parts of the image
  • PixMo-Docs (fine-tuning): synthetic image-question-answer triplets about various kinds of computer-generated charts, tables, diagrams and documents. Code available here.
  • PixMo-Clocks (fine-tuning): virtual watch faces and time annotations
  • PixMo-Count (fine-tuning): diverse images with counting QA pairs

All datasets were constructed without the use of VLMs.

Pixmo and Molmo

Datasets in PixMo (left) and the capabilities they enable in Molmo (right).

  • [2024/09/24] 🔥 Molmo, a new family of open VLMs, has been released. The Molmo family consists of:

  • MolmoE-1B: a mixture of experts model with 1B (active) 7B (total)

  • Molmo-7B-O: our most open 7B model
  • Molmo-7B-D: our best 7B and demo model
  • Molmo-72B: our best 72B model

Installation

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

To install dependencies, run:

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

For training and evaluating MolmoE-1B, please install megablocks by running pip install git+https://github.com/Muennighoff/megablocks.git@olmoe.

Huggingface Models and Logs

The core models in the Molmo family released so far are:

Model Vision Encoder LLM 11-benchmark avg
MolmoE-1B-0924 OpenAI CLIP ViT-L/14@336 OLMoE-1B-7B-0924 68.6
Molmo-7B-O-0924 OLMo-7B-1024-preview 74.6
Molmo-7B-D-0924 Qwen2-7B 77.3
Molmo-72B-0924 Qwen2-72B 81.2

W&B logs: pre-training, fine-tuning

Data Downloading and Setup

Molmo uses huggingface datasets for most data, therefore most data will be stored in the default huggingface cache. See here for how to set it. Some additional data is stored separately in the path set by 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

Data can then be downloaded with:

python3 scripts/download.py all --n_proc 12

Downloading the pixmo datasets requires downloading images from URLs. The download script will do this automatically, but it will take some time. Downloading everything from scratch can take up to a day. More processes can make it faster, but it also increases the risk of getting rate-limited.

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

Some datasets (InfoQa and Scene-Text) require manually downloading the files. The download scripts will throw an error if those files are not found.

Downloading the android control dataset requires additional dependencies since it requires parsing the original tfrecords.

To download a specific dataset pass in the dataset name run:

python3 scripts/download_data.py ChartQa --n_proc 12

Visualizing Data

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

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

Trained Models

We release model weights both after pre-training and after fine-tuning in a format compatible with this codebase. The fine-tuned weights match the ones in the hugging face repos, but have a slightly different format. The config files are backwards-compatible with this repo, but also have a slightly different format.

Model Pretrained Fine-Tuned
MolmoE-1B-0924 pretrained fine-tuned
Molmo-7B-O-0924 pretrained fine-tuned
Molmo-7B-D-0924 pretrained fine-tuned
Molmo-72B-0924 pretrained fine-tuned

To use them, download the file and untar them. Each folder contains the needed config file and model weights. For example:

wget https://storage.googleapis.com/oe-training-public/Molmo-0924/Molmo-7B-D-0924.tar
tar -xf Molmo-7B-D-0924.tar 

Evaluation

Evaluation is done with the launch_scripts/eval_downstream.py script. FSDP can be used to evaluate large models, or for high-resolution processing. 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 pass the task name, or task_name:split:

torchrun --nproc-per-node 8 launch_scripts/eval_downstream.py Molmo-7B-D-0924 text_vqa --save_to_checkpoint_dir

For most tasks, we evaluate with high resolution:

torchrun --nproc-per-node 8 launch_scripts/eval_downstream.py Molmo-7B-D-0924 text_vqa --save_to_checkpoint_dir --high_res --fsdp --device_batch_size=2

The --fsdp flag will use FSDP which is needed for to avoid OOMs when using high resolution.

To evaluate on our default eval set (including the 11 tasks in the paper):

torchrun --nproc-per-node 8 launch_scripts/eval_downstream.py Molmo-7B-D-0924 low-res --save_to_checkpoint_dir
torchrun --nproc-per-node 8 launch_scripts/eval_downstream.py Molmo-7B-D-0924 high-res --save_to_checkpoint_dir --high_res --fsdp --device_batch_size=2

To get test numbers, use low-res-test and high-res-test. Some test numbers will require re-formatting the prediction files and then submitting to test servers.

To evaluate the 72B model with this codebase you will need to run on multiple nodes and might need to set PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True

These scripts will save the metrics and predictions in the save directory. Future calls to the eval script will re-use cached metrics if they exist, to overwrite these cached metrics use the --overwrite flag.

Evaluation with VLMEvalkit

Evaluation of the HF models is also supported via open-compass/VLMEvalkit. Check PR#648 for supported prompts and evaluation settings to reproduce results from the paper. However a few datasets (e.g., PixMo-Count) are not supported.

Pretrained Models for Initialization

Training end-to-end requires downloading the pre-trained models used to initialize Molmo. This can be done with the script scripts/convert_hf_to_molmo.py

For example, to load the Qwen2 LLM and OpenAI CLIP model, run:

python3 scripts/convert_hf_to_molmo.py qwen2_7b
python3 scripts/convert_hf_to_molmo.py openai

The model will be downloaded from huggingface, converted into a compatible format, and then saved into the MOLMO_DATA_DIR directory.

Pre-Training

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

To start a debugging run:

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

To train with the Qwen2 LLM and the CLIP vision encoder:

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

You can use other vision encoders including SigLIP, MetaCLIP and DINOv2 with the option --vision_backbone=model_name.

To run without wandb, use:

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

Multitask Training

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

WANDB_API_KEY=key torchrun --nproc-per-node=8 launch_scripts/train_multitask_model.py 3.2-synthetic /path/to/checkpoint --wandb.name=run_name --wandb.entity=entity --wandb.project=project --save_folder=/path/to/save/folder

Here 3.2-synthetic refers to what training mixture to use and /path/to/checkpoint points to a model checkpoint to start from, typically a dense captioning model.

To launch a debug run:

` torchrun --nproc-per-node=1 launch_scripts/train_multit

Core symbols most depended-on inside this repo

get
called by 119
olmo/data/dataset.py
get_global_rank
called by 39
olmo/torch_util.py
get_world_size
called by 35
olmo/torch_util.py
barrier
called by 34
olmo/torch_util.py
decode
called by 28
olmo/tokenizer.py
load
called by 26
olmo/config.py
reset_parameters
called by 24
olmo/model.py
load_state_dict
called by 18
olmo/checkpoint.py

Shape

Method 639
Function 254
Class 198

Languages

Python100%

Modules by API surface

olmo/model.py120 symbols
olmo/data/academic_datasets.py97 symbols
olmo/checkpoint.py78 symbols
olmo/util.py68 symbols
olmo/config.py66 symbols
olmo/train.py63 symbols
olmo/eval/evaluators.py49 symbols
olmo/beam_search.py48 symbols
olmo/image_vit.py46 symbols
olmo/optim.py45 symbols
olmo/data/pixmo_datasets.py41 symbols
tests/beam_search_test.py36 symbols

For agents

$ claude mcp add molmo \
  -- python -m otcore.mcp_server <graph>

⬇ download graph artifact

Ask about this repo answers extend the page