feat(sft): multimodal (VLM) SFT by hubert-marek · Pull Request #3075 · PrimeIntellect-ai/prime-rl

hubert-marek · 2026-07-17T19:49:53Z

Summary

Rebased replacement for #2945, now targeting main after #3032 merged.

Carry renderer-produced multimodal samples end to end through SFT, including image-safe truncation, mm_token_type_ids, multimodal processor outputs, and cat packing.
Load the model processor, attach it to the renderer, and save it with weight checkpoints.
Add the dense Qwen3.5 composite VLM body (HF vision encoder + custom PrimeRL text model) with MRoPE positions derived from renderer metadata.
Require custom PrimeRL implementations for VLM training and keep vision collectives symmetric across text-only and multimodal microbatches.
Freeze vision encoders for text-only VLM training and fail invalid LoRA configurations with no matching targets.
Enforce the VLM dtype, micro-batch-size, LoRA, and no-CP constraints in SFT config validation.
Update the multimodal nightly config and add dense and MoE VLM SFT examples.

The branch preserves #3032's explicit seq_lens contract: trailing padding is merged into the final sequence segment and no separate padding metadata is transported.

Validation

23 passed — existing SFT dataset suite after the image-safe truncation fix
32 passed, 5 deselected — focused non-GPU VLM model/forward/SFT tests using an import-only ring_flash_attn shim; no attention kernels were exercised
41 passed, 80 deselected — config unit tests excluding external-plugin config loading
Ruff format/check and diff checks passed
Historical multimodal KL run: custom Qwen3.5 VLM, mean mismatch KL 0.0138, step-0 KL 0.0025
Historical HF control: Qwen3-VL, mean mismatch KL 0.1145

Post-rebase CI fixes: the exact examples/vlm_sft_moe/sft.toml config test and an sft --dry-run validation both pass. The nightly RL VLM config now sets impl = "custom" in [trainer.model], and obsolete GPU tests for removed modeling-layer validation were dropped.

Latest main merge: merged main at 445392294; the README overlap retains the current multimodal docs anchor and keeps VLM CP marked unsupported, while main's obsolete debug-config deletion is preserved.

Cleanup: multimodal CUDA transfer is inlined in the singleton SFT collate path; the single-use helper was removed.

API cleanup: [model.vlm] now requires model.impl = "custom". The dense and MoE custom VLM bodies accept token IDs only, share embedding accessors with their text-only bodies, and no longer carry unsupported inputs_embeds fallbacks or redundant modeling-layer validation. Multimodal samples without [model.vlm] still fail at the lazy dataset boundary.

Truncation follow-up: multimodal samples are now causal-shifted before image-safe truncation, keeping placeholder ranges, modality IDs, targets, and processor sidecars on the same model-input boundary.

Review follow-up: the multimodal processor is now loaded from model.name, so overriding tokenizer.name does not discard model-specific image preprocessing. configure_trainable_parameters is limited to LoRA application and frozen-vision discovery; router transforms and EP remain explicit in setup_model. LoRA keeps its existing initial freeze and logging, while EP performs the same conditional re-freeze as main after replacing parameters with DTensors. The VLM-specific vision freeze runs afterward.

Weight-tying follow-up: Qwen3.5 honors the checkpoint's standard tied-weight mapping. Checkpoints that omit lm_head.weight load the shared language embedding through the existing tied-model path without custom loader handling.

LoRA documentation follow-up: freeze_vision_encoder now promises parameter freezing only. The docs explicitly note that the default Qwen3.5 LoRA targets do not match vision modules, without claiming that arbitrary user-supplied target patterns are filtered.

Note

Medium Risk
Touches core training paths (model setup, forward, checkpoints, packing) and distributed VLM collectives; misconfiguration is mostly caught by validators, but multimodal truncation and vision/LoRA interactions need careful review.

Overview
Adds end-to-end multimodal SFT for Qwen3.5 dense and MoE when [model.vlm] is set with model.impl = "custom".

The SFT pipeline now carries renderer-produced mm_kwargs, mm_token_type_ids, and image-safe truncation through packing and the training forward; it loads an AutoProcessor from model.name, wires it into the renderer, and saves it with HF weight checkpoints. Config validation enforces VLM bfloat16 dtypes, micro_batch_size = 1, no CP, and rejects freeze_vision_encoder=false with LoRA; the old blocker that forbade VLM in renderer-only SFT is removed.

