Enable VLM context parallel training#2909
Conversation
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Fail fast for unsupported multimodal packing runtimes and move the VLM context-parallelism incompatibility into config validation. Co-authored-by: Cursor <cursoragent@cursor.com>
Avoid an unnecessary multimodal sidecar copy in sample preparation and keep Qwen3-VL configs opted out of strict multimodal packing by default. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Simplify packed multimodal support detection to the root model flag and align unit fixtures with token-level advantages after the main merge. Co-authored-by: Cursor <cursoragent@cursor.com>
Reject trainer.pack_multimodal with context parallelism regardless of whether an explicit VLM config is present. Co-authored-by: Cursor <cursoragent@cursor.com>
(cherry picked from commit 8e4c8b3)
Two fixes that let context parallelism run for the Qwen3.5 VLM path (the
DeltaNet CP correctness fix landed separately):
1. Vision attention under CP. substitute_ulysses_attn patches the global
flash_attention_2 dispatch to do the Ulysses all-to-all, which is causal-
and sharded-sequence-only. The vision encoder runs *before* CP sharding
(full, replicated patches) with bidirectional attention, so it must not go
through that path. Pin the vision encoder to sdpa under CP so it dispatches
around the patched flash_attention_2 entry.
2. FSDP vs the pre-shard merge. The merge (prepare_vlm_inputs_for_context_
parallel) calls the vision encoder and embed_tokens standalone, before the
root FSDP forward. FSDP2 requires the first forward of an iteration to go
through the root, so a separately-sharded module called here lazy-inits as
its own root and then collides with the real root forward ("FSDP state has
already been lazily initialized"). Keep these pre-shard-merge modules out of
FSDP (full/replicated, materialized by model.to_empty) so the standalone
calls just work — but only when they are FROZEN, since ignored params get no
FSDP gradient reduction. Frozen -> ignore; trainable -> raise a clear error
pointing at LoRA / freeze_vision_encoder / the structural fix. This makes
CP+VLM work out of the box for frozen-embedding recipes (e.g. LoRA) and fails
loudly (instead of silently mis-training) for full fine-tuning.
Validation (2x H200, cp=2 ulysses, random-init Qwen3.5-35B-A3B, fake data):
- full fine-tune (trainable embed) -> raises the clear error, as intended.
- LoRA (embed + vision frozen) -> guard takes the ignore branch, FSDP setup
succeeds, reaches the training loop.
- frozen-embed run -> 6 steps end-to-end with images, finite loss, no crash.
Known follow-ups (out of scope here): full-parameter VLM-CP needs the
structural fix (run the pre-shard merge through the root forward); and full
rl-pipeline (orchestrator + inference) validation at real scale is still
required.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
(cherry picked from commit 7d83615)
Checkpoint current VLM CP cleanup so the branch can return to this state while review continues. Co-authored-by: Cursor <cursoragent@cursor.com>
Use a model-level dummy vision pass for text-only VLM batches so packing can stay generic while CP sharding still operates on the merged language sequence. Co-authored-by: Cursor <cursoragent@cursor.com>
Resolve the removed qwen35 RLM SWE config in favor of main while updating the branch. Co-authored-by: Cursor <cursoragent@cursor.com>
Restore empty trainer table headers for config readability and drop a redundant packing distribution test. Co-authored-by: Cursor <cursoragent@cursor.com>
hubert-marek
requested review from
S1ro1,
hallerite and
samsja
and removed request for
S1ro1,
hallerite and
samsja
![cursor[bot]](https://avatars.githubusercontent.com/in/1210556?s=60&v=4){kind=small}
| else: | ||
| for key in mm_kwargs: | ||
| mm_kwargs[key].data += sample.mm_kwargs[key].data | ||
| mm_kwargs[key].shape[0] += sample.mm_kwargs[key].shape[0] |
There was a problem hiding this comment.
Packed MM kwargs key mismatch
Medium Severity
When merging packed multimodal samples, _materialize_bin only concatenates keys already present in the accumulated mm_kwargs. Extra keys on a later sample are dropped, and a missing key on a later sample causes KeyError during packing.
Reviewed by Cursor Bugbot for commit 24eee16. Configure here.
There was a problem hiding this comment.
Not possible to happen under 1 model / 1 image processor
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 24eee168c9
ℹ️ 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".
| # Separate by modality so each step index has uniform modality across all ranks | ||
| mm_batches = [b for b in micro_batches if _is_multimodal_sample(b)] | ||
| text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)] | ||
| micro_batches = _pad_group_for_distribution(micro_batches, num_train_workers) |
There was a problem hiding this comment.
When the packed list contains separate image and text-only microbatches and num_train_workers > 1, distributing one combined pool can schedule an image microbatch on one rank and a text-only microbatch at the same step on another. For registered HF VLMs that do not have the new dummy-vision path (the custom VLM mapping only covers qwen3_5_moe), the text-only rank skips the FSDP-sharded vision encoder while the image rank enters it, reintroducing the unmatched all-gather/hang that the old mm/text grouping avoided.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
We do it on purpose - VLMs with CP need custom implementation
| return False | ||
| if (first_sample.routed_experts is None) != (sample.routed_experts is None): | ||
| return False | ||
| return True |
There was a problem hiding this comment.
When two Qwen-style multimodal samples fit in one bin, this unconditional acceptance packs them into a single sequence. trainer.model.forward() drops the trainer's reset 2-D position_ids whenever image_grid_thw is present, and only the custom qwen3_5_moe path consumes seq_lens; HF VLM paths such as qwen3_vl recompute MRoPE for the concatenation as one sample, so positions and causal attention cross sample boundaries and corrupt multimodal training.
Useful? React with 👍 / 👎.
Shard language embeddings normally under VLM context parallelism while keeping frozen vision params ignored. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 975b11721f
ℹ️ 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".
| # Separate by modality so each step index has uniform modality across all ranks | ||
| mm_batches = [b for b in micro_batches if _is_multimodal_sample(b)] | ||
| text_batches = [b for b in micro_batches if not _is_multimodal_sample(b)] | ||
| micro_batches = _pad_group_for_distribution(micro_batches, num_train_workers) |
There was a problem hiding this comment.
When VLM training uses an HF/generic VLM (for example the registered gemma4 path loaded through AutoModelForImageTextToText) and a step contains both text-only and image samples, this single mixed group can now be distributed so one data rank runs the vision encoder while another rank gets a text-only forward. Unlike the custom Qwen3.5 path, those HF text-only forwards do not run a dummy vision encoder, so FSDP ranks enter different vision-encoder collectives and can hang—the case the removed modality split was preventing. Keep modality grouping unless the selected model is known to execute a dummy vision path for text-only batches.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
VLM require custom implementations for CP
| if cu_seqlens is not None and int(cu_seqlens[-1].item()) == global_seq_len: | ||
| global_cu_seqlens = cu_seqlens.to(device=device, dtype=torch.int32) | ||
| else: | ||
| global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device) |
There was a problem hiding this comment.
When a Qwen3.5 VLM CP microbatch is all text-only, position_ids stay 2D and Qwen3_5MoeModel.forward recomputes cu_seqlens after sharding, so cu_seqlens[-1] is only the local shard length. This fallback then replaces the packed boundaries with [0, global_seq_len], making FLA's CP convolution/DeltaNet treat multiple packed text samples as one long sequence and leak state across sample boundaries, silently corrupting logprobs. Use the global seq_lens already passed by the packer for this path instead of collapsing to one segment.
Useful? React with 👍 / 👎.
| forward_cp_rank = None | ||
| forward_cp_world_size = None | ||
| forward_cp_style = None | ||
| vlm_cp_enabled = cp_enabled and config.model.vlm is not None |
There was a problem hiding this comment.
With CP enabled, this condition takes the root-level VLM CP path for every [model.vlm] configuration, but only the custom Qwen3.5 VLM forward consumes cp_group/cp_rank and shards the merged embeddings internally. Registered HF/generic VLMs such as gemma4 are loaded through AutoModelForImageTextToText, so they either reject the unexpected CP kwargs or run the full sequence on every CP rank while labels and temperatures are already sharded, producing mismatched/duplicated CP outputs. Restrict this path to models that implement the CP-aware root forward, or keep the normal setup_cp_params path for other VLMs.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
CP on VLMs require custom implementation
| return False | ||
| if (first_sample.routed_experts is None) != (sample.routed_experts is None): | ||
| return False | ||
| return True |
There was a problem hiding this comment.
When a non-Qwen/generic VLM sample has mm_kwargs without image_grid_thw (for example the registered HF gemma4 path), this now allows multiple image prompts to be concatenated into one microbatch. That forward path only sends reset position_ids and no attention mask or seq_lens, so HF causal attention treats the packed tokens and images as one long prompt; later samples can attend to earlier samples and corrupt training targets. Keep multimodal packing limited to models that actually consume packed boundaries, or leave these samples as separate microbatches.
Useful? React with 👍 / 👎.
Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a973064e36
ℹ️ 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".
| VLM models handle text-only samples with a synthetic dummy vision path, so | ||
| microbatches can be distributed through one generic packing path. |
There was a problem hiding this comment.
This assumption only holds for the custom Qwen3.5 VLM path; get_model still selects AutoModelForImageTextToText for registered VLMs without a custom class, such as Gemma/Qwen3-VL. In those HF VLM runs, text-only microbatches skip the FSDP-wrapped vision encoder while image microbatches enter it, and the combined distribution below can put those two modalities at the same micro-step on different ranks, recreating the vision-collective hang the deleted modality grouping prevented. Keep modality grouping for non-custom VLMs or add the dummy vision path there too.
Useful? React with 👍 / 👎.
Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
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 7218a53. Configure here.
| bias=self.conv1d.bias, | ||
| activation=self.activation, | ||
| cp_context=cp_context, | ||
| ) |
There was a problem hiding this comment.
Missing FLA conv import guard
High Severity
When Ulysses context parallelism builds a FLA cp_context, the linear-attention forward always calls fla_causal_conv1d, but that symbol is imported in a separate optional block and can be None. Training then crashes with a TypeError instead of using a fallback or a clear error, blocking VLM CP on environments where FLA’s CP conv module is unavailable.
Reviewed by Cursor Bugbot for commit 7218a53. Configure here.
Co-authored-by: Cursor <cursoragent@cursor.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2015e717aa
ℹ️ 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".
| def _qwen35_vlm_ulysses_equivalence_worker(rank: int, port: int): | ||
| os.environ["MASTER_ADDR"] = "127.0.0.1" | ||
| os.environ["MASTER_PORT"] = str(port) | ||
| dist.init_process_group("gloo", rank=rank, world_size=2) |
There was a problem hiding this comment.
When this GPU test runs, the worker builds the model inputs/logits on CUDA and the Ulysses path uses distributed all-to-all, followed by dist.all_gather, on those CUDA tensors through dist.group.WORLD. A Gloo process group is CPU-only for these collectives, so GPU CI will fail before the equivalence assertion is reached; use NCCL with one GPU per rank, or skip this test when that setup is unavailable.
Useful? React with 👍 / 👎.
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
3 participants
Summary
Enables Qwen3.5 VLM training with Ulysses context parallelism.
FSDP state has already been lazily initialized for model.language_model.embed_tokens) that happened when embedding/vision children were called before the FSDP root.causal_conv1dso conv/recurrent state crosses CP rank boundaries instead of resetting at the shard.seq_lensas the single source for packedcu_seqlensunder CP, so packed text-only microbatches don't leak state across sample boundaries.model.vlm+cp > 1requires a custom PrimeRL VLM implementation (generic HF VLMs remain supported atcp=1).PreTrainedModelPrimeRL.prime_forward_kwargs(...), Qwen3.5 returns theseq_lenskwargs it needs, and generic HF models plus unrelated custom models keep their normal forward signatures.Notes for new VLM models
VLM CP requires a custom implementation: the root forward must, when CP metadata (
cp_group/cp_rank/cp_world_size/cp_style) is present, do embedding + (real-or-dummy) vision merge + position IDs, then shard for CP. Generic HF VLMs don't satisfy this and are gated tocp=1.Models that need trainer-owned packed-boundary metadata should opt in via
prime_forward_kwargs(...)on their PrimeRL top-level model class instead of broadening every custom model forward signature.Verification
uvx ruff==0.13.0 format --check --config=pyproject.toml src/prime_rl/trainer/model.pypassed after the CI formatting failure.uv run --no-sync python -m py_compile src/prime_rl/trainer/model.py tests/unit/train/test_model_forward.py src/prime_rl/trainer/models/*/modeling_*.pypassed for the latest custom-model packed metadata interface.fatal error: cuda/atomic: No such file or directory) — environment issue, not logic.seq_lensunexpected-kwarg failure is gone. Five of six modules passed; onlybenchmark_regressionfailed because measured throughput/step time were much faster than the stored baseline.deps/pydantic-configis missing Python project metadata.Made with Cursor