From 9b2113033b82baf5caa4e14e60ba3cb1a036a916 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Fri, 3 Jul 2026 06:36:23 +0000 Subject: [PATCH 1/2] feat(rl): multimodal context parallelism MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror the SFT VLM CP path in the RL trainer, replacing the blanket NotImplementedError: - MRoPE multimodal batches under ulysses defer sharding to the model: input_ids/position_ids stay global so the vision encoder and image-embed merge see the full sequence; the model root forward shards embeds/positions after merge (and routed_experts with them). Text batches keep the existing pre-shard path. - seq_lens passes with seq_lens_are_global=True under CP — micro-batch boundaries span the pre-shard sequence and pad_to_multiple_of=cp already guarantees shard divisibility (padding is its own seq_lens segment). - The LoRA token-count CP adjustment derives chunk_size from the sharded length even when input_ids stays global. - Config gate relaxed: VLM + CP now requires cp_style='ulysses' and a custom model implementation instead of being rejected outright; resolve_pack_multimodal drops its CP rejection (packed MM rows meet the same divisibility guarantee). Co-Authored-By: Claude Fable 5 --- docs/development.md | 13 ++++++ .../src/prime_rl/configs/trainer.py | 6 ++- src/prime_rl/trainer/rl/train.py | 33 ++++++++++----- tests/unit/test_configs.py | 40 +++++++++++++++++++ 4 files changed, 80 insertions(+), 12 deletions(-) diff --git a/docs/development.md b/docs/development.md index 253c36fb29..06eb135e80 100644 --- a/docs/development.md +++ b/docs/development.md @@ -14,6 +14,7 @@ This page covers workflows for developing on `prime-rl` itself — running the t - [Implement the Modeling Code](#implement-the-modeling-code) - [Register a Mini Preset](#register-a-mini-preset) - [Run the Smoke Test](#run-the-smoke-test) +- [Adding a Custom VLM Implementation](#adding-a-custom-vlm-implementation) ## Test Suite @@ -138,3 +139,15 @@ Before merging a new model, you need to ensure the following: - The small smoke test passes. In the PR that adds the new model, you also need to provide a table covering the KL mismatch across 20 steps on `math` environment with `batch_size=64`. All the entries in the table must lower than 0.015. If this is not met, the PR will not be merged (unless reasonable justification is provided). This is to ensure all our models are consistent and their implementations match the implementations in the inference framework. + +## Adding a Custom VLM Implementation + +VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementation-only: `get_model` rejects models without a custom PrimeRL VLM class at load time. To make a new VLM family trainable, extend a custom text model with a composite VLM body — the Qwen3.5 dense (`models/qwen3_5/`) and MoE (`models/qwen3_5_moe/`) implementations are the reference. The pieces, in dependency order: + +1. **Custom text model first.** The VLM body wraps a custom `*ForCausalLM` (see [Adding a New Model](#adding-a-new-model)), so the text side — including its state-dict conversion and KL-mismatch table — comes first. +2. **Composite VLM body.** A `*VLMModel` that holds the HF vision encoder and the custom text model, with a `prepare_inputs_embeds_and_position_ids` step: embed tokens, run the vision encoder, scatter image embeddings over placeholder tokens, and build MRoPE 3D positions from `mm_token_type_ids` (the renderer owns the token→modality mapping). The unified `*ForCausalLM` dispatches on the config: composite config → VLM path, text config → text path. +3. **Always run the vision encoder.** Text-only micro-batches must feed the encoder dummy pixels and graft the result into the graph with zero contribution (`inputs_embeds + image_embeds.sum() * 0.0`) so FSDP/EP collectives stay symmetric across ranks when the encoder is trainable. +4. **Packed-boundary consumption.** Samples pack into shared rows with per-document boundaries in `seq_lens`; every custom model's `forward()` declares the typed `seq_lens`/`seq_lens_are_global` parameters (the trainer passes them unconditionally) and must honor the boundaries — varlen flash `cu_seqlens`, linear-attention state resets per document, and a loud rejection on attention paths that can't (see the packed-batch guard in any modeling file). Set `self.supports_packed_multimodal_training` on VLM configs once packed rows are handled — RL fails loudly at startup for VLM models without it. +5. **Registration.** Register the composite `model_type` in `_CUSTOM_VLM_MAPPING` (`models/__init__.py`) so `get_model` dispatches to the custom class, and describe the family in `VLM_REGISTRY` (`utils/vlm.py`). +6. **Context parallelism (optional).** Implement `set_context_parallel_attributes` and shard embeds/positions inside the model after the vision merge (the trainers defer CP sharding to the model for MRoPE batches under ulysses). Without this, CP configs are rejected for the family. +7. **Validation.** Same bar as text models: the KL-mismatch table for the text path, plus an SFT run and an RL run on a real multimodal dataset (the `color-codeword` environment is the reference task). diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index 051d8fab22..c6303477e4 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -605,7 +605,11 @@ def vlm_freeze_incompatible_with_lora(self): @model_validator(mode="after") def vlm_incompatible_with_cp(self): if self.model.vlm is not None and self.model.cp > 1: - raise ValueError("VLM models are not supported with context parallelism") + if self.model.cp_style != "ulysses": + raise ValueError("VLM models require cp_style='ulysses' for context parallelism") + # impl="auto" passes here; get_model rejects it at load time if it resolves to hf. + if self.model.impl == "hf": + raise ValueError("VLM context parallelism requires the custom model implementation") return self @model_validator(mode="after") diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 8324625641..616b21acf2 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -432,19 +432,28 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: labels = shift_tensor_left(input_ids) - # VLM + CP is not supported: MRoPE requires global positions but CP shards the sequence - if cp_enabled and mm_kwargs is not None: - raise NotImplementedError("Context parallelism is not supported with VLM/multimodal training") + seq_lens_are_global = False + # MRoPE MM batches keep global input_ids/position_ids: the vision + # encoder and image-embed merge need the full sequence, so the model + # shards after merge (as in SFT). + defer_vlm_cp_to_model = ( + cp_enabled + and mm_kwargs is not None + and "image_grid_thw" in mm_kwargs + and config.model.cp_style == "ulysses" + ) if cp_enabled: - # seq_lens spans the pre-shard sequence; consuming it under CP needs - # global-boundary handling that isn't wired up here yet. - seq_lens = None - input_ids, forward_position_ids = setup_cp_params( - input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style - ) + if not defer_vlm_cp_to_model: + input_ids, forward_position_ids = setup_cp_params( + input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style + ) + else: + forward_position_ids = position_ids + seq_lens_are_global = seq_lens is not None labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size) - if routed_experts is not None: + if routed_experts is not None and not defer_vlm_cp_to_model: + # The model shards routed_experts itself when CP is deferred. routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_size) else: forward_position_ids = position_ids @@ -452,7 +461,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: if config.model.lora: lora_num_tokens = micro_batch["lora_num_tokens"].to("cuda") if cp_enabled: - chunk_size = input_ids.shape[1] + # input_ids stays global when VLM CP defers sharding to the model. + chunk_size = input_ids.shape[1] // cp_size if defer_vlm_cp_to_model else input_ids.shape[1] # Convert to cumsum, adjust for CP chunk, convert back to num_tokens cu_offsets = lora_num_tokens.cumsum(dim=0, dtype=torch.int32) adjusted_cu = torch.clip(cu_offsets - chunk_size * cp_rank, min=0, max=chunk_size) @@ -478,6 +488,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: mm_kwargs=mm_kwargs, mm_token_type_ids=mm_token_type_ids, seq_lens=seq_lens, + seq_lens_are_global=seq_lens_are_global, routed_experts=routed_experts, ) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index dd8be2da1c..7f26453abe 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -280,6 +280,46 @@ def test_trainer_rejects_declared_vlm_cp(): TrainerConfig.model_validate(config) +def test_trainer_allows_vlm_cp_with_ulysses_custom_impl(): + config = TrainerConfig.model_validate( + { + "model": { + "cp": 2, + "cp_style": "ulysses", + "impl": "custom", + "optimization_dtype": "bfloat16", + "reduce_dtype": "bfloat16", + "vlm": { + "vision_encoder_attr": "model.visual", + "language_model_attr": "model.language_model", + }, + }, + } + ) + + assert config.model.vlm is not None + assert config.model.cp == 2 + + +def test_trainer_rejects_vlm_cp_with_hf_impl(): + with pytest.raises(ValidationError, match="custom model implementation"): + TrainerConfig.model_validate( + { + "model": { + "cp": 2, + "cp_style": "ulysses", + "impl": "hf", + "optimization_dtype": "bfloat16", + "reduce_dtype": "bfloat16", + "vlm": { + "vision_encoder_attr": "model.visual", + "language_model_attr": "model.language_model", + }, + }, + } + ) + + def test_trainer_allows_text_only_cp(): config = TrainerConfig.model_validate({"model": {"cp": 2}}) From bd09c630eb2d0f2f62760b60331d00837aeb0d01 Mon Sep 17 00:00:00 2001 From: hubert-marek Date: Thu, 9 Jul 2026 23:44:19 +0000 Subject: [PATCH 2/2] refactor(rl): align multimodal CP with SFT contract Use the pre-shard boundary API, rely on shared CP implementation validation, and derive local LoRA chunk sizes from already-sharded labels. Co-authored-by: Cursor --- docs/development.md | 4 +-- .../src/prime_rl/configs/trainer.py | 10 ++----- src/prime_rl/trainer/rl/train.py | 29 +++++-------------- tests/unit/test_configs.py | 26 ++--------------- 4 files changed, 16 insertions(+), 53 deletions(-) diff --git a/docs/development.md b/docs/development.md index 06eb135e80..6cfbaf2f44 100644 --- a/docs/development.md +++ b/docs/development.md @@ -147,7 +147,7 @@ VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementatio 1. **Custom text model first.** The VLM body wraps a custom `*ForCausalLM` (see [Adding a New Model](#adding-a-new-model)), so the text side — including its state-dict conversion and KL-mismatch table — comes first. 2. **Composite VLM body.** A `*VLMModel` that holds the HF vision encoder and the custom text model, with a `prepare_inputs_embeds_and_position_ids` step: embed tokens, run the vision encoder, scatter image embeddings over placeholder tokens, and build MRoPE 3D positions from `mm_token_type_ids` (the renderer owns the token→modality mapping). The unified `*ForCausalLM` dispatches on the config: composite config → VLM path, text config → text path. 3. **Always run the vision encoder.** Text-only micro-batches must feed the encoder dummy pixels and graft the result into the graph with zero contribution (`inputs_embeds + image_embeds.sum() * 0.0`) so FSDP/EP collectives stay symmetric across ranks when the encoder is trainable. -4. **Packed-boundary consumption.** Samples pack into shared rows with per-document boundaries in `seq_lens`; every custom model's `forward()` declares the typed `seq_lens`/`seq_lens_are_global` parameters (the trainer passes them unconditionally) and must honor the boundaries — varlen flash `cu_seqlens`, linear-attention state resets per document, and a loud rejection on attention paths that can't (see the packed-batch guard in any modeling file). Set `self.supports_packed_multimodal_training` on VLM configs once packed rows are handled — RL fails loudly at startup for VLM models without it. +4. **Packed-boundary consumption.** Samples pack into shared rows with per-document boundaries in `seq_lens`; every custom model's `forward()` declares the typed `seq_lens`/`seq_lens_are_pre_shard` parameters (the trainer passes them unconditionally) and must honor the boundaries — varlen flash `cu_seqlens`, linear-attention state resets per document, and a loud rejection on attention paths that can't (see the packed-batch guard in any modeling file). Set `self.supports_packed_multimodal_training` on VLM configs once packed rows are handled — RL fails loudly at startup for VLM models without it. 5. **Registration.** Register the composite `model_type` in `_CUSTOM_VLM_MAPPING` (`models/__init__.py`) so `get_model` dispatches to the custom class, and describe the family in `VLM_REGISTRY` (`utils/vlm.py`). -6. **Context parallelism (optional).** Implement `set_context_parallel_attributes` and shard embeds/positions inside the model after the vision merge (the trainers defer CP sharding to the model for MRoPE batches under ulysses). Without this, CP configs are rejected for the family. +6. **Context parallelism (optional).** CP-capable VLMs implement `set_context_parallel_attributes` and shard embeds/positions inside the model after the vision merge; the trainers defer sharding to the model for MRoPE batches under ulysses. 7. **Validation.** Same bar as text models: the KL-mismatch table for the text path, plus an SFT run and an RL run on a real multimodal dataset (the `color-codeword` environment is the reference task). diff --git a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py index c6303477e4..5fa522fd2d 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -603,13 +603,9 @@ def vlm_freeze_incompatible_with_lora(self): return self @model_validator(mode="after") - def vlm_incompatible_with_cp(self): - if self.model.vlm is not None and self.model.cp > 1: - if self.model.cp_style != "ulysses": - raise ValueError("VLM models require cp_style='ulysses' for context parallelism") - # impl="auto" passes here; get_model rejects it at load time if it resolves to hf. - if self.model.impl == "hf": - raise ValueError("VLM context parallelism requires the custom model implementation") + def validate_vlm_cp(self): + if self.model.vlm is not None and self.model.cp > 1 and self.model.cp_style != "ulysses": + raise ValueError("VLM models require cp_style='ulysses' for context parallelism") return self @model_validator(mode="after") diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 616b21acf2..d2ec0816b7 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -432,37 +432,24 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: labels = shift_tensor_left(input_ids) - seq_lens_are_global = False - # MRoPE MM batches keep global input_ids/position_ids: the vision - # encoder and image-embed merge need the full sequence, so the model - # shards after merge (as in SFT). - defer_vlm_cp_to_model = ( - cp_enabled - and mm_kwargs is not None - and "image_grid_thw" in mm_kwargs - and config.model.cp_style == "ulysses" - ) + seq_lens_are_pre_shard = False if cp_enabled: + # MRoPE batches must merge image embeddings before sharding. + defer_vlm_cp_to_model = mm_kwargs is not None and "image_grid_thw" in mm_kwargs if not defer_vlm_cp_to_model: - input_ids, forward_position_ids = setup_cp_params( + input_ids, position_ids = setup_cp_params( input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style ) - else: - forward_position_ids = position_ids - seq_lens_are_global = seq_lens is not None + seq_lens_are_pre_shard = seq_lens is not None labels = shard_for_cp(labels, cp_rank=cp_rank, cp_world_size=cp_size) if routed_experts is not None and not defer_vlm_cp_to_model: - # The model shards routed_experts itself when CP is deferred. routed_experts = shard_for_cp(routed_experts, cp_rank=cp_rank, cp_world_size=cp_size) - else: - forward_position_ids = position_ids if config.model.lora: lora_num_tokens = micro_batch["lora_num_tokens"].to("cuda") if cp_enabled: - # input_ids stays global when VLM CP defers sharding to the model. - chunk_size = input_ids.shape[1] // cp_size if defer_vlm_cp_to_model else input_ids.shape[1] + chunk_size = labels.shape[1] # Convert to cumsum, adjust for CP chunk, convert back to num_tokens cu_offsets = lora_num_tokens.cumsum(dim=0, dtype=torch.int32) adjusted_cu = torch.clip(cu_offsets - chunk_size * cp_rank, min=0, max=chunk_size) @@ -482,13 +469,13 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: out = forward( model, input_ids, - forward_position_ids, + position_ids, labels=labels, temperature=temperatures, mm_kwargs=mm_kwargs, mm_token_type_ids=mm_token_type_ids, seq_lens=seq_lens, - seq_lens_are_global=seq_lens_are_global, + seq_lens_are_pre_shard=seq_lens_are_pre_shard, routed_experts=routed_experts, ) diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 7f26453abe..3c6cf39e2d 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -263,7 +263,7 @@ def test_orchestrator_vlm_requires_renderer(): assert config.renderer is not None -def test_trainer_rejects_declared_vlm_cp(): +def test_trainer_rejects_vlm_cp_with_ring(): config = { "model": { "cp": 2, @@ -276,17 +276,16 @@ def test_trainer_rejects_declared_vlm_cp(): }, } - with pytest.raises(ValidationError, match="VLM models.*context parallelism"): + with pytest.raises(ValidationError, match="cp_style='ulysses'"): TrainerConfig.model_validate(config) -def test_trainer_allows_vlm_cp_with_ulysses_custom_impl(): +def test_trainer_allows_vlm_cp_with_ulysses(): config = TrainerConfig.model_validate( { "model": { "cp": 2, "cp_style": "ulysses", - "impl": "custom", "optimization_dtype": "bfloat16", "reduce_dtype": "bfloat16", "vlm": { @@ -301,25 +300,6 @@ def test_trainer_allows_vlm_cp_with_ulysses_custom_impl(): assert config.model.cp == 2 -def test_trainer_rejects_vlm_cp_with_hf_impl(): - with pytest.raises(ValidationError, match="custom model implementation"): - TrainerConfig.model_validate( - { - "model": { - "cp": 2, - "cp_style": "ulysses", - "impl": "hf", - "optimization_dtype": "bfloat16", - "reduce_dtype": "bfloat16", - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", - }, - }, - } - ) - - def test_trainer_allows_text_only_cp(): config = TrainerConfig.model_validate({"model": {"cp": 2}})