refactor(trainer)!: fold multimodal into trainer/vlm.py; one VLM registry; drop vlm attr fields by faresobeid · Pull Request #3861 · PrimeIntellect-ai/prime-rl

faresobeid · 2026-10-04T12:40:05Z

Folds the trainer's multimodal code into one module, keeps one VLM registry, and removes VLM config fields that nothing reads.

What changes

One module. src/prime_rl/multimodal/ (5 files), trainer/multimodal.py, utils/vlm.py and setup_processor (from model.py) become src/prime_rl/trainer/vlm.py: get_vision_encoder, get_language_model, get_layer_prefix, load_image, setup_processor.
Image processing lives on the model. Each VLM class exposes materialize_images(refs, processor); the RL trainer calls model.materialize_images, so a new VLM family brings its own. Qwen3.5 and Qwen3.8 Flash Next each define materialize_qwen_vl_images in their own vision.py (the packages stay self-contained), which runs the Qwen-VL image processor and checks placeholder lengths against vLLM; the family-agnostic data-URL decoding is load_image in trainer/vlm.py.
Deleted: ForwardPolicy, MaterializedMM, the MultimodalAdapter Protocol and its registry, and the Kimi K2.5 / Qwen3-VL image paths. No Kimi K2.5 or Qwen3-VL trainer model exists, so get_model already rejected them.
One VLM registry. _CUSTOM_VLM_MAPPING (trainer/models/__init__.py) is the only one; VLM_REGISTRY and is_vlm_architecture are gone. The vision encoder, language model and NCCL layer prefix come from the model layout: model.model.visual and model.model.language_model, with model.model for text models (Muse Glimmer, Qwen3.8 Flash (qwen4_exp) and GLM-5.3 Flash (glm5_next) resolve the same way; they load no vision tower).
Qwen3.5 VLM owns its positions and CP. The model builds MRoPE positions whenever image_grid_thw is given; forward() always passes position_ids, used only for text-only rows. RL and SFT defer CP to the model for every batch with images. setup_context_parallel refuses cp > 1 for a VLM without cp_context.
Config (breaking). [model.vlm] keeps only freeze_vision_encoder; vision_encoder_attr / language_model_attr and orchestrator.model.vlm are removed (the orchestrator detects images from message content). vlm moves to the trainer ModelConfig and the shared [model.vlm] propagates only to trainer.model.vlm. The VLM bfloat16 optimization_dtype / reduce_dtype validator is removed. Old keys fail with "Extra inputs are not permitted".

Migration: delete vision_encoder_attr / language_model_attr and any [orchestrator.model.vlm] block.

Behavior

Training is unchanged for every trainable VLM (Qwen3.5 dense and MoE): same MRoPE positions, same CP shards, same component lookup. Configs with fp32 dtypes and [model.vlm] are no longer rejected.

Checks

pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train tests/unit/utils after merging main (1810d5b): 651 passed. ruff check / ruff format --check: clean. No dependency changes.
CPU, main vs this branch, tiny Qwen3.5 VLM through the trainer forward(): embeddings, positions and logits reaching the decoder are bit-identical for an image batch at cp=1 and ulysses cp=2, and for a text-only batch.
GPU, 1×H200, VLM SFT Qwen3.5-0.8B, 6 steps on image+text samples at cp=1 and ulysses cp=2: main and this branch match to 1e-4 on every step.

run	loss step 1 → 6	grad norm step 1 → 6
main and branch, cp=1 / cp=2	12.4225 → 12.1251	24.88 → 28.88

RL with images is not GPU-validated here: the nightly multimodal RL config stops on main and on this branch with the orchestrator error that #3949 fixes. #3949 and this PR share no files.

Note

Medium Risk
Breaking VLM config and a broad trainer refactor (FSDP, NCCL broadcast, CP, RL micro-batch prep), though trainable Qwen3.5 paths were validated for numerical parity.

Overview
Breaking config: [model.vlm] now only has freeze_vision_encoder. Remove vision_encoder_attr / language_model_attr and any [orchestrator.model.vlm] block. Shared [model.vlm] propagates to the trainer only. SFT no longer requires optimization_dtype / reduce_dtype bfloat16 for VLM runs.

