feat(rl): multimodal context parallelism

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_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).** 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).

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -603,9 +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:
-            raise ValueError("VLM models are not supported with context parallelism")
+    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")

src/prime_rl/trainer/rl/train.py

@@ -432,27 +432,24 @@ 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_pre_shard = False
 
             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
-                )
+                # 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, position_ids = setup_cp_params(
+                        input_ids, position_ids, cp_rank, cp_size, cp_group, cp_style=config.model.cp_style
+                    )
+                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:
+                if routed_experts is not None and not defer_vlm_cp_to_model:
                     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:
-                    chunk_size = 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)
@@ -472,12 +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_pre_shard=seq_lens_are_pre_shard,
                     routed_experts=routed_experts,
                 )
 

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,10 +276,30 @@ 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():
+    config = TrainerConfig.model_validate(
+        {
+            "model": {
+                "cp": 2,
+                "cp_style": "ulysses",
+                "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_allows_text_only_cp():
     config = TrainerConfig.model_validate({"model": {"cp": 2}})