Add Molmo2 by SangbumChoi · Pull Request #43451 · huggingface/transformers

SangbumChoi · 2026-01-23T14:47:55Z

What does this PR do?

Fixes # (issue)

Before submitting

This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
Did you read the contributor guideline,
Pull Request section?
Was this discussed/approved via a Github issue or the forum? Please add a link
to it if that's the case.
Did you make sure to update the documentation with your changes? Here are the
documentation guidelines, and
here are tips on formatting docstrings.
Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Adds AllenAI Molmo2 multimodal VLM to transformers, supporting: - Molmo2ForConditionalGeneration (image+video+text → text) - Molmo2TextModel / Molmo2TextForCausalLM (text-only) - Molmo2ImageProcessor and Molmo2VideoProcessor - Molmo2Processor Key implementation details: - Uses is_first_iteration (v5 API) for prepare_inputs_for_generation - Custom Molmo2Embedding with embedding + new_embedding parameters - Vision backbone with pooling adapter and multi-layer ViT features - Dynamic full cache support for generation Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

@strict

…odel_prefix - Replace einops.rearrange with native numpy reshape+transpose+reshape - Add @strict decorator to all 4 config classes (Molmo2VitConfig, Molmo2AdapterConfig, Molmo2TextConfig, Molmo2Config) to satisfy TRF010 - Set Molmo2Model.base_model_prefix = "model" (was empty, violating TRF002) - Fix image_mean/image_std mutable shared list (copy constants on init) - Fix test_image_processing: use image_processing_class instead of image_processor_list; skip CHW torch and 4-channel unsupported tests Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

- Re-sort _toctree.yml to place Molmo2 after mllama alphabetically - Add None guard in test_video_processor_from_dict_with_kwargs to skip when fast_video_processing_class is not defined Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Molmo2TextModel is an internal sub-component used by Molmo2Model and Molmo2ForConditionalGeneration and is tested implicitly through those. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

requests is not part of the standard library and caused ImportError in minimal environments (e.g. HuggingFace Jobs). Use urllib.request instead. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Molmo2's processor has several behaviors that are incompatible with the default ProcessorTesterMixin assumptions: - Chat template enforces strict user/assistant alternation (no system role) - Processor inserts BOS token, shifting sequence length by 1 - Image processor patchifies output, so rescale_factor passthrough fails - Video processor requires FPS metadata not provided by base tests - Hub processor_config.json contains auto_map not preserved in save/load Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Add @auto_docstring(checkpoint="allenai/Molmo2-8B") decorator to Molmo2TextConfig and Molmo2Config with custom_args for documenting non-standard parameters. This fixes check_config_docstrings CI check. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

… date Add parameter docstrings to Molmo2TextConfig and Molmo2Config __init__ methods so @strict-wrapped classes pass config docstring CI checks. Update model doc date to 2026-03-28. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Move top-level `import torch` and `import torchvision.transforms` behind `is_torch_available()` / `is_torchvision_available()` guards in both image and video processors to prevent ModuleNotFoundError when torchvision is not installed. Also skip test_kwargs_overrides_default_image_processor_kwargs since Molmo2's patchifying image processor doesn't support rescale_factor passthrough. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Convert all absolute imports (from transformers.xxx) to relative imports (from ...xxx) in image_processing, video_processing, and processing modules to match the convention used by all other in-library models. Remove register_for_auto_class() calls which are only needed for custom hub models and were causing dynamic_module_utils to incorrectly scan local files for relative imports during save_pretrained. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…n_available The processor's top-level imports from image_processing_molmo2 and video_processing_molmo2 pull in PILImageResampling which requires PIL. Guard these imports with is_vision_available() so `from transformers import *` works when only torch is installed (no PIL/torchvision). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

…L imports Move Molmo2ImagesKwargs and Molmo2VideosKwargs definitions directly into processing_molmo2.py instead of importing them from image/video processor modules which require PIL. Also remove Molmo2ImageProcessor/VideoProcessor type hints from __init__ to avoid NameError when vision is unavailable. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

SangbumChoi · 2026-03-29T01:49:56Z

@molbap Hi I am still working on it since I have to make some example visualizer for this and (most of the code is generated by Claude code). However, you can start review this with brief level of code review! cc. @merveenoyan

Add integration tests for Molmo2-8B covering: - Image generation with exact expected text verification - Video QA (penguin identification) - Video pointing (coordinate output) - Multi-image comparison All expected values derived from actual model inference on A10G. Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

