tiny-random-molmo2

A tiny Molmo2 (allenai/MolmoWeb-4B architecture) with random weights, used for the OpenVINO GenAI VLMPipeline and genai-cli export tests. It produces meaningless text and is meant only for testing.

• Text model: 3 layers, hidden size 128, 4 attention heads / 2 KV heads, head_dim 32, full Qwen vocabulary
• Vision (ViT): 3 layers, hidden size 64, 2 heads, head_dim 32; adapter uses the last two ViT layers
• 39.6M parameters, float32
• Tokenizer, processor configs, chat template and remote code (*_molmo2.py) are copied unchanged from allenai/MolmoWeb-4B

Molmo2 is remote code: load it with trust_remote_code=True and transformers>=4.57,<4.58 (einops and torchvision are also required).

Generation script

# uv run --with transformers==4.57.6 --with einops --with torchvision python make_tiny_molmo2.py <MolmoWeb-4B snapshot> <out>
import math
import shutil
import sys
from pathlib import Path

import torch
from transformers import AutoConfig, AutoModelForImageTextToText

src, out = Path(sys.argv[1]), Path(sys.argv[2])
torch.manual_seed(0)
cfg = AutoConfig.from_pretrained(src, trust_remote_code=True)
t, v, a = cfg.text_config, cfg.vit_config, cfg.adapter_config
t.hidden_size, t.num_attention_heads, t.num_key_value_heads, t.head_dim = 128, 4, 2, 32
t.num_hidden_layers, t.intermediate_size = 3, 256
v.hidden_size, v.intermediate_size, v.num_hidden_layers = 64, 128, 3
v.num_attention_heads = v.num_key_value_heads = 2
v.head_dim = 32
a.hidden_size, a.num_attention_heads, a.num_key_value_heads, a.head_dim = 64, 2, 2, 32
a.intermediate_size, a.text_hidden_size, a.vit_layers = 128, t.hidden_size, [-1, -2]

model = AutoModelForImageTextToText.from_config(cfg, trust_remote_code=True).float()
with torch.no_grad():
    for name, p in model.named_parameters():
        if name.endswith("norm.weight") or "ln_f" in name or "_norm" in name:
            continue
        if p.ndim == 2:  # linear / embedding / positional embedding
            std = 1.0 if ("wte" in name or "positional_embedding" in name) else 1.0 / math.sqrt(p.shape[1])
            if name == "lm_head.weight":
                std = 4.0 / math.sqrt(p.shape[1])
            p.normal_(0.0, std)
        elif name.endswith("bias"):
            p.normal_(0.0, 0.02)
model.save_pretrained(out, safe_serialization=True)
for f in src.iterdir():
    if f.is_file() and f.suffix != ".safetensors" and f.name not in ("config.json", "model.safetensors.index.json"):
        shutil.copy(f, out / f.name)