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

docs/advanced.md

@@ -155,10 +155,12 @@ The experimental `shard_fused_on_dim1 = true` shards fused 2-D weights along dim
 
 The built-in VLM registry covers:
 
-| Family | `model_type` | Vision attr | LM attr |
-|---|---|---|---|
-| Qwen3.5 | `qwen3_5` | `model.visual` | `model.language_model` |
-| Qwen3.5-MoE | `qwen3_5_moe` | `model.visual` | `model.language_model` |
+| Family | `model_type` |
+|---|---|
+| Qwen3.5 | `qwen3_5` |
+| Qwen3.5-MoE | `qwen3_5_moe` |
+
+The trainer finds the vision encoder at `model.visual` and the language model at `model.language_model`.
 
 ### Enabling VLM Mode
 
@@ -171,12 +173,10 @@ optimization_dtype = "bfloat16"
 reduce_dtype = "bfloat16"
 
 [model.vlm]
-vision_encoder_attr = "model.visual"
-language_model_attr = "model.language_model"
 # freeze_vision_encoder = true  # default; set false to fine-tune the encoder
 ```
 
-The weight-broadcast key prefix is derived as `{language_model_attr}.layers.` automatically.
+For `rl`, keep `[model.vlm]` in the shared `[model]` block but set the dtypes under `[trainer.model]`.
 
 VLM training requires a registered custom PrimeRL implementation.
 

docs/development.md

@@ -142,9 +142,9 @@ In the PR that adds the new model, you also need to provide a table covering the
 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 unified dense/MoE Qwen3.5 implementation in `models/qwen3_5/` is 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 vision encoder and custom text model, embeds tokens, runs vision, scatters image embeddings over placeholder tokens, and builds 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.
+2. **Composite VLM body.** A `*VLMModel` that holds the vision encoder and custom text model, embeds tokens, runs vision, scatters image embeddings over placeholder tokens, and builds MRoPE 3D positions from `mm_token_type_ids` whenever `image_grid_thw` is given (the renderer owns the token→modality mapping; the trainer always passes plain `position_ids`, which the model uses for text-only rows). Keep the vision encoder at `model.model.visual` and the text model at `model.model.language_model`: FSDP, freezing, activation checkpointing, compile, the MoE runtime and NCCL weight broadcast find them there (`trainer/vlm.py`). 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 `supports_packed_multimodal_training` on the VLM model 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 hold a `CPContext` on all modules which need to be CP-aware and shard embeds/positions inside the model after the vision merge; the trainers defer sharding to the model for MRoPE batches under ulysses.
+5. **Registration.** Register the composite `model_type` in `_CUSTOM_VLM_MAPPING` (`models/__init__.py`) so `get_model` dispatches to the custom class. RL turns raw image refs into vision kwargs in `materialize_images` (`trainer/vlm.py`), which today calls the Qwen image processor; a new family with different processor outputs extends it.
+6. **Context parallelism (optional).** CP-capable VLMs hold a `CPContext` on all modules which need to be CP-aware and shard embeds/positions inside the model after the vision merge; the trainers defer sharding to the model for every batch with images.
 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).

examples/extra/vlm/sft-moe.toml

@@ -12,8 +12,6 @@ ep = 8                      # keep in sync with deployment.num_train_gpus
 name = "auto"
 
 [model.vlm]
-vision_encoder_attr = "model.visual"
-language_model_attr = "model.language_model"
 freeze_vision_encoder = true
 
 [model.lora]

examples/extra/vlm/sft.toml

@@ -11,8 +11,6 @@ reduce_dtype = "bfloat16"
 name = "auto"
 
 [model.vlm]
-vision_encoder_attr = "model.visual"
-language_model_attr = "model.language_model"
 freeze_vision_encoder = true
 
 [model.lora]

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

@@ -505,10 +505,6 @@ def validate_cp_micro_batch_size(self):
     def validate_vlm_constraints(self):
         if self.model.vlm is None:
             return self
-        if self.model.optimization_dtype != "bfloat16" or self.model.reduce_dtype != "bfloat16":
-            raise ValueError(
-                "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference."
-            )
         if self.data.micro_batch_size != 1:
             raise ValueError("VLM SFT requires data.micro_batch_size = 1.")
         if self.val is not None and self.val.data.micro_batch_size != 1:

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

@@ -193,12 +193,6 @@ def resolve_project_dir(self):
 
 
 class VLMConfig(BaseConfig):
-    vision_encoder_attr: str
-    """Dotted attribute path to the vision encoder module (e.g. ``model.visual``)."""
-
-    language_model_attr: str
-    """Dotted attribute path to the language model module (e.g. ``model.language_model``)."""
-
     freeze_vision_encoder: bool = True
     """Freeze the vision encoder parameters during training."""
 
@@ -210,9 +204,6 @@ class BaseModelConfig(BaseConfig):
     trust_remote_code: bool = False
     """Trust remote code when initializing the tokenizer."""
 
-    vlm: "VLMConfig | None" = None
-    """VLM configuration. Setting this enables vision-language model support."""
-
 
 class DynamoConfig(BaseConfig):
     enabled: bool = True

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

@@ -15,6 +15,7 @@
     ResumeConfig,
     TrainerLogConfig,
     TransportConfig,
+    VLMConfig,
     WeightBroadcastConfig,
     ZMQTransportConfig,
 )
@@ -262,6 +263,9 @@ class MoERuntimeConfig(BaseConfig):
 
 
 class ModelConfig(BaseModelConfig):
+    vlm: VLMConfig | None = None
+    """VLM configuration. Setting this enables vision-language model support."""
+
     conversion_dir: Path | None = None
     """Directory for the auto-converted weights (written to a `prime`/`hf` subdirectory). If not set, we write into the model snapshot directory."""
 
@@ -330,6 +334,16 @@ class ModelConfig(BaseModelConfig):
     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")
+    def validate_vlm(self):
+        if self.vlm is None:
+            return self
+        if self.optimization_dtype != "bfloat16" or self.reduce_dtype != "bfloat16":
+            raise ValueError(
+                "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference."
+            )
+        return self
+
     @model_validator(mode="after")
     def validate_cp(self):
         if self.cp > 1 and self.attn not in ["flash_attention_2", "flash_attention_3", "flash_attention_4", "auto"]:

packages/prime-rl-configs/src/prime_rl/utils/validation.py

@@ -67,11 +67,7 @@ def propagate(shared_path: str, *targets: str) -> None:
         "inference.vllm.model",
         "orchestrator.model.name",
     )
-    propagate(
-        "model.vlm",
-        "trainer.model.vlm",
-        "orchestrator.model.vlm",
-    )
+    propagate("model.vlm", "trainer.model.vlm")
 
     # [log]
     propagate("log.level", "trainer.log.level", "orchestrator.log.level", "inference.log.level")

src/prime_rl/orchestrator/batch.py

@@ -369,11 +369,6 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch
     )
 
 
-def _is_multimodal_sample(sample: MicroBatch) -> bool:
-    """Check if a sample contains multimodal data (images)."""
-    return sample.mm_refs is not None
-
-
 @dataclass
 class _MicroBatchBin:
     samples: list[MicroBatch]
@@ -767,8 +762,8 @@ def prepare_batch(
     micro_batches = [pad_micro_batch(micro_batch, pad_to_multiple_of) for micro_batch in micro_batches]
 
     # 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)]
+    mm_batches = [b for b in micro_batches if b.mm_refs is not None]
+    text_batches = [b for b in micro_batches if b.mm_refs is None]
 
     # Pad each group independently so its count is divisible by num_train_workers
     mm_batches = _pad_group_for_distribution(mm_batches, num_train_workers)