zucchini-nlp

Hey @SangbumChoi

Great model to add to Transformers. After reviewing I see that using modular would be much better since a lot of part are copy-paste from different models. I left comments on each class about where it can be copied from. Apart from that, there are a few places where we need to clean up and align API with the rest of VLMs for consistency

If you have q, ping me on slack. I will unsubscribe myself from this PR to not get notif about each commit, so when you want another review ping me again by @

zucchini-nlp

Hey @SangbumChoi

Great model to add to Transformers. After reviewing I see that using modular would be much better since a lot of part are copy-paste from different models. I left comments on each class about where it can be copied from. Apart from that, there are a few places where we need to clean up and align API with the rest of VLMs for consistency

If you have q, ping me on slack. I will unsubscribe myself from this PR to not get notif about each commit, so when you want another review ping me again by @

`token_type_ids_mask_function` only clamped `kv_idx` against the length of `mm_token_type_ids` (which covers just the original prompt), but indexed `token_type_ids[batch_idx, q_idx]` directly. When generation verifies several new tokens in one forward (e.g. assisted/speculative decoding), the query length exceeds the prompt length and `q_idx` runs out of bounds: `IndexError: index N is out of bounds for dimension 1 with size N`. Newly generated positions are always text, never image patches, so clamp the `q_idx` lookup the same way as `kv_idx` and treat out-of-range positions as token-type 0. Fixes test_assisted_decoding_sample and test_assisted_decoding_matches_greedy_search. Fast Molmo2ModelTest is now fully green (129 passed, 119 skipped).

SangbumChoi · 2026-06-29T23:41:34Z

molmo2_finetune.ipynb
molmo2_demo_(5).ipynb

SangbumChoi · 2026-07-02T13:27:38Z

molmo2_finetune.ipynb

SangbumChoi · 2026-07-05T01:50:52Z

@molbap @merveenoyan @guarin gentle ping

molbap

Hey @SangbumChoi , thanks a lot for iterating on this.

The modular implementation is muuuch cleaner. processor is also looking good. There's no more numpy, einsum, tests are leaner, there's not much missing.

The PR has been open for a long time, and we need to merge it before main drifts away, we get new API changes, and so on. What's left is a small, well-understood set of items that have been through two review rounds already, so rather than start a third one I'll finish them myself directly on this branch.

One ask: please don't push to this branch, so we don't overwrite each other. Your commits and authorship stay exactly as they are and if you spot something you'd like changed, review the PR and I'll include it. I'll ping you when it's ready to merge. Thanks again for carrying the model this far!

molbap · 2026-06-30T15:38:44Z

+from transformers import Molmo2ForConditionalGeneration, Molmo2Processor
+import torch
+
+processor = Molmo2Processor.from_pretrained("allenai/Molmo2-8B")
+model = Molmo2ForConditionalGeneration.from_pretrained(
+    "allenai/Molmo2-8B",


these should work with AutoModel/AutoProcessor

molbap · 2026-06-30T15:39:02Z

+from transformers import Molmo2ForConditionalGeneration, Molmo2Processor
+from transformers.video_utils import load_video


same comment

molbap · 2026-06-30T15:40:37Z

            ("mistral3", "PixtralProcessor"),
            ("mm-grounding-dino", "GroundingDinoProcessor"),
            ("modernvbert", "Idefics3Processor"),
+            ("molmo2", "Molmo2Processor"),


not needed here AFAIK

Suggested change

("molmo2", "Molmo2Processor"),

molbap · 2026-06-30T15:48:43Z

+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+    processor_kwargs={"video_metadata": [metadata]},


i don't think we should need to pass explicitly processor_kwargs here... we should be able to send video_metadata directly imo

we dont need to load video manually tbh, video processor does it internally for us

molbap · 2026-06-30T15:53:04Z

+    image_num_pos: int = 577
+    attention_dropout: float = 0.0
+    residual_dropout: float = 0.0
+    float32_attention: bool = True


if you upcast inputs just before the dispatched call then sdpa is computed in whatever dtype the inputs are in, that should be enough

molbap · 2026-06-30T15:53:45Z

+    def __post_init__(self, **kwargs):
+        if self.attn_implementation is not None:
+            kwargs["attn_implementation"] = self.attn_implementation
+        super().__post_init__(**kwargs)


molbap · 2026-06-30T15:54:05Z

+    def __post_init__(self, **kwargs):
+        if self.attn_implementation is not None:
+            kwargs["attn_implementation"] = self.attn_implementation
+        super().__post_init__(**kwargs)


still to remove

molbap · 2026-07-06T12:54:10Z

+    def tearDown(self):
+        cleanup(torch_device, gc_collect=True)


these should call super.teardown

molbap · 2026-07-06T13:00:36Z

+        if self.config._attn_implementation == "eager":
+            key_states = repeat_kv(key_states, self.num_key_value_groups)
+            value_states = repeat_kv(value_states, self.num_key_value_groups)
+            attn_weights = torch.matmul(query_states / math.sqrt(query_states.size(-1)), key_states.transpose(2, 3))
+            if attention_mask is not None:
+                # The image pooling adapter passes a boolean keep-mask (valid patches); exclude the
+                # invalid ones with -inf to match the SDPA path. A float mask is already additive.
+                if attention_mask.dtype == torch.bool:
+                    attn_weights = attn_weights.masked_fill(~attention_mask, torch.finfo(attn_weights.dtype).min)
+                else:
+                    attn_weights = attn_weights + attention_mask
+
+            attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+            attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=self.training)
+            attn_output = torch.matmul(attn_weights.to(value_states.dtype), value_states)
+            attn_output = attn_output.transpose(1, 2).contiguous()
+        elif self.config._attn_implementation == "sdpa":