Architecture cleanup: Deletes prime_rl/multimodal/, utils/vlm.py (VLM_REGISTRY, per-family attr overrides), and ForwardPolicy / MultimodalAdapter (including Kimi paths). VLM helpers live in trainer/vlm.py (encoder/LM lookup, load_image, setup_processor). Vision and LM modules are resolved from a fixed layout: model.model.visual and model.model.language_model.

RL image prep: Each custom VLM exposes materialize_images(refs, processor) (Qwen3.5 / Qwen3.8 implement materialize_qwen_vl_images in their vision.py). The RL data path calls model.materialize_images instead of a global adapter registry.

Forward / CP: Trainer forward() always passes position_ids; Qwen VLMs build MRoPE only when image_grid_thw is set. Context parallelism defers sharding to the model for any batch with mm_kwargs, and setup_context_parallel errors if a VLM lacks cp_context with cp > 1.

Docs and example TOMLs are updated to match.

^{Reviewed by Cursor Bugbot for commit bba4004. Bugbot is set up for automated code reviews on this repo. Configure here.}

Stacked on #3861. Deletes the Kimi K2.5 and Qwen3-VL image paths from `trainer/vlm.py`. `materialize_images(refs, processor)` now always uses the Qwen3.5 image processor (`pixel_values` + `image_grid_thw`, with the placeholder-length check). ## Why - Neither path can run. `get_model` rejects any `[model.vlm]` model whose `model_type` is not in `_CUSTOM_VLM_MAPPING` (`qwen3_5`, `qwen3_5_moe`, `qwen4_exp`). There is no Kimi K2.5 or Qwen3-VL trainer model. - RL also requires `supports_packed_multimodal_training`, which only the Qwen3.5 VLM sets. - A Kimi image path belongs with a future Kimi VLM model class. Re-adding it there is about 15 lines. ## Behavior None for trainable models. Qwen3.5 image batches go through the same processor call and the same length check as before. Kimi/Qwen3-VL `model_type`s used to fail in `get_model`, and they still do. ## Checks - `pytest -m "not gpu" tests/unit/orchestrator tests/unit/test_configs.py tests/unit/train tests/unit/transports`: all pass (440 + `test_cpu_adam.py` 3/3, run with `ninja` on `PATH`) - `ruff check` + `ruff format --check`: clean. - GPU: none needed beyond #3861. The Qwen3.5 code path is unchanged from #3861, whose SFT check covered it. ## Conflicts None beyond #3861's. This PR touches only `trainer/vlm.py`, `rl/train.py` (one call site), `tests/unit/train/test_multimodal.py` and `docs/development.md` (one sentence).  --- > [!NOTE] > **Low Risk** > Dead-code removal and API simplification; the active Qwen3.5 multimodal path is unchanged and Kimi/Qwen3-VL types were already rejected at model load. > > **Overview** > Removes **model-type dispatch** for RL image preprocessing: `materialize_images` now takes only `(refs, processor)` and always runs the **Qwen3.5-style** path (`pixel_values` + `image_grid_thw`, merge-size placeholder length check against vLLM). > > The **Kimi K2.5** branch (`grid_thws`, `preprocess` media dict) and the **`_IMAGE_INPUTS` map** for `qwen3_vl` / `qwen3_vl_moe` are deleted. The RL trainer call site drops `model.config.model_type`. Docs note that new VLM families must **extend** `materialize_images` if their processor outputs differ; unit tests drop the Kimi case and match the slimmer API. > > **No intended change** for currently trainable Qwen3.5 VLMs—the same processor call and length validation as the old Qwen entries. > > <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit ba6ca9b. Bugbot is set up for automated code reviews on this repo. Configure [here](https://www.cursor.com/dashboard/bugbot).</sup>

…stry - Move src/prime_rl/multimodal/, trainer/multimodal.py, utils/vlm.py and setup_processor into trainer/vlm.py. Delete ForwardPolicy, MaterializedMM, the adapter Protocol/registry and the Kimi K2.5 / Qwen3-VL image paths; materialize_images() returns the Qwen3.5 vision kwargs. - One VLM registry: _CUSTOM_VLM_MAPPING. Vision encoder, language model and weight-broadcast layer prefix come from the model layout (model.model.visual, model.model.language_model). - Qwen3.5 VLM builds MRoPE positions whenever image_grid_thw is given; the trainer always passes position_ids, and RL and SFT defer CP to the model for every batch with images. - Drop [model.vlm].vision_encoder_attr / language_model_attr and the unread orchestrator.model.vlm (breaking config change).