Modeling: Registers dense qwen3_5 in the custom VLM map and extends Qwen3_5ForCausalLM / MoE VLM bodies with HF vision + custom text, MRoPE from renderer metadata, and dummy vision passes on text-only steps for distributed collective symmetry. Trainability is centralized in configure_trainable_parameters (LoRA + optional vision freeze after EP).

Docs and CI shift the multimodal story to Qwen3.5 (CP marked unsupported for VLMs); new example TOMLs cover dense and MoE VLM SFT.

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

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 084d489008

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

Renderer-produced multimodal samples flow end to end through SFT: - SFTDataset keeps the renderer's MultiModalData: image-safe truncation (never cuts inside a placeholder run, drops truncated-out images, and skips samples whose EOS was truncated away), mm_token_type_ids kept aligned through EOS append and causal shift, video inputs rejected explicitly. _flatten_mm_items folds per-image processor outputs into model-forward kwargs (pixel_values, image_grid_thw, ...). - Cat packing concatenates mm_kwargs across samples in a pack and represents text-only spans as zeros in mm_token_type_ids; stack bucketing emits multimodal samples solo. Collates move mm tensors to CUDA and keep packed rows batch-size 1. - Local data_files normalization wraps string message content into typed content blocks under a multimodal flag so Arrow can unify text and image rows. - setup_processor loads the AutoProcessor and attaches it to the renderer; VLM SFT fails fast without one; weight checkpoints save the processor alongside the tokenizer. - Dense Qwen3.5 gains the composite VLM body (HF vision encoder + custom text model) with MRoPE positions from mm_token_type_ids and registers in _CUSTOM_VLM_MAPPING; VLM training requires a custom PrimeRL implementation. Both dense and MoE VLM bodies always run the vision encoder (dummy input on text-only microbatches, kept in the backward graph) for FSDP/EP collective symmetry. - Config: VLM SFT requires micro_batch_size=1; unfreezing the vision encoder is incompatible with LoRA; the renderer-vs-VLM rejection is gone (renderers own multimodal tokenization now). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

processor.save_pretrained rewrites tokenizer files from the processor's separately loaded tokenizer, clobbering the configured pad token and any custom chat template. Save the tokenizer last so it wins. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Without [model.vlm] the custom VLM class still loads and runs the vision tower on dummy input every microbatch, so its params carry zero — not None — grads and AdamW weight decay silently shrinks them. Freeze the tower when no VLM training is configured; mixed image/text batches cannot occur without [model.vlm], so collective symmetry is unaffected. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

…l.vlm] The processor auto-loads for VLM checkpoints and mm samples would flow through packing and forward unvalidated (no mbs/attn/CP checks, no freeze policy). Require the explicit [model.vlm] declaration instead. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Enforce the documented VLM dtype contract while reducing misleading conversion, packing, and documentation noise. Co-authored-by: Cursor <cursoragent@cursor.com>

Validate MRoPE inputs only when positions must be computed and rely on the dataset's established token-type alignment invariant. Co-authored-by: Cursor <cursoragent@cursor.com>

Keep the VLM validation and packing changes aligned with the repository formatter used by CI. Co-authored-by: Cursor <cursoragent@cursor.com>

Co-authored-by: Cursor <cursoragent@cursor.com>

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 56c34fcff1

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

S1ro1

would be a bit more skeptical about the llm generated tests also, they dont seem to test a lot and just pollute code
also data part needs checking

S1ro1 · 2026-07-18T23:39:57Z

+        if self._is_vlm:
+            self.model = Qwen3_5VLMModel(config)
+            text_config = config.text_config
+            self._tied_weights_keys = {"lm_head.weight": "model.language_model.embed_tokens.weight"}


I think tying weights breaks things in a lot of cases, do we need it?

chatgpt-codex-connector

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 7b7c64663e

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

Open a pull request for review
Mark a draft as ready
Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

samsja

wondering if it wouldn;'t be easier to just not support lora for vlm, kinda trying to keep the lora path as is for now

samsja

 and apply_ep

 Before: apply_lora → freeze_moe_router → apply_fp32_moe_router →
 freeze_sparse_indexer → apply_ep

 After: apply_lora → apply_fp32_moe_router → apply_ep → freeze_all_except_lora →
 freeze_moe_router → freeze_sparse_indexer

lets keep the otherd the same as before

cursor

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

There are 2 total unresolved issues (including 1 from previous review).

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

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

# Conflicts: # README.md # configs/debug/multimodal.toml

S1ro1

like it!

hubert-marek requested a review from samsja July 17, 2026 19:55

hubert-marek marked this pull request as ready for review July 17, 2026 21:52

cursor Bot reviewed Jul 17, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/__init__.py