this conditional branching, as said, only reason is the upcast, which can done before the call, and casted back after output.

SangbumChoi · 2026-07-06T14:25:59Z

@molbap Got it. From now on I will just leave it as a comment only

molbap · 2026-07-06T14:28:14Z

Thanks @SangbumChoi ! I also am trying to help narrow down AI-assisted PRs scope to boost contribution. I am trying out a "Self-review" skill to run on your end on your PRs, which should cover the recent API changes/alignment and boost a bit. If you have feedback on it it'd be very helpful!
transformers_contribution_skill.md

molbap · 2026-07-21T12:13:33Z

run-slow: molmo2

github-actions · 2026-07-21T12:14:53Z

Workflow Run ⚙️

This comment contains run-slow, running the specified jobs:

models: ["models/molmo2"]
quantizations: []

github-actions · 2026-07-21T12:41:41Z

CI Results

Workflow Run ⚙️

Commit Info

Context	Commit	Description
RUN	506086a8	workflow commit (merge commit)
PR	45bf547b	branch commit (from PR)
main	29985e67	base commit (on `main`)

⚠️ Model CI failed to report results

The test failure analysis could not be completed. Please check the workflow run for details.

zucchini-nlp

Adding my two cents as the model is closer to merge

zucchini-nlp · 2026-07-21T12:09:38Z

+    add_generation_prompt=True,
+    return_tensors="pt",
+    return_dict=True,
+    processor_kwargs={"video_metadata": [metadata]},


we dont need to load video manually tbh, video processor does it internally for us

zucchini-nlp · 2026-07-21T12:21:17Z

+def select_tiling(height: int, width: int, patch_size: int, max_num_crops: int) -> tuple[int, int]:
+    """Select the image tiling in the same height/width order as the original Molmo2 processor."""
+    tilings = []
+    for tile_height in range(1, max_num_crops + 1):
+        for tile_width in range(1, max_num_crops + 1):
+            if tile_height * tile_width <= max_num_crops:
+                tilings.append((tile_height, tile_width))
+    tilings.sort(key=lambda x: (x[0] * x[1], x[0]))
+
+    candidate_resolutions = torch.tensor(tilings, dtype=torch.int32) * patch_size
+    original_size = torch.tensor([height, width], dtype=torch.float32)
+
+    required_scales = candidate_resolutions.to(torch.float32) / original_size
+    required_scale = required_scales.amin(dim=-1, keepdim=True)
+
+    if torch.all(required_scale < 1):
+        return tilings[int(required_scale.argmax())]
+
+    required_scale = torch.where(required_scale < 1.0, 10e9, required_scale)
+    return tilings[int(required_scale.argmin())]


looks same as get_optimal_tiled_canvas in cohere vision

same core min-scale selection, but molmo sorts candidates by (area + height) while cohere enumerates aspect ratios sorted by area only, and returns (w, h) not (h, w). so using it changes the grid (except for near-square imags)

zucchini-nlp · 2026-07-21T13:04:26Z

+    def test_training_gradient_checkpointing_use_reentrant_true(self):
+        pass
+
+    @unittest.skip(reason="VLMs have dynamic control flow in preparing inputs for generation")