… too); one validate_vlm

# Conflicts: # tests/unit/orchestrator/test_batch.py

hubert-marek · 2026-10-21T13:15:14Z

    fused_lm_head_token_chunk_size: int | Literal["disabled"] = 8192
    """Flattened token chunk size for the fused LM head. ``int >= 1`` sets the tokens per LM-head chunk explicitly; ``disabled`` uses the vanilla LM head. In SFT the fused head computes the summed cross-entropy and its gradients chunk by chunk, holding one chunk's full-vocab logits at a time."""

+    @model_validator(mode="after")


Is there a reason we still validate it?

hubert-marek · 2026-10-21T13:27:49Z

+        return image.convert("RGB")
+
+
+def materialize_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]:


Hm I don't really like it - it put's model specific image processing login inside the trainer, e.g. Nemoron-VL will have different logic than Qwen. Can we keep this in model module, and trainer calls the image processing function inside model?

hubert-marek · 2026-10-21T13:35:31Z

-                )
+        if image_grid_thw is not None:
+            # Image inputs get MRoPE positions; text-only inputs keep the caller's positions.
+            if mm_token_type_ids is None:


Do we need this check?

hubert-marek · 2026-10-21T13:39:42Z

+                spatial_merge_size=self.config.vision_config.spatial_merge_size,
+                seq_lens=seq_lens,
+            )
+        elif position_ids is None:


Shoudn't those always be passed rather than reconstructed if not present?

hubert-marek · 2026-10-21T13:49:02Z

            # CP requires the sequence length to be divisible by cp_size. CatDataset
            # pads every pack to a multiple of it; shard_for_cp raises on violations.
-            defer_vlm_cp_to_model = mm_kwargs is not None and "image_grid_thw" in mm_kwargs
+            # The VLM shards image inputs itself, after building their positions.


Maybe we should have some check to verify if custom cp handling is implemented, so it doesn't fail later with some rank hang?

…at cannot shard its inputs

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit fecaf4e. Configure here.}

# Conflicts: # src/prime_rl/utils/vlm.py

# Conflicts: # tests/unit/train/models/test_glm5_next.py

# Conflicts: # src/prime_rl/trainer/model.py

…_vl.py Qwen3.8 Flash Next no longer imports the Qwen3.5 modeling module; both point materialize_images at materialize_qwen_vl_images. The family-agnostic data-URL decoding moves to trainer/vlm.py as load_image. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

… vision.py Model packages stay self-contained, like the duplicated vision towers: Qwen3.5 and Qwen3.8 Flash Next each define materialize_qwen_vl_images next to their vision model, and the top-level models/qwen_vl.py goes away. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

hubert-marek

LGTM assuming tested e2e

hubert-marek

LGTM assuming tested e2e

faresobeid mentioned this pull request Oct 4, 2026

chore(configs)!: drop vlm attr-path fields and the unread orchestrator.model.vlm #3852

Closed

faresobeid marked this pull request as ready for review October 4, 2026 12:42

This was referenced Oct 4, 2026

refactor(trainer): drop the Kimi K2.5 and Qwen3-VL image paths #3862

Merged

fix(qwen3_5): tie lm_head to embed_tokens for tied checkpoints #3863

Closed

fix(qwen3_5): build the FLA CP context once, outside checkpointed/compiled layers #3864

Closed

samsja closed this Oct 4, 2026

faresobeid reopened this Oct 8, 2026

faresobeid force-pushed the refactor/trainer-vlm branch from 2091ecd to a7293b2 Compare October 8, 2026 12:21

faresobeid changed the title ~~refactor(trainer): fold multimodal into trainer/vlm.py; one VLM registry~~ refactor(trainer)!: fold multimodal into trainer/vlm.py; one VLM registry; drop vlm attr fields Oct 8, 2026

faresobeid force-pushed the refactor/trainer-vlm branch 3 times, most recently from 4008f0e to eaf5654 Compare October 12, 2026 18:32

faresobeid added 3 commits October 17, 2026 18:45