Comment thread src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py

chatgpt-codex-connector Bot reviewed Jul 17, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/sft/data.py Outdated

Comment thread src/prime_rl/trainer/lora.py Outdated

Comment thread src/prime_rl/trainer/model.py Outdated

Comment thread src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py

hubert-marek and others added 14 commits July 18, 2026 00:11

Single set trainable modules function

ab61e0d

Cleanup

2f18f3e

Remove validation on shape - we only support bs = 1 for custom attention

7da8f79

fix(sft): tighten VLM constraints and remove dead code

cf31aa7

Enforce the documented VLM dtype contract while reducing misleading conversion, packing, and documentation noise. Co-authored-by: Cursor <cursoragent@cursor.com>

refactor(vlm): narrow multimodal validation

3300532

Validate MRoPE inputs only when positions must be computed and rely on the dataset's established token-type alignment invariant. Co-authored-by: Cursor <cursoragent@cursor.com>

style(sft): apply ruff formatting

f14cb88

Keep the VLM validation and packing changes aligned with the repository formatter used by CI. Co-authored-by: Cursor <cursoragent@cursor.com>

test(sft): align VLM packing with merged padding

e32a1af

Co-authored-by: Cursor <cursoragent@cursor.com>

fix(config): remove obsolete SFT loss field

d6c493b

Co-authored-by: Cursor <cursoragent@cursor.com>

refactor(sft): inline multimodal collate transfer

a62eae3

Co-authored-by: Cursor <cursoragent@cursor.com>

refactor(trainer): require token inputs in forward wrapper

928f22d

Co-authored-by: Cursor <cursoragent@cursor.com>

hubert-marek force-pushed the feat/sft-vlm branch from 084d489 to 928f22d Compare July 18, 2026 00:14

hubert-marek and others added 2 commits July 17, 2026 17:18

Merge branch 'main' into feat/sft-vlm

f8ed6db

fix(trainer): enforce LoRA trainability policy

56c34fc

chatgpt-codex-connector Bot reviewed Jul 18, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/sft/train.py Outdated

hubert-marek requested review from S1ro1 and eligotts July 18, 2026 23:22

S1ro1 reviewed Jul 18, 2026

View reviewed changes

hubert-marek mentioned this pull request Jul 19, 2026

feat(cp): context parallelism for hybrid + VLM SFT #2946

Closed

2 tasks

refactor(vlm): narrow custom model contract

7b7c646

cursor Bot reviewed Jul 19, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/sft/data.py

chatgpt-codex-connector Bot reviewed Jul 19, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/sft/data.py

Comment thread src/prime_rl/trainer/ckpt.py

samsja reviewed Jul 20, 2026

View reviewed changes

Comment thread examples/vlm_sft_moe/sft.toml

samsja reviewed Jul 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/model.py Outdated

samsja reviewed Jul 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/model.py Outdated

samsja reviewed Jul 20, 2026

View reviewed changes

eligotts reviewed Jul 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/lora.py Outdated

fix(sft): align multimodal truncation after shift

c7ad29e

eligotts reviewed Jul 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/model.py Outdated

hubert-marek added 2 commits July 20, 2026 21:59

fix(sft): address VLM processor review

a34246a

fix(trainer): materialize untied custom lm heads

3531688

cursor Bot reviewed Jul 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py

refactor(lora): drop vision module exclusions

c8a54d9

cursor Bot reviewed Jul 20, 2026

View reviewed changes

Comment thread src/prime_rl/trainer/model.py

hubert-marek added 4 commits July 20, 2026 22:39

refactor(trainer): use standard tied weight loading

877aa92

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

fece249

# Conflicts: # README.md # configs/debug/multimodal.toml

refactor(trainer): align LoRA EP freezing with main

719e71a

fix(ci): align VLM config and tests with validation

0c882df

S1ro1 approved these changes Jul 21, 2026

View reviewed changes

hubert-marek merged commit 0af5506 into main Jul 21, 2026
16 of 18 checks passed

hubert-marek mentioned this pull request Jul 21, 2026

feat(cp): context parallelism for hybrid + VLM SFT #3099

Merged

Uh oh!

Conversation

hubert-marek commented Jul 17, 2026 • edited by cursor Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Validation

Uh oh!

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

S1ro1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

S1ro1 Jul 18, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

chatgpt-codex-connector Bot left a comment

Choose a reason for hiding this comment

💡 Codex Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

samsja left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

samsja left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

cursor Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

S1ro1 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

hubert-marek commented Jul 17, 2026 •

edited by cursor Bot

Loading