+1, preparing inputs for generation shouldn't be the issue since other VLMs dont skip the test

zucchini-nlp · 2026-07-21T13:06:56Z

+        input_len = device_inputs["input_ids"].shape[1]
+        generated_text = self.processor.batch_decode(generated_ids[:, input_len:], skip_special_tokens=True)[0]
+        EXPECTED_TEXT = "Penguins appear in the video."
+        self.assertEqual(generated_text.strip(), EXPECTED_TEXT)


quite a lot of slow tests, do we need all of them? 👁️

yeah well the 3 models are quite different... maybe we can skip the 4B ones

github-actions · 2026-07-22T09:16:51Z

[For maintainers] Suggested jobs to run (before merge)

run-slow: auto, molmo2

github-actions · 2026-07-22T09:28:05Z

CI recap

Dashboard: View test results in Grafana
Latest run: 29907256921:2
Result: failure | Jobs: 15 | Tests: 172,855 | Failures: 1 | Duration: 12h 40m

merveenoyan requested a review from molbap February 27, 2026 05:55

SangbumChoi force-pushed the molmo2 branch from b15107c to 3fee343 Compare March 26, 2026 23:28

SangbumChoi and others added 14 commits March 27, 2026 08:56

Merge branch 'main' into molmo2

bc04776

fix(molmo2): add Molmo2TextModel to IGNORE_NON_TESTED

e38b0a3

Molmo2TextModel is an internal sub-component used by Molmo2Model and Molmo2ForConditionalGeneration and is tested implicitly through those. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

fix(molmo2): replace requests with stdlib urllib in video processor

cc06cbe

requests is not part of the standard library and caused ImportError in minimal environments (e.g. HuggingFace Jobs). Use urllib.request instead. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Merge branch 'main' into molmo2

f000173

Merge branch 'main' into molmo2

de8c268

zucchini-nlp self-requested a review March 30, 2026 15:23

zucchini-nlp reviewed Apr 2, 2026

View reviewed changes

SangbumChoi added 5 commits April 4, 2026 09:13

add comments fix

5d9811f

remove safe loading

3d6ec36

fix configuration with large refactor

2604a49

fix tests

e257e17

make formalize to standard way

d75fdbe

SangbumChoi force-pushed the molmo2 branch from 77a8c8c to 99c8582 Compare April 4, 2026 07:47

Merge branch 'main' into molmo2

91b6a78

SangbumChoi force-pushed the molmo2 branch from 99c8582 to 91b6a78 Compare April 4, 2026 07:51

Merge branch 'main' into molmo2

8d83a8a

molbap reviewed Jul 6, 2026

View reviewed changes

molbap and others added 9 commits July 6, 2026 17:06

minor changes

2903861

Merge branch 'main' into molmo2

597327d

simplify causal mapping

0e2a957

fix processor tests

8ca21f0

Merge branch 'main' into molmo2

05ae03e

merge main + conflict

cc7b3e1

update Expectations wrt original implem

1255831

Fix rope for 7B-O

b442aba

fix style

45bf547

zucchini-nlp reviewed Jul 21, 2026

View reviewed changes

molbap added 7 commits July 21, 2026 17:36

address review comments

16ee563

better modular, some more convert chunking

4919f36

skip test for correct reason + compile

2251b5a

forgot that one

3e3fa30

make style & fix-repo (linter issue?)

e3e3510

Merge branch 'main' into molmo2-pr-update

49fb4ce

why did make fix-repo change this 🤔

abbcb95

		from transformers import Molmo2ForConditionalGeneration, Molmo2Processor
		from transformers.video_utils import load_video

Uh oh!

Conversation

SangbumChoi commented Jan 23, 2026 • edited by github-actions Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Before submitting

Who can review?

Uh oh!

SangbumChoi commented Mar 29, 2026

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

SangbumChoi commented Jun 29, 2026

Uh oh!

SangbumChoi commented Jul 2, 2026

Uh oh!

SangbumChoi commented Jul 5, 2026

Uh oh!

molbap left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

SangbumChoi commented Jul 6, 2026

Uh oh!

molbap commented Jul 6, 2026

Uh oh!

molbap commented Jul 21, 2026

Uh oh!

github-actions Bot commented Jul 21, 2026

Uh oh!

github-actions Bot commented Jul 21, 2026

CI Results

Commit Info

Uh oh!

zucchini-nlp left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

SangbumChoi commented Jan 23, 2026 •

edited by github-actions Bot

Loading