fix(tools): import setup_processor from trainer.vlm

c9aa101

fix(configs): enforce VLM bf16 dtypes in the trainer model config (RL…

6113c9a

… too); one validate_vlm

faresobeid force-pushed the refactor/trainer-vlm branch from 6342813 to 6113c9a Compare October 17, 2026 18:50

faresobeid added 3 commits October 19, 2026 08:48

Merge branch 'main' into refactor/trainer-vlm

24447b9

Merge remote-tracking branch 'origin/main' into refactor/trainer-vlm

969672e

Merge remote-tracking branch 'origin/main' into vlm

eaf0ef0

# Conflicts: # tests/unit/orchestrator/test_batch.py

hubert-marek self-requested a review October 21, 2026 13:10

Merge branch 'main' into refactor/trainer-vlm

82a65e9

hubert-marek reviewed Oct 21, 2026

View reviewed changes

hubert-marek requested a review from eligotts October 21, 2026 13:53

faresobeid added 2 commits October 22, 2026 12:59

fix(trainer): positions come from the caller; refuse cp with a VLM th…

a976bd3

…at cannot shard its inputs

Merge remote-tracking branch 'origin/main' into refactor/trainer-vlm

44e9418

faresobeid added 2 commits October 24, 2026 18:34

refactor(configs): drop the VLM bf16 dtype validator

fef8808

refactor(trainer): materialize_images lives on the Qwen3.5 model class

fecaf4e

faresobeid mentioned this pull request Oct 24, 2026

fix(orchestrator): skip branches whose image placeholders do not match their images #3949

Closed

cursor Bot reviewed Oct 24, 2026

View reviewed changes

Comment thread packages/prime-rl-configs/src/prime_rl/configs/sft.py Resolved

faresobeid and others added 6 commits October 26, 2026 11:52

Merge remote-tracking branch 'origin/main' into refactor/trainer-vlm

5130274

# Conflicts: # src/prime_rl/utils/vlm.py

Merge remote-tracking branch 'origin/main' into refactor/trainer-vlm

34388d1

# Conflicts: # tests/unit/train/models/test_glm5_next.py

Merge remote-tracking branch 'origin/main' into refactor/trainer-vlm

f7af0f8

Merge remote-tracking branch 'origin/main' into r8-3861

2fc44f4

# Conflicts: # src/prime_rl/trainer/model.py

hubert-marek self-requested a review October 29, 2026 17:39

hubert-marek approved these changes Oct 29, 2026

View reviewed changes

faresobeid merged commit f98f993 into main Oct 29, 2026
11 checks passed

faresobeid deleted the refactor/trainer-vlm branch October 29, 2026 19:40

faresobeid mentioned this pull request Oct 30, 2026

feat: train on explicit multimodal placeholder positions, not image token ids #4343

Draft

Navigation Menu

Repository navigation

refactor(trainer)!: fold multimodal into trainer/vlm.py; one VLM registry; drop vlm attr fields - #3861

refactor(trainer)!: fold multimodal into trainer/vlm.py; one VLM registry; drop vlm attr fields#3861
faresobeid merged 17 commits into
mainfrom
refactor/trainer-vlm

faresobeid commented Oct 4, 2026 •

edited by cursor Bot

Loading

Uh oh!

hubert-marek Oct 21, 2026

Uh oh!

hubert-marek Oct 21, 2026

Uh oh!

hubert-marek Oct 21, 2026 •

edited

Loading

Uh oh!

hubert-marek Oct 21, 2026

Uh oh!

hubert-marek Oct 21, 2026

Uh oh!

cursor Bot left a comment

Uh oh!

Uh oh!

hubert-marek left a comment

Uh oh!

hubert-marek left a comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		return image.convert("RGB")


		def materialize_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]:

Conversation

faresobeid commented Oct 4, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What changes

Behavior

Checks

Uh oh!

hubert-marek Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

hubert-marek Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

hubert-marek Oct 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hubert-marek Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

hubert-marek Oct 21, 2026

Choose a reason for hiding this comment

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

hubert-marek left a comment

Choose a reason for hiding this comment

Uh oh!

hubert-marek left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

faresobeid commented Oct 4, 2026 •

edited by cursor Bot

Loading

hubert-marek Oct 21, 2026 •

edited

Loading