diff --git a/docs/advanced.md b/docs/advanced.md index ac5485c12d0..102ee2e27c3 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -158,11 +158,13 @@ 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` | -| Qwen3.8 Flash Next | `qwen4_exp` | `model.visual` | `model.language_model` | +| Family | `model_type` | +|---|---| +| Qwen3.5 | `qwen3_5` | +| Qwen3.5-MoE | `qwen3_5_moe` | +| Qwen3.8 Flash Next | `qwen4_exp` | + +The trainer finds the vision encoder at `model.visual` and the language model at `model.language_model`. ### Enabling VLM Mode @@ -173,12 +175,10 @@ Add `[model.vlm]` to your RL or SFT config: name = "Qwen/Qwen3.5-4B" [model.vlm] -vision_encoder_attr = "model.visual" -language_model_attr = "model.language_model" # freeze_vision_encoder = true # default; set false to fine-tune the encoder ``` -The weight-broadcast key prefix is derived as `{language_model_attr}.layers.` automatically. +For `rl`, set `[model.vlm]` in the shared `[model]` block. VLM training requires a registered custom PrimeRL implementation. diff --git a/docs/development.md b/docs/development.md index 4b317186e25..6a190c63dd3 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 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/examples/extra/vlm/rl-qwen3.5.toml b/examples/extra/vlm/rl-qwen3.5.toml index aa9d9612f89..544e1d021ee 100644 --- a/examples/extra/vlm/rl-qwen3.5.toml +++ b/examples/extra/vlm/rl-qwen3.5.toml @@ -11,8 +11,6 @@ name = "your-run-name" name = "Qwen/Qwen3.5-122B-A10B" [model.vlm] -language_model_attr = "model.language_model" -vision_encoder_attr = "model.visual" freeze_vision_encoder = true [deployment] diff --git a/examples/extra/vlm/rl-qwen3.8.toml b/examples/extra/vlm/rl-qwen3.8.toml index 642a8faa8f1..5f5c8b2c738 100644 --- a/examples/extra/vlm/rl-qwen3.8.toml +++ b/examples/extra/vlm/rl-qwen3.8.toml @@ -11,8 +11,6 @@ name = "your-run-name" name = "Qwen/Qwen3.8-Flash-Next" [model.vlm] -language_model_attr = "model.language_model" -vision_encoder_attr = "model.visual" freeze_vision_encoder = true [deployment] 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/sft.py b/packages/prime-rl-configs/src/prime_rl/configs/sft.py index 7e578e5f0e0..99f39ca6a3c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -505,10 +505,6 @@ def validate_cp_micro_batch_size(self): def validate_vlm_constraints(self): if self.model.vlm is None: return self - if self.model.optimization_dtype != "bfloat16" or self.model.reduce_dtype != "bfloat16": - raise ValueError( - "VLM models must use optimization_dtype='bfloat16' and reduce_dtype='bfloat16' to match vLLM inference." - ) if self.data.micro_batch_size != 1: raise ValueError("VLM SFT requires data.micro_batch_size = 1.") if self.val is not None and self.val.data.micro_batch_size != 1: 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 85ce32a1d2c..66b167c3c95 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 52538acee2b..6ab69b2ff58 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -16,6 +16,7 @@ ResumeConfig, TrainerLogConfig, TransportConfig, + VLMConfig, WeightBroadcastConfig, ZMQTransportConfig, ) @@ -261,6 +262,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 d4bebc13cc7..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", "qwen4_exp"}) - 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 c4610659828..10fe51df346 100644 --- a/src/prime_rl/orchestrator/batch.py +++ b/src/prime_rl/orchestrator/batch.py @@ -289,11 +289,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 - - def _has_routing(sample: MicroBatch) -> bool: """Whether the sample's arrays, kept by an env server, include routing rows.""" return any(payload is not None and "routed_experts" in payload.fields for payload in sample.payloads or []) @@ -675,8 +670,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 47969ddc130..8c0c369c1d2 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -36,7 +36,6 @@ MXFP8Config, TokenizerConfig, ) -from prime_rl.multimodal import ForwardPolicy from prime_rl.trainer.activation_checkpointing import get_activation_checkpoint_wrapper from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora, strip_lora_from_state_dict from prime_rl.trainer.models import ( @@ -64,10 +63,10 @@ from prime_rl.trainer.ngram_runtime import configure_ngram_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 @@ -294,8 +290,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( @@ -303,9 +297,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", {}): @@ -365,7 +358,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. " @@ -424,23 +417,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, @@ -490,9 +466,8 @@ 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") @@ -511,7 +486,7 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None: fully_shard(table, mesh=table_mesh, **fsdp_config) table.set_gradient_divide_factor(parallel_dims.fsdp_gradient_divide_factor) - 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 for transformer_block in transformer_layers: @@ -818,7 +793,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: @@ -933,10 +908,7 @@ def setup_model( freeze_all_except_lora(model) 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: @@ -973,7 +945,6 @@ def forward( sampling_mask: Int[Tensor, "batch seq mask"] | None = None, topk_ids: Int[Tensor, "batch seq head"] | 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). @@ -981,6 +952,7 @@ def forward( ) -> PrimeLmOutput: kwargs = { "input_ids": input_ids, + "position_ids": position_ids, "labels": labels, "temperature": temperature, "sampling_mask": sampling_mask, @@ -993,15 +965,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 c245d02c081..8aebd9a8ea1 100644 --- a/src/prime_rl/trainer/models/layers/lm_head.py +++ b/src/prime_rl/trainer/models/layers/lm_head.py @@ -471,9 +471,7 @@ def new_forward( topk_ids: 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 6533d125526..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,7 +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.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 @@ -236,19 +236,15 @@ 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. + 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, + ) return inputs_embeds, position_ids def forward( @@ -308,6 +304,8 @@ def get_input_embeddings(self) -> nn.Embedding: def set_input_embeddings(self, embeddings: nn.Embedding) -> None: self.model.set_input_embeddings(embeddings) + materialize_images = staticmethod(materialize_qwen_vl_images) + 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/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 bc6527ced16..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,7 +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.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 @@ -273,19 +273,15 @@ 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.8 Flash Next image inputs") - position_ids = build_mrope_position_ids( - input_ids=input_ids, - mm_token_type_ids=mm_token_type_ids, - image_grid_thw=image_grid_thw, - spatial_merge_size=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. + position_ids = build_mrope_position_ids( + input_ids=input_ids, + mm_token_type_ids=mm_token_type_ids, + image_grid_thw=image_grid_thw, + spatial_merge_size=vision_config.spatial_merge_size, + seq_lens=seq_lens, + ) return inputs_embeds, position_ids def forward( @@ -349,6 +345,8 @@ def get_input_embeddings(self) -> nn.Embedding: def set_input_embeddings(self, embeddings: nn.Embedding) -> None: self.model.set_input_embeddings(embeddings) + materialize_images = staticmethod(materialize_qwen_vl_images) + def num_active_params(self) -> int: if not self.is_vlm: return super().num_active_params() 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/moe_runtime.py b/src/prime_rl/trainer/moe_runtime.py index d6fedf5cad9..939242eae8d 100644 --- a/src/prime_rl/trainer/moe_runtime.py +++ b/src/prime_rl/trainer/moe_runtime.py @@ -36,8 +36,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 setup_moe_comm_env(config: ModelConfig) -> None: @@ -104,9 +104,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 e5c6979aba2..9de65aa1418 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -1,7 +1,8 @@ import time +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor from pathlib import Path -from typing import Any, NotRequired, TypedDict +from typing import NotRequired, TypedDict import numpy as np import torch @@ -9,8 +10,6 @@ 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.world import get_world from prime_rl.transports.batch import ( BatchReceiver, @@ -57,10 +56,9 @@ class TensorMicroBatch(TypedDict): top_logprobs_values: Float[Tensor, "batch seq head"] | 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 @@ -195,18 +193,15 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: def prepare_micro_batch( - micro_batch: TensorMicroBatch, processor: Any | None, mm_adapter: MultimodalAdapter | None + 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 - 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 materialize_images 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) 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 47b20672cfd..7a6b1c99811 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -48,7 +48,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, @@ -57,11 +56,11 @@ get_global_moe_stats, is_tt_moe_model, setup_model, - setup_processor, take_moe_routing_stats, ) from prime_rl.trainer.moe_runtime import setup_moe_comm_env 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, @@ -161,19 +160,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 - mm_adapter = 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") - mm_adapter = get_multimodal_adapter(model.config.model_type) - prepare = partial(prepare_micro_batch, processor=processor, mm_adapter=mm_adapter) + 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) @@ -461,7 +458,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 @@ -482,9 +478,8 @@ def train(config: TrainerConfig): if cp_enabled: validate_cp_seq_lens(micro_batch["seq_lens"], input_ids.shape[-1]) - 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, position_ids, cp_rank, cp_size) seq_lens_are_pre_shard = True @@ -522,7 +517,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 b6030445bd3..c3b3cecb252 100644 --- a/src/prime_rl/trainer/sft/train.py +++ b/src/prime_rl/trainer/sft/train.py @@ -30,7 +30,6 @@ get_global_moe_stats, is_tt_moe_model, reshard_module, - setup_processor, setup_tokenizer, setup_model, take_moe_routing_stats, @@ -59,6 +58,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 from prime_rl.control.server import ControlServer @@ -259,7 +259,8 @@ def compute_loss(micro_batch: dict) -> tuple[torch.Tensor, torch.Tensor]: # 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. validate_cp_seq_lens(micro_batch["seq_lens"], input_ids.shape[-1]) - defer_vlm_cp_to_model = mm_kwargs is not None and "image_grid_thw" in mm_kwargs + # 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, position_ids, cp_rank, cp_size) seq_lens_are_pre_shard = True diff --git a/src/prime_rl/trainer/vlm.py b/src/prime_rl/trainer/vlm.py new file mode 100644 index 00000000000..a51e0845d7e --- /dev/null +++ b/src/prime_rl/trainer/vlm.py @@ -0,0 +1,49 @@ +"""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 + + +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 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 + + 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 diff --git a/src/prime_rl/transports/weights/nccl.py b/src/prime_rl/transports/weights/nccl.py index 0a48e8e76c5..0f45bddbef7 100644 --- a/src/prime_rl/transports/weights/nccl.py +++ b/src/prime_rl/transports/weights/nccl.py @@ -14,10 +14,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 get_chunked_conversion_ops, resolve_wire_dtype @@ -189,7 +189,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) chunked_ops = get_chunked_conversion_ops(model, state_dict) num_state_dict_to_send = num_layers + 1 + sum(op.shard_count for op in chunked_ops.values()) diff --git a/src/prime_rl/utils/cp.py b/src/prime_rl/utils/cp.py index 045584717fb..2332a3e7a6c 100644 --- a/src/prime_rl/utils/cp.py +++ b/src/prime_rl/utils/cp.py @@ -29,6 +29,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() diff --git a/src/prime_rl/utils/vlm.py b/src/prime_rl/utils/vlm.py deleted file mode 100644 index 7cdf52351cd..00000000000 --- a/src/prime_rl/utils/vlm.py +++ /dev/null @@ -1,125 +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"), - "glm5_next": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), - "muse_glimmer": VLMModelInfo(vision_encoder_attr="model.vision_tower", 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 87cc36ff1c7..4782285ec71 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -2,7 +2,6 @@ import pytest from prime_rl.orchestrator.batch import ( - _is_multimodal_sample, pad_micro_batch, prepare_batch, prepare_sample, @@ -566,7 +565,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_keeps_the_payload_reference(): diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 97402575469..c77ba115aee 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -790,36 +790,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_ac_none_disables_ac_offloading(): config = TrainerModelConfig.model_validate({"ac": None}) assert config.ac is None diff --git a/tests/unit/train/models/qwen3_8_flash_next/test_model.py b/tests/unit/train/models/qwen3_8_flash_next/test_model.py index 4c1e0e8d656..771859becc2 100644 --- a/tests/unit/train/models/qwen3_8_flash_next/test_model.py +++ b/tests/unit/train/models/qwen3_8_flash_next/test_model.py @@ -9,7 +9,7 @@ Qwen3_8FlashNextVisionConfig, ) from prime_rl.trainer.perf import FlopCounter -from prime_rl.utils.vlm import get_language_model, get_vision_encoder +from prime_rl.trainer.vlm import get_language_model, get_vision_encoder @pytest.fixture(params=[False, True], ids=["text", "composite"]) diff --git a/tests/unit/train/models/test_glm5_next.py b/tests/unit/train/models/test_glm5_next.py index 35a14f590ff..49a4001d540 100644 --- a/tests/unit/train/models/test_glm5_next.py +++ b/tests/unit/train/models/test_glm5_next.py @@ -3,7 +3,7 @@ from prime_rl.trainer.models.conversion_ops import apply_hf_to_prime, apply_prime_to_hf from prime_rl.trainer.models.glm5_next import Glm5NextConfig, Glm5NextForCausalLM, Glm5NextTextConfig, sparse_attention -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/models/test_muse_glimmer.py b/tests/unit/train/models/test_muse_glimmer.py index ef8caa546cc..b47fc7eb44a 100644 --- a/tests/unit/train/models/test_muse_glimmer.py +++ b/tests/unit/train/models/test_muse_glimmer.py @@ -2,7 +2,7 @@ from prime_rl.trainer.models import get_custom_vlm_cls, supports_custom_impl from prime_rl.trainer.models.muse_glimmer import MuseGlimmerConfig, MuseGlimmerForCausalLM, MuseGlimmerTextConfig -from prime_rl.utils.vlm import get_language_model, get_layer_prefix +from prime_rl.trainer.vlm import get_language_model, get_layer_prefix def _text_config(**overrides) -> MuseGlimmerTextConfig: @@ -38,7 +38,6 @@ def test_composite_config_roundtrip_keeps_vision_config(): assert restored.vision_config == {"hidden_size": 32} assert supports_custom_impl(restored) assert get_custom_vlm_cls(restored) is MuseGlimmerForCausalLM - assert get_layer_prefix(restored) == "model.language_model.layers." def test_state_dict_uses_checkpoint_names(): @@ -46,6 +45,7 @@ def test_state_dict_uses_checkpoint_names(): config._attn_implementation = "flash_attention_2" model = MuseGlimmerForCausalLM(config) assert len(get_language_model(model).layers) == 1 + assert get_layer_prefix(model) == "model.language_model.layers." layer = "model.language_model.layers.0" assert set(model.state_dict()) == { "lm_head.weight", 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..db5e0371e0a 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.models.qwen3_5.vision import materialize_qwen_vl_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_qwen_vl_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_qwen_vl_images(_refs(1), processor) diff --git a/tools/convert_dcp_to_bf16.py b/tools/convert_dcp_to_bf16.py index 4298771f75e..5977544ec6c 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