From 770b2e65b7db6fce6c3bf2144eebd434e7bcfd96 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Thu, 8 Oct 2026 12:20:18 +0000 Subject: [PATCH 1/8] refactor(trainer)!: fold multimodal into trainer/vlm.py; one VLM registry - Move src/prime_rl/multimodal/, trainer/multimodal.py, utils/vlm.py and setup_processor into trainer/vlm.py. Delete ForwardPolicy, MaterializedMM, the adapter Protocol/registry and the Kimi K2.5 / Qwen3-VL image paths; materialize_images() returns the Qwen3.5 vision kwargs. - One VLM registry: _CUSTOM_VLM_MAPPING. Vision encoder, language model and weight-broadcast layer prefix come from the model layout (model.model.visual, model.model.language_model). - Qwen3.5 VLM builds MRoPE positions whenever image_grid_thw is given; the trainer always passes position_ids, and RL and SFT defer CP to the model for every batch with images. - Drop [model.vlm].vision_encoder_attr / language_model_attr and the unread orchestrator.model.vlm (breaking config change). --- .../ci/nightly/multimodal_color_codeword.toml | 2 - docs/advanced.md | 14 +- docs/development.md | 6 +- examples/extra/vlm/sft-moe.toml | 2 - examples/extra/vlm/sft.toml | 2 - .../src/prime_rl/configs/shared.py | 9 -- .../src/prime_rl/configs/trainer.py | 4 + .../src/prime_rl/utils/validation.py | 6 +- src/prime_rl/multimodal/__init__.py | 9 -- src/prime_rl/multimodal/base.py | 40 ------ src/prime_rl/multimodal/kimi_k25.py | 33 ----- src/prime_rl/multimodal/qwen_vl.py | 35 ----- src/prime_rl/multimodal/registry.py | 17 --- src/prime_rl/orchestrator/batch.py | 9 +- src/prime_rl/trainer/model.py | 57 ++------ src/prime_rl/trainer/models/layers/lm_head.py | 4 +- .../models/qwen3_5/modeling_qwen3_5.py | 26 ++-- src/prime_rl/trainer/moe_runtime.py | 6 +- src/prime_rl/trainer/multimodal.py | 26 ---- src/prime_rl/trainer/rl/data.py | 17 +-- src/prime_rl/trainer/rl/train.py | 14 +- src/prime_rl/trainer/sft/train.py | 7 +- src/prime_rl/trainer/vlm.py | 69 ++++++++++ src/prime_rl/transports/weights/nccl.py | 4 +- src/prime_rl/utils/vlm.py | 123 ------------------ tests/unit/orchestrator/test_batch.py | 3 +- tests/unit/test_configs.py | 35 +---- .../train/models/test_qwen3_8_flash_next.py | 2 +- tests/unit/train/test_model_forward.py | 87 ------------- tests/unit/train/test_multimodal.py | 32 +---- 30 files changed, 132 insertions(+), 568 deletions(-) delete mode 100644 src/prime_rl/multimodal/__init__.py delete mode 100644 src/prime_rl/multimodal/base.py delete mode 100644 src/prime_rl/multimodal/kimi_k25.py delete mode 100644 src/prime_rl/multimodal/qwen_vl.py delete mode 100644 src/prime_rl/multimodal/registry.py delete mode 100644 src/prime_rl/trainer/multimodal.py create mode 100644 src/prime_rl/trainer/vlm.py delete mode 100644 src/prime_rl/utils/vlm.py delete mode 100644 tests/unit/train/test_model_forward.py diff --git a/configs/ci/nightly/multimodal_color_codeword.toml b/configs/ci/nightly/multimodal_color_codeword.toml index 46999258bbd..9e2050d3711 100644 --- a/configs/ci/nightly/multimodal_color_codeword.toml +++ b/configs/ci/nightly/multimodal_color_codeword.toml @@ -10,8 +10,6 @@ num_infer_gpus = 4 name = "Qwen/Qwen3.5-4B" [model.vlm] -vision_encoder_attr = "model.visual" -language_model_attr = "model.language_model" [orchestrator] batch_size = 256 diff --git a/docs/advanced.md b/docs/advanced.md index 4d11cb52c44..d6d34999c9e 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -142,10 +142,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 @@ -158,13 +160,9 @@ 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. - VLM training requires a registered custom PrimeRL implementation. ### Limitations diff --git a/docs/development.md b/docs/development.md index 98eb4a3fec5..16af61a541b 100644 --- a/docs/development.md +++ b/docs/development.md @@ -143,9 +143,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). diff --git a/examples/extra/vlm/sft-moe.toml b/examples/extra/vlm/sft-moe.toml index 1fedb448b51..1c5027abe65 100644 --- a/examples/extra/vlm/sft-moe.toml +++ b/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] diff --git a/examples/extra/vlm/sft.toml b/examples/extra/vlm/sft.toml index a0152f16d42..c98a56daeb6 100644 --- a/examples/extra/vlm/sft.toml +++ b/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] diff --git a/packages/prime-rl-configs/src/prime_rl/configs/shared.py b/packages/prime-rl-configs/src/prime_rl/configs/shared.py index 1ce168817a9..d2bed9e5ffd 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/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 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 3db4fc4a015..15a30307205 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -15,6 +15,7 @@ ResumeConfig, TrainerLogConfig, TransportConfig, + VLMConfig, WeightBroadcastConfig, ZMQTransportConfig, ) @@ -232,6 +233,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.""" diff --git a/packages/prime-rl-configs/src/prime_rl/utils/validation.py b/packages/prime-rl-configs/src/prime_rl/utils/validation.py index 41f2b610d44..6f168290246 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/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") diff --git a/src/prime_rl/multimodal/__init__.py b/src/prime_rl/multimodal/__init__.py deleted file mode 100644 index 7bc9ff129ea..00000000000 --- a/src/prime_rl/multimodal/__init__.py +++ /dev/null @@ -1,9 +0,0 @@ -from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, MultimodalAdapter -from prime_rl.multimodal.registry import get_multimodal_adapter - -__all__ = [ - "ForwardPolicy", - "MaterializedMM", - "MultimodalAdapter", - "get_multimodal_adapter", -] diff --git a/src/prime_rl/multimodal/base.py b/src/prime_rl/multimodal/base.py deleted file mode 100644 index 54fbd089e61..00000000000 --- a/src/prime_rl/multimodal/base.py +++ /dev/null @@ -1,40 +0,0 @@ -from __future__ import annotations - -from dataclasses import dataclass -from typing import Any, Protocol - -import torch -from PIL.Image import Image - - -@dataclass(frozen=True) -class ForwardPolicy: - pass_position_ids: bool = True - requires_mm_token_type_ids: bool = False - defer_context_parallelism: bool = False - - -@dataclass(frozen=True) -class MaterializedMM: - kwargs: dict[str, torch.Tensor] - forward_policy: ForwardPolicy - - -class MultimodalAdapter(Protocol): - model_types: frozenset[str] - forward_policy: ForwardPolicy - - def materialize( - self, - image_processor: Any, - images: list[Image], - placeholder_lengths: list[int], - ) -> MaterializedMM: ... - - -def required_tensors(values: Any, keys: tuple[str, ...]) -> dict[str, torch.Tensor]: - data = dict(values) - missing = [key for key in keys if key not in data] - if missing: - raise ValueError(f"Image processor did not return {', '.join(missing)}") - return {key: torch.as_tensor(data[key]).contiguous() for key in keys} diff --git a/src/prime_rl/multimodal/kimi_k25.py b/src/prime_rl/multimodal/kimi_k25.py deleted file mode 100644 index d4efa501cdd..00000000000 --- a/src/prime_rl/multimodal/kimi_k25.py +++ /dev/null @@ -1,33 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from PIL.Image import Image - -from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, required_tensors - - -class KimiK25Adapter: - model_types = frozenset({"kimi_k25"}) - forward_policy = ForwardPolicy() - - def materialize( - self, - image_processor: Any, - images: list[Image], - placeholder_lengths: list[int], - ) -> MaterializedMM: - preprocess = getattr(image_processor, "preprocess", None) - if preprocess is None: - raise ValueError("Kimi image processor is missing preprocess") - media = [{"type": "image", "image": image} for image in images] - kwargs = required_tensors( - preprocess(media, return_tensors="pt"), - ("pixel_values", "grid_thws"), - ) - lengths = [1] * len(kwargs["grid_thws"].reshape(-1, 3)) - if lengths != placeholder_lengths: - raise ValueError( - f"Kimi image placeholder lengths differ from vLLM: expected {placeholder_lengths}, got {lengths}" - ) - return MaterializedMM(kwargs=kwargs, forward_policy=self.forward_policy) diff --git a/src/prime_rl/multimodal/qwen_vl.py b/src/prime_rl/multimodal/qwen_vl.py deleted file mode 100644 index 4cf3c7a78ea..00000000000 --- a/src/prime_rl/multimodal/qwen_vl.py +++ /dev/null @@ -1,35 +0,0 @@ -from __future__ import annotations - -from typing import Any - -from PIL.Image import Image - -from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, required_tensors - - -class QwenVLAdapter: - model_types = frozenset({"qwen3_vl", "qwen3_vl_moe", "qwen3_5", "qwen3_5_moe"}) - forward_policy = ForwardPolicy( - pass_position_ids=False, - requires_mm_token_type_ids=True, - defer_context_parallelism=True, - ) - - def materialize( - self, - image_processor: Any, - images: list[Image], - placeholder_lengths: list[int], - ) -> MaterializedMM: - kwargs = required_tensors( - image_processor(images=images, return_tensors="pt"), - ("pixel_values", "image_grid_thw"), - ) - merge_size = int(image_processor.merge_size) - # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. - lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] - if lengths != placeholder_lengths: - raise ValueError( - f"Qwen image placeholder lengths differ from vLLM: expected {placeholder_lengths}, got {lengths}" - ) - return MaterializedMM(kwargs=kwargs, forward_policy=self.forward_policy) diff --git a/src/prime_rl/multimodal/registry.py b/src/prime_rl/multimodal/registry.py deleted file mode 100644 index d86bc96bd2e..00000000000 --- a/src/prime_rl/multimodal/registry.py +++ /dev/null @@ -1,17 +0,0 @@ -from __future__ import annotations - -from prime_rl.multimodal.base import MultimodalAdapter -from prime_rl.multimodal.kimi_k25 import KimiK25Adapter -from prime_rl.multimodal.qwen_vl import QwenVLAdapter - -_ADAPTERS = (QwenVLAdapter(), KimiK25Adapter()) -_BY_MODEL_TYPE: dict[str, MultimodalAdapter] = { - model_type: adapter for adapter in _ADAPTERS for model_type in adapter.model_types -} - - -def get_multimodal_adapter(model_type: str) -> MultimodalAdapter: - try: - return _BY_MODEL_TYPE[model_type] - except KeyError as exc: - raise NotImplementedError(f"Raw image training is not implemented for model type {model_type!r}") from exc diff --git a/src/prime_rl/orchestrator/batch.py b/src/prime_rl/orchestrator/batch.py index 1417ed2bdea..32795d68bec 100644 --- a/src/prime_rl/orchestrator/batch.py +++ b/src/prime_rl/orchestrator/batch.py @@ -466,11 +466,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] @@ -842,8 +837,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) diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 14944637e08..649a9fced2d 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -35,7 +35,6 @@ MXFP8Config, TokenizerConfig, ) -from prime_rl.multimodal import ForwardPolicy from prime_rl.trainer.activation_checkpointing import get_activation_checkpoint_wrapper from prime_rl.trainer.distributed.embedding_parallel import EmbeddingParallel from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict @@ -64,10 +63,10 @@ from prime_rl.trainer.moe_runtime import configure_moe_runtime from prime_rl.trainer.parallel_dims import ParallelDims from prime_rl.trainer.utils import get_dcp_group +from prime_rl.trainer.vlm import get_language_model, get_vision_encoder from prime_rl.trainer.world import get_world, wait_for_master from prime_rl.utils.logger import get_logger from prime_rl.utils.utils import format_time -from prime_rl.utils.vlm import get_language_model, get_vision_encoder, is_vlm_architecture from prime_rl.utils.weights import ( load_state_dict, load_state_dict_keys, @@ -118,11 +117,8 @@ def pre_download_model(model_name: str, *, skip_weights: bool = False) -> None: torch._dynamo.config.cache_size_limit = 64 # default: 8 -def freeze_vision_encoder(model: nn.Module, override_attr: str | None = None) -> None: +def freeze_vision_encoder(vision_encoder: nn.Module) -> None: logger = get_logger() - vision_encoder = get_vision_encoder(model, override=override_attr) - if vision_encoder is None: - raise ValueError("Could not find vision encoder to freeze") num_frozen = 0 for param in vision_encoder.parameters(): param.requires_grad = False @@ -338,8 +334,6 @@ def get_model( f"Loading model config (name={config.name}, attn={config.attn}, trust_remote_code={config.trust_remote_code})" ) - is_vlm_training = config.vlm is not None - model_config = cast( PretrainedConfig, AutoConfig.from_pretrained( @@ -347,9 +341,8 @@ def get_model( ), ) model_config.use_cache = False - is_vlm_arch = is_vlm_architecture(model_config) - if is_vlm_training: + if config.vlm is not None: logger.info(f"Detected vision-language model: {config.name}") for subconfig_key in getattr(model_config, "sub_configs", {}): @@ -409,7 +402,7 @@ def get_model( ) target_config.num_hidden_layers = num_hidden_layers - custom_vlm_cls = get_custom_vlm_cls(model_config) if is_vlm_arch else None + custom_vlm_cls = get_custom_vlm_cls(model_config) if custom_vlm_cls is None and not supports_custom_impl(model_config): raise ValueError( f"{model_config.model_type!r} has no PrimeRL model implementation. " @@ -473,23 +466,6 @@ def setup_tokenizer(config: TokenizerConfig) -> PreTrainedTokenizer: return tokenizer -def setup_processor(config: ModelConfig): - """Load an ``AutoProcessor`` for VLM models. Returns ``None`` for text-only models.""" - from transformers import AutoProcessor - - logger = get_logger() - try: - processor = AutoProcessor.from_pretrained(config.name, trust_remote_code=config.trust_remote_code) - except (ValueError, OSError, KeyError) as e: - logger.debug(f"No AutoProcessor available for {config.name} ({type(e).__name__}); treating as text-only.") - return None - if not (getattr(processor, "image_processor", None) or getattr(processor, "video_processor", None)): - logger.debug(f"AutoProcessor for {config.name} has no image/video processor; treating as text-only.") - return None - logger.info(f"Loaded multimodal processor: {type(processor).__name__}") - return processor - - def _expert_shard_placement_fn( experts: nn.Module, expert_mesh_info: FSDPMeshInfo, @@ -539,16 +515,15 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None: expert_mesh_info = _get_mesh_info(dp_mod_ep_mesh) assert isinstance(expert_mesh_info, FSDPMeshInfo) - is_vlm_training = config.vlm is not None - if is_vlm_training: - vision_encoder = get_vision_encoder(model, override=config.vlm.vision_encoder_attr) + if config.vlm is not None: + vision_encoder = get_vision_encoder(model) if vision_encoder is None: raise ValueError(f"VLM model {config.name} has no recognized vision encoder") fully_shard(vision_encoder, mesh=hsdp_mesh, **fsdp_config) get_logger().info(f"Applied FSDP to vision encoder (frozen={config.vlm.freeze_vision_encoder})") - language_model = get_language_model(model, override=config.vlm.language_model_attr if is_vlm_training else None) + language_model = get_language_model(model) transformer_layers = language_model.layers fullgraph = config.compile is not None and config.compile.fullgraph @@ -870,7 +845,7 @@ def configure_trainable_parameters(model: nn.Module, config: ModelConfig) -> nn. """Apply LoRA and identify any vision encoder that must remain frozen.""" frozen_vision_encoder = None if config.vlm is not None and config.vlm.freeze_vision_encoder: - frozen_vision_encoder = get_vision_encoder(model, override=config.vlm.vision_encoder_attr) + frozen_vision_encoder = get_vision_encoder(model) elif config.vlm is None: frozen_vision_encoder = get_vision_encoder(model) if frozen_vision_encoder is not None: @@ -1005,10 +980,7 @@ def setup_model( freeze_all_except_lora_and_specified(model, config.lora) if frozen_vision_encoder is not None: - freeze_vision_encoder( - model, - override_attr=config.vlm.vision_encoder_attr if config.vlm is not None else None, - ) + freeze_vision_encoder(frozen_vision_encoder) # the right order is AC -> Compile -> FSDP if config.ac is not None: @@ -1044,7 +1016,6 @@ def forward( routed_experts: Int[Tensor, "batch seq layers topk"] | None = None, sampling_mask: Int[Tensor, "batch seq mask"] | None = None, mm_kwargs: dict[str, Tensor] | None = None, - mm_forward_policy: ForwardPolicy | None = None, mm_token_type_ids: Int[Tensor, "batch seq"] | None = None, # True when seq_lens holds the full pre-CP-shard document boundaries # (kept global because documents can straddle the shard cut). @@ -1052,6 +1023,7 @@ def forward( ) -> PrimeLmOutput: kwargs = { "input_ids": input_ids, + "position_ids": position_ids, "labels": labels, "temperature": temperature, "sampling_mask": sampling_mask, @@ -1061,15 +1033,6 @@ def forward( kwargs.update(mm_kwargs) if mm_token_type_ids is not None: kwargs["mm_token_type_ids"] = mm_token_type_ids - # SFT still uses its existing eager processor path and does not provide - # an adapter policy yet, so preserve its current kwargs-based behavior. - policy = mm_forward_policy or ForwardPolicy(pass_position_ids="image_grid_thw" not in mm_kwargs) - if policy.requires_mm_token_type_ids and mm_token_type_ids is None: - raise ValueError("Multimodal forward policy requires mm_token_type_ids") - if policy.pass_position_ids: - kwargs["position_ids"] = position_ids - else: - kwargs["position_ids"] = position_ids kwargs["seq_lens"] = seq_lens kwargs["seq_lens_are_pre_shard"] = seq_lens_are_pre_shard diff --git a/src/prime_rl/trainer/models/layers/lm_head.py b/src/prime_rl/trainer/models/layers/lm_head.py index bfd1aa01d9b..a27f939012a 100644 --- a/src/prime_rl/trainer/models/layers/lm_head.py +++ b/src/prime_rl/trainer/models/layers/lm_head.py @@ -413,9 +413,7 @@ def new_forward( sampling_mask: torch.Tensor | None = None, **kwargs: object, ) -> PrimeLmOutput: - # For VLM with images, don't create position_ids - let model compute MRoPE internally - is_multimodal = kwargs.get("pixel_values") is not None - if position_ids is None and not is_multimodal: + if position_ids is None: reference_tensor = input_ids if input_ids is not None else inputs_embeds position_ids = torch.arange(reference_tensor.shape[1], device=reference_tensor.device).unsqueeze(0) model_kwargs = {"input_ids": input_ids, "position_ids": position_ids, **kwargs} diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py index ab269415c45..a92bd74fda1 100644 --- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py +++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py @@ -250,19 +250,19 @@ def prepare_inputs( # Every rank must retain the vision graph so FSDP collectives stay symmetric. inputs_embeds = inputs_embeds + image_embeds.sum() * 0.0 - if position_ids is None: - if image_grid_thw is None: - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) - else: - if mm_token_type_ids is None: - raise ValueError("mm_token_type_ids are required with Qwen3.5 image inputs") - position_ids = build_qwen3_5_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=self.config.vision_config.spatial_merge_size, - seq_lens=seq_lens, - ) + if image_grid_thw is not None: + # Image inputs get MRoPE positions; text-only inputs keep the caller's positions. + if mm_token_type_ids is None: + raise ValueError("mm_token_type_ids are required with Qwen3.5 image inputs") + position_ids = build_qwen3_5_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=self.config.vision_config.spatial_merge_size, + seq_lens=seq_lens, + ) + elif position_ids is None: + position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) return inputs_embeds, position_ids def forward( diff --git a/src/prime_rl/trainer/moe_runtime.py b/src/prime_rl/trainer/moe_runtime.py index fa1dc2bdc35..3d04c71d722 100644 --- a/src/prime_rl/trainer/moe_runtime.py +++ b/src/prime_rl/trainer/moe_runtime.py @@ -27,8 +27,8 @@ ) from prime_rl.trainer.models.layers.moe import MoE from prime_rl.trainer.parallel_dims import ParallelDims +from prime_rl.trainer.vlm import get_language_model from prime_rl.utils.logger import get_logger -from prime_rl.utils.vlm import get_language_model def _resolve_expert_compute(config: ModelConfig) -> ExpertCompute: @@ -68,9 +68,7 @@ def configure_moe_runtime(model: nn.Module, config: ModelConfig, parallel_dims: selected_moes = set(moe_layers) if config.moe.compute.apply_to != "all": - language_model = get_language_model( - model, override=config.vlm.language_model_attr if config.vlm is not None else None - ) + language_model = get_language_model(model) selected_layers = config.moe.compute.resolve_layers(len(language_model.layers)) selected_moes = { module diff --git a/src/prime_rl/trainer/multimodal.py b/src/prime_rl/trainer/multimodal.py deleted file mode 100644 index 04431c24da0..00000000000 --- a/src/prime_rl/trainer/multimodal.py +++ /dev/null @@ -1,26 +0,0 @@ -from __future__ import annotations - -import base64 -from io import BytesIO -from typing import Any - -from PIL import Image - -from prime_rl.multimodal import MaterializedMM, MultimodalAdapter -from prime_rl.transports.batch import MMRefs - - -def _load_image(data_url: str) -> Image.Image: - header, separator, payload = data_url.partition(",") - if not separator or not header.startswith("data:image/") or ";base64" not in header: - raise ValueError("Multimodal training requires base64 data image URLs") - with Image.open(BytesIO(base64.b64decode(payload, validate=True))) as image: - return image.convert("RGB") - - -def materialize_mm_refs(refs: MMRefs, processor: Any, adapter: MultimodalAdapter) -> MaterializedMM: - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - raise ValueError("Multimodal samples require a model image processor") - images = [_load_image(ref.url) for ref in refs.images] - return adapter.materialize(image_processor, images, [ref.length for ref in refs.images]) diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 45ae121dea8..a181f4991f2 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -7,8 +7,7 @@ from torch import Tensor from prime_rl.configs.trainer import FakeDataLoaderConfig -from prime_rl.multimodal import ForwardPolicy, MultimodalAdapter -from prime_rl.trainer.multimodal import materialize_mm_refs +from prime_rl.trainer.vlm import materialize_images from prime_rl.trainer.world import get_world from prime_rl.transports.batch import ( BatchReceiver, @@ -50,10 +49,9 @@ class TensorMicroBatch(TypedDict): sampling_mask: Int[Tensor, "batch seq mask"] | None # Multimodal inputs. The data loader sets mm_refs (undecoded image references), and prepare_micro_batch - # replaces them with mm_kwargs (CPU tensors for the vision encoder) and mm_forward_policy. + # replaces them with mm_kwargs (CPU tensors for the vision encoder). mm_refs: MMRefs | None mm_kwargs: NotRequired[dict[str, Tensor] | None] - mm_forward_policy: NotRequired[ForwardPolicy | None] # mm_token_type_ids: token type per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: Int[Tensor, "batch seq"] | None @@ -177,19 +175,14 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: } -def prepare_micro_batch( - micro_batch: TensorMicroBatch, processor: Any | None, mm_adapter: MultimodalAdapter | None -) -> TensorMicroBatch: +def prepare_micro_batch(micro_batch: TensorMicroBatch, processor: Any | None) -> TensorMicroBatch: """Prepare a micro batch for the training loop. Currently only multimodal samples need preparation.""" micro_batch["mm_kwargs"] = None - micro_batch["mm_forward_policy"] = None mm_refs = micro_batch.get("mm_refs") if mm_refs is not None: - if processor is None or mm_adapter is None: + if processor is None: raise ValueError("Received multimodal samples but [model.vlm] is not set") - materialized = materialize_mm_refs(mm_refs, processor, mm_adapter) - micro_batch["mm_kwargs"] = materialized.kwargs - micro_batch["mm_forward_policy"] = materialized.forward_policy + micro_batch["mm_kwargs"] = materialize_images(mm_refs, processor) micro_batch["mm_refs"] = None return micro_batch diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 16dba8ee9ae..36d9b05bc14 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -37,7 +37,6 @@ shift_tensor_left, shift_tensor_right, ) -from prime_rl.multimodal import get_multimodal_adapter from prime_rl.trainer.rl.annotations import AnnotationWriter from prime_rl.trainer.model import ( forward, @@ -46,9 +45,9 @@ get_global_moe_stats, is_tt_moe_model, setup_model, - setup_processor, ) from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep +from prime_rl.trainer.vlm import setup_processor from prime_rl.trainer.perf import get_perf_counter from prime_rl.trainer.utils import ( GarbageCollection, @@ -144,13 +143,11 @@ def train(config: TrainerConfig): logger.debug(f"Initialized model in {format_time(time.perf_counter() - t0)}") processor = None - mm_adapter = None if config.model.vlm is not None: processor = setup_processor(config.model) if processor is None: raise ValueError("Multimodal training requires a model image processor") - mm_adapter = get_multimodal_adapter(model.config.model_type) - prepare = partial(prepare_micro_batch, processor=processor, mm_adapter=mm_adapter) + prepare = partial(prepare_micro_batch, processor=processor) micro_batch_workers = WorkerPool(config.data.num_workers) if config.model.vlm is not None and not getattr(model, "supports_packed_multimodal_training", False): @@ -381,7 +378,6 @@ def train(config: TrainerConfig): mm_kwargs = micro_batch.pop("mm_kwargs") if mm_kwargs is not None: mm_kwargs = {key: value.to("cuda") for key, value in mm_kwargs.items()} - mm_forward_policy = micro_batch.pop("mm_forward_policy") mm_token_type_ids = ( micro_batch["mm_token_type_ids"].to("cuda") if micro_batch.get("mm_token_type_ids") is not None @@ -399,9 +395,8 @@ def train(config: TrainerConfig): seq_lens_are_pre_shard = False if cp_enabled: - defer_vlm_cp_to_model = bool( - mm_forward_policy is not None and mm_forward_policy.defer_context_parallelism - ) + # The VLM shards image inputs itself, after building their positions. + defer_vlm_cp_to_model = mm_kwargs is not None if not defer_vlm_cp_to_model: input_ids, position_ids = setup_cp_params( input_ids, @@ -454,7 +449,6 @@ def train(config: TrainerConfig): labels=labels, temperature=temperatures, mm_kwargs=mm_kwargs, - mm_forward_policy=mm_forward_policy, mm_token_type_ids=mm_token_type_ids, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, diff --git a/src/prime_rl/trainer/sft/train.py b/src/prime_rl/trainer/sft/train.py index a49e770281e..52f33aaea81 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -32,7 +32,6 @@ get_load_balance_stats, is_tt_moe_model, reshard_module, - setup_processor, setup_tokenizer, setup_model, ) @@ -58,6 +57,7 @@ setup_full_cpu_optimizer_offload, setup_torch_distributed, ) +from prime_rl.trainer.vlm import setup_processor from prime_rl.trainer.world import get_world from prime_rl.control.controller import Controller, broadcast_pause from prime_rl.control.server import ControlServer @@ -259,9 +259,8 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: if cp_enabled: # CP requires the sequence length to be divisible by cp_size. CatDataset # pads every pack to a multiple of it; shard_for_cp raises on violations. - defer_vlm_cp_to_model = ( - mm_kwargs is not None and "image_grid_thw" in mm_kwargs and config.model.cp_style == "ulysses" - ) + # The VLM shards image inputs itself, after building their positions. + defer_vlm_cp_to_model = mm_kwargs is not None if not defer_vlm_cp_to_model: input_ids, position_ids = setup_cp_params( input_ids, diff --git a/src/prime_rl/trainer/vlm.py b/src/prime_rl/trainer/vlm.py new file mode 100644 index 00000000000..c03d14cd735 --- /dev/null +++ b/src/prime_rl/trainer/vlm.py @@ -0,0 +1,69 @@ +"""VLMs are the ``_CUSTOM_VLM_MAPPING`` models. They keep the vision encoder at ``model.model.visual`` +and the decoder at ``model.model.language_model``; text models keep the decoder at ``model.model``.""" + +import base64 +from io import BytesIO +from typing import Any + +import torch +import torch.nn as nn +from PIL import Image + +from prime_rl.configs.trainer import ModelConfig +from prime_rl.transports.batch import MMRefs +from prime_rl.utils.logger import get_logger + + +def get_vision_encoder(model: nn.Module) -> nn.Module | None: + return getattr(model.model, "visual", None) + + +def get_language_model(model: nn.Module) -> nn.Module: + return getattr(model.model, "language_model", model.model) + + +def get_layer_prefix(model: nn.Module) -> str: + """Weight key prefix of the decoder layers.""" + return "model.language_model.layers." if hasattr(model.model, "language_model") else "model.layers." + + +def setup_processor(config: ModelConfig): + """Load an ``AutoProcessor`` for VLM models. Returns ``None`` for text-only models.""" + from transformers import AutoProcessor + + logger = get_logger() + try: + processor = AutoProcessor.from_pretrained(config.name, trust_remote_code=config.trust_remote_code) + except (ValueError, OSError, KeyError) as e: + logger.debug(f"No AutoProcessor available for {config.name} ({type(e).__name__}); treating as text-only.") + return None + if not (getattr(processor, "image_processor", None) or getattr(processor, "video_processor", None)): + logger.debug(f"AutoProcessor for {config.name} has no image/video processor; treating as text-only.") + return None + logger.info(f"Loaded multimodal processor: {type(processor).__name__}") + return processor + + +def _load_image(data_url: str) -> Image.Image: + header, separator, payload = data_url.partition(",") + if not separator or not header.startswith("data:image/") or ";base64" not in header: + raise ValueError("Multimodal training requires base64 data image URLs") + with Image.open(BytesIO(base64.b64decode(payload, validate=True))) as image: + return image.convert("RGB") + + +def materialize_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: + """Decode the image refs of a micro batch into the Qwen3.5 vision forward kwargs.""" + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError("Multimodal samples require a model image processor") + images = [_load_image(ref.url) for ref in refs.images] + processed = image_processor(images=images, return_tensors="pt") + kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} + merge_size = int(image_processor.merge_size) + # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. + lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] + expected = [ref.length for ref in refs.images] + if lengths != expected: + raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") + return kwargs diff --git a/src/prime_rl/transports/weights/nccl.py b/src/prime_rl/transports/weights/nccl.py index 8ef57378b46..3537c8c86fe 100644 --- a/src/prime_rl/transports/weights/nccl.py +++ b/src/prime_rl/transports/weights/nccl.py @@ -12,10 +12,10 @@ from prime_rl.configs.shared import NCCLWeightBroadcastConfig from prime_rl.trainer.models import PreTrainedModelPrimeRL from prime_rl.trainer.utils import get_world +from prime_rl.trainer.vlm import get_layer_prefix from prime_rl.transports.weights.base import WeightReceiver, WeightSender from prime_rl.utils.logger import get_logger from prime_rl.utils.nccl import disable_nccl_p2p_if_unavailable, iter_tensor_buckets -from prime_rl.utils.vlm import get_layer_prefix from prime_rl.utils.weights import resolve_wire_dtype @@ -151,7 +151,7 @@ def __init__( def send(self, model: nn.Module) -> None: """Broadcast the state dict of a model into the inference pool using NCCL.""" state_dict = model.state_dict() - layer_prefix = get_layer_prefix(model.config) + layer_prefix = get_layer_prefix(model) num_layers = get_max_layer_num(state_dict, layer_prefix) num_state_dict_to_send = num_layers + 1 # we send all layer plus the remaining weights diff --git a/src/prime_rl/utils/vlm.py b/src/prime_rl/utils/vlm.py deleted file mode 100644 index 0f5e5baf24f..00000000000 --- a/src/prime_rl/utils/vlm.py +++ /dev/null @@ -1,123 +0,0 @@ -"""Vision-Language Model (VLM) support utilities. - -Central registry for VLM model families. All model-specific knowledge -lives here. Add new VLM families by extending VLM_REGISTRY. - -For custom models not in the registry, set overrides in config: - [model.vlm] - vision_encoder_attr = "model.my_vision" - language_model_attr = "model.my_lm" -""" - -from dataclasses import dataclass - -import torch.nn as nn -from transformers.configuration_utils import PretrainedConfig - - -@dataclass(frozen=True) -class VLMModelInfo: - """Per-model-family VLM architecture metadata.""" - - vision_encoder_attr: str - language_model_attr: str - - -# Central registry: model_type -> architecture info. -VLM_REGISTRY: dict[str, VLMModelInfo] = { - "qwen3_5": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), - "qwen3_5_moe": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), - "qwen4_exp": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), -} - -# Text-only default -DEFAULT_LAYER_PREFIX = "model.layers." - - -# --------------------------------------------------------------------------- -# Model component access -# --------------------------------------------------------------------------- - - -def get_vision_encoder(model: nn.Module, override: str | None = None) -> nn.Module | None: - """Get the vision encoder module. - - Checks: config override -> registry. Returns None if not found. - Raises ValueError on a bad config override. - """ - if override is not None: - result = _resolve_attr(model, override) - if result is None: - raise ValueError(f"vlm.vision_encoder_attr='{override}' does not resolve on this model") - return result - - info = _get_model_info(model) - if info is not None: - return _resolve_attr(model, info.vision_encoder_attr) - - return None - - -def get_language_model(model: nn.Module, override: str | None = None) -> nn.Module: - """Get the language model module (the part with transformer layers). - - Checks: config override -> registry -> model.model (text-only default). - Raises ValueError on a bad config override. - """ - if override is not None: - result = _resolve_attr(model, override) - if result is None: - raise ValueError(f"vlm.language_model_attr='{override}' does not resolve on this model") - return result - - info = _get_model_info(model) - if info is not None: - result = _resolve_attr(model, info.language_model_attr) - if result is not None: - return result - - # Text-only models: language model is directly at model.model - return model.model - - -def is_vlm_architecture(model_config: PretrainedConfig) -> bool: - """Check if the model config belongs to a known VLM architecture.""" - return _get_model_info_from_config(model_config) is not None - - -def get_layer_prefix(model_config: PretrainedConfig, override: str | None = None) -> str: - """Return the weight key prefix for language model layers. - - Derived from language_model_attr + '.layers.' for registered VLMs, - or 'model.layers.' for text-only / unknown models. - """ - if override is not None: - return override - info = _get_model_info_from_config(model_config) - if info is not None: - return info.language_model_attr + ".layers." - return DEFAULT_LAYER_PREFIX - - -# --------------------------------------------------------------------------- -# Internal -# --------------------------------------------------------------------------- - - -def _get_model_info(model: nn.Module) -> VLMModelInfo | None: - model_type = getattr(getattr(model, "config", None), "model_type", None) - return VLM_REGISTRY.get(model_type) if model_type else None - - -def _get_model_info_from_config(model_config: PretrainedConfig) -> VLMModelInfo | None: - model_type = getattr(model_config, "model_type", None) - return VLM_REGISTRY.get(model_type) if model_type else None - - -def _resolve_attr(obj, dotted_path: str): - """Resolve a dotted attribute path like 'model.visual' on an object.""" - for part in dotted_path.split("."): - obj = getattr(obj, part, None) - if obj is None: - return None - return obj diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index d498d9bb1ba..92d20dd5a22 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -4,7 +4,6 @@ import pytest from prime_rl.orchestrator.batch import ( - _is_multimodal_sample, build_bin_cost, pad_micro_batch, prepare_batch, @@ -525,7 +524,7 @@ def make_mm_sample(token: int) -> TrainingSample: batches = _flatten_batches(batches_per_gpu) assert len(batches) == 2 - assert all(_has_loss_tokens(batch) and _is_multimodal_sample(batch) for batch in batches) + assert all(_has_loss_tokens(batch) and batch.mm_refs is not None for batch in batches) def test_prepare_sample_none_routed_experts(): diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 8ce6d6b566a..e560308f2b2 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -731,36 +731,6 @@ def test_disaggregated_inference_inherits_deployment_gpus_per_node(): RLConfig.model_validate(config) -def test_orchestrator_vlm_requires_renderer(): - with pytest.raises(ValidationError, match="renderer"): - OrchestratorConfig.model_validate( - { - "model": { - "name": "Qwen/Qwen3-VL-4B-Instruct", - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", - }, - }, - "renderer": None, - } - ) - - config = OrchestratorConfig.model_validate( - { - "model": { - "name": "Qwen/Qwen3-VL-4B-Instruct", - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", - }, - }, - } - ) - - assert config.renderer is not None - - def test_trainer_rejects_vlm_cp_with_ring(): config = { "model": { @@ -768,10 +738,7 @@ def test_trainer_rejects_vlm_cp_with_ring(): "cp_style": "ring", "optimization_dtype": "bfloat16", "reduce_dtype": "bfloat16", - "vlm": { - "vision_encoder_attr": "model.visual", - "language_model_attr": "model.language_model", - }, + "vlm": {}, }, } diff --git a/tests/unit/train/models/test_qwen3_8_flash_next.py b/tests/unit/train/models/test_qwen3_8_flash_next.py index 0d16e3053da..20915528883 100644 --- a/tests/unit/train/models/test_qwen3_8_flash_next.py +++ b/tests/unit/train/models/test_qwen3_8_flash_next.py @@ -7,7 +7,7 @@ Qwen3_8FlashNextForCausalLM, Qwen3_8FlashNextTextConfig, ) -from prime_rl.utils.vlm import get_language_model +from prime_rl.trainer.vlm import get_language_model @pytest.fixture(params=[False, True], ids=["text", "composite"]) diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py deleted file mode 100644 index 853eef3ebc9..00000000000 --- a/tests/unit/train/test_model_forward.py +++ /dev/null @@ -1,87 +0,0 @@ -from types import SimpleNamespace - -import torch -import torch.nn as nn - -from prime_rl.multimodal.qwen_vl import QwenVLAdapter -from prime_rl.trainer.model import forward - - -class _CaptureModel(nn.Module): - def __init__(self, config: SimpleNamespace): - super().__init__() - self.config = config - self.kwargs = None - - def forward(self, **kwargs): - self.kwargs = kwargs - input_ids = kwargs["input_ids"] - return {"logits": torch.zeros(*input_ids.shape, 4)} - - -def test_forward_passes_renderer_mm_token_type_ids_through(): - """``forward()`` forwards renderer-supplied ``mm_token_type_ids`` - verbatim — the trainer no longer auto-computes from the model - config, since the renderer is the source of truth.""" - model = _CaptureModel(SimpleNamespace(model_type="qwen3_vl")) - input_ids = torch.tensor([[1, 10, 10, 2, 20]]) - position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) - pixel_values = torch.ones(2, 3) - image_grid_thw = torch.tensor([[1, 1, 2]]) - mm_token_type_ids = torch.tensor([[0, 1, 1, 0, 2]]) - - forward( - model, - input_ids, - position_ids, - seq_lens=torch.tensor([input_ids.shape[1]]), - mm_kwargs={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, - mm_forward_policy=QwenVLAdapter.forward_policy, - mm_token_type_ids=mm_token_type_ids, - ) - - assert model.kwargs is not None - assert "position_ids" not in model.kwargs - torch.testing.assert_close(model.kwargs["pixel_values"], pixel_values) - torch.testing.assert_close(model.kwargs["image_grid_thw"], image_grid_thw) - torch.testing.assert_close(model.kwargs["mm_token_type_ids"], mm_token_type_ids) - - -def test_forward_omits_mm_token_type_ids_when_renderer_does_not_supply(): - """When the renderer doesn't ship ``mm_token_type_ids`` (text-only - or a family without modality markers), ``forward()`` doesn't - fabricate one.""" - model = _CaptureModel(SimpleNamespace(model_type="qwen3_vl")) - input_ids = torch.tensor([[1, 10, 10, 2]]) - position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) - - forward( - model, - input_ids, - position_ids, - seq_lens=torch.tensor([input_ids.shape[1]]), - mm_kwargs={"pixel_values": torch.ones(2, 3), "image_grid_thw": torch.tensor([[1, 1, 2]])}, - ) - - assert model.kwargs is not None - assert "position_ids" not in model.kwargs - assert "mm_token_type_ids" not in model.kwargs - - -def test_forward_keeps_position_ids_for_non_mrope_vlm(): - """Non-MRoPE VLM families (no ``image_grid_thw``) keep the trainer's - pre-computed ``position_ids``.""" - model = _CaptureModel(SimpleNamespace(model_type="gemma3")) - input_ids = torch.tensor([[1, 10, 10, 2]]) - position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) - - forward( - model, - input_ids, - position_ids, - seq_lens=torch.tensor([input_ids.shape[1]]), - mm_kwargs={"pixel_values": torch.ones(2, 3)}, - ) - - assert model.kwargs is not None - torch.testing.assert_close(model.kwargs["position_ids"], position_ids) diff --git a/tests/unit/train/test_multimodal.py b/tests/unit/train/test_multimodal.py index a7dd5694ee6..2686ec8716b 100644 --- a/tests/unit/train/test_multimodal.py +++ b/tests/unit/train/test_multimodal.py @@ -3,10 +3,7 @@ import pytest import torch -from prime_rl.multimodal import get_multimodal_adapter -from prime_rl.multimodal.kimi_k25 import KimiK25Adapter -from prime_rl.multimodal.qwen_vl import QwenVLAdapter -from prime_rl.trainer.multimodal import materialize_mm_refs +from prime_rl.trainer.vlm import materialize_images from prime_rl.transports.batch import MMImageRef, MMRefs _IMAGE_URL = ( @@ -18,7 +15,7 @@ def _refs(length: int) -> MMRefs: return MMRefs(images=[MMImageRef(url=_IMAGE_URL, offset=1, length=length)]) -def test_qwen_adapter_materializes_and_validates_expansion(): +def test_qwen_images_materialize_and_validate_expansion(): class ImageProcessor: merge_size = 1 @@ -30,27 +27,6 @@ def __call__(self, *, images, return_tensors): } processor = SimpleNamespace(image_processor=ImageProcessor()) - adapter = get_multimodal_adapter("qwen3_vl") - materialized = materialize_mm_refs(_refs(2), processor, adapter) - - assert set(materialized.kwargs) == {"pixel_values", "image_grid_thw"} - assert materialized.forward_policy == QwenVLAdapter.forward_policy + assert set(materialize_images(_refs(2), processor)) == {"pixel_values", "image_grid_thw"} with pytest.raises(ValueError, match="placeholder lengths differ"): - materialize_mm_refs(_refs(1), processor, adapter) - - -def test_kimi_adapter_materializes_sparse_image_position(): - class ImageProcessor: - def preprocess(self, media, *, return_tensors): - assert len(media) == 1 and media[0]["type"] == "image" - assert media[0]["image"].mode == "RGB" and return_tensors == "pt" - return { - "pixel_values": torch.ones(4, 3), - "grid_thws": torch.tensor([[1, 2, 2]]), - } - - processor = SimpleNamespace(image_processor=ImageProcessor()) - materialized = materialize_mm_refs(_refs(1), processor, get_multimodal_adapter("kimi_k25")) - - assert set(materialized.kwargs) == {"pixel_values", "grid_thws"} - assert materialized.forward_policy == KimiK25Adapter.forward_policy + materialize_images(_refs(1), processor) From c9aa101b47154511a410dde2cdba97ce54d12c12 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sat, 10 Oct 2026 16:12:11 +0000 Subject: [PATCH 2/8] fix(tools): import setup_processor from trainer.vlm --- tools/convert_dcp_to_bf16.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tools/convert_dcp_to_bf16.py b/tools/convert_dcp_to_bf16.py index a42bbf77029..69876a61005 100644 --- a/tools/convert_dcp_to_bf16.py +++ b/tools/convert_dcp_to_bf16.py @@ -33,9 +33,10 @@ from prime_rl.configs.trainer import ModelConfig, MoERuntimeConfig, TokenizerConfig from prime_rl.trainer.ckpt import AppState -from prime_rl.trainer.model import setup_model, setup_processor, setup_tokenizer +from prime_rl.trainer.model import setup_model, setup_tokenizer from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep from prime_rl.trainer.utils import setup_torch_distributed +from prime_rl.trainer.vlm import setup_processor from prime_rl.trainer.world import get_world from prime_rl.utils.logger import get_logger, setup_logger from prime_rl.utils.pathing import get_config_dir From 6113c9a39ec1b356f37ca809ce8b30ea7e2e8bfc Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Mon, 12 Oct 2026 18:32:44 +0000 Subject: [PATCH 3/8] fix(configs): enforce VLM bf16 dtypes in the trainer model config (RL too); one validate_vlm --- docs/advanced.md | 2 ++ packages/prime-rl-configs/src/prime_rl/configs/sft.py | 4 ---- .../prime-rl-configs/src/prime_rl/configs/trainer.py | 10 ++++++++-- tests/unit/test_configs.py | 5 +++++ 4 files changed, 15 insertions(+), 6 deletions(-) diff --git a/docs/advanced.md b/docs/advanced.md index d6d34999c9e..7014a8c82da 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -163,6 +163,8 @@ reduce_dtype = "bfloat16" # freeze_vision_encoder = true # default; set false to fine-tune the encoder ``` +For `rl`, keep `[model.vlm]` in the shared `[model]` block but set the dtypes under `[trainer.model]` (see `configs/ci/nightly/multimodal_color_codeword.toml`). + VLM training requires a registered custom PrimeRL implementation. ### Limitations diff --git a/packages/prime-rl-configs/src/prime_rl/configs/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 660d33864e1..259b2906d5b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -598,10 +598,6 @@ def vlm_freeze_incompatible_with_lora(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: 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 15a30307205..ec3e20def4a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -311,9 +311,15 @@ class ModelConfig(BaseModelConfig): """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 vlm_cp_requires_ulysses(self): - if self.vlm is not None and self.cp > 1 and self.cp_style not in ("auto", "ulysses"): + def validate_vlm(self): + if self.vlm is None: + return self + if self.cp > 1 and self.cp_style not in ("auto", "ulysses"): raise ValueError("VLM models require cp_style='ulysses' for context parallelism") + 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") diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index e560308f2b2..f991a19acc8 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -746,6 +746,11 @@ def test_trainer_rejects_vlm_cp_with_ring(): TrainerConfig.model_validate(config) +def test_trainer_vlm_requires_bf16(): + with pytest.raises(ValidationError, match="optimization_dtype='bfloat16'"): + TrainerConfig.model_validate({"model": {"vlm": {}}}) + + def test_ac_none_disables_ac_offloading(): config = TrainerModelConfig.model_validate({"ac": None}) assert config.ac is None From a976bd3e51beae40857a5384c190767aad39315a Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Thu, 22 Oct 2026 12:59:49 +0000 Subject: [PATCH 4/8] fix(trainer): positions come from the caller; refuse cp with a VLM that cannot shard its inputs --- src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py | 4 ---- src/prime_rl/utils/cp.py | 3 +++ 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py index 780bf3a29cc..9f7557d1de2 100644 --- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py +++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py @@ -238,8 +238,6 @@ def prepare_inputs( if image_grid_thw is not None: # Image inputs get MRoPE positions; text-only inputs keep the caller's positions. - if mm_token_type_ids is None: - raise ValueError("mm_token_type_ids are required with Qwen3.5 image inputs") position_ids = build_qwen3_5_mrope_position_ids( input_ids=input_ids, mm_token_type_ids=mm_token_type_ids, @@ -247,8 +245,6 @@ def prepare_inputs( spatial_merge_size=self.config.vision_config.spatial_merge_size, seq_lens=seq_lens, ) - elif position_ids is None: - position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device).unsqueeze(0) return inputs_embeds, position_ids def forward( diff --git a/src/prime_rl/utils/cp.py b/src/prime_rl/utils/cp.py index 147ce2725f3..588b97632f7 100644 --- a/src/prime_rl/utils/cp.py +++ b/src/prime_rl/utils/cp.py @@ -30,6 +30,9 @@ def cp_enabled(self) -> bool: def setup_context_parallel(model: nn.Module, config: ModelConfig, parallel_dims: ParallelDims) -> None: + # The trainer leaves multimodal inputs unsharded; the VLM shards them after building their positions. + if config.vlm is not None and not hasattr(model.model, "cp_context"): + raise ValueError(f"{type(model).__name__} does not shard multimodal inputs for context parallelism; set cp=1") cp_group = parallel_dims.world_mesh["cp"].get_group() cp_rank = parallel_dims.world_mesh["cp"].get_local_rank() From fef88085dd7e490a05efc59ab036a28a00b3decb Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:34:26 +0000 Subject: [PATCH 5/8] refactor(configs): drop the VLM bf16 dtype validator --- docs/advanced.md | 7 ++----- .../prime-rl-configs/src/prime_rl/configs/trainer.py | 10 ---------- tests/unit/test_configs.py | 5 ----- 3 files changed, 2 insertions(+), 20 deletions(-) diff --git a/docs/advanced.md b/docs/advanced.md index dfdf8cb0955..5845a63608f 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -165,26 +165,23 @@ The trainer finds the vision encoder at `model.visual` and the language model at ### Enabling VLM Mode -Add `[model.vlm]` and bfloat16 dtypes: +Add `[model.vlm]`: ```toml [model] name = "Qwen/Qwen3.5-4B" -optimization_dtype = "bfloat16" -reduce_dtype = "bfloat16" [model.vlm] # freeze_vision_encoder = true # default; set false to fine-tune the encoder ``` -For `rl`, keep `[model.vlm]` in the shared `[model]` block but set the dtypes under `[trainer.model]`. +For `rl`, set `[model.vlm]` in the shared `[model]` block. VLM training requires a registered custom PrimeRL implementation. ### Limitations - **Vision encoder frozen by default.** The default LoRA targets do not match Qwen3.5 vision modules. Set `freeze_vision_encoder = false` to fine-tune the encoder; this is incompatible with LoRA because LoRA freezes all non-adapter parameters. -- **bfloat16 mandatory.** The trainer config validator refuses any other `optimization_dtype` / `reduce_dtype` for VLMs — vLLM serves VLMs in bfloat16 and a mismatch breaks the importance ratio. - **Higher KL mismatch with multi-image inputs.** Expect noisier `mismatch_kl` than text-only; this is from minor numerical differences between the trainer's and vLLM's image processing. - **Images aren't logged to monitors.** Sample logging captures the prompt text but not the actual images. 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 7735fb67010..907c3553623 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -334,16 +334,6 @@ 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"]: diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 2dd659b39b5..c77ba115aee 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -790,11 +790,6 @@ def test_disaggregated_inference_inherits_deployment_gpus_per_node(): RLConfig.model_validate(config) -def test_trainer_vlm_requires_bf16(): - with pytest.raises(ValidationError, match="optimization_dtype='bfloat16'"): - TrainerConfig.model_validate({"model": {"vlm": {}}}) - - def test_ac_none_disables_ac_offloading(): config = TrainerModelConfig.model_validate({"ac": None}) assert config.ac is None From fecaf4ea334bd078899b2491bb5d8f9158e2e7d4 Mon Sep 17 00:00:00 2001 From: faresobeid <111092724+faresobeid@users.noreply.github.com> Date: Sat, 24 Oct 2026 18:40:41 +0000 Subject: [PATCH 6/8] refactor(trainer): materialize_images lives on the Qwen3.5 model class --- docs/development.md | 2 +- .../models/qwen3_5/modeling_qwen3_5.py | 31 ++++++++++++++++++ src/prime_rl/trainer/rl/data.py | 12 ++++--- src/prime_rl/trainer/rl/train.py | 10 +++--- src/prime_rl/trainer/vlm.py | 32 ------------------- tests/unit/train/test_multimodal.py | 6 ++-- 6 files changed, 47 insertions(+), 46 deletions(-) diff --git a/docs/development.md b/docs/development.md index 775544bab99..d136c31f1b8 100644 --- a/docs/development.md +++ b/docs/development.md @@ -145,6 +145,6 @@ VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementatio 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. 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. +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 with the model class's `materialize_images(refs, processor)` (see `Qwen3_5ForCausalLM`); a new family implements its own. 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). diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py index 8e42173dc2d..513e3176570 100644 --- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py +++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py @@ -1,4 +1,9 @@ +import base64 +from io import BytesIO +from typing import Any + import torch +from PIL import Image from torch import Tensor, nn from transformers.modeling_outputs import BaseModelOutput @@ -23,6 +28,7 @@ build_qwen3_5_mrope_position_ids, ) from prime_rl.trainer.models.qwen3_5.vision import Qwen3_5VisionModel +from prime_rl.transports.batch import MMRefs from prime_rl.utils.cp import CPContext, shard_for_cp, shard_position_ids_for_cp from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens @@ -284,6 +290,14 @@ def forward( ) +def _load_image(data_url: str) -> Image.Image: + header, separator, payload = data_url.partition(",") + if not separator or not header.startswith("data:image/") or ";base64" not in header: + raise ValueError("Multimodal training requires base64 data image URLs") + with Image.open(BytesIO(base64.b64decode(payload, validate=True))) as image: + return image.convert("RGB") + + class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel): def __init__(self, config) -> None: super().__init__(config) @@ -304,6 +318,23 @@ def get_input_embeddings(self) -> nn.Embedding: def set_input_embeddings(self, embeddings: nn.Embedding) -> None: self.model.set_input_embeddings(embeddings) + @staticmethod + def materialize_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: + """Decode the image refs of a micro batch into the vision forward kwargs.""" + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError("Multimodal samples require a model image processor") + images = [_load_image(ref.url) for ref in refs.images] + processed = image_processor(images=images, return_tensors="pt") + kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} + merge_size = int(image_processor.merge_size) + # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. + lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] + expected = [ref.length for ref in refs.images] + if lengths != expected: + raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") + return kwargs + def num_active_params(self) -> int: # The vision encoder runs per image, not per text token. if not self.is_vlm: diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 4fc5f594cef..6eade8bb1f3 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -1,5 +1,6 @@ +from collections.abc import Callable from pathlib import Path -from typing import Any, NotRequired, TypedDict +from typing import NotRequired, TypedDict import numpy as np import torch @@ -7,7 +8,6 @@ from torch import Tensor from prime_rl.configs.trainer import FakeDataLoaderConfig -from prime_rl.trainer.vlm import materialize_images from prime_rl.trainer.world import get_world from prime_rl.transports.batch import ( BatchReceiver, @@ -183,14 +183,16 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: } -def prepare_micro_batch(micro_batch: TensorMicroBatch, processor: Any | None) -> TensorMicroBatch: +def prepare_micro_batch( + micro_batch: TensorMicroBatch, materialize_images: Callable[[MMRefs], dict[str, Tensor]] | None +) -> TensorMicroBatch: """Prepare a micro batch for the training loop. Currently only multimodal samples need preparation.""" micro_batch["mm_kwargs"] = None mm_refs = micro_batch.get("mm_refs") if mm_refs is not None: - if processor is None: + if materialize_images is None: raise ValueError("Received multimodal samples but [model.vlm] is not set") - micro_batch["mm_kwargs"] = materialize_images(mm_refs, processor) + micro_batch["mm_kwargs"] = materialize_images(mm_refs) micro_batch["mm_refs"] = None return micro_batch diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 795d842a4c6..65860ad03a3 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -146,17 +146,17 @@ def train(config: TrainerConfig): model = setup_model(config.model, parallel_dims, loading_from_ckpt_later) logger.debug(f"Initialized model in {format_time(time.perf_counter() - t0)}") - processor = None + materialize_images = None if config.model.vlm is not None: + if not getattr(model, "supports_packed_multimodal_training", False): + raise ValueError("Packed multimodal training requires model support") processor = setup_processor(config.model) if processor is None: raise ValueError("Multimodal training requires a model image processor") - prepare = partial(prepare_micro_batch, processor=processor) + materialize_images = partial(model.materialize_images, processor=processor) + prepare = partial(prepare_micro_batch, materialize_images=materialize_images) micro_batch_workers = WorkerPool(config.data.num_workers) - if config.model.vlm is not None and not getattr(model, "supports_packed_multimodal_training", False): - raise ValueError("Packed multimodal training requires model support") - # Set up the loss function for the RL loss type (ce / ref_kl are fixed) logger.info(f"Initializing loss function ({config.loss})") rl_loss_fn = setup_rl_loss_fn(config.loss) diff --git a/src/prime_rl/trainer/vlm.py b/src/prime_rl/trainer/vlm.py index c03d14cd735..c3ef0c78339 100644 --- a/src/prime_rl/trainer/vlm.py +++ b/src/prime_rl/trainer/vlm.py @@ -1,16 +1,9 @@ """VLMs are the ``_CUSTOM_VLM_MAPPING`` models. They keep the vision encoder at ``model.model.visual`` and the decoder at ``model.model.language_model``; text models keep the decoder at ``model.model``.""" -import base64 -from io import BytesIO -from typing import Any - -import torch import torch.nn as nn -from PIL import Image from prime_rl.configs.trainer import ModelConfig -from prime_rl.transports.batch import MMRefs from prime_rl.utils.logger import get_logger @@ -42,28 +35,3 @@ def setup_processor(config: ModelConfig): return None logger.info(f"Loaded multimodal processor: {type(processor).__name__}") return processor - - -def _load_image(data_url: str) -> Image.Image: - header, separator, payload = data_url.partition(",") - if not separator or not header.startswith("data:image/") or ";base64" not in header: - raise ValueError("Multimodal training requires base64 data image URLs") - with Image.open(BytesIO(base64.b64decode(payload, validate=True))) as image: - return image.convert("RGB") - - -def materialize_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: - """Decode the image refs of a micro batch into the Qwen3.5 vision forward kwargs.""" - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - raise ValueError("Multimodal samples require a model image processor") - images = [_load_image(ref.url) for ref in refs.images] - processed = image_processor(images=images, return_tensors="pt") - kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} - merge_size = int(image_processor.merge_size) - # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. - lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] - expected = [ref.length for ref in refs.images] - if lengths != expected: - raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") - return kwargs diff --git a/tests/unit/train/test_multimodal.py b/tests/unit/train/test_multimodal.py index 2686ec8716b..4d2e20f41d8 100644 --- a/tests/unit/train/test_multimodal.py +++ b/tests/unit/train/test_multimodal.py @@ -3,7 +3,7 @@ import pytest import torch -from prime_rl.trainer.vlm import materialize_images +from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM from prime_rl.transports.batch import MMImageRef, MMRefs _IMAGE_URL = ( @@ -27,6 +27,6 @@ def __call__(self, *, images, return_tensors): } processor = SimpleNamespace(image_processor=ImageProcessor()) - assert set(materialize_images(_refs(2), processor)) == {"pixel_values", "image_grid_thw"} + assert set(Qwen3_5ForCausalLM.materialize_images(_refs(2), processor)) == {"pixel_values", "image_grid_thw"} with pytest.raises(ValueError, match="placeholder lengths differ"): - materialize_images(_refs(1), processor) + Qwen3_5ForCausalLM.materialize_images(_refs(1), processor) From 4028409f4082ffedf48f926b8336c6b1533cd900 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hubert=20Pysk=C5=82o?= Date: Thu, 29 Oct 2026 18:35:06 +0100 Subject: [PATCH 7/8] refactor(trainer): share Qwen-VL image materialization in models/qwen_vl.py Qwen3.8 Flash Next no longer imports the Qwen3.5 modeling module; both point materialize_images at materialize_qwen_vl_images. The family-agnostic data-URL decoding moves to trainer/vlm.py as load_image. Co-Authored-By: Claude Opus 5.5 --- docs/development.md | 2 +- .../models/qwen3_5/modeling_qwen3_5.py | 32 ++----------------- .../modeling_qwen3_8_flash_next.py | 5 ++- src/prime_rl/trainer/models/qwen_vl.py | 23 +++++++++++++ src/prime_rl/trainer/vlm.py | 12 +++++++ tests/unit/train/test_multimodal.py | 6 ++-- 6 files changed, 43 insertions(+), 37 deletions(-) create mode 100644 src/prime_rl/trainer/models/qwen_vl.py diff --git a/docs/development.md b/docs/development.md index a17917fb908..8201eeae3b5 100644 --- a/docs/development.md +++ b/docs/development.md @@ -146,6 +146,6 @@ VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementatio 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. RL turns raw image refs into vision kwargs with the model class's `materialize_images(refs, processor)` (see `Qwen3_5ForCausalLM`); a new family implements its own. +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 with the model class's `materialize_images(refs, processor)`; Qwen3.5 and Qwen3.8 Flash Next share `models/qwen_vl.py`, and a new family implements its own. 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). diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py index 513e3176570..e9141846ca4 100644 --- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py +++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py @@ -1,9 +1,4 @@ -import base64 -from io import BytesIO -from typing import Any - import torch -from PIL import Image from torch import Tensor, nn from transformers.modeling_outputs import BaseModelOutput @@ -28,7 +23,7 @@ build_qwen3_5_mrope_position_ids, ) from prime_rl.trainer.models.qwen3_5.vision import Qwen3_5VisionModel -from prime_rl.transports.batch import MMRefs +from prime_rl.trainer.models.qwen_vl import materialize_qwen_vl_images from prime_rl.utils.cp import CPContext, shard_for_cp, shard_position_ids_for_cp from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens @@ -290,14 +285,6 @@ def forward( ) -def _load_image(data_url: str) -> Image.Image: - header, separator, payload = data_url.partition(",") - if not separator or not header.startswith("data:image/") or ";base64" not in header: - raise ValueError("Multimodal training requires base64 data image URLs") - with Image.open(BytesIO(base64.b64decode(payload, validate=True))) as image: - return image.convert("RGB") - - class Qwen3_5ForCausalLM(Qwen3_5PreTrainedModel): def __init__(self, config) -> None: super().__init__(config) @@ -318,22 +305,7 @@ def get_input_embeddings(self) -> nn.Embedding: def set_input_embeddings(self, embeddings: nn.Embedding) -> None: self.model.set_input_embeddings(embeddings) - @staticmethod - def materialize_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: - """Decode the image refs of a micro batch into the vision forward kwargs.""" - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - raise ValueError("Multimodal samples require a model image processor") - images = [_load_image(ref.url) for ref in refs.images] - processed = image_processor(images=images, return_tensors="pt") - kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} - merge_size = int(image_processor.merge_size) - # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. - lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] - expected = [ref.length for ref in refs.images] - if lengths != expected: - raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") - return kwargs + materialize_images = staticmethod(materialize_qwen_vl_images) def num_active_params(self) -> int: # The vision encoder runs per image, not per text token. diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py index 46c011833a2..f3ce7dd5ff2 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py @@ -9,7 +9,6 @@ from prime_rl.trainer.models.layers.moe import MoE from prime_rl.trainer.models.layers.ngram import NGramEmbeddings, NGramTable from prime_rl.trainer.models.layers.norms import ZeroCenteredRMSNorm -from prime_rl.trainer.models.qwen3_5.modeling_qwen3_5 import Qwen3_5ForCausalLM from prime_rl.trainer.models.qwen3_8_flash_next.attention import IndexedGatedAttention from prime_rl.trainer.models.qwen3_8_flash_next.configuration_qwen3_8_flash_next import ( Qwen3_8FlashNextConfig, @@ -26,6 +25,7 @@ from prime_rl.trainer.models.qwen3_8_flash_next.position_learning import PositionLearningEnhancement from prime_rl.trainer.models.qwen3_8_flash_next.rotary_embedding import RotaryEmbedding, build_mrope_position_ids from prime_rl.trainer.models.qwen3_8_flash_next.vision import Qwen3_8FlashNextVisionModel +from prime_rl.trainer.models.qwen_vl import materialize_qwen_vl_images from prime_rl.utils.cp import CPContext, shard_for_cp, shard_position_ids_for_cp from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens @@ -346,8 +346,7 @@ def get_input_embeddings(self) -> nn.Embedding: def set_input_embeddings(self, embeddings: nn.Embedding) -> None: self.model.set_input_embeddings(embeddings) - # Same Qwen-VL image processor and placeholder count as Qwen3.5. - materialize_images = staticmethod(Qwen3_5ForCausalLM.materialize_images) + materialize_images = staticmethod(materialize_qwen_vl_images) def num_active_params(self) -> int: if not self.is_vlm: diff --git a/src/prime_rl/trainer/models/qwen_vl.py b/src/prime_rl/trainer/models/qwen_vl.py new file mode 100644 index 00000000000..91a4a3268a4 --- /dev/null +++ b/src/prime_rl/trainer/models/qwen_vl.py @@ -0,0 +1,23 @@ +from typing import Any + +import torch + +from prime_rl.trainer.vlm import load_image +from prime_rl.transports.batch import MMRefs + + +def materialize_qwen_vl_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: + """Decode the image refs of a micro batch into Qwen-VL vision forward kwargs.""" + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError("Multimodal samples require a model image processor") + images = [load_image(ref.url) for ref in refs.images] + processed = image_processor(images=images, return_tensors="pt") + kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} + merge_size = int(image_processor.merge_size) + # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. + lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] + expected = [ref.length for ref in refs.images] + if lengths != expected: + raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") + return kwargs diff --git a/src/prime_rl/trainer/vlm.py b/src/prime_rl/trainer/vlm.py index c3ef0c78339..a51e0845d7e 100644 --- a/src/prime_rl/trainer/vlm.py +++ b/src/prime_rl/trainer/vlm.py @@ -1,7 +1,11 @@ """VLMs are the ``_CUSTOM_VLM_MAPPING`` models. They keep the vision encoder at ``model.model.visual`` and the decoder at ``model.model.language_model``; text models keep the decoder at ``model.model``.""" +import base64 +from io import BytesIO + import torch.nn as nn +from PIL import Image from prime_rl.configs.trainer import ModelConfig from prime_rl.utils.logger import get_logger @@ -20,6 +24,14 @@ def get_layer_prefix(model: nn.Module) -> str: return "model.language_model.layers." if hasattr(model.model, "language_model") else "model.layers." +def load_image(data_url: str) -> Image.Image: + header, separator, payload = data_url.partition(",") + if not separator or not header.startswith("data:image/") or ";base64" not in header: + raise ValueError("Multimodal training requires base64 data image URLs") + with Image.open(BytesIO(base64.b64decode(payload, validate=True))) as image: + return image.convert("RGB") + + def setup_processor(config: ModelConfig): """Load an ``AutoProcessor`` for VLM models. Returns ``None`` for text-only models.""" from transformers import AutoProcessor diff --git a/tests/unit/train/test_multimodal.py b/tests/unit/train/test_multimodal.py index 4d2e20f41d8..089b297becb 100644 --- a/tests/unit/train/test_multimodal.py +++ b/tests/unit/train/test_multimodal.py @@ -3,7 +3,7 @@ import pytest import torch -from prime_rl.trainer.models.qwen3_5 import Qwen3_5ForCausalLM +from prime_rl.trainer.models.qwen_vl import materialize_qwen_vl_images from prime_rl.transports.batch import MMImageRef, MMRefs _IMAGE_URL = ( @@ -27,6 +27,6 @@ def __call__(self, *, images, return_tensors): } processor = SimpleNamespace(image_processor=ImageProcessor()) - assert set(Qwen3_5ForCausalLM.materialize_images(_refs(2), processor)) == {"pixel_values", "image_grid_thw"} + assert set(materialize_qwen_vl_images(_refs(2), processor)) == {"pixel_values", "image_grid_thw"} with pytest.raises(ValueError, match="placeholder lengths differ"): - Qwen3_5ForCausalLM.materialize_images(_refs(1), processor) + materialize_qwen_vl_images(_refs(1), processor) From bba40041de2751cc3239e14ca91888cab52c7fe3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Hubert=20Pysk=C5=82o?= Date: Thu, 29 Oct 2026 18:37:47 +0100 Subject: [PATCH 8/8] refactor(trainer): keep Qwen-VL image materialization in each model's vision.py Model packages stay self-contained, like the duplicated vision towers: Qwen3.5 and Qwen3.8 Flash Next each define materialize_qwen_vl_images next to their vision model, and the top-level models/qwen_vl.py goes away. Co-Authored-By: Claude Opus 5.5 --- docs/development.md | 2 +- .../models/qwen3_5/modeling_qwen3_5.py | 3 +-- src/prime_rl/trainer/models/qwen3_5/vision.py | 22 +++++++++++++++++- .../modeling_qwen3_8_flash_next.py | 3 +-- .../models/qwen3_8_flash_next/vision.py | 22 +++++++++++++++++- src/prime_rl/trainer/models/qwen_vl.py | 23 ------------------- tests/unit/train/test_multimodal.py | 2 +- 7 files changed, 46 insertions(+), 31 deletions(-) delete mode 100644 src/prime_rl/trainer/models/qwen_vl.py diff --git a/docs/development.md b/docs/development.md index 8201eeae3b5..6a190c63dd3 100644 --- a/docs/development.md +++ b/docs/development.md @@ -146,6 +146,6 @@ VLM training (any run with `[model.vlm]` set, SFT or RL) is custom-implementatio 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. RL turns raw image refs into vision kwargs with the model class's `materialize_images(refs, processor)`; Qwen3.5 and Qwen3.8 Flash Next share `models/qwen_vl.py`, and a new family implements its own. +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 with the model class's `materialize_images(refs, processor)`; Qwen3.5 and Qwen3.8 Flash Next each define it in their `vision.py`, and a new family implements its own. 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). diff --git a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py index e9141846ca4..5c6dcbe7eaa 100644 --- a/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py +++ b/src/prime_rl/trainer/models/qwen3_5/modeling_qwen3_5.py @@ -22,8 +22,7 @@ Qwen3_5RotaryEmbedding, build_qwen3_5_mrope_position_ids, ) -from prime_rl.trainer.models.qwen3_5.vision import Qwen3_5VisionModel -from prime_rl.trainer.models.qwen_vl import materialize_qwen_vl_images +from prime_rl.trainer.models.qwen3_5.vision import Qwen3_5VisionModel, materialize_qwen_vl_images from prime_rl.utils.cp import CPContext, shard_for_cp, shard_position_ids_for_cp from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens diff --git a/src/prime_rl/trainer/models/qwen3_5/vision.py b/src/prime_rl/trainer/models/qwen3_5/vision.py index 149903b8591..0c69f407b6a 100644 --- a/src/prime_rl/trainer/models/qwen3_5/vision.py +++ b/src/prime_rl/trainer/models/qwen3_5/vision.py @@ -1,4 +1,5 @@ import math +from typing import Any import torch import torch.nn.functional as F @@ -12,6 +13,8 @@ ) from prime_rl.trainer.models.layers.rotary_emb import rotate_half from prime_rl.trainer.models.qwen3_5.configuration_qwen3_5 import Qwen3_5VisionConfig +from prime_rl.trainer.vlm import load_image +from prime_rl.transports.batch import MMRefs class Qwen3_5VisionRotaryEmbedding(nn.Module): @@ -260,4 +263,21 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> BaseMod ) -__all__ = ["Qwen3_5VisionModel"] +def materialize_qwen_vl_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: + """Decode the image refs of a micro batch into Qwen-VL vision forward kwargs.""" + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError("Multimodal samples require a model image processor") + images = [load_image(ref.url) for ref in refs.images] + processed = image_processor(images=images, return_tensors="pt") + kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} + merge_size = int(image_processor.merge_size) + # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. + lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] + expected = [ref.length for ref in refs.images] + if lengths != expected: + raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") + return kwargs + + +__all__ = ["Qwen3_5VisionModel", "materialize_qwen_vl_images"] diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py index f3ce7dd5ff2..5261811ac5a 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/modeling_qwen3_8_flash_next.py @@ -24,8 +24,7 @@ from prime_rl.trainer.models.qwen3_8_flash_next.ngram_hasher import Qwen3_8FlashNextNGramHasher from prime_rl.trainer.models.qwen3_8_flash_next.position_learning import PositionLearningEnhancement from prime_rl.trainer.models.qwen3_8_flash_next.rotary_embedding import RotaryEmbedding, build_mrope_position_ids -from prime_rl.trainer.models.qwen3_8_flash_next.vision import Qwen3_8FlashNextVisionModel -from prime_rl.trainer.models.qwen_vl import materialize_qwen_vl_images +from prime_rl.trainer.models.qwen3_8_flash_next.vision import Qwen3_8FlashNextVisionModel, materialize_qwen_vl_images from prime_rl.utils.cp import CPContext, shard_for_cp, shard_position_ids_for_cp from prime_rl.utils.sequence import get_cu_seqlens_from_seq_lens diff --git a/src/prime_rl/trainer/models/qwen3_8_flash_next/vision.py b/src/prime_rl/trainer/models/qwen3_8_flash_next/vision.py index a54b6099284..0e2ce3660fe 100644 --- a/src/prime_rl/trainer/models/qwen3_8_flash_next/vision.py +++ b/src/prime_rl/trainer/models/qwen3_8_flash_next/vision.py @@ -1,4 +1,5 @@ import math +from typing import Any import torch import torch.nn.functional as F @@ -12,6 +13,8 @@ ) from prime_rl.trainer.models.layers.rotary_emb import rotate_half from prime_rl.trainer.models.qwen3_8_flash_next.configuration_qwen3_8_flash_next import Qwen3_8FlashNextVisionConfig +from prime_rl.trainer.vlm import load_image +from prime_rl.transports.batch import MMRefs class VisionRotaryEmbedding(nn.Module): @@ -271,4 +274,21 @@ def forward(self, pixel_values: torch.Tensor, grid_thw: torch.Tensor) -> BaseMod ) -__all__ = ["Qwen3_8FlashNextVisionModel"] +def materialize_qwen_vl_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: + """Decode the image refs of a micro batch into Qwen-VL vision forward kwargs.""" + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError("Multimodal samples require a model image processor") + images = [load_image(ref.url) for ref in refs.images] + processed = image_processor(images=images, return_tensors="pt") + kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} + merge_size = int(image_processor.merge_size) + # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. + lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] + expected = [ref.length for ref in refs.images] + if lengths != expected: + raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") + return kwargs + + +__all__ = ["Qwen3_8FlashNextVisionModel", "materialize_qwen_vl_images"] diff --git a/src/prime_rl/trainer/models/qwen_vl.py b/src/prime_rl/trainer/models/qwen_vl.py deleted file mode 100644 index 91a4a3268a4..00000000000 --- a/src/prime_rl/trainer/models/qwen_vl.py +++ /dev/null @@ -1,23 +0,0 @@ -from typing import Any - -import torch - -from prime_rl.trainer.vlm import load_image -from prime_rl.transports.batch import MMRefs - - -def materialize_qwen_vl_images(refs: MMRefs, processor: Any) -> dict[str, torch.Tensor]: - """Decode the image refs of a micro batch into Qwen-VL vision forward kwargs.""" - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - raise ValueError("Multimodal samples require a model image processor") - images = [load_image(ref.url) for ref in refs.images] - processed = image_processor(images=images, return_tensors="pt") - kwargs = {key: torch.as_tensor(processed[key]).contiguous() for key in ("pixel_values", "image_grid_thw")} - merge_size = int(image_processor.merge_size) - # HF Qwen-VL / renderer pad count: T*H*W / merge_size^2. - lengths = [int(grid.prod()) // (merge_size * merge_size) for grid in kwargs["image_grid_thw"].reshape(-1, 3)] - expected = [ref.length for ref in refs.images] - if lengths != expected: - raise ValueError(f"Image placeholder lengths differ from vLLM: expected {expected}, got {lengths}") - return kwargs diff --git a/tests/unit/train/test_multimodal.py b/tests/unit/train/test_multimodal.py index 089b297becb..db5e0371e0a 100644 --- a/tests/unit/train/test_multimodal.py +++ b/tests/unit/train/test_multimodal.py @@ -3,7 +3,7 @@ import pytest import torch -from prime_rl.trainer.models.qwen_vl import materialize_qwen_vl_images +from prime_rl.trainer.models.qwen3_5.vision import materialize_qwen_vl_images from prime_rl.transports.batch import MMImageRef, MMRefs _IMAGE_URL = (