From 034e83e6147b3f4e17be9883fec2e0ed60742823 Mon Sep 17 00:00:00 2001 From: eligotts Date: Wed, 24 Jun 2026 01:11:43 +0000 Subject: [PATCH 01/52] Support v1 raw multimodal offload --- deps/renderers | 2 +- deps/verifiers | 2 +- .../src/prime_rl/configs/trainer.py | 4 + src/prime_rl/entrypoints/rl.py | 8 +- src/prime_rl/inference/vllm/serving_tokens.py | 196 +++++++++++++-- src/prime_rl/multimodal/__init__.py | 12 + src/prime_rl/multimodal/adapters/__init__.py | 4 + src/prime_rl/multimodal/adapters/base.py | 52 ++++ src/prime_rl/multimodal/adapters/kimi_k25.py | 162 ++++++++++++ src/prime_rl/multimodal/adapters/qwen_vl.py | 165 ++++++++++++ src/prime_rl/multimodal/registry.py | 17 ++ src/prime_rl/multimodal/schema.py | 66 +++++ src/prime_rl/orchestrator/orchestrator.py | 8 +- src/prime_rl/orchestrator/trajectories.py | 44 +--- .../templates/multi_node_rl.sbatch.j2 | 3 + src/prime_rl/trainer/batch.py | 77 +----- src/prime_rl/trainer/model.py | 24 +- .../models/glm4_moe/modeling_glm4_moe.py | 6 +- src/prime_rl/trainer/rl/data.py | 87 ++++++- src/prime_rl/trainer/rl/train.py | 13 +- src/prime_rl/transport/types.py | 29 ++- src/prime_rl/utils/mm.py | 237 ++++++++++++++++++ src/prime_rl/utils/run_assets.py | 24 ++ tests/unit/inference/test_serving_tokens.py | 62 +++++ tests/unit/orchestrator/test_batch.py | 90 ++++--- tests/unit/orchestrator/test_qwen3_vl_e2e.py | 44 ++-- tests/unit/test_configs.py | 5 + tests/unit/train/test_model_forward.py | 27 +- tests/unit/utils/test_mm.py | 146 +++++++++++ 29 files changed, 1393 insertions(+), 223 deletions(-) create mode 100644 src/prime_rl/multimodal/__init__.py create mode 100644 src/prime_rl/multimodal/adapters/__init__.py create mode 100644 src/prime_rl/multimodal/adapters/base.py create mode 100644 src/prime_rl/multimodal/adapters/kimi_k25.py create mode 100644 src/prime_rl/multimodal/adapters/qwen_vl.py create mode 100644 src/prime_rl/multimodal/registry.py create mode 100644 src/prime_rl/multimodal/schema.py create mode 100644 src/prime_rl/utils/mm.py create mode 100644 src/prime_rl/utils/run_assets.py create mode 100644 tests/unit/utils/test_mm.py diff --git a/deps/renderers b/deps/renderers index 1933293372..eaa07bb86b 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 193329337245d614c42a7ad6d9205542f9cc4743 +Subproject commit eaa07bb86b66eddbcbd4924d746725fbdc8ef658 diff --git a/deps/verifiers b/deps/verifiers index 20abcd9163..de37650b9b 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 20abcd91636c5457475e3a990c254824c3cc0fc3 +Subproject commit de37650b9b49ffd5541bd97cf4b0cd426c96c160 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 f1522d4efe..297db330be 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -19,6 +19,7 @@ AttnImplementation: TypeAlias = Literal["eager", "sdpa", "flash_attention_2", "flash_attention_3", "fa4"] EPCommBackend: TypeAlias = Literal["torch", "deepep"] +MissingMMImagePolicy: TypeAlias = Literal["error", "placeholder_zero_loss"] # User-facing name -> internal name. Users set `flash_attention_4` in configs, # which gets rewritten to `fa4` before pydantic validation. @@ -565,6 +566,9 @@ class TrainerConfig(BaseConfig): max_concurrent_runs: int = Field(1, ge=1) """Maximum number of concurrent runs to allow. If 1, only one run may run at a time.""" + missing_mm_image_policy: MissingMMImagePolicy = "placeholder_zero_loss" + """Policy when raw multimodal image files disappear before trainer materialization. ``placeholder_zero_loss`` warns, synthesizes zero-valued image tensors with the original descriptor geometry, and masks out the affected microbatch loss; ``error`` preserves fail-fast behavior.""" + enable_token_export: bool = False """Opt-in per-token JSONL export for rollout debugging. When enabled, writes token ids and aligned trainer metrics after each forward pass.""" diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 533febcaf7..7a7e50f21e 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -26,6 +26,7 @@ validate_output_dir, ) from prime_rl.utils.process import cleanup_processes, cleanup_threads, monitor_process, set_proc_title +from prime_rl.utils.run_assets import run_asset_env RL_TOML = "rl.toml" RL_SBATCH = "rl.sbatch" @@ -116,6 +117,7 @@ def rl_local(config: RLConfig): "WANDB_SHARED_MODE": "1", "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } + shared_run_asset_env = run_asset_env(config.orchestrator.output_dir) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.student.client.is_elastic: @@ -161,7 +163,7 @@ def sigterm_handler(signum, frame): inference_process = Popen( inference_cmd, env={ - **os.environ, + **shared_run_asset_env, "CUDA_VISIBLE_DEVICES": ",".join(map(str, infer_gpu_ids)), }, stdout=log_file, @@ -205,7 +207,7 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ - **os.environ, + **shared_run_asset_env, **wandb_shared_env, "WANDB_SHARED_LABEL": "orchestrator", "LOGURU_FORCE_COLORS": "1", @@ -251,7 +253,7 @@ def sigterm_handler(signum, frame): trainer_process = Popen( trainer_cmd, env={ - **os.environ, + **shared_run_asset_env, **wandb_shared_env, "WANDB_SHARED_LABEL": "trainer", "CUDA_VISIBLE_DEVICES": ",".join(map(str, trainer_gpu_ids)), diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 52c9a76b5d..a3b35df9e4 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -4,7 +4,8 @@ ``vllm.entrypoints.serve.disagg.serving.ServingTokens`` that already covers prefix-cache salting, lora dispatch, multimodal features, prompt logprobs, priority, ``data_parallel_rank`` header routing and server-side ``max_tokens`` -defaulting. We subclass it for the bits still missing from the upstream handler: +defaulting. We subclass it for prime-RL behavior that is still missing or +customized: 1. ``data_parallel_rank`` routing — read from the ``X-data-parallel-rank`` header and forwarded to ``engine_client.generate``. Upstream ``ServingTokens`` @@ -15,7 +16,12 @@ decisions, surface them as base64 raw-byte payloads without requiring a vLLM source fork. -3. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies +3. Raw image refs for multimodal rollouts — renderers send lightweight raw + descriptor refs for new images and ``None`` for cache-only prior images. + This handler materializes refs through multimodal adapters and turns cache + misses into a structured retryable error. + +4. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies this itself (via ``GenerateRequest.is_sampling_param_provided`` + ``get_max_tokens``); we keep an equivalent guard so callers that omit ``max_tokens`` don't truncate at vLLM's 16-token ``SamplingParams`` default. @@ -26,8 +32,13 @@ from __future__ import annotations +import asyncio +import hashlib from collections.abc import AsyncGenerator, AsyncIterable -from functools import cached_property +from dataclasses import dataclass +from functools import cached_property, lru_cache +from http import HTTPStatus +from io import BytesIO from typing import Any from fastapi import Request @@ -48,6 +59,15 @@ from vllm.sampling_params import RequestOutputKind, SamplingParams from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture +from prime_rl.multimodal.registry import get_multimodal_adapter +from prime_rl.multimodal.schema import RawMMItem + + +@dataclass +class _MMImageRefError(Exception): + message: str + err_type: str = "invalid_mm_image_ref" + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST class PrimeRlGenerateResponseChoice(GenerateResponseChoice): @@ -147,8 +167,133 @@ async def _client_set_max_tokens(raw_request: Request | None) -> bool: return isinstance(sp, dict) and "max_tokens" in sp +def _is_missing_mm_cache_error(exc: BaseException) -> bool: + return "Expected a cached item for mm_hash=" in str(exc) + + +def _cache_only_mm_hashes(features: Any) -> list[str]: + missing: list[str] = [] + kwargs_data = features.kwargs_data or {} + for modality, hashes in features.mm_hashes.items(): + items = kwargs_data.get(modality) + if items is None: + missing.extend(f"{modality}:{mm_hash}" for mm_hash in hashes) + continue + for idx, mm_hash in enumerate(hashes): + if idx >= len(items) or items[idx] is None: + missing.append(f"{modality}:{mm_hash}") + return missing + + +def _missing_mm_cache_message(features: Any, exc: BaseException) -> str: + hashes = _cache_only_mm_hashes(features) + if not hashes: + return str(exc) + joined = ", ".join(hashes[:8]) + suffix = "" if len(hashes) <= 8 else f", ... (+{len(hashes) - 8} more)" + return f"vLLM multimodal cache miss for {joined}{suffix}" + + +@lru_cache(maxsize=8) +def _load_image_processor(model_name: str, trust_remote_code: bool): + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=trust_remote_code) + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError(f"{model_name!r} does not expose an image_processor") + return image_processor + + +def _materialize_raw_image_ref_sync( + raw_ref: str, + *, + expected_modality: str, + expected_hash: str, + expected_placeholder_length: int | None, + processor_model_name: str, + trust_remote_code: bool, +): + from PIL import Image + from renderers.mm_store import raw_image_path, split_mmraw_ref + + try: + ref = split_mmraw_ref(raw_ref) + if ref.modality != expected_modality: + raise ValueError(f"Expected modality {expected_modality!r}, got {ref.modality!r}") + if ref.mm_hash != expected_hash: + raise ValueError(f"Expected image hash {expected_hash}, got {ref.mm_hash}") + + raw = raw_image_path(run_id=ref.run_id, raw_image_id=ref.raw_image_id).read_bytes() + actual_hash = hashlib.sha256(raw).hexdigest()[:32] + if actual_hash != ref.mm_hash: + raise ValueError(f"Raw image hash mismatch: expected {ref.mm_hash}, got {actual_hash}") + + image_processor = _load_image_processor(processor_model_name, trust_remote_code) + item = RawMMItem( + modality=ref.modality, + family=ref.family, + layout_fingerprint=ref.fingerprint, + payload=dict(ref.payload), + raw_ref=raw_ref, + ) + adapter = get_multimodal_adapter(ref.family) + image = Image.open(BytesIO(raw)).convert("RGB") + return adapter.materialize_for_vllm( + image_processor, + item, + image, + expected_placeholder_length, + ) + except Exception as exc: + raise _MMImageRefError(str(exc)) from exc + + +async def _decode_raw_mm_kwargs( + features: Any, + *, + processor_model_name: str, + trust_remote_code: bool, +) -> dict[str, list[Any | None]]: + from renderers.mm_store import IMAGE_REF_PREFIX + + kwargs_data = features.kwargs_data or {} + mm_kwargs: dict[str, list[Any | None]] = {} + for modality, hashes in features.mm_hashes.items(): + placeholders = features.mm_placeholders.get(modality, []) + items = kwargs_data.get(modality) + if items is None: + mm_kwargs[modality] = [None] * len(hashes) + continue + if len(items) != len(hashes): + raise _MMImageRefError( + f"Multimodal kwargs/hash length mismatch for {modality}: {len(items)} != {len(hashes)}" + ) + decoded: list[Any | None] = [] + for idx, item in enumerate(items): + if item is None: + decoded.append(None) + continue + if not isinstance(item, str) or not item.startswith(f"{IMAGE_REF_PREFIX}:"): + raise _MMImageRefError("v1 multimodal inference accepts raw descriptor refs only") + placeholder_length = placeholders[idx].length if idx < len(placeholders) else None + decoded.append( + await asyncio.to_thread( + _materialize_raw_image_ref_sync, + item, + expected_modality=modality, + expected_hash=hashes[idx], + expected_placeholder_length=placeholder_length, + processor_model_name=processor_model_name, + trust_remote_code=trust_remote_code, + ) + ) + mm_kwargs[modality] = decoded + return mm_kwargs + + class PrimeRlServingTokens(ServingTokens): - """ServingTokens + DP-rank routing + compact routed experts + max_tokens defaulting.""" + """ServingTokens + DP-rank routing + routed experts + raw image refs + max_tokens defaulting.""" @cached_property def _max_tokens_defaults(self) -> tuple[dict, int | None]: @@ -198,26 +343,28 @@ async def serve_tokens( raw_request.state.request_metadata = request_metadata # Build the engine input — features-aware (MM) or text-only fallback. - # Identical to upstream so we keep tracking it. if features := request.features: - from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item from vllm.inputs import mm_input - from vllm.multimodal.inputs import ( - MultiModalKwargsItem, - PlaceholderRange, - ) + from vllm.multimodal.inputs import PlaceholderRange mm_placeholders = { modality: [PlaceholderRange(offset=p.offset, length=p.length) for p in ranges] for modality, ranges in features.mm_placeholders.items() } - mm_kwargs: dict[str, list[MultiModalKwargsItem | None]] = {} - if features.kwargs_data is not None: - for modality, items in features.kwargs_data.items(): - mm_kwargs[modality] = [decode_mm_kwargs_item(item) if item is not None else None for item in items] - else: - for modality, hashes in features.mm_hashes.items(): - mm_kwargs[modality] = [None] * len(hashes) + processor_model_name = getattr(self.model_config, "model", None) or model_name + trust_remote_code = bool(getattr(self.model_config, "trust_remote_code", False)) + try: + mm_kwargs = await _decode_raw_mm_kwargs( + features, + processor_model_name=processor_model_name, + trust_remote_code=trust_remote_code, + ) + except _MMImageRefError as exc: + return self.create_error_response( + message=exc.message, + err_type=exc.err_type, + status_code=exc.status_code, + ) engine_input = mm_input( prompt_token_ids=request.token_ids, mm_kwargs=mm_kwargs, # type: ignore[arg-type] @@ -330,9 +477,18 @@ async def serve_tokens_full_generator( # type: ignore[override] final_capture = _FinalOutputCapture(result_generator) result_generator = final_capture - response = await super().serve_tokens_full_generator( - request, result_generator, request_id, model_name, request_metadata - ) + try: + response = await super().serve_tokens_full_generator( + request, result_generator, request_id, model_name, request_metadata + ) + except AssertionError as exc: + if request.features is not None and _is_missing_mm_cache_error(exc): + return self.create_error_response( + message=_missing_mm_cache_message(request.features, exc), + err_type="missing_mm_cache_item", + status_code=HTTPStatus.CONFLICT, + ) + raise if not isinstance(response, GenerateResponse): return response diff --git a/src/prime_rl/multimodal/__init__.py b/src/prime_rl/multimodal/__init__.py new file mode 100644 index 0000000000..961b12ada0 --- /dev/null +++ b/src/prime_rl/multimodal/__init__.py @@ -0,0 +1,12 @@ +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM, MultimodalAdapter +from prime_rl.multimodal.registry import get_multimodal_adapter +from prime_rl.multimodal.schema import RawMMItem, parse_raw_mm_item + +__all__ = [ + "ForwardPolicy", + "MaterializedMM", + "MultimodalAdapter", + "RawMMItem", + "get_multimodal_adapter", + "parse_raw_mm_item", +] diff --git a/src/prime_rl/multimodal/adapters/__init__.py b/src/prime_rl/multimodal/adapters/__init__.py new file mode 100644 index 0000000000..453f5c39f7 --- /dev/null +++ b/src/prime_rl/multimodal/adapters/__init__.py @@ -0,0 +1,4 @@ +from prime_rl.multimodal.adapters.kimi_k25 import KimiK25Adapter +from prime_rl.multimodal.adapters.qwen_vl import QwenVLAdapter + +__all__ = ["KimiK25Adapter", "QwenVLAdapter"] diff --git a/src/prime_rl/multimodal/adapters/base.py b/src/prime_rl/multimodal/adapters/base.py new file mode 100644 index 0000000000..86571509fb --- /dev/null +++ b/src/prime_rl/multimodal/adapters/base.py @@ -0,0 +1,52 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Protocol + +if TYPE_CHECKING: + import torch + from PIL.Image import Image + + from prime_rl.multimodal.schema import RawMMItem + + +@dataclass(frozen=True) +class ForwardPolicy: + pass_position_ids_with_mm: bool = True + requires_mm_token_type_ids: bool = False + + +@dataclass(frozen=True) +class MaterializedMM: + kwargs: dict[str, "torch.Tensor"] + forward_policy: ForwardPolicy + + +class MultimodalAdapter(Protocol): + family: str + forward_policy: ForwardPolicy + + def validate_item(self, item: "RawMMItem") -> None: ... + + def processor_fingerprint(self, image_processor: Any) -> str: ... + + def materialize_for_trainer( + self, + image_processor: Any, + items: list["RawMMItem"], + images: list["Image"], + ) -> MaterializedMM: ... + + def materialize_for_vllm( + self, + image_processor: Any, + item: "RawMMItem", + image: "Image", + expected_placeholder_length: int | None, + ) -> Any: ... + + def synthesize_placeholder( + self, + image_processor: Any, + items: list["RawMMItem"], + ) -> MaterializedMM | None: ... diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py new file mode 100644 index 0000000000..e04f4132cf --- /dev/null +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import math +from typing import Any + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RawMMItem + +KIMI_K25_DEFAULTS = { + "patch_size": 14, + "merge_kernel_size": 2, + "in_patch_limit": 16384, + "patch_limit_on_one_side": 512, + "fixed_output_tokens": None, + "image_mean": [0.5, 0.5, 0.5], + "image_std": [0.5, 0.5, 0.5], +} + + +def _tensorize(value: Any): + import torch + + if isinstance(value, torch.Tensor): + return value.contiguous() + return torch.as_tensor(value).contiguous() + + +def _cfg_value(image_processor: Any, name: str) -> Any: + for source in ( + image_processor, + getattr(image_processor, "media_proc_cfg", None), + getattr(image_processor, "config", None), + ): + if source is None: + continue + if isinstance(source, dict) and name in source: + return source[name] + value = getattr(source, name, None) + if value is not None: + return value + return KIMI_K25_DEFAULTS[name] + + +def _grid_payload(item: RawMMItem) -> list[int]: + grid = item.payload.get("grid_thws") + if grid is None: + raise ValueError("Kimi raw descriptor payload is missing grid_thws") + if len(grid) == 1 and isinstance(grid[0], list): + grid = grid[0] + if not isinstance(grid, list | tuple) or len(grid) != 3: + raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") + out = [int(v) for v in grid] + if any(v <= 0 for v in out): + raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") + return out + + +def _process_images(image_processor: Any, images: list[Any], *, return_tensors: str): + medias = [{"type": "image", "image": image} for image in images] + preprocess = getattr(image_processor, "preprocess", None) + if preprocess is None: + raise ValueError("Kimi image processor is missing preprocess") + return preprocess(medias, return_tensors=return_tensors) + + +class KimiK25Adapter: + family = "kimi_k25" + forward_policy = ForwardPolicy(pass_position_ids_with_mm=True) + + def validate_item(self, item: RawMMItem) -> None: + if item.family != self.family: + raise ValueError(f"Kimi adapter cannot handle family {item.family!r}") + _grid_payload(item) + + def processor_fingerprint(self, image_processor: Any) -> str: + from renderers.mm_store import image_layout_fingerprint + + fixed_output_tokens = _cfg_value(image_processor, "fixed_output_tokens") + return image_layout_fingerprint( + family=self.family, + patch_size=int(_cfg_value(image_processor, "patch_size")), + merge_kernel_size=int(_cfg_value(image_processor, "merge_kernel_size")), + in_patch_limit=int(_cfg_value(image_processor, "in_patch_limit")), + patch_limit_on_one_side=int(_cfg_value(image_processor, "patch_limit_on_one_side")), + fixed_output_tokens=None if fixed_output_tokens is None else int(fixed_output_tokens), + image_mean=list(_cfg_value(image_processor, "image_mean")), + image_std=list(_cfg_value(image_processor, "image_std")), + ) + + def materialize_for_trainer( + self, + image_processor: Any, + items: list[RawMMItem], + images: list[Any], + ) -> MaterializedMM: + for item in items: + self.validate_item(item) + processed = _process_images(image_processor, images, return_tensors="pt") + tensors = {str(k): _tensorize(v) for k, v in dict(processed).items()} + if "grid_thws" not in tensors: + raise ValueError("Kimi processor did not return grid_thws") + actual_grids = tensors["grid_thws"].reshape(-1, 3).tolist() + for idx, item in enumerate(items): + expected = _grid_payload(item) + if actual_grids[idx] != expected: + raise ValueError(f"Kimi grid mismatch at index {idx}: expected {expected}, got {actual_grids[idx]}") + return MaterializedMM(kwargs=tensors, forward_policy=self.forward_policy) + + def materialize_for_vllm( + self, + image_processor: Any, + item: RawMMItem, + image: Any, + expected_placeholder_length: int | None, + ) -> Any: + from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems + + self.validate_item(item) + actual_fingerprint = self.processor_fingerprint(image_processor) + if actual_fingerprint != item.layout_fingerprint: + raise ValueError( + f"Image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" + ) + hf_inputs = _process_images(image_processor, [image], return_tensors="pt") + tensors = {str(k): _tensorize(v) for k, v in dict(hf_inputs).items()} + expected_grid = _grid_payload(item) + actual_grid = tensors["grid_thws"].reshape(-1, 3).tolist()[0] + if actual_grid != expected_grid: + raise ValueError(f"Kimi grid mismatch: expected {expected_grid}, got {actual_grid}") + if expected_placeholder_length is not None and expected_placeholder_length != 1: + raise ValueError(f"Kimi image placeholder length mismatch: expected {expected_placeholder_length}, got 1") + grid_sizes = tensors["grid_thws"].reshape(-1, 3).prod(-1) + config_by_key = { + "pixel_values": MultiModalFieldConfig.flat_from_sizes("vision_chunk", grid_sizes), + "grid_thws": MultiModalFieldConfig.batched("vision_chunk"), + } + return MultiModalKwargsItems.from_hf_inputs(tensors, config_by_key)["vision_chunk"][0] + + def synthesize_placeholder( + self, + image_processor: Any, + items: list[RawMMItem], + ) -> MaterializedMM | None: + if not items: + return None + import torch + + patch_size = int(_cfg_value(image_processor, "patch_size")) + grids: list[list[int]] = [] + pixel_values: list[torch.Tensor] = [] + for item in items: + self.validate_item(item) + grid = _grid_payload(item) + grids.append(grid) + pixel_values.append(torch.zeros((math.prod(grid), 3, patch_size, patch_size), dtype=torch.float32)) + return MaterializedMM( + kwargs={ + "pixel_values": torch.cat(pixel_values, dim=0).contiguous(), + "grid_thws": torch.tensor(grids, dtype=torch.long), + }, + forward_policy=self.forward_policy, + ) diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py new file mode 100644 index 0000000000..59551ee327 --- /dev/null +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -0,0 +1,165 @@ +from __future__ import annotations + +import math +from typing import Any + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RawMMItem + + +def _processor_value(processor: Any, name: str, *, size_key: str | None = None) -> int: + value = getattr(processor, name, None) + if value is None and size_key is not None: + size = getattr(processor, "size", None) + get_size_value = getattr(size, "get", None) + if callable(get_size_value): + value = get_size_value(size_key) + if value is None: + raise ValueError(f"Image processor is missing {name}") + return int(value) + + +def _tensorize(value: Any): + import torch + + if isinstance(value, torch.Tensor): + return value.contiguous() + return torch.as_tensor(value).contiguous() + + +def _grid_payload(item: RawMMItem) -> list[int]: + grid = item.payload.get("image_grid_thw") + if grid is None: + raise ValueError("Qwen raw descriptor payload is missing image_grid_thw") + if len(grid) == 1 and isinstance(grid[0], list): + grid = grid[0] + if not isinstance(grid, list | tuple) or len(grid) != 3: + raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") + out = [int(v) for v in grid] + if any(v <= 0 for v in out): + raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") + return out + + +def _patch_area(patch_size: Any) -> int: + if isinstance(patch_size, list | tuple): + return math.prod(int(dim) for dim in patch_size) + size = int(patch_size) + return size * size + + +def _temporal_patch_extent(temporal_patch_size: Any) -> int: + if isinstance(temporal_patch_size, list | tuple): + return math.prod(int(dim) for dim in temporal_patch_size) + return int(temporal_patch_size) + + +class QwenVLAdapter: + family = "qwen_vl" + forward_policy = ForwardPolicy( + pass_position_ids_with_mm=False, + requires_mm_token_type_ids=True, + ) + + def validate_item(self, item: RawMMItem) -> None: + if item.family != self.family: + raise ValueError(f"Qwen adapter cannot handle family {item.family!r}") + _grid_payload(item) + + def processor_fingerprint(self, image_processor: Any) -> str: + from renderers.mm_store import image_layout_fingerprint + + return image_layout_fingerprint( + family=self.family, + patch_size=_processor_value(image_processor, "patch_size"), + merge_size=_processor_value(image_processor, "merge_size"), + temporal_patch_size=_processor_value(image_processor, "temporal_patch_size"), + min_pixels=_processor_value(image_processor, "min_pixels", size_key="shortest_edge"), + max_pixels=_processor_value(image_processor, "max_pixels", size_key="longest_edge"), + ) + + def materialize_for_trainer( + self, + image_processor: Any, + items: list[RawMMItem], + images: list[Any], + ) -> MaterializedMM: + for item in items: + self.validate_item(item) + processed = image_processor(images=images, return_tensors="pt") + tensors = {str(k): _tensorize(v) for k, v in dict(processed).items()} + if "image_grid_thw" not in tensors: + raise ValueError("Qwen processor did not return image_grid_thw") + actual_grids = tensors["image_grid_thw"].tolist() + for idx, item in enumerate(items): + expected = _grid_payload(item) + if actual_grids[idx] != expected: + raise ValueError(f"Image grid mismatch at index {idx}: expected {expected}, got {actual_grids[idx]}") + return MaterializedMM(kwargs=tensors, forward_policy=self.forward_policy) + + def materialize_for_vllm( + self, + image_processor: Any, + item: RawMMItem, + image: Any, + expected_placeholder_length: int | None, + ) -> Any: + from vllm.model_executor.models.qwen2_vl import _create_qwen2vl_field_factory + from vllm.multimodal.inputs import MultiModalKwargsItems + + self.validate_item(item) + actual_fingerprint = self.processor_fingerprint(image_processor) + if actual_fingerprint != item.layout_fingerprint: + raise ValueError( + f"Image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" + ) + hf_inputs = image_processor(images=[image], return_tensors="pt") + merge_size = _processor_value(image_processor, "merge_size") + config_by_key = _create_qwen2vl_field_factory(merge_size)(hf_inputs) + mm_item = MultiModalKwargsItems.from_hf_inputs(hf_inputs, config_by_key)["image"][0] + expected_grid = _grid_payload(item) + actual_grid = mm_item["image_grid_thw"].data.tolist() + if actual_grid != expected_grid: + raise ValueError(f"Image grid mismatch: expected {expected_grid}, got {actual_grid}") + num_image_tokens = int(expected_grid[0] * expected_grid[1] * expected_grid[2] // (merge_size * merge_size)) + if expected_placeholder_length is not None and expected_placeholder_length != num_image_tokens: + raise ValueError( + f"Image placeholder length mismatch: expected {expected_placeholder_length}, got {num_image_tokens}" + ) + return mm_item + + def placeholder_feature_dim(self, image_processor: Any) -> int: + patch_size = getattr(image_processor, "patch_size", None) + temporal_patch_size = getattr(image_processor, "temporal_patch_size", None) + image_mean = getattr(image_processor, "image_mean", None) + channels = len(image_mean) if image_mean is not None else getattr(image_processor, "num_channels", 3) + if patch_size is None or temporal_patch_size is None: + raise ValueError( + "Cannot synthesize raw image placeholders without image processor patch_size and temporal_patch_size" + ) + return int(channels) * _temporal_patch_extent(temporal_patch_size) * _patch_area(patch_size) + + def synthesize_placeholder( + self, + image_processor: Any, + items: list[RawMMItem], + ) -> MaterializedMM | None: + if not items: + return None + import torch + + feature_dim = self.placeholder_feature_dim(image_processor) + pixel_values: list[torch.Tensor] = [] + image_grid_thw: list[list[int]] = [] + for idx, item in enumerate(items): + self.validate_item(item) + grid = _grid_payload(item) + pixel_values.append(torch.zeros((math.prod(grid), feature_dim), dtype=torch.float32)) + image_grid_thw.append(grid) + return MaterializedMM( + kwargs={ + "pixel_values": torch.cat(pixel_values, dim=0).contiguous(), + "image_grid_thw": torch.tensor(image_grid_thw, dtype=torch.long), + }, + forward_policy=self.forward_policy, + ) diff --git a/src/prime_rl/multimodal/registry.py b/src/prime_rl/multimodal/registry.py new file mode 100644 index 0000000000..88ccab849f --- /dev/null +++ b/src/prime_rl/multimodal/registry.py @@ -0,0 +1,17 @@ +from __future__ import annotations + +from prime_rl.multimodal.adapters.base import MultimodalAdapter +from prime_rl.multimodal.adapters.kimi_k25 import KimiK25Adapter +from prime_rl.multimodal.adapters.qwen_vl import QwenVLAdapter + +_ADAPTERS: dict[str, MultimodalAdapter] = { + QwenVLAdapter.family: QwenVLAdapter(), + KimiK25Adapter.family: KimiK25Adapter(), +} + + +def get_multimodal_adapter(family: str) -> MultimodalAdapter: + try: + return _ADAPTERS[family] + except KeyError as exc: + raise NotImplementedError(f"No multimodal adapter registered for family {family!r}") from exc diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py new file mode 100644 index 0000000000..4cc04f3eaa --- /dev/null +++ b/src/prime_rl/multimodal/schema.py @@ -0,0 +1,66 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +RAW_MM_ITEM_KIND = "prime_raw_mm_item" +RAW_MM_ITEM_VERSION = 1 +PROCESSED_MM_KEYS = {"pixel_values", "image_embeds", "image_features"} + + +@dataclass(frozen=True) +class RawMMItem: + modality: str + family: str + layout_fingerprint: str + payload: dict[str, Any] + raw_ref: str | None = None + vllm_modality: str | None = None + + +def contains_processed_payload_key(value: Any) -> bool: + if isinstance(value, Mapping): + return bool(PROCESSED_MM_KEYS.intersection(value)) or any( + contains_processed_payload_key(v) for v in value.values() + ) + if isinstance(value, list | tuple): + return any(contains_processed_payload_key(v) for v in value) + return False + + +def parse_raw_mm_item(value: Any) -> RawMMItem: + if not isinstance(value, Mapping): + raise TypeError(f"v1 multimodal sidecars must be raw descriptor dicts, got {type(value).__name__}") + if contains_processed_payload_key(value): + raise TypeError("v1 multimodal sidecars must not carry processed multimodal payloads") + if value.get("kind") != RAW_MM_ITEM_KIND: + raise ValueError("raw multimodal descriptor is missing the common envelope kind") + if int(value.get("version", -1)) != RAW_MM_ITEM_VERSION: + raise ValueError(f"unsupported raw multimodal descriptor version: {value.get('version')!r}") + modality = value.get("modality") + family = value.get("family") + layout_fingerprint = value.get("layout_fingerprint") + payload = value.get("payload") + if not isinstance(modality, str) or not modality: + raise ValueError("raw multimodal descriptor is missing modality") + if not isinstance(family, str) or not family: + raise ValueError("raw multimodal descriptor is missing family") + if not isinstance(layout_fingerprint, str) or not layout_fingerprint: + raise ValueError("raw multimodal descriptor is missing layout_fingerprint") + if not isinstance(payload, Mapping): + raise ValueError("raw multimodal descriptor payload must be a dict") + raw_ref = value.get("raw_ref") + if raw_ref is not None and not isinstance(raw_ref, str): + raise ValueError("raw multimodal descriptor raw_ref must be a string when present") + vllm_modality = value.get("vllm_modality") + if vllm_modality is not None and not isinstance(vllm_modality, str): + raise ValueError("raw multimodal descriptor vllm_modality must be a string when present") + return RawMMItem( + modality=modality, + family=family, + layout_fingerprint=layout_fingerprint, + payload={str(k): v for k, v in payload.items()}, + raw_ref=raw_ref, + vllm_modality=vllm_modality, + ) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 1344596cfd..2281eb8683 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -101,11 +101,9 @@ # resumed when the watcher advances ``policy.version``. TARGET_LAG = 1 -# Drop the per-node training tensors when dumping a Trace to disk (rollout jsonl / wandb -# tables): the multimodal `mm_kwargs` carrier and the router-replay `routed_experts` array are -# training inputs, not part of the rollout record, and would bloat every line. They also can't -# round-trip the json dump — their `__nd__` carriers hold raw bytes. `__all__` applies the -# exclude to every node in the list. +# Drop per-node training sidecars when dumping a Trace to disk (rollout jsonl / wandb +# tables): multimodal descriptors and router replay are training inputs, not part of the +# rollout record. `__all__` applies the exclude to every node in the list. ROLLOUT_DUMP_EXCLUDE = {"nodes": {"__all__": {"multi_modal_data", "routed_experts"}}} diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index ff9d3ef3e2..771d88db40 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -9,8 +9,8 @@ Training is renderer-only across every mode (RL/OPD student, SFT teacher), so every node always carries its tokens — no backfill needed. For multimodal rollouts the branch also carries -the images it introduced (`branch.multi_modal_data`), rebuilt here into the flat `mm_kwargs` / -`mm_token_type_ids` the trainer forwards. +raw image refs and renderer descriptors (`branch.multi_modal_data`), preserved here as +`mm_refs` for trainer-side materialization. """ from __future__ import annotations @@ -19,32 +19,9 @@ import verifiers.v1 as vf from prime_rl.transport import TrainingSample -from prime_rl.transport.types import EncodedTensor, RoutedExperts +from prime_rl.transport.types import MMRefs, RoutedExperts from prime_rl.utils.logger import get_logger - - -def _to_numpy(val) -> np.ndarray: - """A renderer mm item value (torch tensor or numpy array) -> a contiguous numpy array.""" - if hasattr(val, "detach"): # torch tensor - val = val.detach().cpu().numpy() - return np.ascontiguousarray(val) - - -def _encode_mm_kwargs(mm_items: dict[str, list[dict]]) -> dict[str, EncodedTensor] | None: - """Concatenate the branch's per-image renderer items into the flat `mm_kwargs` the trainer - forwards — one `EncodedTensor` per kwarg key (e.g. `pixel_values`, `image_grid_thw`), images - cat'd along dim 0 in branch token order. Model-agnostic: the keys are whatever the processor - emits. Returns None when there are no items.""" - bins: dict[str, list[np.ndarray]] = {} - for items in mm_items.values(): # per modality - for item in items: # per image - for key, val in item.items(): - bins.setdefault(key, []).append(_to_numpy(val)) - encoded: dict[str, EncodedTensor] = {} - for key, arrs in bins.items(): - arr = np.concatenate(arrs, axis=0) - encoded[key] = EncodedTensor(dtype=str(arr.dtype), shape=list(arr.shape), data=arr.tobytes()) - return encoded or None +from prime_rl.utils.mm import build_mm_refs def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExperts | None: @@ -76,9 +53,10 @@ def trace_to_samples( `branch.sampled_mask` / `branch.logprobs`), so a sample carries it directly: `mask` marks the trainable (model-sampled) tokens, the context tokens between completions stay masked out. On a rollout error the whole completion is masked out. A branch carrying images also - gets `mm_kwargs` (the concatenated pixel tensors) and `mm_token_type_ids` (the renderer's - `mm_token_type_id_map` applied to the branch tokens). Branches with no sampled tokens - (e.g. an openai client carrying none) yield nothing. + gets `mm_refs` (raw image URIs + JSON-safe renderer metadata) and + `mm_token_type_ids` (the renderer's `mm_token_type_id_map` applied to the branch + tokens). Branches with no sampled tokens (e.g. an openai client carrying none) yield + nothing. """ has_error = trace.has_error samples: list[TrainingSample] = [] @@ -87,11 +65,11 @@ def trace_to_samples( if not any(mask): continue token_ids = branch.token_ids - mm_kwargs: dict[str, EncodedTensor] | None = None + mm_refs: MMRefs | None = None mm_token_type_ids: list[int] | None = None mmd = branch.multi_modal_data if mmd is not None: - mm_kwargs = _encode_mm_kwargs(mmd.mm_items) + mm_refs = build_mm_refs(mmd, branch.messages) mapping = mm_token_type_ids_mapping or {} mm_token_type_ids = [mapping.get(t, 0) for t in token_ids] samples.append( @@ -103,7 +81,7 @@ def trace_to_samples( teacher_logprobs=None, advantage=None, env_name=env_name, - mm_kwargs=mm_kwargs, + mm_refs=mm_refs, mm_token_type_ids=mm_token_type_ids, routed_experts=_encode_routed_experts(branch.routed_experts, len(token_ids)), ) diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 67ce7e8f41..33c9e7ee29 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -59,6 +59,9 @@ export PROJECT_DIR={{ project_dir }} export CONFIG_DIR={{ config_dir }} export OUTPUT_DIR={{ output_dir }} export ORCHESTRATOR_OUTPUT_DIR={{ orchestrator_output_dir }} +if [ -z "${VF_RENDERER_IMAGE_OFFLOAD_DIR:-}" ] && [ -z "${PRIME_RL_RUN_DIR:-}" ] && [ -z "${RUN_ID:-}" ]; then + export PRIME_RL_RUN_DIR="$ORCHESTRATOR_OUTPUT_DIR" +fi mkdir -p $OUTPUT_DIR/logs/trainer $OUTPUT_DIR/logs/inference rm -f $OUTPUT_DIR/logs/inference/*.log ln -sfn trainer/node_0.log $OUTPUT_DIR/logs/trainer.log diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 6da2ca3f8d..cdb1b083f8 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -2,10 +2,8 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass -import numpy as np - from prime_rl.trainer.utils import balanced_partition -from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import MicroBatch, RoutedExperts, TrainingSample ROUTED_EXPERTS_DTYPE_ITEMSIZE = { "uint8": 1, @@ -43,58 +41,6 @@ def _pad_routed_experts(micro_batch: MicroBatch, padding_size: int) -> None: routed_experts.shape[0] += padding_size -def _slice_encoded(tensor: EncodedTensor, n_rows: int) -> EncodedTensor: - """First `n_rows` rows of a dim-0-stacked encoded tensor (e.g. pixel_values, image_grid_thw).""" - row = int(np.prod(tensor.shape[1:])) if len(tensor.shape) > 1 else 1 - itemsize = np.dtype(tensor.dtype).itemsize - return EncodedTensor( - dtype=tensor.dtype, - shape=[n_rows, *tensor.shape[1:]], - data=tensor.data[: n_rows * row * itemsize], - ) - - -def _truncate_mm( - mm_token_type_ids: list[int], mm_kwargs: dict[str, EncodedTensor], seq_len: int -) -> tuple[int, dict[str, EncodedTensor] | None]: - """Truncating a sample must not split an image's placeholder block, else the surviving image - token count no longer matches the image embeddings in `mm_kwargs`. Returns the cut point - (<= seq_len, never inside an image block) and `mm_kwargs` sliced to the images whose - placeholders fully survive (None if no image survives).""" - grid = np.frombuffer(bytearray(mm_kwargs["image_grid_thw"].data), dtype=mm_kwargs["image_grid_thw"].dtype).reshape( - mm_kwargs["image_grid_thw"].shape - ) - patches_per_image = [int(g.prod()) for g in grid] - total_patches = mm_kwargs["pixel_values"].shape[0] - total_tokens = sum(1 for t in mm_token_type_ids if t) - ppt = total_patches // total_tokens if total_tokens else 1 # patches per token (merge^2) - tokens_per_image = [p // ppt for p in patches_per_image] - - surviving = sum(1 for t in mm_token_type_ids[:seq_len] if t) - kept = acc = 0 - for n in tokens_per_image: - if acc + n > surviving: - break - acc += n - kept += 1 - if acc == surviving: - cut = seq_len # surviving image tokens are exactly `kept` whole images - else: - # `surviving` lands inside image `kept`; cut to its first placeholder, dropping it. - seen, cut = 0, seq_len - for i, t in enumerate(mm_token_type_ids): - if t: - seen += 1 - if seen == acc + 1: - cut = i - break - if not kept: - return cut, None - kept_patches = sum(patches_per_image[:kept]) - sliced = {k: _slice_encoded(v, kept if k == "image_grid_thw" else kept_patches) for k, v in mm_kwargs.items()} - return cut, sliced - - def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch: """ Prepare a problem for sequence packing training. @@ -108,7 +54,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch rewards = [reward] * len(input_ids) position_ids = list(range(len(input_ids))) mm_token_type_ids = training_example.mm_token_type_ids - mm_kwargs = training_example.mm_kwargs + if training_example.mm_kwargs is not None: + raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") + mm_refs = copy.deepcopy(training_example.mm_refs) assert training_example.env_name != "all", "env_name='all' is reserved for aggregate metric keys" env_names = [training_example.env_name] * len(input_ids) @@ -123,11 +71,12 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ) if len(input_ids) > seq_len: - # Multimodal: never split an image's placeholder block — cut to a whole-image boundary - # and slice mm_kwargs to match, so image-token count == image-embedding count. + if mm_refs is not None: + raise ValueError( + "Multimodal samples cannot be truncated after raw image offload. " + f"Got {len(input_ids)} tokens with seq_len={seq_len}." + ) cut = seq_len - if mm_token_type_ids is not None and mm_kwargs is not None: - cut, mm_kwargs = _truncate_mm(mm_token_type_ids, mm_kwargs, seq_len) input_ids = input_ids[:cut] loss_mask = loss_mask[:cut] inference_logprobs = inference_logprobs[:cut] @@ -182,14 +131,14 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch routed_experts=routed_experts, mm_token_type_ids=mm_token_type_ids, env_names=env_names, - mm_kwargs=mm_kwargs, + mm_refs=mm_refs, training_mode=training_example.training_mode, ) def _is_multimodal_sample(sample: MicroBatch) -> bool: """Check if a sample contains multimodal data (images).""" - return sample.mm_kwargs is not None + return sample.mm_refs is not None or sample.mm_kwargs is not None @dataclass @@ -307,7 +256,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: routed_experts=routed_experts, mm_token_type_ids=mm_token_type_ids, env_names=env_names, - mm_kwargs=first_sample.mm_kwargs if _is_multimodal_sample(first_sample) else None, + mm_refs=first_sample.mm_refs if _is_multimodal_sample(first_sample) else None, training_mode=first_sample.training_mode, ) @@ -341,7 +290,7 @@ def packed_samples_into_micro_bs( We follow the First Fit Decreasing algorithm to pack the samples into bins and minimize potential padding while never truncating. With per-token temperatures, samples can be packed together regardless of their temperature values. - NOTE: Multimodal samples (with mm_kwargs) are NOT packed together as they have variable-sized + NOTE: Multimodal samples are NOT packed together as they have variable-sized vision data that doesn't pack well. Each multimodal sample becomes its own micro batch. """ # Sort by (lora_idx, -length) for packing efficiency diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 2cdf8bbef8..48d529b1c7 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -27,6 +27,7 @@ from transformers.utils.import_utils import is_flash_attn_3_available from prime_rl.configs.trainer import ActivationCheckpointConfig, CompileConfig, ModelConfig, TokenizerConfig +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.distributed import DeepEPExpertParallel from prime_rl.trainer.lora import apply_lora_to_model, freeze_all_except_lora_and_specified, strip_lora_from_state_dict from prime_rl.trainer.models import ( @@ -1171,13 +1172,12 @@ def forward( labels: Int[Tensor, "batch seq"] | None = None, temperature: Tensor | None = None, routed_experts: Int[Tensor, "batch seq layers topk"] | None = None, - # Generic multimodal kwargs (e.g. {"pixel_values": ..., - # "image_grid_thw": ...} for Qwen3-VL; just {"pixel_values": ...} - # for Gemma3). Passed straight through to ``model(**kwargs)`` so - # the model's HF forward signature is the schema. ``mm_token_type_ids`` - # is split out because it's prime-rl-computed (from token ids), - # not a renderer/processor output. + # Generic multimodal kwargs materialized by the trainer's model processor. + # Passed straight through to ``model(**kwargs)`` so the model's HF forward + # signature is the schema. ``mm_token_type_ids`` is split out because it's + # prime-rl-computed from token ids. mm_kwargs: dict[str, Tensor] | None = None, + mm_forward_policy: ForwardPolicy | None = None, mm_token_type_ids: Int[Tensor, "batch seq"] | None = None, ) -> PrimeLmOutput: # Build kwargs for model forward @@ -1189,16 +1189,14 @@ def forward( if mm_kwargs: # Forward the per-model multimodal tensors verbatim, plus the - # renderer-supplied ``mm_token_type_ids`` (renderer owns the - # token→modality mapping via ``mm_token_type_id_map``). + # token→modality map derived from renderer token ids. kwargs.update(mm_kwargs) if mm_token_type_ids is not None: kwargs["mm_token_type_ids"] = mm_token_type_ids - # ``position_ids`` for MRoPE families: Qwen3-VL's HF forward - # recomputes 3D positions from ``image_grid_thw`` and breaks if - # given the trainer's pre-computed 1D ``position_ids``. Detect - # via the mm_kwargs shape so we don't enumerate model_types. - if "image_grid_thw" not in mm_kwargs: + policy = mm_forward_policy or ForwardPolicy() + 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_with_mm: kwargs["position_ids"] = position_ids else: kwargs["position_ids"] = position_ids diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index d16102e954..710659fd5d 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -227,11 +227,7 @@ def forward( position_embeddings = self.rotary_emb(hidden_states, position_ids) for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): - # .contiguous(): the per-layer slice is a strided view of [tokens, layers, topk] - # (dim-1 stride = layers*topk); torch.compile's MoE kernel asserts a contiguous input. - routed_experts_layer = ( - routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None - ) + routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None hidden_states = decoder_layer( hidden_states, position_embeddings=position_embeddings, diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 7c3189bc4d..73c8a942dd 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -1,3 +1,4 @@ +import time from collections.abc import Callable, Sequence from pathlib import Path from typing import TypedDict @@ -7,7 +8,8 @@ from torch import Tensor from transformers.tokenization_utils import PreTrainedTokenizer -from prime_rl.configs.trainer import FakeDataLoaderConfig +from prime_rl.configs.trainer import FakeDataLoaderConfig, MissingMMImagePolicy +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.rl.packer import BasePacker, setup_packer from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.trainer.world import get_world @@ -17,6 +19,8 @@ TransportConfig, setup_micro_batch_receiver, ) +from prime_rl.utils.logger import get_logger +from prime_rl.utils.mm import RawImageMaterializer, missing_file_uris class TensorMicroBatch(TypedDict): @@ -40,12 +44,9 @@ class TensorMicroBatch(TypedDict): # MoE router replay routed_experts: Int[Tensor, "batch seq layers topk"] | None - # Generic multimodal kwargs — flat dict matching the model's forward - # signature (e.g. ``{"pixel_values": ..., "image_grid_thw": ...}`` for - # Qwen3-VL; ``{"pixel_values": ...}`` for Gemma3-VL). The trainer - # ``**`` -unpacks this into the forward call, so any HF VLM whose - # processor and forward agree on kwarg names works out of the box. + # Generic multimodal kwargs materialized by the trainer's model processor. mm_kwargs: dict[str, Tensor] | None + mm_forward_policy: 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 @@ -132,6 +133,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc "lora_num_tokens": lora_num_tokens, "routed_experts": None, "mm_kwargs": None, + "mm_forward_policy": None, "mm_token_type_ids": None, "training_mode": "rl", "run_id": None, @@ -163,6 +165,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: "lora_num_tokens": lora_num_tokens, "routed_experts": None, "mm_kwargs": None, + "mm_forward_policy": None, "mm_token_type_ids": None, "training_mode": "rl", "run_id": None, @@ -183,6 +186,9 @@ def __init__( tokenizer: PreTrainedTokenizer, bin_cost: Callable[[Sequence[int]], int], config: TransportConfig, + model_name: str, + model_trust_remote_code: bool, + missing_mm_image_policy: MissingMMImagePolicy = "placeholder_zero_loss", ): self.world = get_world() @@ -202,6 +208,11 @@ def __init__( self.multi_run_manager = get_multi_run_manager() self.receiver: MicroBatchReceiver = setup_micro_batch_receiver(output_dir, dp_rank, start_step, config) + self.mm_materializer = RawImageMaterializer(model_name, trust_remote_code=model_trust_remote_code) + self.missing_mm_image_policy = missing_mm_image_policy + self.last_mm_materialize_time = 0.0 + self.last_mm_images_materialized = 0 + self.last_mm_images_placeholdered = 0 def wait_for_batch(self) -> None: if self.world.is_master: @@ -215,6 +226,9 @@ def wait_for_batch(self) -> None: def get_batch(self) -> list[TensorMicroBatch]: micro_batches = self.receiver.receive() + self.last_mm_materialize_time = 0.0 + self.last_mm_images_materialized = 0 + self.last_mm_images_placeholdered = 0 return [self._micro_batch_to_tensor(mb) for mb in micro_batches] def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: @@ -223,14 +237,60 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: micro_batch.lora_num_tokens = [0] * self.multi_run_manager.max_runs micro_batch.lora_num_tokens[0] = len(micro_batch.input_ids) mm_kwargs: dict[str, Tensor] | None = None + mm_forward_policy: ForwardPolicy | None = None if micro_batch.mm_kwargs: - # Each value is an EncodedTensor (dtype, shape, raw bytes). - # No batch dim — the orchestrator concatenates per-image along - # dim=0 generically, matching what each HF VLM's forward expects. - mm_kwargs = { - key: torch.frombuffer(bytearray(payload.data), dtype=_torch_dtype(payload.dtype)).reshape(payload.shape) - for key, payload in micro_batch.mm_kwargs.items() - } + raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") + if micro_batch.mm_refs is not None: + materialize_start = time.perf_counter() + try: + materialized = self.mm_materializer.materialize(micro_batch.mm_refs) + if materialized is not None: + mm_kwargs = materialized.kwargs + mm_forward_policy = materialized.forward_policy + self.last_mm_materialize_time += time.perf_counter() - materialize_start + self.last_mm_images_materialized += len(micro_batch.mm_refs.uris) + except FileNotFoundError as exc: + self.last_mm_materialize_time += time.perf_counter() - materialize_start + run_idx = next((i for i, n in enumerate(micro_batch.lora_num_tokens or []) if n > 0), None) + if self.missing_mm_image_policy == "error": + get_logger().error( + f"raw image materialization failed (run_idx={run_idx}, run_id={micro_batch.run_id}, " + f"run_step={micro_batch.run_step}, uris={micro_batch.mm_refs.uris}): {exc!r}" + ) + raise + + placeholder_start = time.perf_counter() + try: + materialized = self.mm_materializer.synthesize_placeholder(micro_batch.mm_refs) + if materialized is not None: + mm_kwargs = materialized.kwargs + mm_forward_policy = materialized.forward_policy + except Exception as placeholder_exc: + get_logger().error( + f"raw image placeholder synthesis failed after missing image " + f"(run_idx={run_idx}, run_id={micro_batch.run_id}, run_step={micro_batch.run_step}, " + f"uris={micro_batch.mm_refs.uris}): {placeholder_exc!r}" + ) + raise placeholder_exc from exc + + self.last_mm_materialize_time += time.perf_counter() - placeholder_start + self.last_mm_images_placeholdered += len(micro_batch.mm_refs.uris) + micro_batch.loss_mask = [False] * len(micro_batch.loss_mask) + micro_batch.advantages = [0.0] * len(micro_batch.advantages) + missing_uris = missing_file_uris(micro_batch.mm_refs.uris) + get_logger().warning( + "raw image materialization missing image(s); using zero-loss placeholder " + f"(run_idx={run_idx}, run_id={micro_batch.run_id}, run_step={micro_batch.run_step}, " + f"missing_uris={missing_uris or ['']}, " + f"uris={micro_batch.mm_refs.uris})" + ) + except Exception as exc: + run_idx = next((i for i, n in enumerate(micro_batch.lora_num_tokens or []) if n > 0), None) + get_logger().error( + f"raw image materialization failed (run_idx={run_idx}, run_id={micro_batch.run_id}, " + f"run_step={micro_batch.run_step}, uris={micro_batch.mm_refs.uris}): {exc!r}" + ) + raise routed_experts = None packed_routed_experts = micro_batch.routed_experts if packed_routed_experts is not None: @@ -260,6 +320,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: sequence_lengths=micro_batch.sequence_lengths, lora_num_tokens=torch.tensor(micro_batch.lora_num_tokens, dtype=torch.int32), mm_kwargs=mm_kwargs, + mm_forward_policy=mm_forward_policy, mm_token_type_ids=torch.tensor(micro_batch.mm_token_type_ids, dtype=torch.long).unsqueeze(0) if micro_batch.mm_token_type_ids is not None else None, diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 3e4aa34a0f..32fbcf06bb 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -239,6 +239,9 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: tokenizer, build_bin_cost(model.config), config.rollout_transport, + config.model.name, + config.model.trust_remote_code, + missing_mm_image_policy=config.missing_mm_image_policy, ) token_exporter = setup_token_exporter(config, parallel_dims, world, logger) @@ -390,12 +393,10 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: # we could've gotten routed experts from the inference server, but we didn't enable router replay routed_experts = None - # Multimodal kwargs are an opaque per-model dict (e.g. - # {"pixel_values": ..., "image_grid_thw": ...} for Qwen3-VL, - # just {"pixel_values": ...} for Gemma3-VL) — we move every - # tensor to CUDA and let the model's forward sort them. + # Multimodal kwargs are materialized by the trainer model processor. mm_kwargs_raw = micro_batch.get("mm_kwargs") mm_kwargs = {k: v.to("cuda") for k, v in mm_kwargs_raw.items()} if mm_kwargs_raw else None + mm_forward_policy = micro_batch.get("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 @@ -445,6 +446,7 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: labels=labels, temperature=temperatures, mm_kwargs=mm_kwargs, + mm_forward_policy=mm_forward_policy, mm_token_type_ids=mm_token_type_ids, routed_experts=routed_experts, ) @@ -646,9 +648,12 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: "time/step": step_time, "time/wait_for_batch": wait_for_batch_time, "time/load_data": load_data_time, + "time/mm_materialize": getattr(dataloader, "last_mm_materialize_time", 0.0), "time/broadcast_weights": broadcast_weights_time, "time/save_ckpt": save_ckpt_time, "time/forward_backward": forward_backward_time, + "mm/images_materialized": getattr(dataloader, "last_mm_images_materialized", 0), + "mm/images_placeholdered": getattr(dataloader, "last_mm_images_placeholdered", 0), "step": progress.step, } monitor.log(time_metrics, step=progress.step) diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index 12aba3b101..790275a121 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -22,6 +22,19 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru dtype: str +class MMRefs(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): + """Raw multimodal sidecar references for one sample. + + ``descriptor`` carries JSON-safe renderer metadata (hashes, grids, placeholder + layout). ``uris`` carries the raw image files that the trainer materializes + with its own processor. Processed tensors are intentionally not part of this + transport. + """ + + descriptor: dict + uris: list[str] + + # Orchestrator -> Packer class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): """A single training example — one branch of a rollout as a flat token sequence. @@ -40,16 +53,12 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr advantage: float | None = None reward: float | None = None - # Generic multimodal kwargs: flat dict keyed by the kwarg names the - # model's forward expects (e.g. {"pixel_values": ..., "image_grid_thw": - # ...} for Qwen3-VL; just {"pixel_values": ...} for Gemma3). The - # orchestrator batches per-image renderer items by torch.cat along - # dim=0 generically — no model-specific knowledge in prime-rl. The - # trainer ``**`` -unpacks this into the model forward, so any VLM - # whose HF processor / forward agree on kwarg names works without - # touching this transport. + # Legacy eager multimodal payloads are rejected by the v1 raw-image-ref + # path. Keep the field so old batches fail at a clear boundary. mm_kwargs: dict[str, EncodedTensor] | None = None + mm_refs: MMRefs | None = None + routed_experts: RoutedExperts | None = None # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) @@ -84,8 +93,10 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): lora_num_tokens: list[int] | None = None routed_experts: RoutedExperts | None = None - # See TrainingSample.mm_kwargs. + # Legacy eager multimodal payloads are rejected by the v1 raw-image-ref + # path. Keep the field so old batches fail at a clear boundary. mm_kwargs: dict[str, EncodedTensor] | None = None + mm_refs: MMRefs | None = None # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) mm_token_type_ids: list[int] | None = None diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py new file mode 100644 index 0000000000..8e495118a3 --- /dev/null +++ b/src/prime_rl/utils/mm.py @@ -0,0 +1,237 @@ +from __future__ import annotations + +import hashlib +from collections.abc import Iterable, Mapping +from io import BytesIO +from pathlib import Path +from typing import Any +from urllib.parse import unquote, urlparse + +from prime_rl.multimodal.adapters.base import MaterializedMM +from prime_rl.multimodal.registry import get_multimodal_adapter +from prime_rl.multimodal.schema import ( + RawMMItem, + contains_processed_payload_key, + parse_raw_mm_item, +) +from prime_rl.transport.types import MMRefs + +IMAGE_MODALITY = "image" +SUPPORTED_MODALITIES = {IMAGE_MODALITY} +PROCESSED_MM_KEYS = {"pixel_values", "image_embeds", "image_features"} + + +def _field(value: Any, name: str, default: Any = None) -> Any: + if isinstance(value, Mapping): + return value.get(name, default) + return getattr(value, name, default) + + +def file_uri_to_path(uri: str) -> Path: + parsed = urlparse(uri) + if parsed.scheme != "file": + raise ValueError(f"Raw multimodal image refs must be file:// URIs, got {uri!r}") + if parsed.netloc not in ("", "localhost"): + raise ValueError(f"file:// multimodal refs must be local paths, got {uri!r}") + return Path(unquote(parsed.path)) + + +def missing_file_uris(uris: Iterable[str]) -> list[str]: + """Return missing local ``file://`` image refs; non-file refs are ignored.""" + missing: list[str] = [] + for uri in uris: + parsed = urlparse(uri) + if parsed.scheme != "file": + continue + if not Path(unquote(parsed.path)).exists(): + missing.append(uri) + return missing + + +def image_file_uris_from_messages(messages: Iterable[Any]) -> list[str]: + uris: list[str] = [] + for message in messages: + content = _field(message, "content") + if not isinstance(content, list): + continue + for part in content: + if _field(part, "type") != "image_url": + continue + image_url = _field(part, "image_url") + url = image_url if isinstance(image_url, str) else _field(image_url, "url") + if isinstance(url, str): + uris.append(url) + return uris + + +def _normalize_json_value(value: Any, path: str) -> Any: + if value is None or isinstance(value, str | int | float | bool): + return value + if isinstance(value, tuple): + return [_normalize_json_value(v, f"{path}[]") for v in value] + if isinstance(value, list): + return [_normalize_json_value(v, f"{path}[]") for v in value] + if isinstance(value, Mapping): + return {str(k): _normalize_json_value(v, f"{path}.{k}") for k, v in value.items()} + raise TypeError( + f"v1 multimodal sidecars must be JSON-safe raw image descriptors; {path} has unsupported {type(value).__name__}" + ) + + +def validate_raw_mm_item(item: Mapping[str, Any]) -> dict[str, Any]: + if contains_processed_payload_key(item): + raise TypeError( + "v1 multimodal sidecars must be raw image descriptors, not processed payloads " + f"({', '.join(sorted(PROCESSED_MM_KEYS))})" + ) + normalized = {str(k): _normalize_json_value(v, str(k)) for k, v in item.items()} + parse_raw_mm_item(normalized) + return normalized + + +def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: + unsupported = sorted( + modality for modality, items in mm_items.items() if items and modality not in SUPPORTED_MODALITIES + ) + if unsupported: + raise NotImplementedError( + "v1 multimodal training currently supports raw image refs only; " + f"unsupported modalities: {', '.join(unsupported)}" + ) + + +def _placeholder_dict(placeholder: Any) -> dict[str, int]: + return { + "offset": int(_field(placeholder, "offset")), + "length": int(_field(placeholder, "length")), + } + + +def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | None: + mm_items = _field(multi_modal_data, "mm_items", None) + if not mm_items: + return None + _validate_modalities(mm_items) + + image_items = [validate_raw_mm_item(item) for item in mm_items.get(IMAGE_MODALITY, [])] + if not image_items: + return None + + mm_hashes = _field(multi_modal_data, "mm_hashes", {}) or {} + image_hashes = list(mm_hashes.get(IMAGE_MODALITY, [])) + if len(image_hashes) != len(image_items): + raise ValueError( + "Raw image descriptor/hash mismatch: " + f"{len(image_items)} image descriptors but {len(image_hashes)} image hashes" + ) + + mm_placeholders = _field(multi_modal_data, "mm_placeholders", {}) or {} + image_placeholders = [_placeholder_dict(p) for p in mm_placeholders.get(IMAGE_MODALITY, [])] + if image_placeholders and len(image_placeholders) != len(image_items): + raise ValueError( + "Raw image placeholder/descriptor mismatch: " + f"{len(image_placeholders)} placeholders but {len(image_items)} image descriptors" + ) + + uris = image_file_uris_from_messages(messages) + if len(uris) != len(image_items): + raise ValueError( + "Raw image URI/descriptor mismatch: " + f"{len(uris)} file refs in messages but {len(image_items)} image descriptors" + ) + for uri in uris: + file_uri_to_path(uri) + + return MMRefs( + descriptor={ + "mm_items": {IMAGE_MODALITY: image_items}, + "mm_hashes": {IMAGE_MODALITY: image_hashes}, + "mm_placeholders": {IMAGE_MODALITY: image_placeholders}, + }, + uris=uris, + ) + + +def sha256_32(data: bytes) -> str: + return hashlib.sha256(data).hexdigest()[:32] + + +def _single_family_adapter(items: list[RawMMItem]): + families = {item.family for item in items} + if len(families) != 1: + raise ValueError(f"Raw multimodal refs must use exactly one adapter family, got {sorted(families)}") + return get_multimodal_adapter(next(iter(families))) + + +class RawImageMaterializer: + """Materialize raw image refs with the trainer model's HF image processor.""" + + def __init__(self, model_name: str, *, trust_remote_code: bool): + self.model_name = model_name + self.trust_remote_code = trust_remote_code + self._image_processor = None + + @property + def image_processor(self): + if self._image_processor is None: + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(self.model_name, trust_remote_code=self.trust_remote_code) + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError(f"{self.model_name!r} does not expose an image_processor") + self._image_processor = image_processor + return self._image_processor + + def materialize(self, refs: MMRefs | None) -> MaterializedMM | None: + if refs is None: + return None + + image_item_dicts = refs.descriptor.get("mm_items", {}).get(IMAGE_MODALITY, []) + image_hashes = refs.descriptor.get("mm_hashes", {}).get(IMAGE_MODALITY, []) + if not image_item_dicts: + return None + if len(refs.uris) != len(image_item_dicts) or len(image_hashes) != len(image_item_dicts): + raise ValueError( + "Raw image refs must have matching URI, descriptor, and hash counts " + f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" + ) + image_items = [parse_raw_mm_item(validate_raw_mm_item(item)) for item in image_item_dicts] + adapter = _single_family_adapter(image_items) + actual_fingerprint = adapter.processor_fingerprint(self.image_processor) + for item in image_items: + if item.layout_fingerprint != actual_fingerprint: + raise ValueError( + "Raw image layout fingerprint mismatch: " + f"expected {item.layout_fingerprint}, got {actual_fingerprint}" + ) + + from PIL import Image + + images = [] + for uri, expected_hash in zip(refs.uris, image_hashes, strict=True): + raw = file_uri_to_path(uri).read_bytes() + actual_hash = sha256_32(raw) + if actual_hash != expected_hash: + raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") + images.append(Image.open(BytesIO(raw)).convert("RGB")) + + return adapter.materialize_for_trainer(self.image_processor, image_items, images) + + def synthesize_placeholder(self, refs: MMRefs | None) -> MaterializedMM | None: + """Build zero-valued multimodal tensors via the owning adapter.""" + if refs is None: + return None + + image_item_dicts = refs.descriptor.get("mm_items", {}).get(IMAGE_MODALITY, []) + image_hashes = refs.descriptor.get("mm_hashes", {}).get(IMAGE_MODALITY, []) + if not image_item_dicts: + return None + if len(refs.uris) != len(image_item_dicts) or len(image_hashes) != len(image_item_dicts): + raise ValueError( + "Raw image refs must have matching URI, descriptor, and hash counts " + f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" + ) + image_items = [parse_raw_mm_item(validate_raw_mm_item(item)) for item in image_item_dicts] + adapter = _single_family_adapter(image_items) + return adapter.synthesize_placeholder(self.image_processor, image_items) diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py new file mode 100644 index 0000000000..ce9539581e --- /dev/null +++ b/src/prime_rl/utils/run_assets.py @@ -0,0 +1,24 @@ +from __future__ import annotations + +import os +from collections.abc import Mapping +from pathlib import Path + +IMAGE_OFFLOAD_DIR_ENV = "VF_RENDERER_IMAGE_OFFLOAD_DIR" +RUN_DIR_ENV = "PRIME_RL_RUN_DIR" +RUN_ID_ENV = "RUN_ID" + + +def run_asset_env(output_dir: Path, base: Mapping[str, str] | None = None) -> dict[str, str]: + """Resolve the environment used by subprocesses that share run image assets. + + Hosted runs resolve ``RUN_ID`` to ``/data/outputs/run_${RUN_ID}`` inside + renderers. Local launches without hosted env vars use the RL output dir as + the explicit run dir. + """ + + env = dict(os.environ if base is None else base) + if env.get(IMAGE_OFFLOAD_DIR_ENV) or env.get(RUN_DIR_ENV) or env.get(RUN_ID_ENV): + return env + env[RUN_DIR_ENV] = str(output_dir.resolve()) + return env diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 2929ff20d7..cc569f33c7 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -11,6 +11,7 @@ from __future__ import annotations import asyncio +import hashlib import numpy as np import pybase64 @@ -267,3 +268,64 @@ def test_client_set_max_tokens_assumes_set_when_body_unreadable(): # non-dict body → can't tell, don't override. assert asyncio.run(_client_set_max_tokens(_FakeRawRequest([1, 2, 3]))) is True + + +def test_materialize_raw_image_ref_uses_generic_family_payload(tmp_path, monkeypatch): + from PIL import Image + from renderers.mm_store import raw_mm_ref + + from prime_rl.inference.vllm import serving_tokens + + image_dir = tmp_path / "run_serving" / "assets" / "images" + image_dir.mkdir(parents=True) + image_path = image_dir / "image.png" + Image.new("RGB", (8, 6), color=(32, 64, 128)).save(image_path) + monkeypatch.setenv("VF_RENDERER_IMAGE_OFFLOAD_DIR", str(image_dir)) + + mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] + fingerprint = "f" * 32 + raw_ref = raw_mm_ref( + run_id="serving", + family="test_family", + fingerprint=fingerprint, + modality="image", + mm_hash=mm_hash, + raw_image_id=image_path.name, + payload={"adapter_owned": [1, 2, 3]}, + ) + processor = object() + captured = {} + + class _Adapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + captured["image_processor"] = image_processor + captured["item"] = item + captured["image_size"] = image.size + captured["expected_placeholder_length"] = expected_placeholder_length + return {"materialized": True} + + def _get_adapter(family): + captured["family"] = family + return _Adapter() + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: processor) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", _get_adapter) + + out = serving_tokens._materialize_raw_image_ref_sync( + raw_ref, + expected_modality="image", + expected_hash=mm_hash, + expected_placeholder_length=7, + processor_model_name="model", + trust_remote_code=True, + ) + + assert out == {"materialized": True} + assert captured["family"] == "test_family" + assert captured["image_processor"] is processor + assert captured["image_size"] == (8, 6) + assert captured["expected_placeholder_length"] == 7 + item = captured["item"] + assert item.family == "test_family" + assert item.layout_fingerprint == fingerprint + assert item.payload == {"adapter_owned": [1, 2, 3]} diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 58bf154a79..14c7e8fa6e 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -5,7 +5,7 @@ from prime_rl.trainer.batch import pad_micro_batch, prepare_batch, prepare_sample from prime_rl.trainer.utils import build_bin_cost -from prime_rl.transport.types import EncodedTensor, MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import MicroBatch, MMRefs, RoutedExperts, TrainingSample def _routed_experts(data, dtype=np.uint8): @@ -331,41 +331,73 @@ def test_prepare_sample_truncates_routed_experts(): assert micro_batch.env_names == ["test-env"] * 3 -def _encoded(arr) -> EncodedTensor: - a = np.asarray(arr) - return EncodedTensor(data=a.tobytes(), shape=list(a.shape), dtype=str(a.dtype)) +def _mm_refs() -> MMRefs: + return MMRefs( + descriptor={ + "mm_items": { + "image": [ + { + "kind": "prime_raw_mm_item", + "version": 1, + "modality": "image", + "family": "qwen_vl", + "layout_fingerprint": "f" * 32, + "payload": {"image_grid_thw": [[1, 1, 1]]}, + } + ] + }, + "mm_hashes": {"image": ["a" * 32]}, + }, + uris=["file:///tmp/image.png"], + ) + + +def test_prepare_sample_preserves_raw_mm_refs(): + sample = TrainingSample( + token_ids=[10, 11, 12], + mask=[False, False, True], + logprobs=[0.0] * 3, + temperatures=[1.0] * 3, + advantage=1.0, + env_name="test-env", + mm_token_type_ids=[0, 1, 0], + mm_refs=_mm_refs(), + ) + + mb = prepare_sample(sample, seq_len=8) + assert mb.mm_refs == sample.mm_refs + assert mb.mm_token_type_ids == [0, 1, 0] + + +def test_prepare_sample_rejects_overlong_raw_mm_refs(): + sample = TrainingSample( + token_ids=[10, 11, 12, 13], + mask=[False, False, True, True], + logprobs=[0.0] * 4, + temperatures=[1.0] * 4, + advantage=1.0, + env_name="test-env", + mm_token_type_ids=[0, 1, 1, 0], + mm_refs=_mm_refs(), + ) + + with pytest.raises(ValueError, match="Multimodal samples cannot be truncated"): + prepare_sample(sample, seq_len=3) -def test_prepare_sample_truncates_mm_at_image_boundary(): - """Truncation never splits an image's placeholder block: it cuts to a whole-image boundary - and slices mm_kwargs to match, so image-token count stays == image-embedding count.""" - # Two 2-token images (patches-per-token = 1): image-pad at indices 1,2 (img0) and 4,5 (img1). - mm_token_type_ids = [0, 1, 1, 0, 1, 1, 0] - pixel_values = np.array([[1.0], [1.0], [2.0], [2.0]], dtype=np.float32) # img0=1.0, img1=2.0 - grid = np.array([[1, 2, 1], [1, 2, 1]], dtype=np.int64) +def test_prepare_sample_rejects_processed_mm_kwargs(): sample = TrainingSample( - token_ids=[10, 11, 12, 13, 14, 15, 16], - mask=[False, False, False, False, False, True, True], - logprobs=[0.0] * 7, - temperatures=[1.0] * 7, + token_ids=[10, 11, 12], + mask=[False, False, True], + logprobs=[0.0] * 3, + temperatures=[1.0] * 3, advantage=1.0, env_name="test-env", - mm_token_type_ids=mm_token_type_ids, - mm_kwargs={"pixel_values": _encoded(pixel_values), "image_grid_thw": _encoded(grid)}, + mm_kwargs={}, ) - # seq_len=5 falls inside img1 (one of its two placeholders survives) -> drop img1 entirely. - mb = prepare_sample(sample, seq_len=5) - assert len(mb.input_ids) == 4 # cut back to img1's first placeholder (index 4) - assert len(mb.mm_token_type_ids) == len(mb.input_ids) - n_placeholders = sum(1 for t in mb.mm_token_type_ids if t) - assert n_placeholders == 2 # only img0's two placeholders remain - # No mismatch: placeholders == image embeddings, and only img0's pixels are kept. - assert mb.mm_kwargs["pixel_values"].shape == [2, 1] - assert mb.mm_kwargs["image_grid_thw"].shape == [1, 3] - kept = np.frombuffer(bytearray(mb.mm_kwargs["pixel_values"].data), dtype=np.float32) - assert kept.tolist() == [1.0, 1.0] - assert n_placeholders == mb.mm_kwargs["pixel_values"].shape[0] # ppt == 1 here + with pytest.raises(ValueError, match="Processed multimodal mm_kwargs are unsupported"): + prepare_sample(sample, seq_len=8) def test_prepare_sample_none_routed_experts(): diff --git a/tests/unit/orchestrator/test_qwen3_vl_e2e.py b/tests/unit/orchestrator/test_qwen3_vl_e2e.py index ffbdc45457..0369bfa35d 100644 --- a/tests/unit/orchestrator/test_qwen3_vl_e2e.py +++ b/tests/unit/orchestrator/test_qwen3_vl_e2e.py @@ -76,19 +76,19 @@ async def post(self, path, *, cast_to=dict, body=None, options=None): return httpx.Response(200, content=json.dumps(payload).encode()) -def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm(): +def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm(tmp_path, monkeypatch): """Walk a Qwen3-VL multimodal turn through the renderer client and verify the resulting ``/inference/v1/generate`` body has a valid ``features`` payload that: 1. parses through vLLM's ``GenerateRequest`` pydantic model, - 2. decodes back to ``MultiModalKwargsItem`` instances carrying - ``pixel_values`` + ``image_grid_thw`` of the right shapes, + 2. carries a raw image ref instead of processed image tensors, 3. has placeholder ranges that exactly cover the ``<|image_pad|>`` runs in the prompt token sequence. """ from PIL import Image from renderers.base import load_tokenizer + from renderers.mm_store import IMAGE_REF_PREFIX, split_raw_mm_ref from renderers.qwen3_vl import Qwen3VLRenderer from transformers import AutoProcessor from verifiers.clients.renderer_client import RendererClient @@ -96,7 +96,6 @@ def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm() ClientConfig, UserMessage, ) - from vllm.entrypoints.serve.disagg.mm_serde import decode_mm_kwargs_item from vllm.entrypoints.serve.disagg.protocol import GenerateRequest # ── Build a real Qwen3VLRenderer with a real processor. ───────────── @@ -116,15 +115,19 @@ def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm() rc.logger = MagicMock() # ── Build a verifiers-shaped user message with an image. ──────────── - img = Image.new("RGB", (224, 224), color=(64, 128, 255)) + run_dir = tmp_path / "run_e2e" + image_dir = run_dir / "assets" / "images" + image_dir.mkdir(parents=True) + image_path = image_dir / "sample.png" + Image.new("RGB", (224, 224), color=(64, 128, 255)).save(image_path) + monkeypatch.setenv("PRIME_RL_RUN_DIR", str(run_dir)) + # The renderer accepts the OpenAI ``image_url`` content-part shape — # the same shape verifiers' UserMessage carries through. user = UserMessage( content=[ {"type": "text", "text": "What's in this picture?"}, - # Embed the PIL image directly. The verifiers→renderer message - # converter forwards content unchanged for our purposes. - {"type": "image", "image": img}, + {"type": "image_url", "image_url": {"url": image_path.as_uri()}}, ] ) @@ -169,19 +172,20 @@ def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm() f"placeholder span ({ph.offset}, {ph.length}) does not cover image_pad tokens; slice={pad_slice[:8]}..." ) - # ── kwargs_data decodes to MultiModalKwargsItem with the right keys. ─ + # ── kwargs_data carries a raw image ref for vLLM to materialize. ───── assert gen_req.features.kwargs_data is not None - encoded_items = gen_req.features.kwargs_data["image"] - assert len(encoded_items) == 1 - item = decode_mm_kwargs_item(encoded_items[0]) - assert set(item.keys()) == {"pixel_values", "image_grid_thw"} - - # The image_grid_thw must match what the HF processor would have - # produced for the same PIL image — strongest signal that the engine - # sees the same image features the trainer will. - direct_proc_out = processor.image_processor(images=[img], return_tensors="pt") - expected_grid = direct_proc_out["image_grid_thw"][0].tolist() - assert item["image_grid_thw"].data.tolist() == expected_grid + ref_items = gen_req.features.kwargs_data["image"] + assert len(ref_items) == 1 + ref_item = ref_items[0] + assert isinstance(ref_item, str) + assert ref_item.startswith(f"{IMAGE_REF_PREFIX}:") + ref = split_raw_mm_ref(ref_item) + assert ref.run_id == "e2e" + assert ref.modality == "image" + assert ref.raw_image_id == image_path.name + assert ref.mm_hash == gen_req.features.mm_hashes["image"][0] + persisted = response["multi_modal_data"].mm_items["image"][0] + assert ref.payload["image_grid_thw"] == persisted["payload"]["image_grid_thw"] # ── Response parsed through renderer's parse_response. ────────────── assert response["completion_ids"] == [50, 60, 151645] diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index 51c458fced..d9d7a21d1e 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -172,6 +172,11 @@ def test_trainer_enable_token_export_cli_flag(): assert cli(TrainerConfig, args=["--enable-token-export"]).enable_token_export +def test_trainer_missing_mm_image_policy_cli_override(): + assert cli(TrainerConfig, args=[]).missing_mm_image_policy == "placeholder_zero_loss" + assert cli(TrainerConfig, args=["--missing-mm-image-policy", "error"]).missing_mm_image_policy == "error" + + def test_single_node_auto_inference_client_dp_rank_count_matches_local_dp(): config = RLConfig.model_validate( { diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py index 7baf7760a8..22d492003c 100644 --- a/tests/unit/train/test_model_forward.py +++ b/tests/unit/train/test_model_forward.py @@ -3,6 +3,7 @@ import torch import torch.nn as nn +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.model import forward @@ -34,11 +35,11 @@ def test_forward_passes_renderer_mm_token_type_ids_through(): input_ids, position_ids, mm_kwargs={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), mm_token_type_ids=mm_token_type_ids, ) assert model.kwargs is not None - # MRoPE families (image_grid_thw present) get position_ids stripped. 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) @@ -58,6 +59,7 @@ def test_forward_omits_mm_token_type_ids_when_renderer_does_not_supply(): input_ids, position_ids, mm_kwargs={"pixel_values": torch.ones(2, 3), "image_grid_thw": torch.tensor([[1, 1, 2]])}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), ) assert model.kwargs is not None @@ -66,8 +68,7 @@ def test_forward_omits_mm_token_type_ids_when_renderer_does_not_supply(): 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``.""" + """Families whose adapter asks for position_ids keep the trainer's values.""" model = _CaptureModel(SimpleNamespace(model_type="gemma3")) input_ids = torch.tensor([[1, 10, 10, 2]]) position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) @@ -77,7 +78,27 @@ def test_forward_keeps_position_ids_for_non_mrope_vlm(): input_ids, position_ids, mm_kwargs={"pixel_values": torch.ones(2, 3)}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=True), ) assert model.kwargs is not None torch.testing.assert_close(model.kwargs["position_ids"], position_ids) + + +def test_forward_policy_can_require_mm_token_type_ids(): + 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) + + try: + forward( + model, + input_ids, + position_ids, + mm_kwargs={"pixel_values": torch.ones(2, 3)}, + mm_forward_policy=ForwardPolicy(requires_mm_token_type_ids=True), + ) + except ValueError as exc: + assert "mm_token_type_ids" in str(exc) + else: + raise AssertionError("forward should require mm_token_type_ids when policy says so") diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py new file mode 100644 index 0000000000..a50307904e --- /dev/null +++ b/tests/unit/utils/test_mm.py @@ -0,0 +1,146 @@ +from types import SimpleNamespace + +import pytest +import torch + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND, RAW_MM_ITEM_VERSION +from prime_rl.trainer.rl.data import DataLoader +from prime_rl.transport.types import MicroBatch, MMRefs +from prime_rl.utils.mm import RawImageMaterializer, build_mm_refs, missing_file_uris + + +class _ImageProcessor: + patch_size = 2 + temporal_patch_size = 2 + image_mean = [0.5, 0.5, 0.5] + + +class _MissingMaterializer: + def materialize(self, refs): + raise FileNotFoundError("missing image") + + def synthesize_placeholder(self, refs): + return MaterializedMM( + kwargs={ + "pixel_values": torch.zeros((1, 24), dtype=torch.float32), + "image_grid_thw": torch.tensor([[1, 1, 1]], dtype=torch.long), + }, + forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), + ) + + +def _qwen_item(grid): + return { + "kind": RAW_MM_ITEM_KIND, + "version": RAW_MM_ITEM_VERSION, + "modality": "image", + "family": "qwen_vl", + "layout_fingerprint": "f" * 32, + "payload": {"image_grid_thw": grid}, + } + + +def _refs(uri: str = "file:///tmp/missing-image.png") -> MMRefs: + return MMRefs( + descriptor={ + "mm_items": {"image": [_qwen_item([[1, 1, 1]])]}, + "mm_hashes": {"image": ["a" * 32]}, + }, + uris=[uri], + ) + + +def _loader(policy: str = "placeholder_zero_loss") -> DataLoader: + loader = object.__new__(DataLoader) + loader.multi_run_manager = SimpleNamespace(max_runs=1) + loader.mm_materializer = _MissingMaterializer() + loader.missing_mm_image_policy = policy + loader.last_mm_materialize_time = 0.0 + loader.last_mm_images_materialized = 0 + loader.last_mm_images_placeholdered = 0 + return loader + + +def _micro_batch() -> MicroBatch: + return MicroBatch( + input_ids=[10, 11, 12], + loss_mask=[False, True, True], + advantages=[1.5, 1.5, 1.5], + inference_logprobs=[0.0, -0.1, -0.2], + position_ids=[0, 1, 2], + sequence_lengths=[3], + temperatures=[1.0, 1.0, 1.0], + env_names=["env", "env", "env"], + lora_num_tokens=[3], + mm_refs=_refs(), + mm_token_type_ids=[0, 1, 0], + ) + + +def test_missing_file_uris_reports_missing_local_refs(tmp_path): + existing = tmp_path / "image.png" + existing.write_bytes(b"image") + + assert missing_file_uris( + [existing.as_uri(), (tmp_path / "missing.png").as_uri(), "https://example.test/i.png"] + ) == [(tmp_path / "missing.png").as_uri()] + + +def test_build_mm_refs_rejects_legacy_descriptor_without_raw_envelope(tmp_path): + image_path = tmp_path / "image.png" + image_path.write_bytes(b"image") + multi_modal_data = SimpleNamespace( + mm_items={"image": [{"image_grid_thw": [[1, 1, 1]]}]}, + mm_hashes={"image": ["a" * 32]}, + mm_placeholders={"image": []}, + ) + messages = [ + { + "content": [ + { + "type": "image_url", + "image_url": {"url": image_path.as_uri()}, + } + ] + } + ] + + with pytest.raises(ValueError, match="common envelope"): + build_mm_refs(multi_modal_data, messages) + + +def test_raw_image_materializer_synthesizes_qwen_placeholder_from_descriptor(): + materializer = RawImageMaterializer("unused", trust_remote_code=False) + materializer._image_processor = _ImageProcessor() + + mm_kwargs = materializer.synthesize_placeholder( + MMRefs( + descriptor={ + "mm_items": {"image": [_qwen_item([[1, 2, 3]])]}, + "mm_hashes": {"image": ["b" * 32]}, + }, + uris=["file:///tmp/missing-image.png"], + ) + ) + + assert mm_kwargs is not None + assert mm_kwargs.kwargs["pixel_values"].shape == (6, 24) + assert not bool(mm_kwargs.kwargs["pixel_values"].any()) + assert mm_kwargs.kwargs["image_grid_thw"].tolist() == [[1, 2, 3]] + + +def test_dataloader_uses_zero_loss_placeholder_for_missing_raw_image(): + tensor_batch = _loader()._micro_batch_to_tensor(_micro_batch()) + + assert tensor_batch["mm_kwargs"] is not None + assert tensor_batch["mm_kwargs"]["pixel_values"].shape == (1, 24) + assert tensor_batch["mm_forward_policy"] == ForwardPolicy(pass_position_ids_with_mm=False) + assert tensor_batch["loss_mask"].tolist() == [[False, False, False]] + assert tensor_batch["advantages"].tolist() == [[0.0, 0.0, 0.0]] + assert tensor_batch["mm_token_type_ids"].tolist() == [[0, 1, 0]] + + +def test_dataloader_can_fail_fast_on_missing_raw_image(): + with pytest.raises(FileNotFoundError): + _loader(policy="error")._micro_batch_to_tensor(_micro_batch()) From b7903f6c77afc7049ee87d25a889ff0414fcbe04 Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 25 Jun 2026 06:46:40 +0000 Subject: [PATCH 02/52] Update raw image renderer pin --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index eaa07bb86b..a8f43867bd 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit eaa07bb86b66eddbcbd4924d746725fbdc8ef658 +Subproject commit a8f43867bd02078f876c9c777bc7a6803814ad71 From 6c0157af6fc7b5d7c957b4c2e0816a5e8272149a Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sat, 27 Jun 2026 00:18:43 +0000 Subject: [PATCH 03/52] Simplify v1 raw multimodal serving: always materialize refs /inference/v1/generate materializes every raw ref (no cache-only None branch); an unresolved ref is a hard error, not a silent None. Removes the cache-miss 409 path + helpers (_cache_only_mm_hashes/_is_missing_mm_cache_error/_missing_mm_cache_message). Bumps renderers/verifiers submodules to the matching cleanup commits. Deployment-agnostic; mm_hash encoder cache still skips re-encode. Co-Authored-By: Claude Opus 4.8 (1M context) --- deps/renderers | 2 +- deps/verifiers | 2 +- src/prime_rl/inference/vllm/serving_tokens.py | 50 ++++--------------- 3 files changed, 13 insertions(+), 41 deletions(-) diff --git a/deps/renderers b/deps/renderers index a8f43867bd..b5167c938c 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit a8f43867bd02078f876c9c777bc7a6803814ad71 +Subproject commit b5167c938ce2d9f618f925fcf6d72f0979b0d56f diff --git a/deps/verifiers b/deps/verifiers index de37650b9b..e6b13dc2b6 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit de37650b9b49ffd5541bd97cf4b0cd426c96c160 +Subproject commit e6b13dc2b64899c5178b6c2677d118d588b36b9f diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index a3b35df9e4..2be310f583 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -16,10 +16,10 @@ decisions, surface them as base64 raw-byte payloads without requiring a vLLM source fork. -3. Raw image refs for multimodal rollouts — renderers send lightweight raw - descriptor refs for new images and ``None`` for cache-only prior images. - This handler materializes refs through multimodal adapters and turns cache - misses into a structured retryable error. +3. Raw image refs for multimodal rollouts — renderers send a lightweight raw + descriptor ref at every image slot (current and prior turns alike). This + handler materializes every ref through multimodal adapters; there is no + cache-only ``None`` path, so an unresolved ref is a hard error, not a retry. 4. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies this itself (via ``GenerateRequest.is_sampling_param_provided`` + @@ -167,31 +167,10 @@ async def _client_set_max_tokens(raw_request: Request | None) -> bool: return isinstance(sp, dict) and "max_tokens" in sp -def _is_missing_mm_cache_error(exc: BaseException) -> bool: - return "Expected a cached item for mm_hash=" in str(exc) -def _cache_only_mm_hashes(features: Any) -> list[str]: - missing: list[str] = [] - kwargs_data = features.kwargs_data or {} - for modality, hashes in features.mm_hashes.items(): - items = kwargs_data.get(modality) - if items is None: - missing.extend(f"{modality}:{mm_hash}" for mm_hash in hashes) - continue - for idx, mm_hash in enumerate(hashes): - if idx >= len(items) or items[idx] is None: - missing.append(f"{modality}:{mm_hash}") - return missing -def _missing_mm_cache_message(features: Any, exc: BaseException) -> str: - hashes = _cache_only_mm_hashes(features) - if not hashes: - return str(exc) - joined = ", ".join(hashes[:8]) - suffix = "" if len(hashes) <= 8 else f", ... (+{len(hashes) - 8} more)" - return f"vLLM multimodal cache miss for {joined}{suffix}" @lru_cache(maxsize=8) @@ -263,8 +242,10 @@ async def _decode_raw_mm_kwargs( placeholders = features.mm_placeholders.get(modality, []) items = kwargs_data.get(modality) if items is None: - mm_kwargs[modality] = [None] * len(hashes) - continue + raise _MMImageRefError( + "v1 raw multimodal: modality %r arrived with no ref payload; every image must " + "carry a raw ref (cache-only None entries are no longer supported)" % modality + ) if len(items) != len(hashes): raise _MMImageRefError( f"Multimodal kwargs/hash length mismatch for {modality}: {len(items)} != {len(hashes)}" @@ -477,18 +458,9 @@ async def serve_tokens_full_generator( # type: ignore[override] final_capture = _FinalOutputCapture(result_generator) result_generator = final_capture - try: - response = await super().serve_tokens_full_generator( - request, result_generator, request_id, model_name, request_metadata - ) - except AssertionError as exc: - if request.features is not None and _is_missing_mm_cache_error(exc): - return self.create_error_response( - message=_missing_mm_cache_message(request.features, exc), - err_type="missing_mm_cache_item", - status_code=HTTPStatus.CONFLICT, - ) - raise + response = await super().serve_tokens_full_generator( + request, result_generator, request_id, model_name, request_metadata + ) if not isinstance(response, GenerateResponse): return response From 5ec8148e3c6bac1abade490fca1b05f6da0f74cf Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sat, 27 Jun 2026 01:47:40 +0000 Subject: [PATCH 04/52] serving: use canonical split_raw_mm_ref + bump renderers pin renderers.mm_store dropped the split_mmraw_ref backcompat alias; use split_raw_mm_ref. Bump renderers submodule pin to the cleanup commit. Co-Authored-By: Claude Opus 4.8 (1M context) --- deps/renderers | 2 +- src/prime_rl/inference/vllm/serving_tokens.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/deps/renderers b/deps/renderers index b5167c938c..b404f80c6f 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit b5167c938ce2d9f618f925fcf6d72f0979b0d56f +Subproject commit b404f80c6f2939f8bf4aa1209e8e6d93241be1c4 diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 2be310f583..290df7a1d5 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -194,10 +194,10 @@ def _materialize_raw_image_ref_sync( trust_remote_code: bool, ): from PIL import Image - from renderers.mm_store import raw_image_path, split_mmraw_ref + from renderers.mm_store import raw_image_path, split_raw_mm_ref try: - ref = split_mmraw_ref(raw_ref) + ref = split_raw_mm_ref(raw_ref) if ref.modality != expected_modality: raise ValueError(f"Expected modality {expected_modality!r}, got {ref.modality!r}") if ref.mm_hash != expected_hash: From bcc93c2c65e77a286e029773b2da829adf0461d6 Mon Sep 17 00:00:00 2001 From: S1ro1 Date: Sat, 27 Jun 2026 01:54:52 +0000 Subject: [PATCH 05/52] Bump renderers pin (drop orphaned image_cache_max) Co-Authored-By: Claude Opus 4.8 (1M context) --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index b404f80c6f..f8ca35415a 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit b404f80c6f2939f8bf4aa1209e8e6d93241be1c4 +Subproject commit f8ca35415a6e6d89a6cee6d9e9127429bb9d68b6 From 75b9e7c6c9a835cffdc2d4016d58713ce1fc648f Mon Sep 17 00:00:00 2001 From: eligotts Date: Sun, 28 Jun 2026 04:30:22 +0000 Subject: [PATCH 06/52] Bump raw image dependency pins --- deps/renderers | 2 +- deps/verifiers | 2 +- uv.lock | 14 -------------- 3 files changed, 2 insertions(+), 16 deletions(-) diff --git a/deps/renderers b/deps/renderers index f8ca35415a..8fcd0c7675 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit f8ca35415a6e6d89a6cee6d9e9127429bb9d68b6 +Subproject commit 8fcd0c76755013d2d3aba6d0d60129f8d4cc04d3 diff --git a/deps/verifiers b/deps/verifiers index e6b13dc2b6..9430999d0d 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit e6b13dc2b64899c5178b6c2677d118d588b36b9f +Subproject commit 9430999d0d07235e7fe3565f6c064edde932ed0d diff --git a/uv.lock b/uv.lock index cd859007f3..b96af7e2cb 100644 --- a/uv.lock +++ b/uv.lock @@ -1112,18 +1112,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/29/11/867a59ae2d233ce81da373513cabaca2abf03f0180b28e85cf3ae2d5f7f6/fastcore-1.13.6-py3-none-any.whl", hash = "sha256:59264df5bf003fcf44602b7418c6178b37d3c31ae87d8ac106139f716e73711f", size = 105794, upload-time = "2026-06-18T06:49:06.499Z" }, ] -[[package]] -name = "fastokens" -version = "0.2.0" -source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/14/8e/7e88ec1d48db5a6e8d8d44318ce285e38c04b81508bdc2a60e17045a116f/fastokens-0.2.0.tar.gz", hash = "sha256:ef0e175de5c8cb1b616b3210d75dce1fab78e35fc02f77f03f7847d4678be686", size = 675822, upload-time = "2026-05-17T10:32:55.642Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/b4/54/e0e4318ee1ad0b5196df72cf93615bba0b81f7869d659a44ccc475969151/fastokens-0.2.0-cp39-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:160253f8d30747cf66e7ed895c513e16f7b173dd9e644fa641e2eecbd43a616a", size = 3303534, upload-time = "2026-05-17T10:32:37.462Z" }, - { url = "https://files.pythonhosted.org/packages/64/44/bfff90e4b1a43c17edf7305dafbd56dc992bbe832cc08da78f1f50104c2d/fastokens-0.2.0-cp39-abi3-manylinux_2_28_aarch64.whl", hash = "sha256:b61b9fe5b41e0bb36ad86e7551dc53293c9833909ef07b1cdbaa2055b06c3b3e", size = 3254096, upload-time = "2026-05-17T10:32:28.489Z" }, - { url = "https://files.pythonhosted.org/packages/05/bf/1cad7f0e8d03f5f5b2b417cda8859e4d968d2eebdca0cd336b23d7dbbdbb/fastokens-0.2.0-cp39-abi3-musllinux_1_2_aarch64.whl", hash = "sha256:01b9bdba818d7b2c67d57d9917faf7a1dad32ece0734440130de94ad768b819f", size = 3336689, upload-time = "2026-05-17T10:32:46.21Z" }, - { url = "https://files.pythonhosted.org/packages/97/d7/f5fb2564e16b1f5733e05c41b090f95a3fe767f6b888ba7d864193bc5447/fastokens-0.2.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:d068bc50082ad67d5d542847075f1f7b8d10f703274e56e241312f18b4d9e772", size = 3598064, upload-time = "2026-05-17T10:32:54.109Z" }, -] - [[package]] name = "fastsafetensors" version = "0.3.2" @@ -4359,7 +4347,6 @@ wheels = [ name = "renderers" source = { editable = "deps/renderers" } dependencies = [ - { name = "fastokens", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "jinja2", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "numpy", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, { name = "openai", marker = "(platform_machine == 'aarch64' and sys_platform == 'linux') or (platform_machine == 'x86_64' and sys_platform == 'linux')" }, @@ -4371,7 +4358,6 @@ dependencies = [ [package.metadata] requires-dist = [ - { name = "fastokens", specifier = ">=0.2.0" }, { name = "jinja2" }, { name = "numpy" }, { name = "openai", specifier = ">=1.108.1" }, From 99b590a31d16f560ad1fef73b15cfefcb8c8f664 Mon Sep 17 00:00:00 2001 From: eligotts Date: Sun, 28 Jun 2026 07:03:46 +0000 Subject: [PATCH 07/52] feat: support inline multimodal image storage --- deps/renderers | 2 +- deps/verifiers | 2 +- .../src/prime_rl/configs/inference.py | 5 +- .../src/prime_rl/configs/orchestrator.py | 4 + .../src/prime_rl/configs/rl.py | 4 + .../src/prime_rl/configs/shared.py | 14 ++ .../src/prime_rl/configs/trainer.py | 4 + .../src/prime_rl/utils/validation.py | 1 + src/prime_rl/entrypoints/orchestrator.py | 2 + src/prime_rl/entrypoints/rl.py | 2 +- src/prime_rl/entrypoints/trainer.py | 2 + src/prime_rl/inference/server.py | 3 + src/prime_rl/inference/vllm/serving_tokens.py | 168 +++++++++++------- src/prime_rl/multimodal/adapters/kimi_k25.py | 91 ++++++---- src/prime_rl/multimodal/schema.py | 74 ++++---- src/prime_rl/orchestrator/orchestrator.py | 7 +- .../models/glm4_moe/modeling_glm4_moe.py | 6 +- src/prime_rl/trainer/rl/data.py | 111 ++++++------ src/prime_rl/trainer/rl/train.py | 3 + src/prime_rl/utils/mm.py | 138 ++++++++------ src/prime_rl/utils/run_assets.py | 69 ++++++- tests/unit/inference/test_serving_tokens.py | 95 +++++++++- tests/unit/utils/test_mm.py | 24 +++ 23 files changed, 572 insertions(+), 259 deletions(-) diff --git a/deps/renderers b/deps/renderers index 8fcd0c7675..e97c812bd1 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 8fcd0c76755013d2d3aba6d0d60129f8d4cc04d3 +Subproject commit e97c812bd16ca3631eb4f68af1ddcf1d4e5a5248 diff --git a/deps/verifiers b/deps/verifiers index 9430999d0d..4a7b37aca8 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 9430999d0d07235e7fe3565f6c064edde932ed0d +Subproject commit 4a7b37aca8035be424b4077e91bb7b94dc445f15 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index fef17c62f1..3207cab3eb 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -5,7 +5,7 @@ from pydantic import Field, model_validator from pydantic_config import BaseConfig -from prime_rl.configs.shared import BaseModelConfig, LogConfig, SlurmConfig +from prime_rl.configs.shared import BaseModelConfig, LogConfig, MultimodalConfig, SlurmConfig from prime_rl.utils.config import find_package_resource, rgetattr, rsetattr from prime_rl.utils.parsers import resolve_reasoning_parser, resolve_tool_call_parser @@ -435,6 +435,9 @@ class InferenceConfig(BaseConfig): experimental: InferenceExperimentalConfig = InferenceExperimentalConfig() + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal storage policy shared with trainer and orchestrator.""" + @model_validator(mode="after") def validate_multi_node_requires_slurm(self): if self.deployment.type in ("multi_node", "disaggregated") and self.slurm is None: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index e7e173e57f..ea842d48f7 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -16,6 +16,7 @@ FileSystemTransportConfig, HeartbeatConfig, LogConfig, + MultimodalConfig, PrimeMonitorConfig, TransportConfig, WandbWithExtrasConfig, @@ -561,6 +562,9 @@ class OrchestratorConfig(BaseConfig): experimental: OrchestratorExperimentalConfig = OrchestratorExperimentalConfig() + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal storage policy shared with trainer and inference.""" + @model_validator(mode="before") @classmethod def _env_to_train(cls, data: Any) -> Any: diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index 079af2620d..ce8c178ba6 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -16,6 +16,7 @@ OrchestratorConfig, ) from prime_rl.configs.shared import ( + MultimodalConfig, SlurmConfig, VLMConfig, ) @@ -219,6 +220,9 @@ class RLConfig(BaseConfig): weight_broadcast: SharedWeightBroadcastConfig | None = None + multimodal: MultimodalConfig = MultimodalConfig() + """Shared raw multimodal storage policy. Propagated to trainer, orchestrator, and inference.""" + bench: bool = False """Benchmark mode. Sets trainer and orchestrator to benchmark mode and, when set, suffixes the W&B project with ``-bench``.""" 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 e0e17c5586..3f614166f1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -60,6 +60,20 @@ def resolve_project_dir(self): ServerType = Literal["vllm", "openai"] +MultimodalImageStorage = Literal["offload", "inline"] + + +class MultimodalImageConfig(BaseConfig): + storage: MultimodalImageStorage = "offload" + """How raw image bytes are transported. ``offload`` writes data images to run assets; ``inline`` keeps base64 data-image URIs in the multimodal payload.""" + + offload_dir: Path | None = None + """Directory for offloaded image assets. Supports environment expansion such as ``/data/outputs/run_${RUN_ID}/assets/images``. When unset, prime-rl resolves a run-scoped default.""" + + +class MultimodalConfig(BaseConfig): + images: MultimodalImageConfig = MultimodalImageConfig() + """Raw multimodal image storage configuration.""" class VLMConfig(BaseConfig): 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 70aa51b0c4..227e08e231 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -9,6 +9,7 @@ FileSystemTransportConfig, HeartbeatConfig, MetricsServerConfig, + MultimodalConfig, TrainerLogConfig, TransportConfig, WandbConfig, @@ -569,6 +570,9 @@ class TrainerConfig(BaseConfig): missing_mm_image_policy: MissingMMImagePolicy = "placeholder_zero_loss" """Policy when raw multimodal image files disappear before trainer materialization. ``placeholder_zero_loss`` warns, synthesizes zero-valued image tensors with the original descriptor geometry, and masks out the affected microbatch loss; ``error`` preserves fail-fast behavior.""" + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal storage policy shared with orchestrator and inference.""" + enable_token_export: bool = False """Opt-in per-token JSONL export for rollout debugging. When enabled, writes token ids and aligned trainer metrics after each forward pass.""" 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 d3feef7894..709c475569 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -126,6 +126,7 @@ def propagate(shared_path: str, *targets: str) -> None: # Top-level scalars. propagate("max_steps", "trainer.max_steps", "orchestrator.max_steps") propagate("seq_len", "trainer.model.seq_len", "orchestrator.seq_len") + propagate("multimodal", "trainer.multimodal", "orchestrator.multimodal", "inference.multimodal") # [slurm] → inference: a multi-node RL run drives its inference deployment under # the same SLURM allocation, so the nested inference inherits [slurm]. This is diff --git a/src/prime_rl/entrypoints/orchestrator.py b/src/prime_rl/entrypoints/orchestrator.py index 48a8a354f8..1bdc92ff8b 100644 --- a/src/prime_rl/entrypoints/orchestrator.py +++ b/src/prime_rl/entrypoints/orchestrator.py @@ -14,11 +14,13 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title +from prime_rl.utils.run_assets import configure_run_asset_env def main(): set_proc_title("Orchestrator") config = cli(OrchestratorConfig) + configure_run_asset_env(config.output_dir, config.multimodal) from prime_rl.orchestrator.orchestrator import run_orchestrator asyncio.run(run_orchestrator(config)) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index ac876e8d17..482a4e5c61 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -118,7 +118,7 @@ def rl_local(config: RLConfig): "WANDB_SHARED_MODE": "1", "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } - shared_run_asset_env = run_asset_env(config.orchestrator.output_dir) + shared_run_asset_env = run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.model.client.is_elastic: diff --git a/src/prime_rl/entrypoints/trainer.py b/src/prime_rl/entrypoints/trainer.py index 0d89f58fba..096966934a 100644 --- a/src/prime_rl/entrypoints/trainer.py +++ b/src/prime_rl/entrypoints/trainer.py @@ -12,11 +12,13 @@ from prime_rl.configs.trainer import TrainerConfig from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title +from prime_rl.utils.run_assets import configure_run_asset_env def main(): set_proc_title("Trainer") config = cli(TrainerConfig) + configure_run_asset_env(config.output_dir, config.multimodal) from prime_rl.trainer.rl.train import train train(config) diff --git a/src/prime_rl/inference/server.py b/src/prime_rl/inference/server.py index 4429146c5f..d16c599e49 100644 --- a/src/prime_rl/inference/server.py +++ b/src/prime_rl/inference/server.py @@ -3,11 +3,14 @@ from prime_rl.configs.inference import InferenceConfig from prime_rl.utils.config import cli +from prime_rl.utils.run_assets import configure_run_asset_env def setup_vllm_env(config: InferenceConfig): """Set vLLM environment variables based on config. Must be called before importing vLLM.""" + configure_run_asset_env(config.output_dir, config.multimodal) + # spawn is more robust in vLLM nightlies and Qwen3-VL (fork can deadlock with multithreaded processes) os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index db072991c1..44b273a9f2 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -18,8 +18,8 @@ 3. Raw image refs for multimodal rollouts — renderers send a lightweight raw descriptor ref at every image slot (current and prior turns alike). This - handler materializes every ref through multimodal adapters; there is no - cache-only ``None`` path, so an unresolved ref is a hard error, not a retry. + handler materializes every ref through multimodal adapters before vLLM sees + the request. 4. Server-side ``max_tokens`` defaulting — upstream ``ServingTokens`` now applies this itself (via ``GenerateRequest.is_sampling_param_provided`` + @@ -33,6 +33,7 @@ from __future__ import annotations import asyncio +import base64 import hashlib from collections.abc import AsyncGenerator, AsyncIterable from dataclasses import dataclass @@ -178,48 +179,106 @@ def _load_image_processor(model_name: str, trust_remote_code: bool): return image_processor +def _parse_raw_image_ref(raw_ref: str, *, feature_modality: str, mm_hash: str): + from renderers.mm_store import split_raw_mm_ref + + try: + ref = split_raw_mm_ref(raw_ref) + except ValueError as exc: + raise _MMImageRefError(str(exc)) from exc + + if ref.modality != feature_modality: + raise _MMImageRefError(f"Expected feature modality {feature_modality!r}, got {ref.modality!r}") + if ref.mm_hash != mm_hash: + raise _MMImageRefError(f"Expected image hash {mm_hash}, got {ref.mm_hash}") + return ref + + +def _read_verified_raw_image(ref) -> bytes: + from renderers.mm_store import raw_image_path + + label = ref.raw_image_id or ref.raw_uri or "raw image" + if ref.raw_image_id is not None: + try: + raw = raw_image_path(run_id=ref.run_id, raw_image_id=ref.raw_image_id).read_bytes() + except OSError as exc: + raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_id!r}: {exc}") from exc + elif isinstance(ref.raw_uri, str) and ref.raw_uri.startswith("data:image/"): + marker = ";base64," + if marker not in ref.raw_uri: + raise _MMImageRefError("Inline raw image refs must use base64 data:image URIs") + raw = base64.b64decode(ref.raw_uri.split(marker, 1)[1]) + else: + raise _MMImageRefError(f"Unsupported raw image source {label!r}") + + actual_hash = hashlib.sha256(raw).hexdigest()[:32] + if actual_hash != ref.mm_hash: + raise _MMImageRefError(f"Raw image hash mismatch: expected {ref.mm_hash}, got {actual_hash}") + return raw + + +def _decode_raw_image(raw: bytes, *, raw_image_id: str): + from PIL import Image + + try: + return Image.open(BytesIO(raw)).convert("RGB") + except OSError as exc: + raise _MMImageRefError(f"Unable to decode raw image asset {raw_image_id!r}: {exc}") from exc + + def _materialize_raw_image_ref_sync( raw_ref: str, *, - expected_modality: str, - expected_hash: str, + feature_modality: str, + mm_hash: str, expected_placeholder_length: int | None, processor_model_name: str, trust_remote_code: bool, ): - from PIL import Image - from renderers.mm_store import raw_image_path, split_raw_mm_ref + ref = _parse_raw_image_ref(raw_ref, feature_modality=feature_modality, mm_hash=mm_hash) + raw = _read_verified_raw_image(ref) + image = _decode_raw_image(raw, raw_image_id=ref.raw_image_id or "inline image") + image_processor = _load_image_processor(processor_model_name, trust_remote_code) + item = RawMMItem( + modality=ref.modality, + family=ref.family, + layout_fingerprint=ref.fingerprint, + payload=dict(ref.payload), + raw_ref=raw_ref, + ) + adapter = get_multimodal_adapter(ref.family) + return adapter.materialize_for_vllm( + image_processor, + item, + image, + expected_placeholder_length, + ) - try: - ref = split_raw_mm_ref(raw_ref) - if ref.modality != expected_modality: - raise ValueError(f"Expected modality {expected_modality!r}, got {ref.modality!r}") - if ref.mm_hash != expected_hash: - raise ValueError(f"Expected image hash {expected_hash}, got {ref.mm_hash}") - - raw = raw_image_path(run_id=ref.run_id, raw_image_id=ref.raw_image_id).read_bytes() - actual_hash = hashlib.sha256(raw).hexdigest()[:32] - if actual_hash != ref.mm_hash: - raise ValueError(f"Raw image hash mismatch: expected {ref.mm_hash}, got {actual_hash}") - - image_processor = _load_image_processor(processor_model_name, trust_remote_code) - item = RawMMItem( - modality=ref.modality, - family=ref.family, - layout_fingerprint=ref.fingerprint, - payload=dict(ref.payload), - raw_ref=raw_ref, + +def _raw_ref_payloads_for_feature( + features: Any, feature_modality: str, hashes: list[str] +) -> tuple[list[str], list[Any]]: + from renderers.mm_store import is_raw_mm_ref + + kwargs_data = features.kwargs_data + if kwargs_data is None or feature_modality not in kwargs_data: + raise _MMImageRefError(f"v1 raw multimodal: modality {feature_modality!r} arrived with no raw refs") + + raw_refs = kwargs_data[feature_modality] + placeholders = features.mm_placeholders.get(feature_modality) + if placeholders is None: + raise _MMImageRefError(f"v1 raw multimodal: modality {feature_modality!r} arrived with no placeholders") + if len(raw_refs) != len(hashes): + raise _MMImageRefError( + f"Multimodal kwargs/hash length mismatch for {feature_modality}: {len(raw_refs)} != {len(hashes)}" ) - adapter = get_multimodal_adapter(ref.family) - image = Image.open(BytesIO(raw)).convert("RGB") - return adapter.materialize_for_vllm( - image_processor, - item, - image, - expected_placeholder_length, + if len(placeholders) != len(hashes): + raise _MMImageRefError( + f"Multimodal placeholder/hash length mismatch for {feature_modality}: {len(placeholders)} != {len(hashes)}" ) - except Exception as exc: - raise _MMImageRefError(str(exc)) from exc + if not all(is_raw_mm_ref(item) for item in raw_refs): + raise _MMImageRefError("v1 multimodal inference accepts raw descriptor refs only") + return raw_refs, placeholders async def _decode_raw_mm_kwargs( @@ -227,43 +286,24 @@ async def _decode_raw_mm_kwargs( *, processor_model_name: str, trust_remote_code: bool, -) -> dict[str, list[Any | None]]: - from renderers.mm_store import IMAGE_REF_PREFIX - - kwargs_data = features.kwargs_data or {} - mm_kwargs: dict[str, list[Any | None]] = {} - for modality, hashes in features.mm_hashes.items(): - placeholders = features.mm_placeholders.get(modality, []) - items = kwargs_data.get(modality) - if items is None: - raise _MMImageRefError( - "v1 raw multimodal: modality %r arrived with no ref payload; every image must " - "carry a raw ref (cache-only None entries are no longer supported)" % modality - ) - if len(items) != len(hashes): - raise _MMImageRefError( - f"Multimodal kwargs/hash length mismatch for {modality}: {len(items)} != {len(hashes)}" - ) - decoded: list[Any | None] = [] - for idx, item in enumerate(items): - if item is None: - decoded.append(None) - continue - if not isinstance(item, str) or not item.startswith(f"{IMAGE_REF_PREFIX}:"): - raise _MMImageRefError("v1 multimodal inference accepts raw descriptor refs only") - placeholder_length = placeholders[idx].length if idx < len(placeholders) else None +) -> dict[str, list[Any]]: + mm_kwargs: dict[str, list[Any]] = {} + for feature_modality, hashes in features.mm_hashes.items(): + raw_refs, placeholders = _raw_ref_payloads_for_feature(features, feature_modality, hashes) + decoded: list[Any] = [] + for idx, raw_ref in enumerate(raw_refs): decoded.append( await asyncio.to_thread( _materialize_raw_image_ref_sync, - item, - expected_modality=modality, - expected_hash=hashes[idx], - expected_placeholder_length=placeholder_length, + raw_ref, + feature_modality=feature_modality, + mm_hash=hashes[idx], + expected_placeholder_length=placeholders[idx].length, processor_model_name=processor_model_name, trust_remote_code=trust_remote_code, ) ) - mm_kwargs[modality] = decoded + mm_kwargs[feature_modality] = decoded return mm_kwargs diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index e04f4132cf..12ebe00b06 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -1,21 +1,14 @@ from __future__ import annotations import math +from collections.abc import Mapping from typing import Any +from renderers.image_layout_specs import KIMI_K25_IMAGE_LAYOUT, KimiK25ImageLayoutSpec + from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem -KIMI_K25_DEFAULTS = { - "patch_size": 14, - "merge_kernel_size": 2, - "in_patch_limit": 16384, - "patch_limit_on_one_side": 512, - "fixed_output_tokens": None, - "image_mean": [0.5, 0.5, 0.5], - "image_std": [0.5, 0.5, 0.5], -} - def _tensorize(value: Any): import torch @@ -25,20 +18,48 @@ def _tensorize(value: Any): return torch.as_tensor(value).contiguous() -def _cfg_value(image_processor: Any, name: str) -> Any: - for source in ( - image_processor, - getattr(image_processor, "media_proc_cfg", None), - getattr(image_processor, "config", None), - ): - if source is None: - continue - if isinstance(source, dict) and name in source: - return source[name] - value = getattr(source, name, None) - if value is not None: - return value - return KIMI_K25_DEFAULTS[name] +def _media_proc_cfg(image_processor: Any) -> Mapping[str, Any]: + cfg = getattr(image_processor, "media_proc_cfg", None) + if not isinstance(cfg, Mapping): + raise ValueError("Kimi image processor must expose media_proc_cfg") + return cfg + + +def _required_cfg(cfg: Mapping[str, Any], name: str) -> Any: + if name not in cfg: + raise ValueError(f"Kimi image processor media_proc_cfg is missing {name!r}") + return cfg[name] + + +def _optional_int(value: Any) -> int | None: + if value is None: + return None + return int(value) + + +def _float_triple(value: Any, *, name: str) -> tuple[float, float, float]: + if not isinstance(value, list | tuple) or len(value) != 3: + raise ValueError(f"Kimi image processor media_proc_cfg[{name!r}] must be a length-3 sequence") + return (float(value[0]), float(value[1]), float(value[2])) + + +def _processor_layout(image_processor: Any) -> KimiK25ImageLayoutSpec: + cfg = _media_proc_cfg(image_processor) + layout = KimiK25ImageLayoutSpec( + patch_size=int(_required_cfg(cfg, "patch_size")), + merge_kernel_size=int(_required_cfg(cfg, "merge_kernel_size")), + in_patch_limit=int(_required_cfg(cfg, "in_patch_limit")), + patch_limit_on_one_side=int(_required_cfg(cfg, "patch_limit_on_one_side")), + fixed_output_tokens=_optional_int(_required_cfg(cfg, "fixed_output_tokens")), + image_mean=_float_triple(_required_cfg(cfg, "image_mean"), name="image_mean"), + image_std=_float_triple(_required_cfg(cfg, "image_std"), name="image_std"), + ) + if layout != KIMI_K25_IMAGE_LAYOUT: + raise ValueError( + "Kimi image processor layout does not match renderer layout contract: " + f"expected {KIMI_K25_IMAGE_LAYOUT}, got {layout}" + ) + return layout def _grid_payload(item: RawMMItem) -> list[int]: @@ -75,16 +96,16 @@ def validate_item(self, item: RawMMItem) -> None: def processor_fingerprint(self, image_processor: Any) -> str: from renderers.mm_store import image_layout_fingerprint - fixed_output_tokens = _cfg_value(image_processor, "fixed_output_tokens") + layout = _processor_layout(image_processor) return image_layout_fingerprint( family=self.family, - patch_size=int(_cfg_value(image_processor, "patch_size")), - merge_kernel_size=int(_cfg_value(image_processor, "merge_kernel_size")), - in_patch_limit=int(_cfg_value(image_processor, "in_patch_limit")), - patch_limit_on_one_side=int(_cfg_value(image_processor, "patch_limit_on_one_side")), - fixed_output_tokens=None if fixed_output_tokens is None else int(fixed_output_tokens), - image_mean=list(_cfg_value(image_processor, "image_mean")), - image_std=list(_cfg_value(image_processor, "image_std")), + patch_size=layout.patch_size, + merge_kernel_size=layout.merge_kernel_size, + in_patch_limit=layout.in_patch_limit, + patch_limit_on_one_side=layout.patch_limit_on_one_side, + fixed_output_tokens=layout.fixed_output_tokens, + image_mean=list(layout.image_mean), + image_std=list(layout.image_std), ) def materialize_for_trainer( @@ -145,14 +166,16 @@ def synthesize_placeholder( return None import torch - patch_size = int(_cfg_value(image_processor, "patch_size")) + layout = _processor_layout(image_processor) grids: list[list[int]] = [] pixel_values: list[torch.Tensor] = [] for item in items: self.validate_item(item) grid = _grid_payload(item) grids.append(grid) - pixel_values.append(torch.zeros((math.prod(grid), 3, patch_size, patch_size), dtype=torch.float32)) + pixel_values.append( + torch.zeros((math.prod(grid), 3, layout.patch_size, layout.patch_size), dtype=torch.float32) + ) return MaterializedMM( kwargs={ "pixel_values": torch.cat(pixel_values, dim=0).contiguous(), diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index 4cc04f3eaa..a5ef15aa37 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -4,9 +4,9 @@ from dataclasses import dataclass from typing import Any -RAW_MM_ITEM_KIND = "prime_raw_mm_item" -RAW_MM_ITEM_VERSION = 1 -PROCESSED_MM_KEYS = {"pixel_values", "image_embeds", "image_features"} +from renderers.mm_store import RAW_MM_ITEM_KIND, RAW_MM_ITEM_VERSION + +PROCESSED_MM_KEYS = frozenset({"pixel_values", "image_embeds", "image_features"}) @dataclass(frozen=True) @@ -29,38 +29,50 @@ def contains_processed_payload_key(value: Any) -> bool: return False -def parse_raw_mm_item(value: Any) -> RawMMItem: - if not isinstance(value, Mapping): - raise TypeError(f"v1 multimodal sidecars must be raw descriptor dicts, got {type(value).__name__}") +def _descriptor_mapping(value: Any) -> Mapping[str, Any]: + if isinstance(value, Mapping): + return value + raise TypeError(f"v1 multimodal sidecars must be raw descriptor dicts, got {type(value).__name__}") + + +def _required_str(value: Mapping[str, Any], field: str) -> str: + item = value.get(field) + if isinstance(item, str) and item: + return item + raise ValueError(f"raw multimodal descriptor is missing {field}") + + +def _optional_str(value: Mapping[str, Any], field: str) -> str | None: + item = value.get(field) + if item is None or isinstance(item, str): + return item + raise ValueError(f"raw multimodal descriptor {field} must be a string when present") + + +def _payload(value: Mapping[str, Any]) -> dict[str, Any]: + payload = value.get("payload") + if not isinstance(payload, Mapping): + raise ValueError("raw multimodal descriptor payload must be a dict") + return {str(k): v for k, v in payload.items()} + + +def _validate_envelope(value: Mapping[str, Any]) -> None: if contains_processed_payload_key(value): raise TypeError("v1 multimodal sidecars must not carry processed multimodal payloads") if value.get("kind") != RAW_MM_ITEM_KIND: raise ValueError("raw multimodal descriptor is missing the common envelope kind") - if int(value.get("version", -1)) != RAW_MM_ITEM_VERSION: + if value.get("version") != RAW_MM_ITEM_VERSION: raise ValueError(f"unsupported raw multimodal descriptor version: {value.get('version')!r}") - modality = value.get("modality") - family = value.get("family") - layout_fingerprint = value.get("layout_fingerprint") - payload = value.get("payload") - if not isinstance(modality, str) or not modality: - raise ValueError("raw multimodal descriptor is missing modality") - if not isinstance(family, str) or not family: - raise ValueError("raw multimodal descriptor is missing family") - if not isinstance(layout_fingerprint, str) or not layout_fingerprint: - raise ValueError("raw multimodal descriptor is missing layout_fingerprint") - if not isinstance(payload, Mapping): - raise ValueError("raw multimodal descriptor payload must be a dict") - raw_ref = value.get("raw_ref") - if raw_ref is not None and not isinstance(raw_ref, str): - raise ValueError("raw multimodal descriptor raw_ref must be a string when present") - vllm_modality = value.get("vllm_modality") - if vllm_modality is not None and not isinstance(vllm_modality, str): - raise ValueError("raw multimodal descriptor vllm_modality must be a string when present") + + +def parse_raw_mm_item(value: Any) -> RawMMItem: + descriptor = _descriptor_mapping(value) + _validate_envelope(descriptor) return RawMMItem( - modality=modality, - family=family, - layout_fingerprint=layout_fingerprint, - payload={str(k): v for k, v in payload.items()}, - raw_ref=raw_ref, - vllm_modality=vllm_modality, + modality=_required_str(descriptor, "modality"), + family=_required_str(descriptor, "family"), + layout_fingerprint=_required_str(descriptor, "layout_fingerprint"), + payload=_payload(descriptor), + raw_ref=_optional_str(descriptor, "raw_ref"), + vllm_modality=_optional_str(descriptor, "vllm_modality"), ) diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index 569fe6d1d9..b3f876dc1b 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -76,6 +76,7 @@ from prime_rl.utils.logger import format_time, get_logger, setup_logger from prime_rl.utils.monitor import setup_monitor from prime_rl.utils.pathing import get_log_dir, get_rollout_dir, get_step_path +from prime_rl.utils.run_assets import configure_run_asset_env from prime_rl.utils.usage_reporter import UsageReporter from prime_rl.utils.utils import ( clean_exit, @@ -100,11 +101,6 @@ # resumed when the watcher advances ``policy.version``. TARGET_LAG = 1 -# Drop per-node training sidecars when dumping a Trace to disk (rollout jsonl / wandb -# tables): multimodal descriptors and router replay are training inputs, not part of the -# rollout record. `__all__` applies the exclude to every node in the list. -ROLLOUT_DUMP_EXCLUDE = {"nodes": {"__all__": {"multi_modal_data", "routed_experts"}}} - class Orchestrator: # Set in ``__init__`` @@ -845,6 +841,7 @@ async def run_orchestrator(config: OrchestratorConfig) -> None: """Top-level entrypoint. Wrapped in ``@clean_exit`` so wandb is flushed on exit (success or crash); keeps that out of the class. """ + configure_run_asset_env(config.output_dir, config.multimodal) await Orchestrator(config).start() diff --git a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py index 710659fd5d..d16102e954 100644 --- a/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py +++ b/src/prime_rl/trainer/models/glm4_moe/modeling_glm4_moe.py @@ -227,7 +227,11 @@ def forward( position_embeddings = self.rotary_emb(hidden_states, position_ids) for layer_idx, decoder_layer in enumerate(self.layers[: self.config.num_hidden_layers]): - routed_experts_layer = routed_experts[:, :, layer_idx, :] if routed_experts is not None else None + # .contiguous(): the per-layer slice is a strided view of [tokens, layers, topk] + # (dim-1 stride = layers*topk); torch.compile's MoE kernel asserts a contiguous input. + routed_experts_layer = ( + routed_experts[:, :, layer_idx, :].contiguous() if routed_experts is not None else None + ) hidden_states = decoder_layer( hidden_states, position_embeddings=position_embeddings, diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 4558d2b713..ed1b429526 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -9,7 +9,7 @@ from transformers.tokenization_utils import PreTrainedTokenizer from prime_rl.configs.trainer import FakeDataLoaderConfig, MissingMMImagePolicy -from prime_rl.multimodal.adapters.base import ForwardPolicy +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.trainer.rl.packer import BasePacker, setup_packer from prime_rl.trainer.runs import get_multi_run_manager from prime_rl.trainer.world import get_world @@ -19,6 +19,7 @@ TransportConfig, setup_micro_batch_receiver, ) +from prime_rl.transport.types import MMRefs from prime_rl.utils.logger import get_logger from prime_rl.utils.mm import RawImageMaterializer, missing_file_uris @@ -60,6 +61,9 @@ class TensorMicroBatch(TypedDict): run_step: int | None +MaterializedMMParts = tuple[dict[str, Tensor] | None, ForwardPolicy | None] + + class FakeDataLoader: def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: int): self.world = get_world() @@ -213,9 +217,7 @@ def __init__( self.receiver: MicroBatchReceiver = setup_micro_batch_receiver(output_dir, dp_rank, start_step, config) self.mm_materializer = RawImageMaterializer(model_name, trust_remote_code=model_trust_remote_code) self.missing_mm_image_policy = missing_mm_image_policy - self.last_mm_materialize_time = 0.0 - self.last_mm_images_materialized = 0 - self.last_mm_images_placeholdered = 0 + self._reset_mm_stats() def wait_for_batch(self) -> None: if self.world.is_master: @@ -229,10 +231,58 @@ def wait_for_batch(self) -> None: def get_batch(self) -> list[TensorMicroBatch]: micro_batches = self.receiver.receive() + self._reset_mm_stats() + return [self._micro_batch_to_tensor(mb) for mb in micro_batches] + + def _reset_mm_stats(self) -> None: self.last_mm_materialize_time = 0.0 self.last_mm_images_materialized = 0 self.last_mm_images_placeholdered = 0 - return [self._micro_batch_to_tensor(mb) for mb in micro_batches] + + @staticmethod + def _materialized_mm_parts(materialized: MaterializedMM | None) -> MaterializedMMParts: + if materialized is None: + return None, None + return materialized.kwargs, materialized.forward_policy + + @staticmethod + def _mm_run_context(micro_batch: MicroBatch) -> str: + run_idx = next((i for i, n in enumerate(micro_batch.lora_num_tokens or []) if n > 0), None) + return f"run_idx={run_idx}, run_id={micro_batch.run_id}, run_step={micro_batch.run_step}" + + def _materialize_mm_refs(self, micro_batch: MicroBatch, refs: MMRefs) -> MaterializedMMParts: + materialize_start = time.perf_counter() + try: + materialized = self.mm_materializer.materialize(refs) + except FileNotFoundError as exc: + self.last_mm_materialize_time += time.perf_counter() - materialize_start + if self.missing_mm_image_policy == "error": + get_logger().error( + f"raw image materialization failed ({self._mm_run_context(micro_batch)}, uris={refs.uris}): {exc!r}" + ) + raise + return self._synthesize_missing_mm_placeholder(micro_batch, refs) + + self.last_mm_materialize_time += time.perf_counter() - materialize_start + self.last_mm_images_materialized += len(refs.uris) + return self._materialized_mm_parts(materialized) + + def _synthesize_missing_mm_placeholder(self, micro_batch: MicroBatch, refs: MMRefs) -> MaterializedMMParts: + placeholder_start = time.perf_counter() + materialized = self.mm_materializer.synthesize_placeholder(refs) + self.last_mm_materialize_time += time.perf_counter() - placeholder_start + self.last_mm_images_placeholdered += len(refs.uris) + micro_batch.loss_mask = [False] * len(micro_batch.loss_mask) + micro_batch.advantages = [0.0] * len(micro_batch.advantages) + + missing_uris = missing_file_uris(refs.uris) + get_logger().warning( + "raw image materialization missing image(s); using zero-loss placeholder " + f"({self._mm_run_context(micro_batch)}, " + f"missing_uris={missing_uris or ['']}, " + f"uris={refs.uris})" + ) + return self._materialized_mm_parts(materialized) def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: """Convert a MicroBatch (msgspec struct with lists) to a TensorMicroBatch (dict with tensors).""" @@ -244,56 +294,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: if micro_batch.mm_kwargs: raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") if micro_batch.mm_refs is not None: - materialize_start = time.perf_counter() - try: - materialized = self.mm_materializer.materialize(micro_batch.mm_refs) - if materialized is not None: - mm_kwargs = materialized.kwargs - mm_forward_policy = materialized.forward_policy - self.last_mm_materialize_time += time.perf_counter() - materialize_start - self.last_mm_images_materialized += len(micro_batch.mm_refs.uris) - except FileNotFoundError as exc: - self.last_mm_materialize_time += time.perf_counter() - materialize_start - run_idx = next((i for i, n in enumerate(micro_batch.lora_num_tokens or []) if n > 0), None) - if self.missing_mm_image_policy == "error": - get_logger().error( - f"raw image materialization failed (run_idx={run_idx}, run_id={micro_batch.run_id}, " - f"run_step={micro_batch.run_step}, uris={micro_batch.mm_refs.uris}): {exc!r}" - ) - raise - - placeholder_start = time.perf_counter() - try: - materialized = self.mm_materializer.synthesize_placeholder(micro_batch.mm_refs) - if materialized is not None: - mm_kwargs = materialized.kwargs - mm_forward_policy = materialized.forward_policy - except Exception as placeholder_exc: - get_logger().error( - f"raw image placeholder synthesis failed after missing image " - f"(run_idx={run_idx}, run_id={micro_batch.run_id}, run_step={micro_batch.run_step}, " - f"uris={micro_batch.mm_refs.uris}): {placeholder_exc!r}" - ) - raise placeholder_exc from exc - - self.last_mm_materialize_time += time.perf_counter() - placeholder_start - self.last_mm_images_placeholdered += len(micro_batch.mm_refs.uris) - micro_batch.loss_mask = [False] * len(micro_batch.loss_mask) - micro_batch.advantages = [0.0] * len(micro_batch.advantages) - missing_uris = missing_file_uris(micro_batch.mm_refs.uris) - get_logger().warning( - "raw image materialization missing image(s); using zero-loss placeholder " - f"(run_idx={run_idx}, run_id={micro_batch.run_id}, run_step={micro_batch.run_step}, " - f"missing_uris={missing_uris or ['']}, " - f"uris={micro_batch.mm_refs.uris})" - ) - except Exception as exc: - run_idx = next((i for i, n in enumerate(micro_batch.lora_num_tokens or []) if n > 0), None) - get_logger().error( - f"raw image materialization failed (run_idx={run_idx}, run_id={micro_batch.run_id}, " - f"run_step={micro_batch.run_step}, uris={micro_batch.mm_refs.uris}): {exc!r}" - ) - raise + mm_kwargs, mm_forward_policy = self._materialize_mm_refs(micro_batch, micro_batch.mm_refs) routed_experts = None packed_routed_experts = micro_batch.routed_experts if packed_routed_experts is not None: diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 11d7ff3669..fa93188b37 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -65,6 +65,7 @@ from prime_rl.utils.monitor import setup_monitor from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title +from prime_rl.utils.run_assets import configure_run_asset_env from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step, to_col_format from ring_flash_attn import substitute_hf_flash_attn from torchtitan.distributed.utils import clip_grad_norm_ @@ -72,6 +73,8 @@ @clean_exit def train(config: TrainerConfig): + configure_run_asset_env(config.output_dir, config.multimodal) + # Setup world and logger world = get_world() logger = setup_logger( diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index 8e495118a3..9012b6b8af 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -1,5 +1,6 @@ from __future__ import annotations +import base64 import hashlib from collections.abc import Iterable, Mapping from io import BytesIO @@ -7,9 +8,10 @@ from typing import Any from urllib.parse import unquote, urlparse -from prime_rl.multimodal.adapters.base import MaterializedMM +from prime_rl.multimodal.adapters.base import MaterializedMM, MultimodalAdapter from prime_rl.multimodal.registry import get_multimodal_adapter from prime_rl.multimodal.schema import ( + PROCESSED_MM_KEYS, RawMMItem, contains_processed_payload_key, parse_raw_mm_item, @@ -18,7 +20,6 @@ IMAGE_MODALITY = "image" SUPPORTED_MODALITIES = {IMAGE_MODALITY} -PROCESSED_MM_KEYS = {"pixel_values", "image_embeds", "image_features"} def _field(value: Any, name: str, default: Any = None) -> Any: @@ -36,19 +37,41 @@ def file_uri_to_path(uri: str) -> Path: return Path(unquote(parsed.path)) +def data_image_uri_to_bytes(uri: str) -> bytes: + if not uri.startswith("data:image/"): + raise ValueError(f"Expected data:image URI, got {uri!r}") + marker = ";base64," + if marker not in uri: + raise ValueError("data:image URI must use base64 encoding") + _, b64 = uri.split(marker, 1) + return base64.b64decode(b64) + + +def image_uri_to_bytes(uri: str) -> bytes: + if uri.startswith("data:image/"): + return data_image_uri_to_bytes(uri) + return file_uri_to_path(uri).read_bytes() + + +def validate_image_uri(uri: str) -> None: + if uri.startswith("data:image/"): + data_image_uri_to_bytes(uri) + return + file_uri_to_path(uri) + + def missing_file_uris(uris: Iterable[str]) -> list[str]: """Return missing local ``file://`` image refs; non-file refs are ignored.""" missing: list[str] = [] for uri in uris: - parsed = urlparse(uri) - if parsed.scheme != "file": + if urlparse(uri).scheme != "file": continue - if not Path(unquote(parsed.path)).exists(): + if not file_uri_to_path(uri).exists(): missing.append(uri) return missing -def image_file_uris_from_messages(messages: Iterable[Any]) -> list[str]: +def image_uris_from_messages(messages: Iterable[Any]) -> list[str]: uris: list[str] = [] for message in messages: content = _field(message, "content") @@ -133,14 +156,14 @@ def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | No f"{len(image_placeholders)} placeholders but {len(image_items)} image descriptors" ) - uris = image_file_uris_from_messages(messages) + uris = image_uris_from_messages(messages) if len(uris) != len(image_items): raise ValueError( "Raw image URI/descriptor mismatch: " - f"{len(uris)} file refs in messages but {len(image_items)} image descriptors" + f"{len(uris)} image refs in messages but {len(image_items)} image descriptors" ) for uri in uris: - file_uri_to_path(uri) + validate_image_uri(uri) return MMRefs( descriptor={ @@ -156,13 +179,49 @@ def sha256_32(data: bytes) -> str: return hashlib.sha256(data).hexdigest()[:32] -def _single_family_adapter(items: list[RawMMItem]): +def _parse_image_refs(refs: MMRefs) -> tuple[list[RawMMItem], list[str]]: + image_item_dicts = refs.descriptor.get("mm_items", {}).get(IMAGE_MODALITY, []) + image_hashes = list(refs.descriptor.get("mm_hashes", {}).get(IMAGE_MODALITY, [])) + if not image_item_dicts: + return [], [] + if len(refs.uris) != len(image_item_dicts) or len(image_hashes) != len(image_item_dicts): + raise ValueError( + "Raw image refs must have matching URI, descriptor, and hash counts " + f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" + ) + return [parse_raw_mm_item(validate_raw_mm_item(item)) for item in image_item_dicts], image_hashes + + +def _single_family_adapter(items: list[RawMMItem]) -> MultimodalAdapter: families = {item.family for item in items} if len(families) != 1: raise ValueError(f"Raw multimodal refs must use exactly one adapter family, got {sorted(families)}") return get_multimodal_adapter(next(iter(families))) +def _validate_processor_layout(adapter: MultimodalAdapter, image_processor: Any, image_items: list[RawMMItem]) -> None: + actual_fingerprint = adapter.processor_fingerprint(image_processor) + for item in image_items: + if item.layout_fingerprint != actual_fingerprint: + raise ValueError( + f"Raw image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" + ) + + +def _load_verified_images(uris: list[str], expected_hashes: list[str]) -> list[Any]: + from PIL import Image + + images = [] + for uri, expected_hash in zip(uris, expected_hashes, strict=True): + raw = image_uri_to_bytes(uri) + actual_hash = sha256_32(raw) + if actual_hash != expected_hash: + raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") + with Image.open(BytesIO(raw)) as image: + images.append(image.convert("RGB")) + return images + + class RawImageMaterializer: """Materialize raw image refs with the trainer model's HF image processor.""" @@ -183,55 +242,22 @@ def image_processor(self): self._image_processor = image_processor return self._image_processor - def materialize(self, refs: MMRefs | None) -> MaterializedMM | None: - if refs is None: + def materialize(self, refs: MMRefs) -> MaterializedMM | None: + image_items, image_hashes = _parse_image_refs(refs) + if not image_items: return None - image_item_dicts = refs.descriptor.get("mm_items", {}).get(IMAGE_MODALITY, []) - image_hashes = refs.descriptor.get("mm_hashes", {}).get(IMAGE_MODALITY, []) - if not image_item_dicts: - return None - if len(refs.uris) != len(image_item_dicts) or len(image_hashes) != len(image_item_dicts): - raise ValueError( - "Raw image refs must have matching URI, descriptor, and hash counts " - f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" - ) - image_items = [parse_raw_mm_item(validate_raw_mm_item(item)) for item in image_item_dicts] + image_processor = self.image_processor adapter = _single_family_adapter(image_items) - actual_fingerprint = adapter.processor_fingerprint(self.image_processor) - for item in image_items: - if item.layout_fingerprint != actual_fingerprint: - raise ValueError( - "Raw image layout fingerprint mismatch: " - f"expected {item.layout_fingerprint}, got {actual_fingerprint}" - ) - - from PIL import Image - - images = [] - for uri, expected_hash in zip(refs.uris, image_hashes, strict=True): - raw = file_uri_to_path(uri).read_bytes() - actual_hash = sha256_32(raw) - if actual_hash != expected_hash: - raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") - images.append(Image.open(BytesIO(raw)).convert("RGB")) - - return adapter.materialize_for_trainer(self.image_processor, image_items, images) - - def synthesize_placeholder(self, refs: MMRefs | None) -> MaterializedMM | None: - """Build zero-valued multimodal tensors via the owning adapter.""" - if refs is None: - return None + _validate_processor_layout(adapter, image_processor, image_items) + images = _load_verified_images(refs.uris, image_hashes) + return adapter.materialize_for_trainer(image_processor, image_items, images) - image_item_dicts = refs.descriptor.get("mm_items", {}).get(IMAGE_MODALITY, []) - image_hashes = refs.descriptor.get("mm_hashes", {}).get(IMAGE_MODALITY, []) - if not image_item_dicts: + def synthesize_placeholder(self, refs: MMRefs) -> MaterializedMM | None: + """Build zero-valued multimodal tensors via the owning adapter.""" + image_items, _ = _parse_image_refs(refs) + if not image_items: return None - if len(refs.uris) != len(image_item_dicts) or len(image_hashes) != len(image_item_dicts): - raise ValueError( - "Raw image refs must have matching URI, descriptor, and hash counts " - f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" - ) - image_items = [parse_raw_mm_item(validate_raw_mm_item(item)) for item in image_item_dicts] + image_processor = self.image_processor adapter = _single_family_adapter(image_items) - return adapter.synthesize_placeholder(self.image_processor, image_items) + return adapter.synthesize_placeholder(image_processor, image_items) diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py index ce9539581e..a77ea2b3ac 100644 --- a/src/prime_rl/utils/run_assets.py +++ b/src/prime_rl/utils/run_assets.py @@ -4,21 +4,76 @@ from collections.abc import Mapping from pathlib import Path +from prime_rl.configs.shared import MultimodalConfig + IMAGE_OFFLOAD_DIR_ENV = "VF_RENDERER_IMAGE_OFFLOAD_DIR" +IMAGE_STORAGE_ENV = "PRIME_RL_MM_IMAGE_STORAGE" RUN_DIR_ENV = "PRIME_RL_RUN_DIR" RUN_ID_ENV = "RUN_ID" +IMAGE_STORAGE_OFFLOAD = "offload" +IMAGE_STORAGE_INLINE = "inline" +RUN_OUTPUT_ROOT = Path("/data/outputs") +IMAGE_ASSET_SUBDIR = Path("assets/images") + + +def _expand_path(path: Path, env: Mapping[str, str]) -> Path: + expanded = os.path.expanduser(str(path)) + for key, value in env.items(): + expanded = expanded.replace(f"${{{key}}}", value).replace(f"${key}", value) + return Path(expanded).resolve() + + +def _run_id_dir(env: Mapping[str, str]) -> Path | None: + raw_run_id = env.get(RUN_ID_ENV, "").strip() + if not raw_run_id: + return None + run_id = raw_run_id.removeprefix("run_") + return RUN_OUTPUT_ROOT / f"run_{run_id}" + + +def resolve_image_offload_dir( + output_dir: Path, + multimodal: MultimodalConfig, + env: Mapping[str, str], +) -> Path: + explicit = multimodal.images.offload_dir + if explicit is not None: + return _expand_path(explicit, env) + hosted_run_dir = _run_id_dir(env) + if hosted_run_dir is not None: + return (hosted_run_dir / IMAGE_ASSET_SUBDIR).resolve() + return (output_dir.resolve() / IMAGE_ASSET_SUBDIR).resolve() -def run_asset_env(output_dir: Path, base: Mapping[str, str] | None = None) -> dict[str, str]: + +def run_asset_env( + output_dir: Path, + multimodal: MultimodalConfig | None = None, + base: Mapping[str, str] | None = None, +) -> dict[str, str]: """Resolve the environment used by subprocesses that share run image assets. - Hosted runs resolve ``RUN_ID`` to ``/data/outputs/run_${RUN_ID}`` inside - renderers. Local launches without hosted env vars use the RL output dir as - the explicit run dir. + Prime-RL config owns the multimodal image policy. Env vars are only the + transport used by verifiers/renderers running in subprocesses. """ env = dict(os.environ if base is None else base) - if env.get(IMAGE_OFFLOAD_DIR_ENV) or env.get(RUN_DIR_ENV) or env.get(RUN_ID_ENV): - return env - env[RUN_DIR_ENV] = str(output_dir.resolve()) + config = multimodal or MultimodalConfig() + storage = config.images.storage + env[IMAGE_STORAGE_ENV] = storage + + if not env.get(RUN_ID_ENV): + env[RUN_DIR_ENV] = str(output_dir.resolve()) + + if storage == IMAGE_STORAGE_OFFLOAD: + env[IMAGE_OFFLOAD_DIR_ENV] = str(resolve_image_offload_dir(output_dir, config, env)) + elif storage == IMAGE_STORAGE_INLINE: + env.pop(IMAGE_OFFLOAD_DIR_ENV, None) + else: + raise ValueError(f"Unknown multimodal image storage mode: {storage!r}") + return env + + +def configure_run_asset_env(output_dir: Path, multimodal: MultimodalConfig) -> None: + os.environ.update(run_asset_env(output_dir, multimodal=multimodal)) diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index cc569f33c7..8e77f92281 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -11,10 +11,13 @@ from __future__ import annotations import asyncio +import base64 import hashlib +from types import SimpleNamespace import numpy as np import pybase64 +import pytest from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.serve.disagg.protocol import GenerateResponse, GenerateResponseChoice @@ -25,8 +28,10 @@ PrimeRlServingTokens, _build_usage, _client_set_max_tokens, + _decode_raw_mm_kwargs, _FinalOutputCapture, _GenerateRoutedExpertsCapture, + _MMImageRefError, ) @@ -313,8 +318,8 @@ def _get_adapter(family): out = serving_tokens._materialize_raw_image_ref_sync( raw_ref, - expected_modality="image", - expected_hash=mm_hash, + feature_modality="image", + mm_hash=mm_hash, expected_placeholder_length=7, processor_model_name="model", trust_remote_code=True, @@ -329,3 +334,89 @@ def _get_adapter(family): assert item.family == "test_family" assert item.layout_fingerprint == fingerprint assert item.payload == {"adapter_owned": [1, 2, 3]} + + +def test_materialize_raw_image_ref_accepts_inline_data_uri(tmp_path, monkeypatch): + from PIL import Image + from renderers.mm_store import raw_mm_ref + + from prime_rl.inference.vllm import serving_tokens + + image_path = tmp_path / "image.png" + Image.new("RGB", (8, 6), color=(16, 32, 48)).save(image_path) + raw = image_path.read_bytes() + data_uri = f"data:image/png;base64,{base64.b64encode(raw).decode('ascii')}" + + mm_hash = hashlib.sha256(raw).hexdigest()[:32] + fingerprint = "f" * 32 + raw_ref = raw_mm_ref( + run_id="serving", + family="test_family", + fingerprint=fingerprint, + modality="image", + mm_hash=mm_hash, + raw_uri=data_uri, + payload={"adapter_owned": [4, 5, 6]}, + ) + processor = object() + captured = {} + + class _Adapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + captured["image_processor"] = image_processor + captured["item"] = item + captured["image_size"] = image.size + captured["expected_placeholder_length"] = expected_placeholder_length + return {"materialized": True} + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: processor) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _Adapter()) + + out = serving_tokens._materialize_raw_image_ref_sync( + raw_ref, + feature_modality="image", + mm_hash=mm_hash, + expected_placeholder_length=7, + processor_model_name="model", + trust_remote_code=True, + ) + + assert out == {"materialized": True} + assert captured["image_processor"] is processor + assert captured["image_size"] == (8, 6) + assert captured["expected_placeholder_length"] == 7 + assert captured["item"].payload == {"adapter_owned": [4, 5, 6]} + + +def test_decode_raw_mm_kwargs_rejects_none_items(): + features = SimpleNamespace( + mm_hashes={"image": ["a" * 32]}, + mm_placeholders={"image": [SimpleNamespace(length=1)]}, + kwargs_data={"image": [None]}, + ) + + with pytest.raises(_MMImageRefError, match="raw descriptor refs"): + asyncio.run( + _decode_raw_mm_kwargs( + features, + processor_model_name="model", + trust_remote_code=True, + ) + ) + + +def test_decode_raw_mm_kwargs_requires_parallel_placeholders(): + features = SimpleNamespace( + mm_hashes={"image": ["a" * 32]}, + mm_placeholders={"image": []}, + kwargs_data={"image": ["mmraw:v2:anything"]}, + ) + + with pytest.raises(_MMImageRefError, match="placeholder/hash length mismatch"): + asyncio.run( + _decode_raw_mm_kwargs( + features, + processor_model_name="model", + trust_remote_code=True, + ) + ) diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index a50307904e..a3604cf01c 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -110,6 +110,30 @@ def test_build_mm_refs_rejects_legacy_descriptor_without_raw_envelope(tmp_path): build_mm_refs(multi_modal_data, messages) +def test_build_mm_refs_accepts_inline_data_uri(): + data_uri = "data:image/png;base64,aGVsbG8=" + multi_modal_data = SimpleNamespace( + mm_items={"image": [_qwen_item([[1, 1, 1]])]}, + mm_hashes={"image": ["a" * 32]}, + mm_placeholders={"image": []}, + ) + messages = [ + { + "content": [ + { + "type": "image_url", + "image_url": {"url": data_uri}, + } + ] + } + ] + + refs = build_mm_refs(multi_modal_data, messages) + + assert refs is not None + assert refs.uris == [data_uri] + + def test_raw_image_materializer_synthesizes_qwen_placeholder_from_descriptor(): materializer = RawImageMaterializer("unused", trust_remote_code=False) materializer._image_processor = _ImageProcessor() From 3518763575080b06ff340da18c6d815ca4fc6f8d Mon Sep 17 00:00:00 2001 From: eligotts Date: Sun, 28 Jun 2026 20:58:30 +0000 Subject: [PATCH 08/52] fix: preserve launcher image asset env --- src/prime_rl/utils/run_assets.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py index a77ea2b3ac..0efc7e57a2 100644 --- a/src/prime_rl/utils/run_assets.py +++ b/src/prime_rl/utils/run_assets.py @@ -43,6 +43,9 @@ def resolve_image_offload_dir( hosted_run_dir = _run_id_dir(env) if hosted_run_dir is not None: return (hosted_run_dir / IMAGE_ASSET_SUBDIR).resolve() + run_dir = env.get(RUN_DIR_ENV, "").strip() + if run_dir: + return (Path(run_dir).resolve() / IMAGE_ASSET_SUBDIR).resolve() return (output_dir.resolve() / IMAGE_ASSET_SUBDIR).resolve() @@ -62,11 +65,12 @@ def run_asset_env( storage = config.images.storage env[IMAGE_STORAGE_ENV] = storage - if not env.get(RUN_ID_ENV): + if not env.get(RUN_ID_ENV) and not env.get(RUN_DIR_ENV): env[RUN_DIR_ENV] = str(output_dir.resolve()) if storage == IMAGE_STORAGE_OFFLOAD: - env[IMAGE_OFFLOAD_DIR_ENV] = str(resolve_image_offload_dir(output_dir, config, env)) + if not env.get(IMAGE_OFFLOAD_DIR_ENV): + env[IMAGE_OFFLOAD_DIR_ENV] = str(resolve_image_offload_dir(output_dir, config, env)) elif storage == IMAGE_STORAGE_INLINE: env.pop(IMAGE_OFFLOAD_DIR_ENV, None) else: From f37ceac864d649befce38ef09b6a3c0c06de6022 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 29 Jun 2026 06:16:27 +0000 Subject: [PATCH 09/52] Simplify raw multimodal offload path --- .../src/prime_rl/configs/inference.py | 2 +- .../src/prime_rl/configs/orchestrator.py | 2 +- .../src/prime_rl/configs/rl.py | 2 +- .../src/prime_rl/configs/shared.py | 11 +- .../src/prime_rl/configs/trainer.py | 2 +- src/prime_rl/entrypoints/rl.py | 6 + src/prime_rl/entrypoints/trainer.py | 2 - src/prime_rl/inference/vllm/serving_tokens.py | 29 +-- src/prime_rl/multimodal/adapters/base.py | 2 +- src/prime_rl/multimodal/adapters/kimi_k25.py | 16 +- src/prime_rl/multimodal/adapters/qwen_vl.py | 16 +- src/prime_rl/multimodal/schema.py | 2 + .../templates/multi_node_rl.sbatch.j2 | 13 +- src/prime_rl/trainer/batch.py | 3 +- src/prime_rl/trainer/model.py | 2 +- src/prime_rl/trainer/rl/data.py | 2 +- src/prime_rl/trainer/rl/train.py | 3 - src/prime_rl/transport/types.py | 10 +- src/prime_rl/utils/mm.py | 44 +--- src/prime_rl/utils/run_assets.py | 31 +-- tests/unit/inference/test_serving_tokens.py | 72 +------ tests/unit/orchestrator/test_batch.py | 18 +- tests/unit/orchestrator/test_qwen3_vl_e2e.py | 194 ------------------ tests/unit/test_configs.py | 5 - tests/unit/train/test_model_forward.py | 7 +- tests/unit/utils/test_mm.py | 77 +------ 26 files changed, 80 insertions(+), 493 deletions(-) delete mode 100644 tests/unit/orchestrator/test_qwen3_vl_e2e.py diff --git a/packages/prime-rl-configs/src/prime_rl/configs/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 3207cab3eb..f03bf1d081 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/inference.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/inference.py @@ -436,7 +436,7 @@ class InferenceConfig(BaseConfig): experimental: InferenceExperimentalConfig = InferenceExperimentalConfig() multimodal: MultimodalConfig = MultimodalConfig() - """Raw multimodal storage policy shared with trainer and orchestrator.""" + """Raw multimodal image offload settings shared with trainer and orchestrator.""" @model_validator(mode="after") def validate_multi_node_requires_slurm(self): diff --git a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py index ea842d48f7..17223517b0 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -563,7 +563,7 @@ class OrchestratorConfig(BaseConfig): experimental: OrchestratorExperimentalConfig = OrchestratorExperimentalConfig() multimodal: MultimodalConfig = MultimodalConfig() - """Raw multimodal storage policy shared with trainer and inference.""" + """Raw multimodal image offload settings shared with trainer and inference.""" @model_validator(mode="before") @classmethod diff --git a/packages/prime-rl-configs/src/prime_rl/configs/rl.py b/packages/prime-rl-configs/src/prime_rl/configs/rl.py index ce8c178ba6..82e158babd 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -221,7 +221,7 @@ class RLConfig(BaseConfig): weight_broadcast: SharedWeightBroadcastConfig | None = None multimodal: MultimodalConfig = MultimodalConfig() - """Shared raw multimodal storage policy. Propagated to trainer, orchestrator, and inference.""" + """Shared raw multimodal image offload settings. Propagated to trainer, orchestrator, and inference.""" bench: bool = False """Benchmark mode. Sets trainer and orchestrator to benchmark mode and, when set, suffixes the W&B project with ``-bench``.""" 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 3f614166f1..6c9b564168 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -60,22 +60,13 @@ def resolve_project_dir(self): ServerType = Literal["vllm", "openai"] -MultimodalImageStorage = Literal["offload", "inline"] -class MultimodalImageConfig(BaseConfig): - storage: MultimodalImageStorage = "offload" - """How raw image bytes are transported. ``offload`` writes data images to run assets; ``inline`` keeps base64 data-image URIs in the multimodal payload.""" - +class MultimodalConfig(BaseConfig): offload_dir: Path | None = None """Directory for offloaded image assets. Supports environment expansion such as ``/data/outputs/run_${RUN_ID}/assets/images``. When unset, prime-rl resolves a run-scoped default.""" -class MultimodalConfig(BaseConfig): - images: MultimodalImageConfig = MultimodalImageConfig() - """Raw multimodal image storage configuration.""" - - class VLMConfig(BaseConfig): vision_encoder_attr: str """Dotted attribute path to the vision encoder module (e.g. ``model.visual``).""" 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 227e08e231..79fe151b2a 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -571,7 +571,7 @@ class TrainerConfig(BaseConfig): """Policy when raw multimodal image files disappear before trainer materialization. ``placeholder_zero_loss`` warns, synthesizes zero-valued image tensors with the original descriptor geometry, and masks out the affected microbatch loss; ``error`` preserves fail-fast behavior.""" multimodal: MultimodalConfig = MultimodalConfig() - """Raw multimodal storage policy shared with orchestrator and inference.""" + """Raw multimodal image offload settings shared with orchestrator and inference.""" enable_token_export: bool = False """Opt-in per-token JSONL export for rollout debugging. When enabled, writes token ids and aligned trainer metrics after each forward pass.""" diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 482a4e5c61..4ebf5fa61f 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -357,12 +357,16 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> kv_offload_disk_path=str(offload.disk.path) if (is_mooncake and offload.disk is not None) else "", kv_offload_device_name=offload.device_name if is_mooncake else "", ) + image_offload_dir = ( + os.path.expanduser(str(config.multimodal.offload_dir)) if config.multimodal.offload_dir is not None else "" + ) if config.deployment.type == "single_node": script = template.render( **config.slurm.template_vars, config_path=config_dir / RL_TOML, output_dir=config.output_dir, + image_offload_dir=image_offload_dir, gpus_per_node=config.deployment.gpus_per_node, ) elif config.inference is not None and config.inference.deployment.type == "disaggregated": @@ -374,6 +378,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> config_dir=config_dir, output_dir=config.output_dir, orchestrator_output_dir=config.orchestrator.output_dir, + image_offload_dir=image_offload_dir, num_train_nodes=config.deployment.num_train_nodes, num_infer_nodes=infer_deploy.num_nodes * config.deployment.num_infer_replicas, nodes_per_infer_replica=infer_deploy.num_nodes, @@ -406,6 +411,7 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> config_dir=config_dir, # TODO: should prob have each subconfig path separately output_dir=config.output_dir, orchestrator_output_dir=config.orchestrator.output_dir, + image_offload_dir=image_offload_dir, num_train_nodes=config.deployment.num_train_nodes, num_infer_nodes=config.deployment.total_infer_nodes, nodes_per_infer_replica=config.deployment.num_infer_nodes, diff --git a/src/prime_rl/entrypoints/trainer.py b/src/prime_rl/entrypoints/trainer.py index 096966934a..0d89f58fba 100644 --- a/src/prime_rl/entrypoints/trainer.py +++ b/src/prime_rl/entrypoints/trainer.py @@ -12,13 +12,11 @@ from prime_rl.configs.trainer import TrainerConfig from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title -from prime_rl.utils.run_assets import configure_run_asset_env def main(): set_proc_title("Trainer") config = cli(TrainerConfig) - configure_run_asset_env(config.output_dir, config.multimodal) from prime_rl.trainer.rl.train import train train(config) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 44b273a9f2..e29a5c0998 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -33,7 +33,6 @@ from __future__ import annotations import asyncio -import base64 import hashlib from collections.abc import AsyncGenerator, AsyncIterable from dataclasses import dataclass @@ -197,19 +196,10 @@ def _parse_raw_image_ref(raw_ref: str, *, feature_modality: str, mm_hash: str): def _read_verified_raw_image(ref) -> bytes: from renderers.mm_store import raw_image_path - label = ref.raw_image_id or ref.raw_uri or "raw image" - if ref.raw_image_id is not None: - try: - raw = raw_image_path(run_id=ref.run_id, raw_image_id=ref.raw_image_id).read_bytes() - except OSError as exc: - raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_id!r}: {exc}") from exc - elif isinstance(ref.raw_uri, str) and ref.raw_uri.startswith("data:image/"): - marker = ";base64," - if marker not in ref.raw_uri: - raise _MMImageRefError("Inline raw image refs must use base64 data:image URIs") - raw = base64.b64decode(ref.raw_uri.split(marker, 1)[1]) - else: - raise _MMImageRefError(f"Unsupported raw image source {label!r}") + try: + raw = raw_image_path(raw_image_id=ref.raw_image_id).read_bytes() + except OSError as exc: + raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_id!r}: {exc}") from exc actual_hash = hashlib.sha256(raw).hexdigest()[:32] if actual_hash != ref.mm_hash: @@ -231,18 +221,19 @@ def _materialize_raw_image_ref_sync( *, feature_modality: str, mm_hash: str, - expected_placeholder_length: int | None, + expected_placeholder_length: int, processor_model_name: str, trust_remote_code: bool, ): ref = _parse_raw_image_ref(raw_ref, feature_modality=feature_modality, mm_hash=mm_hash) raw = _read_verified_raw_image(ref) - image = _decode_raw_image(raw, raw_image_id=ref.raw_image_id or "inline image") + image = _decode_raw_image(raw, raw_image_id=ref.raw_image_id) image_processor = _load_image_processor(processor_model_name, trust_remote_code) item = RawMMItem( modality=ref.modality, family=ref.family, layout_fingerprint=ref.fingerprint, + raw_image_id=ref.raw_image_id, payload=dict(ref.payload), raw_ref=raw_ref, ) @@ -291,14 +282,14 @@ async def _decode_raw_mm_kwargs( for feature_modality, hashes in features.mm_hashes.items(): raw_refs, placeholders = _raw_ref_payloads_for_feature(features, feature_modality, hashes) decoded: list[Any] = [] - for idx, raw_ref in enumerate(raw_refs): + for raw_ref, mm_hash, placeholder in zip(raw_refs, hashes, placeholders, strict=True): decoded.append( await asyncio.to_thread( _materialize_raw_image_ref_sync, raw_ref, feature_modality=feature_modality, - mm_hash=hashes[idx], - expected_placeholder_length=placeholders[idx].length, + mm_hash=mm_hash, + expected_placeholder_length=placeholder.length, processor_model_name=processor_model_name, trust_remote_code=trust_remote_code, ) diff --git a/src/prime_rl/multimodal/adapters/base.py b/src/prime_rl/multimodal/adapters/base.py index 86571509fb..b746b9f90a 100644 --- a/src/prime_rl/multimodal/adapters/base.py +++ b/src/prime_rl/multimodal/adapters/base.py @@ -42,7 +42,7 @@ def materialize_for_vllm( image_processor: Any, item: "RawMMItem", image: "Image", - expected_placeholder_length: int | None, + expected_placeholder_length: int, ) -> Any: ... def synthesize_placeholder( diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index 12ebe00b06..39923a2c9f 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -4,7 +4,7 @@ from collections.abc import Mapping from typing import Any -from renderers.image_layout_specs import KIMI_K25_IMAGE_LAYOUT, KimiK25ImageLayoutSpec +from renderers.kimi_k25 import KIMI_K25_IMAGE_LAYOUT, KimiK25ImageLayoutSpec from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem @@ -66,7 +66,9 @@ def _grid_payload(item: RawMMItem) -> list[int]: grid = item.payload.get("grid_thws") if grid is None: raise ValueError("Kimi raw descriptor payload is missing grid_thws") - if len(grid) == 1 and isinstance(grid[0], list): + if not isinstance(grid, list | tuple): + raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") + if len(grid) == 1 and isinstance(grid[0], list | tuple): grid = grid[0] if not isinstance(grid, list | tuple) or len(grid) != 3: raise ValueError(f"Invalid Kimi grid_thws: {grid!r}") @@ -121,10 +123,10 @@ def materialize_for_trainer( if "grid_thws" not in tensors: raise ValueError("Kimi processor did not return grid_thws") actual_grids = tensors["grid_thws"].reshape(-1, 3).tolist() - for idx, item in enumerate(items): + for idx, (item, actual_grid) in enumerate(zip(items, actual_grids, strict=True)): expected = _grid_payload(item) - if actual_grids[idx] != expected: - raise ValueError(f"Kimi grid mismatch at index {idx}: expected {expected}, got {actual_grids[idx]}") + if actual_grid != expected: + raise ValueError(f"Kimi grid mismatch at index {idx}: expected {expected}, got {actual_grid}") return MaterializedMM(kwargs=tensors, forward_policy=self.forward_policy) def materialize_for_vllm( @@ -132,7 +134,7 @@ def materialize_for_vllm( image_processor: Any, item: RawMMItem, image: Any, - expected_placeholder_length: int | None, + expected_placeholder_length: int, ) -> Any: from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems @@ -148,7 +150,7 @@ def materialize_for_vllm( actual_grid = tensors["grid_thws"].reshape(-1, 3).tolist()[0] if actual_grid != expected_grid: raise ValueError(f"Kimi grid mismatch: expected {expected_grid}, got {actual_grid}") - if expected_placeholder_length is not None and expected_placeholder_length != 1: + if expected_placeholder_length != 1: raise ValueError(f"Kimi image placeholder length mismatch: expected {expected_placeholder_length}, got 1") grid_sizes = tensors["grid_thws"].reshape(-1, 3).prod(-1) config_by_key = { diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py index 59551ee327..6adf9c0bfa 100644 --- a/src/prime_rl/multimodal/adapters/qwen_vl.py +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -31,7 +31,9 @@ def _grid_payload(item: RawMMItem) -> list[int]: grid = item.payload.get("image_grid_thw") if grid is None: raise ValueError("Qwen raw descriptor payload is missing image_grid_thw") - if len(grid) == 1 and isinstance(grid[0], list): + if not isinstance(grid, list | tuple): + raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") + if len(grid) == 1 and isinstance(grid[0], list | tuple): grid = grid[0] if not isinstance(grid, list | tuple) or len(grid) != 3: raise ValueError(f"Invalid Qwen image_grid_thw: {grid!r}") @@ -91,10 +93,10 @@ def materialize_for_trainer( if "image_grid_thw" not in tensors: raise ValueError("Qwen processor did not return image_grid_thw") actual_grids = tensors["image_grid_thw"].tolist() - for idx, item in enumerate(items): + for idx, (item, actual_grid) in enumerate(zip(items, actual_grids, strict=True)): expected = _grid_payload(item) - if actual_grids[idx] != expected: - raise ValueError(f"Image grid mismatch at index {idx}: expected {expected}, got {actual_grids[idx]}") + if actual_grid != expected: + raise ValueError(f"Image grid mismatch at index {idx}: expected {expected}, got {actual_grid}") return MaterializedMM(kwargs=tensors, forward_policy=self.forward_policy) def materialize_for_vllm( @@ -102,7 +104,7 @@ def materialize_for_vllm( image_processor: Any, item: RawMMItem, image: Any, - expected_placeholder_length: int | None, + expected_placeholder_length: int, ) -> Any: from vllm.model_executor.models.qwen2_vl import _create_qwen2vl_field_factory from vllm.multimodal.inputs import MultiModalKwargsItems @@ -122,7 +124,7 @@ def materialize_for_vllm( if actual_grid != expected_grid: raise ValueError(f"Image grid mismatch: expected {expected_grid}, got {actual_grid}") num_image_tokens = int(expected_grid[0] * expected_grid[1] * expected_grid[2] // (merge_size * merge_size)) - if expected_placeholder_length is not None and expected_placeholder_length != num_image_tokens: + if expected_placeholder_length != num_image_tokens: raise ValueError( f"Image placeholder length mismatch: expected {expected_placeholder_length}, got {num_image_tokens}" ) @@ -151,7 +153,7 @@ def synthesize_placeholder( feature_dim = self.placeholder_feature_dim(image_processor) pixel_values: list[torch.Tensor] = [] image_grid_thw: list[list[int]] = [] - for idx, item in enumerate(items): + for item in items: self.validate_item(item) grid = _grid_payload(item) pixel_values.append(torch.zeros((math.prod(grid), feature_dim), dtype=torch.float32)) diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index a5ef15aa37..73cc7be7de 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -14,6 +14,7 @@ class RawMMItem: modality: str family: str layout_fingerprint: str + raw_image_id: str payload: dict[str, Any] raw_ref: str | None = None vllm_modality: str | None = None @@ -72,6 +73,7 @@ def parse_raw_mm_item(value: Any) -> RawMMItem: modality=_required_str(descriptor, "modality"), family=_required_str(descriptor, "family"), layout_fingerprint=_required_str(descriptor, "layout_fingerprint"), + raw_image_id=_required_str(descriptor, "raw_image_id"), payload=_payload(descriptor), raw_ref=_optional_str(descriptor, "raw_ref"), vllm_modality=_optional_str(descriptor, "vllm_modality"), diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index 33c9e7ee29..23a33f71a9 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -59,9 +59,18 @@ export PROJECT_DIR={{ project_dir }} export CONFIG_DIR={{ config_dir }} export OUTPUT_DIR={{ output_dir }} export ORCHESTRATOR_OUTPUT_DIR={{ orchestrator_output_dir }} -if [ -z "${VF_RENDERER_IMAGE_OFFLOAD_DIR:-}" ] && [ -z "${PRIME_RL_RUN_DIR:-}" ] && [ -z "${RUN_ID:-}" ]; then - export PRIME_RL_RUN_DIR="$ORCHESTRATOR_OUTPUT_DIR" +{% if image_offload_dir %} +export VF_RENDERER_IMAGE_OFFLOAD_DIR="{{ image_offload_dir }}" +{% else %} +if [ -z "${VF_RENDERER_IMAGE_OFFLOAD_DIR:-}" ]; then + if [ -n "${RUN_ID:-}" ]; then + RUN_ASSET_ID="${RUN_ID#run_}" + export VF_RENDERER_IMAGE_OFFLOAD_DIR="/data/outputs/run_${RUN_ASSET_ID}/assets/images" + else + export VF_RENDERER_IMAGE_OFFLOAD_DIR="$ORCHESTRATOR_OUTPUT_DIR/assets/images" + fi fi +{% endif %} mkdir -p $OUTPUT_DIR/logs/trainer $OUTPUT_DIR/logs/inference rm -f $OUTPUT_DIR/logs/inference/*.log ln -sfn trainer/node_0.log $OUTPUT_DIR/logs/trainer.log diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 18e3afbec8..85c46a3480 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -168,7 +168,7 @@ 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 or sample.mm_kwargs is not None + return sample.mm_refs is not None @dataclass @@ -291,7 +291,6 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: mm_token_type_ids=mm_token_type_ids, env_names=env_names, mm_refs=first_sample.mm_refs if _is_multimodal_sample(first_sample) else None, - mm_kwargs=first_sample.mm_kwargs if _is_multimodal_sample(first_sample) else None, rl_weights=streams["rl_weights"], ce_weights=streams["ce_weights"], ref_kl_weights=streams["ref_kl_weights"], diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index 48d529b1c7..7dec82785f 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -1187,7 +1187,7 @@ def forward( "temperature": temperature, } - if mm_kwargs: + if mm_kwargs is not None: # Forward the per-model multimodal tensors verbatim, plus the # token→modality map derived from renderer token ids. kwargs.update(mm_kwargs) diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index ed1b429526..2c9e9287e9 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -291,7 +291,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: micro_batch.lora_num_tokens[0] = len(micro_batch.input_ids) mm_kwargs: dict[str, Tensor] | None = None mm_forward_policy: ForwardPolicy | None = None - if micro_batch.mm_kwargs: + if micro_batch.mm_kwargs is not None: raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") if micro_batch.mm_refs is not None: mm_kwargs, mm_forward_policy = self._materialize_mm_refs(micro_batch, micro_batch.mm_refs) diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index fa93188b37..11d7ff3669 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -65,7 +65,6 @@ from prime_rl.utils.monitor import setup_monitor from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title -from prime_rl.utils.run_assets import configure_run_asset_env from prime_rl.utils.utils import clean_exit, resolve_latest_ckpt_step, to_col_format from ring_flash_attn import substitute_hf_flash_attn from torchtitan.distributed.utils import clip_grad_norm_ @@ -73,8 +72,6 @@ @clean_exit def train(config: TrainerConfig): - configure_run_asset_env(config.output_dir, config.multimodal) - # Setup world and logger world = get_world() logger = setup_logger( diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index 6724fe254c..5b8b03358a 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -21,8 +21,8 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru class MMRefs(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): """Raw multimodal sidecar references for one sample. - ``descriptor`` carries JSON-safe renderer metadata (hashes, grids, placeholder - layout). ``uris`` carries the raw image files that the trainer materializes + ``descriptor`` carries JSON-safe renderer metadata (hashes and adapter + payloads). ``uris`` carries the raw image files that the trainer materializes with its own processor. Processed tensors are intentionally not part of this transport. """ @@ -47,8 +47,7 @@ class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tr env_name: str ref_logprobs: list[float] | None = None # reference-model logprobs (ref_kl component) - # Legacy eager multimodal payloads are rejected by the v1 raw-image-ref - # path. Keep the field so old batches fail at a clear boundary. + # Processed multimodal payloads are rejected by the v1 raw-image-ref path. mm_kwargs: dict[str, EncodedTensor] | None = None mm_refs: MMRefs | None = None @@ -102,8 +101,7 @@ class MicroBatch(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): lora_num_tokens: list[int] | None = None routed_experts: RoutedExperts | None = None - # Legacy eager multimodal payloads are rejected by the v1 raw-image-ref - # path. Keep the field so old batches fail at a clear boundary. + # Processed multimodal payloads are rejected by the v1 raw-image-ref path. mm_kwargs: dict[str, EncodedTensor] | None = None mm_refs: MMRefs | None = None # mm_token_type_ids: token type ids per token [batch seq], int64 (0=text, 1=image, 2=video) diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index 9012b6b8af..a25610609e 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -1,6 +1,5 @@ from __future__ import annotations -import base64 import hashlib from collections.abc import Iterable, Mapping from io import BytesIO @@ -37,29 +36,6 @@ def file_uri_to_path(uri: str) -> Path: return Path(unquote(parsed.path)) -def data_image_uri_to_bytes(uri: str) -> bytes: - if not uri.startswith("data:image/"): - raise ValueError(f"Expected data:image URI, got {uri!r}") - marker = ";base64," - if marker not in uri: - raise ValueError("data:image URI must use base64 encoding") - _, b64 = uri.split(marker, 1) - return base64.b64decode(b64) - - -def image_uri_to_bytes(uri: str) -> bytes: - if uri.startswith("data:image/"): - return data_image_uri_to_bytes(uri) - return file_uri_to_path(uri).read_bytes() - - -def validate_image_uri(uri: str) -> None: - if uri.startswith("data:image/"): - data_image_uri_to_bytes(uri) - return - file_uri_to_path(uri) - - def missing_file_uris(uris: Iterable[str]) -> list[str]: """Return missing local ``file://`` image refs; non-file refs are ignored.""" missing: list[str] = [] @@ -123,13 +99,6 @@ def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: ) -def _placeholder_dict(placeholder: Any) -> dict[str, int]: - return { - "offset": int(_field(placeholder, "offset")), - "length": int(_field(placeholder, "length")), - } - - def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | None: mm_items = _field(multi_modal_data, "mm_items", None) if not mm_items: @@ -148,14 +117,6 @@ def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | No f"{len(image_items)} image descriptors but {len(image_hashes)} image hashes" ) - mm_placeholders = _field(multi_modal_data, "mm_placeholders", {}) or {} - image_placeholders = [_placeholder_dict(p) for p in mm_placeholders.get(IMAGE_MODALITY, [])] - if image_placeholders and len(image_placeholders) != len(image_items): - raise ValueError( - "Raw image placeholder/descriptor mismatch: " - f"{len(image_placeholders)} placeholders but {len(image_items)} image descriptors" - ) - uris = image_uris_from_messages(messages) if len(uris) != len(image_items): raise ValueError( @@ -163,13 +124,12 @@ def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | No f"{len(uris)} image refs in messages but {len(image_items)} image descriptors" ) for uri in uris: - validate_image_uri(uri) + file_uri_to_path(uri) return MMRefs( descriptor={ "mm_items": {IMAGE_MODALITY: image_items}, "mm_hashes": {IMAGE_MODALITY: image_hashes}, - "mm_placeholders": {IMAGE_MODALITY: image_placeholders}, }, uris=uris, ) @@ -213,7 +173,7 @@ def _load_verified_images(uris: list[str], expected_hashes: list[str]) -> list[A images = [] for uri, expected_hash in zip(uris, expected_hashes, strict=True): - raw = image_uri_to_bytes(uri) + raw = file_uri_to_path(uri).read_bytes() actual_hash = sha256_32(raw) if actual_hash != expected_hash: raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py index 0efc7e57a2..9202ad47f1 100644 --- a/src/prime_rl/utils/run_assets.py +++ b/src/prime_rl/utils/run_assets.py @@ -3,24 +3,19 @@ import os from collections.abc import Mapping from pathlib import Path +from string import Template from prime_rl.configs.shared import MultimodalConfig IMAGE_OFFLOAD_DIR_ENV = "VF_RENDERER_IMAGE_OFFLOAD_DIR" -IMAGE_STORAGE_ENV = "PRIME_RL_MM_IMAGE_STORAGE" -RUN_DIR_ENV = "PRIME_RL_RUN_DIR" RUN_ID_ENV = "RUN_ID" -IMAGE_STORAGE_OFFLOAD = "offload" -IMAGE_STORAGE_INLINE = "inline" RUN_OUTPUT_ROOT = Path("/data/outputs") IMAGE_ASSET_SUBDIR = Path("assets/images") def _expand_path(path: Path, env: Mapping[str, str]) -> Path: - expanded = os.path.expanduser(str(path)) - for key, value in env.items(): - expanded = expanded.replace(f"${{{key}}}", value).replace(f"${key}", value) + expanded = Template(os.path.expanduser(str(path))).safe_substitute(env) return Path(expanded).resolve() @@ -37,15 +32,12 @@ def resolve_image_offload_dir( multimodal: MultimodalConfig, env: Mapping[str, str], ) -> Path: - explicit = multimodal.images.offload_dir + explicit = multimodal.offload_dir if explicit is not None: return _expand_path(explicit, env) hosted_run_dir = _run_id_dir(env) if hosted_run_dir is not None: return (hosted_run_dir / IMAGE_ASSET_SUBDIR).resolve() - run_dir = env.get(RUN_DIR_ENV, "").strip() - if run_dir: - return (Path(run_dir).resolve() / IMAGE_ASSET_SUBDIR).resolve() return (output_dir.resolve() / IMAGE_ASSET_SUBDIR).resolve() @@ -56,28 +48,19 @@ def run_asset_env( ) -> dict[str, str]: """Resolve the environment used by subprocesses that share run image assets. - Prime-RL config owns the multimodal image policy. Env vars are only the + Prime-RL config owns the multimodal image offload path. Env vars are only the transport used by verifiers/renderers running in subprocesses. """ env = dict(os.environ if base is None else base) config = multimodal or MultimodalConfig() - storage = config.images.storage - env[IMAGE_STORAGE_ENV] = storage - if not env.get(RUN_ID_ENV) and not env.get(RUN_DIR_ENV): - env[RUN_DIR_ENV] = str(output_dir.resolve()) - - if storage == IMAGE_STORAGE_OFFLOAD: - if not env.get(IMAGE_OFFLOAD_DIR_ENV): - env[IMAGE_OFFLOAD_DIR_ENV] = str(resolve_image_offload_dir(output_dir, config, env)) - elif storage == IMAGE_STORAGE_INLINE: - env.pop(IMAGE_OFFLOAD_DIR_ENV, None) - else: - raise ValueError(f"Unknown multimodal image storage mode: {storage!r}") + env[IMAGE_OFFLOAD_DIR_ENV] = str(resolve_image_offload_dir(output_dir, config, env)) return env def configure_run_asset_env(output_dir: Path, multimodal: MultimodalConfig) -> None: + if os.environ.get(IMAGE_OFFLOAD_DIR_ENV): + return os.environ.update(run_asset_env(output_dir, multimodal=multimodal)) diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 8e77f92281..d3314c8cd7 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -11,7 +11,6 @@ from __future__ import annotations import asyncio -import base64 import hashlib from types import SimpleNamespace @@ -290,7 +289,6 @@ def test_materialize_raw_image_ref_uses_generic_family_payload(tmp_path, monkeyp mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] fingerprint = "f" * 32 raw_ref = raw_mm_ref( - run_id="serving", family="test_family", fingerprint=fingerprint, modality="image", @@ -333,61 +331,10 @@ def _get_adapter(family): item = captured["item"] assert item.family == "test_family" assert item.layout_fingerprint == fingerprint + assert item.raw_image_id == image_path.name assert item.payload == {"adapter_owned": [1, 2, 3]} -def test_materialize_raw_image_ref_accepts_inline_data_uri(tmp_path, monkeypatch): - from PIL import Image - from renderers.mm_store import raw_mm_ref - - from prime_rl.inference.vllm import serving_tokens - - image_path = tmp_path / "image.png" - Image.new("RGB", (8, 6), color=(16, 32, 48)).save(image_path) - raw = image_path.read_bytes() - data_uri = f"data:image/png;base64,{base64.b64encode(raw).decode('ascii')}" - - mm_hash = hashlib.sha256(raw).hexdigest()[:32] - fingerprint = "f" * 32 - raw_ref = raw_mm_ref( - run_id="serving", - family="test_family", - fingerprint=fingerprint, - modality="image", - mm_hash=mm_hash, - raw_uri=data_uri, - payload={"adapter_owned": [4, 5, 6]}, - ) - processor = object() - captured = {} - - class _Adapter: - def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): - captured["image_processor"] = image_processor - captured["item"] = item - captured["image_size"] = image.size - captured["expected_placeholder_length"] = expected_placeholder_length - return {"materialized": True} - - monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: processor) - monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _Adapter()) - - out = serving_tokens._materialize_raw_image_ref_sync( - raw_ref, - feature_modality="image", - mm_hash=mm_hash, - expected_placeholder_length=7, - processor_model_name="model", - trust_remote_code=True, - ) - - assert out == {"materialized": True} - assert captured["image_processor"] is processor - assert captured["image_size"] == (8, 6) - assert captured["expected_placeholder_length"] == 7 - assert captured["item"].payload == {"adapter_owned": [4, 5, 6]} - - def test_decode_raw_mm_kwargs_rejects_none_items(): features = SimpleNamespace( mm_hashes={"image": ["a" * 32]}, @@ -403,20 +350,3 @@ def test_decode_raw_mm_kwargs_rejects_none_items(): trust_remote_code=True, ) ) - - -def test_decode_raw_mm_kwargs_requires_parallel_placeholders(): - features = SimpleNamespace( - mm_hashes={"image": ["a" * 32]}, - mm_placeholders={"image": []}, - kwargs_data={"image": ["mmraw:v2:anything"]}, - ) - - with pytest.raises(_MMImageRefError, match="placeholder/hash length mismatch"): - asyncio.run( - _decode_raw_mm_kwargs( - features, - processor_model_name="model", - trust_remote_code=True, - ) - ) diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 0b7f05cf44..ad22aaac20 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -419,6 +419,7 @@ def _mm_refs() -> MMRefs: "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, + "raw_image_id": "image.png", "payload": {"image_grid_thw": [[1, 1, 1]]}, } ] @@ -429,23 +430,6 @@ def _mm_refs() -> MMRefs: ) -def test_prepare_sample_preserves_raw_mm_refs(): - sample = TrainingSample( - token_ids=[10, 11, 12], - mask=[False, False, True], - logprobs=[0.0] * 3, - temperatures=[1.0] * 3, - advantages=[0.0, 0.0, 1.0], - env_name="test-env", - mm_token_type_ids=[0, 1, 0], - mm_refs=_mm_refs(), - ) - - mb = prepare_sample(sample, seq_len=8) - assert mb.mm_refs == sample.mm_refs - assert mb.mm_token_type_ids == [0, 1, 0] - - def test_prepare_sample_rejects_overlong_raw_mm_refs(): sample = TrainingSample( token_ids=[10, 11, 12, 13], diff --git a/tests/unit/orchestrator/test_qwen3_vl_e2e.py b/tests/unit/orchestrator/test_qwen3_vl_e2e.py deleted file mode 100644 index 0369bfa35d..0000000000 --- a/tests/unit/orchestrator/test_qwen3_vl_e2e.py +++ /dev/null @@ -1,194 +0,0 @@ -"""End-to-end integration test for the Qwen3-VL renderer path. - -Walks a multimodal request through the full client stack — RendererClient -→ renderers.client.generate → /inference/v1/generate features payload — -with the HTTP layer mocked, and verifies that vLLM can deserialize the -features back into engine inputs identical to what its own server-side -processor would have produced for the same messages. - -This is the strongest end-to-end check we can run without a GPU. The -remaining missing piece (vLLM actually consuming the engine input, -sampling tokens, and returning them) is exercised in real rollouts. -""" - -from __future__ import annotations - -import asyncio -import json -from pathlib import Path -from typing import Any -from unittest.mock import MagicMock - -import httpx -import pytest - -_HF_CACHE = Path("~/.cache/huggingface/hub").expanduser() -_MODEL = "Qwen/Qwen3-VL-4B-Instruct" - - -def _model_cached() -> bool: - safe = "models--" + _MODEL.replace("/", "--") - snapshots = _HF_CACHE / safe / "snapshots" - if not snapshots.is_dir(): - return False - return any(p.is_dir() for p in snapshots.iterdir()) - - -pytestmark = pytest.mark.skipif( - not _model_cached(), - reason=f"{_MODEL}: HF snapshot not cached locally", -) - - -class _FakeOpenAI: - """Minimal AsyncOpenAI stand-in that captures POST bodies. - - The renderer client calls ``client.post(absolute_url, body=...)``; - we capture the body for assertions and return a canned generate - response so the parse-side of the flow runs. - """ - - def __init__(self): - self.calls: list[dict[str, Any]] = [] - self.base_url = "http://fake-host:8000/v1" - - async def post(self, path, *, cast_to=dict, body=None, options=None): - self.calls.append({"path": path, "body": body, "options": options}) - # Reply with two sampled tokens + <|im_end|>. The renderer's - # parse_response slices the content tokens. - payload = { - "request_id": "qwen-vl-e2e", - "choices": [ - { - "index": 0, - "token_ids": [50, 60, 151645], - "logprobs": { - "content": [ - {"token": "t1", "logprob": -0.1}, - {"token": "t2", "logprob": -0.2}, - {"token": "t3", "logprob": -0.3}, - ] - }, - "finish_reason": "stop", - }, - ], - } - return httpx.Response(200, content=json.dumps(payload).encode()) - - -def test_renderer_client_qwen3_vl_e2e_features_payload_roundtrips_through_vllm(tmp_path, monkeypatch): - """Walk a Qwen3-VL multimodal turn through the renderer client and - verify the resulting ``/inference/v1/generate`` body has a valid - ``features`` payload that: - - 1. parses through vLLM's ``GenerateRequest`` pydantic model, - 2. carries a raw image ref instead of processed image tensors, - 3. has placeholder ranges that exactly cover the ``<|image_pad|>`` - runs in the prompt token sequence. - """ - from PIL import Image - from renderers.base import load_tokenizer - from renderers.mm_store import IMAGE_REF_PREFIX, split_raw_mm_ref - from renderers.qwen3_vl import Qwen3VLRenderer - from transformers import AutoProcessor - from verifiers.clients.renderer_client import RendererClient - from verifiers.types import ( - ClientConfig, - UserMessage, - ) - from vllm.entrypoints.serve.disagg.protocol import GenerateRequest - - # ── Build a real Qwen3VLRenderer with a real processor. ───────────── - tokenizer = load_tokenizer(_MODEL) - processor = AutoProcessor.from_pretrained(_MODEL) - renderer = Qwen3VLRenderer(tokenizer, processor=processor) - - image_pad_id = tokenizer.convert_tokens_to_ids("<|image_pad|>") - - # ── Manually wire a RendererClient bypassing the pool factory. ────── - client_cfg = ClientConfig(client_type="renderer", base_url="http://fake-host:8000/v1") - rc = object.__new__(RendererClient) - rc._config = client_cfg - rc._renderer = renderer - rc._pool_size = 1 - rc._client = _FakeOpenAI() - rc.logger = MagicMock() - - # ── Build a verifiers-shaped user message with an image. ──────────── - run_dir = tmp_path / "run_e2e" - image_dir = run_dir / "assets" / "images" - image_dir.mkdir(parents=True) - image_path = image_dir / "sample.png" - Image.new("RGB", (224, 224), color=(64, 128, 255)).save(image_path) - monkeypatch.setenv("PRIME_RL_RUN_DIR", str(run_dir)) - - # The renderer accepts the OpenAI ``image_url`` content-part shape — - # the same shape verifiers' UserMessage carries through. - user = UserMessage( - content=[ - {"type": "text", "text": "What's in this picture?"}, - {"type": "image_url", "image_url": {"url": image_path.as_uri()}}, - ] - ) - - # to_native_prompt converts to renderer-shaped messages. - prompt, _ = asyncio.run(rc.to_native_prompt([user])) - sampling = {"max_tokens": 16} - - response = asyncio.run( - rc.get_native_response( - prompt=prompt, - model=_MODEL, - sampling_args=sampling, - tools=None, - ) - ) - - # ── The HTTP body should carry a features payload. ────────────────── - fake = rc.client - assert isinstance(fake, _FakeOpenAI) - assert len(fake.calls) == 1 - body = fake.calls[0]["body"] - assert "features" in body, "RendererClient should ship features for image content" - features = body["features"] - - # ── Pydantic-roundtrip through vLLM's GenerateRequest model. ──────── - gen_req = GenerateRequest( - token_ids=body["token_ids"], - features=features, - sampling_params=body["sampling_params"], - ) - assert gen_req.features is not None - assert "image" in gen_req.features.mm_hashes - assert len(gen_req.features.mm_hashes["image"]) == 1 - - # ── Placeholder anchoring: the offset/length in features must land - # exactly on a run of <|image_pad|> ids in the prompt. ─────────── - placeholders = gen_req.features.mm_placeholders["image"] - assert len(placeholders) == 1 - ph = placeholders[0] - pad_slice = body["token_ids"][ph.offset : ph.offset + ph.length] - assert all(t == image_pad_id for t in pad_slice), ( - f"placeholder span ({ph.offset}, {ph.length}) does not cover image_pad tokens; slice={pad_slice[:8]}..." - ) - - # ── kwargs_data carries a raw image ref for vLLM to materialize. ───── - assert gen_req.features.kwargs_data is not None - ref_items = gen_req.features.kwargs_data["image"] - assert len(ref_items) == 1 - ref_item = ref_items[0] - assert isinstance(ref_item, str) - assert ref_item.startswith(f"{IMAGE_REF_PREFIX}:") - ref = split_raw_mm_ref(ref_item) - assert ref.run_id == "e2e" - assert ref.modality == "image" - assert ref.raw_image_id == image_path.name - assert ref.mm_hash == gen_req.features.mm_hashes["image"][0] - persisted = response["multi_modal_data"].mm_items["image"][0] - assert ref.payload["image_grid_thw"] == persisted["payload"]["image_grid_thw"] - - # ── Response parsed through renderer's parse_response. ────────────── - assert response["completion_ids"] == [50, 60, 151645] - # multi_modal_data surfaces on the result so the caller can persist it. - assert response["multi_modal_data"] is not None - assert len(response["multi_modal_data"].mm_items["image"]) == 1 diff --git a/tests/unit/test_configs.py b/tests/unit/test_configs.py index e91002ef6b..49243a9f98 100644 --- a/tests/unit/test_configs.py +++ b/tests/unit/test_configs.py @@ -191,11 +191,6 @@ def test_trainer_enable_token_export_cli_flag(): assert cli(TrainerConfig, args=["--enable-token-export"]).enable_token_export -def test_trainer_missing_mm_image_policy_cli_override(): - assert cli(TrainerConfig, args=[]).missing_mm_image_policy == "placeholder_zero_loss" - assert cli(TrainerConfig, args=["--missing-mm-image-policy", "error"]).missing_mm_image_policy == "error" - - def test_single_node_auto_inference_client_dp_rank_count_matches_local_dp(): config = RLConfig.model_validate( { diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py index 22d492003c..b49a08bf99 100644 --- a/tests/unit/train/test_model_forward.py +++ b/tests/unit/train/test_model_forward.py @@ -1,5 +1,6 @@ from types import SimpleNamespace +import pytest import torch import torch.nn as nn @@ -90,7 +91,7 @@ def test_forward_policy_can_require_mm_token_type_ids(): input_ids = torch.tensor([[1, 10, 10, 2]]) position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) - try: + with pytest.raises(ValueError, match="mm_token_type_ids"): forward( model, input_ids, @@ -98,7 +99,3 @@ def test_forward_policy_can_require_mm_token_type_ids(): mm_kwargs={"pixel_values": torch.ones(2, 3)}, mm_forward_policy=ForwardPolicy(requires_mm_token_type_ids=True), ) - except ValueError as exc: - assert "mm_token_type_ids" in str(exc) - else: - raise AssertionError("forward should require mm_token_type_ids when policy says so") diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index a3604cf01c..3d4b812bc6 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -1,19 +1,12 @@ from types import SimpleNamespace -import pytest import torch from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND, RAW_MM_ITEM_VERSION from prime_rl.trainer.rl.data import DataLoader from prime_rl.transport.types import MicroBatch, MMRefs -from prime_rl.utils.mm import RawImageMaterializer, build_mm_refs, missing_file_uris - - -class _ImageProcessor: - patch_size = 2 - temporal_patch_size = 2 - image_mean = [0.5, 0.5, 0.5] +from prime_rl.utils.mm import build_mm_refs class _MissingMaterializer: @@ -37,6 +30,7 @@ def _qwen_item(grid): "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, + "raw_image_id": "image.png", "payload": {"image_grid_thw": grid}, } @@ -51,11 +45,11 @@ def _refs(uri: str = "file:///tmp/missing-image.png") -> MMRefs: ) -def _loader(policy: str = "placeholder_zero_loss") -> DataLoader: +def _loader() -> DataLoader: loader = object.__new__(DataLoader) loader.multi_run_manager = SimpleNamespace(max_runs=1) loader.mm_materializer = _MissingMaterializer() - loader.missing_mm_image_policy = policy + loader.missing_mm_image_policy = "placeholder_zero_loss" loader.last_mm_materialize_time = 0.0 loader.last_mm_images_materialized = 0 loader.last_mm_images_placeholdered = 0 @@ -78,51 +72,19 @@ def _micro_batch() -> MicroBatch: ) -def test_missing_file_uris_reports_missing_local_refs(tmp_path): - existing = tmp_path / "image.png" - existing.write_bytes(b"image") - - assert missing_file_uris( - [existing.as_uri(), (tmp_path / "missing.png").as_uri(), "https://example.test/i.png"] - ) == [(tmp_path / "missing.png").as_uri()] - - -def test_build_mm_refs_rejects_legacy_descriptor_without_raw_envelope(tmp_path): +def test_build_mm_refs_accepts_offloaded_file_uri(tmp_path): image_path = tmp_path / "image.png" image_path.write_bytes(b"image") - multi_modal_data = SimpleNamespace( - mm_items={"image": [{"image_grid_thw": [[1, 1, 1]]}]}, - mm_hashes={"image": ["a" * 32]}, - mm_placeholders={"image": []}, - ) - messages = [ - { - "content": [ - { - "type": "image_url", - "image_url": {"url": image_path.as_uri()}, - } - ] - } - ] - - with pytest.raises(ValueError, match="common envelope"): - build_mm_refs(multi_modal_data, messages) - - -def test_build_mm_refs_accepts_inline_data_uri(): - data_uri = "data:image/png;base64,aGVsbG8=" multi_modal_data = SimpleNamespace( mm_items={"image": [_qwen_item([[1, 1, 1]])]}, mm_hashes={"image": ["a" * 32]}, - mm_placeholders={"image": []}, ) messages = [ { "content": [ { "type": "image_url", - "image_url": {"url": data_uri}, + "image_url": {"url": image_path.as_uri()}, } ] } @@ -131,27 +93,7 @@ def test_build_mm_refs_accepts_inline_data_uri(): refs = build_mm_refs(multi_modal_data, messages) assert refs is not None - assert refs.uris == [data_uri] - - -def test_raw_image_materializer_synthesizes_qwen_placeholder_from_descriptor(): - materializer = RawImageMaterializer("unused", trust_remote_code=False) - materializer._image_processor = _ImageProcessor() - - mm_kwargs = materializer.synthesize_placeholder( - MMRefs( - descriptor={ - "mm_items": {"image": [_qwen_item([[1, 2, 3]])]}, - "mm_hashes": {"image": ["b" * 32]}, - }, - uris=["file:///tmp/missing-image.png"], - ) - ) - - assert mm_kwargs is not None - assert mm_kwargs.kwargs["pixel_values"].shape == (6, 24) - assert not bool(mm_kwargs.kwargs["pixel_values"].any()) - assert mm_kwargs.kwargs["image_grid_thw"].tolist() == [[1, 2, 3]] + assert refs.uris == [image_path.as_uri()] def test_dataloader_uses_zero_loss_placeholder_for_missing_raw_image(): @@ -163,8 +105,3 @@ def test_dataloader_uses_zero_loss_placeholder_for_missing_raw_image(): assert tensor_batch["loss_mask"].tolist() == [[False, False, False]] assert tensor_batch["advantages"].tolist() == [[0.0, 0.0, 0.0]] assert tensor_batch["mm_token_type_ids"].tolist() == [[0, 1, 0]] - - -def test_dataloader_can_fail_fast_on_missing_raw_image(): - with pytest.raises(FileNotFoundError): - _loader(policy="error")._micro_batch_to_tensor(_micro_batch()) From 80dc59f35c7048996c5a19b47d48e93959c6cfe8 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 29 Jun 2026 06:39:09 +0000 Subject: [PATCH 10/52] Simplify raw multimodal validation --- src/prime_rl/multimodal/schema.py | 14 ------- src/prime_rl/utils/mm.py | 43 +++++---------------- tests/unit/inference/test_serving_tokens.py | 21 ---------- tests/unit/orchestrator/test_batch.py | 15 ------- 4 files changed, 10 insertions(+), 83 deletions(-) diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index 73cc7be7de..4b3267acc0 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -6,8 +6,6 @@ from renderers.mm_store import RAW_MM_ITEM_KIND, RAW_MM_ITEM_VERSION -PROCESSED_MM_KEYS = frozenset({"pixel_values", "image_embeds", "image_features"}) - @dataclass(frozen=True) class RawMMItem: @@ -20,16 +18,6 @@ class RawMMItem: vllm_modality: str | None = None -def contains_processed_payload_key(value: Any) -> bool: - if isinstance(value, Mapping): - return bool(PROCESSED_MM_KEYS.intersection(value)) or any( - contains_processed_payload_key(v) for v in value.values() - ) - if isinstance(value, list | tuple): - return any(contains_processed_payload_key(v) for v in value) - return False - - def _descriptor_mapping(value: Any) -> Mapping[str, Any]: if isinstance(value, Mapping): return value @@ -58,8 +46,6 @@ def _payload(value: Mapping[str, Any]) -> dict[str, Any]: def _validate_envelope(value: Mapping[str, Any]) -> None: - if contains_processed_payload_key(value): - raise TypeError("v1 multimodal sidecars must not carry processed multimodal payloads") if value.get("kind") != RAW_MM_ITEM_KIND: raise ValueError("raw multimodal descriptor is missing the common envelope kind") if value.get("version") != RAW_MM_ITEM_VERSION: diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index a25610609e..c3c73af761 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -9,12 +9,7 @@ from prime_rl.multimodal.adapters.base import MaterializedMM, MultimodalAdapter from prime_rl.multimodal.registry import get_multimodal_adapter -from prime_rl.multimodal.schema import ( - PROCESSED_MM_KEYS, - RawMMItem, - contains_processed_payload_key, - parse_raw_mm_item, -) +from prime_rl.multimodal.schema import RawMMItem, parse_raw_mm_item from prime_rl.transport.types import MMRefs IMAGE_MODALITY = "image" @@ -63,31 +58,6 @@ def image_uris_from_messages(messages: Iterable[Any]) -> list[str]: return uris -def _normalize_json_value(value: Any, path: str) -> Any: - if value is None or isinstance(value, str | int | float | bool): - return value - if isinstance(value, tuple): - return [_normalize_json_value(v, f"{path}[]") for v in value] - if isinstance(value, list): - return [_normalize_json_value(v, f"{path}[]") for v in value] - if isinstance(value, Mapping): - return {str(k): _normalize_json_value(v, f"{path}.{k}") for k, v in value.items()} - raise TypeError( - f"v1 multimodal sidecars must be JSON-safe raw image descriptors; {path} has unsupported {type(value).__name__}" - ) - - -def validate_raw_mm_item(item: Mapping[str, Any]) -> dict[str, Any]: - if contains_processed_payload_key(item): - raise TypeError( - "v1 multimodal sidecars must be raw image descriptors, not processed payloads " - f"({', '.join(sorted(PROCESSED_MM_KEYS))})" - ) - normalized = {str(k): _normalize_json_value(v, str(k)) for k, v in item.items()} - parse_raw_mm_item(normalized) - return normalized - - def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: unsupported = sorted( modality for modality, items in mm_items.items() if items and modality not in SUPPORTED_MODALITIES @@ -99,13 +69,20 @@ def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: ) +def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: + item_dicts = [dict(item) for item in items] + for item in item_dicts: + parse_raw_mm_item(item) + return item_dicts + + def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | None: mm_items = _field(multi_modal_data, "mm_items", None) if not mm_items: return None _validate_modalities(mm_items) - image_items = [validate_raw_mm_item(item) for item in mm_items.get(IMAGE_MODALITY, [])] + image_items = _raw_item_dicts(mm_items.get(IMAGE_MODALITY, [])) if not image_items: return None @@ -149,7 +126,7 @@ def _parse_image_refs(refs: MMRefs) -> tuple[list[RawMMItem], list[str]]: "Raw image refs must have matching URI, descriptor, and hash counts " f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" ) - return [parse_raw_mm_item(validate_raw_mm_item(item)) for item in image_item_dicts], image_hashes + return [parse_raw_mm_item(item) for item in image_item_dicts], image_hashes def _single_family_adapter(items: list[RawMMItem]) -> MultimodalAdapter: diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index d3314c8cd7..fa8954fd67 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -12,11 +12,9 @@ import asyncio import hashlib -from types import SimpleNamespace import numpy as np import pybase64 -import pytest from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.serve.disagg.protocol import GenerateResponse, GenerateResponseChoice @@ -27,10 +25,8 @@ PrimeRlServingTokens, _build_usage, _client_set_max_tokens, - _decode_raw_mm_kwargs, _FinalOutputCapture, _GenerateRoutedExpertsCapture, - _MMImageRefError, ) @@ -333,20 +329,3 @@ def _get_adapter(family): assert item.layout_fingerprint == fingerprint assert item.raw_image_id == image_path.name assert item.payload == {"adapter_owned": [1, 2, 3]} - - -def test_decode_raw_mm_kwargs_rejects_none_items(): - features = SimpleNamespace( - mm_hashes={"image": ["a" * 32]}, - mm_placeholders={"image": [SimpleNamespace(length=1)]}, - kwargs_data={"image": [None]}, - ) - - with pytest.raises(_MMImageRefError, match="raw descriptor refs"): - asyncio.run( - _decode_raw_mm_kwargs( - features, - processor_model_name="model", - trust_remote_code=True, - ) - ) diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index ad22aaac20..a5c89bcb11 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -446,21 +446,6 @@ def test_prepare_sample_rejects_overlong_raw_mm_refs(): prepare_sample(sample, seq_len=3) -def test_prepare_sample_rejects_processed_mm_kwargs(): - sample = TrainingSample( - token_ids=[10, 11, 12], - mask=[False, False, True], - logprobs=[0.0] * 3, - temperatures=[1.0] * 3, - advantages=[0.0, 0.0, 1.0], - env_name="test-env", - mm_kwargs={}, - ) - - with pytest.raises(ValueError, match="Processed multimodal mm_kwargs are unsupported"): - prepare_sample(sample, seq_len=8) - - def test_prepare_sample_none_routed_experts(): """When routed_experts is None, micro_batch.routed_experts is None.""" sample = TrainingSample( From bf2bc8e314d200be11f7325d7d59230439e943ee Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 29 Jun 2026 16:58:34 +0000 Subject: [PATCH 11/52] Clarify raw image asset env wiring --- deps/renderers | 2 +- src/prime_rl/entrypoints/orchestrator.py | 4 ++-- src/prime_rl/entrypoints/rl.py | 4 ++-- src/prime_rl/inference/server.py | 4 ++-- src/prime_rl/orchestrator/orchestrator.py | 4 ++-- src/prime_rl/utils/run_assets.py | 8 +++++--- 6 files changed, 14 insertions(+), 12 deletions(-) diff --git a/deps/renderers b/deps/renderers index 673b790916..af84c192df 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 673b790916ec70efe89526a66b3fa2eed0b88bd6 +Subproject commit af84c192df5a8cbaa5fa51b1ee5c85fdd8d8e5e2 diff --git a/src/prime_rl/entrypoints/orchestrator.py b/src/prime_rl/entrypoints/orchestrator.py index 1bdc92ff8b..842ddb01f5 100644 --- a/src/prime_rl/entrypoints/orchestrator.py +++ b/src/prime_rl/entrypoints/orchestrator.py @@ -14,13 +14,13 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title -from prime_rl.utils.run_assets import configure_run_asset_env +from prime_rl.utils.run_assets import apply_run_asset_env def main(): set_proc_title("Orchestrator") config = cli(OrchestratorConfig) - configure_run_asset_env(config.output_dir, config.multimodal) + apply_run_asset_env(config.output_dir, config.multimodal) from prime_rl.orchestrator.orchestrator import run_orchestrator asyncio.run(run_orchestrator(config)) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 04ff72b240..bc754876eb 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -35,7 +35,7 @@ monitor_process, set_proc_title, ) -from prime_rl.utils.run_assets import run_asset_env +from prime_rl.utils.run_assets import build_run_asset_env RL_TOML = "rl.toml" RL_SBATCH = "rl.sbatch" @@ -126,7 +126,7 @@ def rl_local(config: RLConfig): "WANDB_SHARED_MODE": "1", "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } - shared_run_asset_env = run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal) + shared_run_asset_env = build_run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.model.client.is_elastic: diff --git a/src/prime_rl/inference/server.py b/src/prime_rl/inference/server.py index d16c599e49..0970b6c291 100644 --- a/src/prime_rl/inference/server.py +++ b/src/prime_rl/inference/server.py @@ -3,13 +3,13 @@ from prime_rl.configs.inference import InferenceConfig from prime_rl.utils.config import cli -from prime_rl.utils.run_assets import configure_run_asset_env +from prime_rl.utils.run_assets import apply_run_asset_env def setup_vllm_env(config: InferenceConfig): """Set vLLM environment variables based on config. Must be called before importing vLLM.""" - configure_run_asset_env(config.output_dir, config.multimodal) + apply_run_asset_env(config.output_dir, config.multimodal) # spawn is more robust in vLLM nightlies and Qwen3-VL (fork can deadlock with multithreaded processes) os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") diff --git a/src/prime_rl/orchestrator/orchestrator.py b/src/prime_rl/orchestrator/orchestrator.py index b3f876dc1b..f8d6e60aee 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -76,7 +76,7 @@ from prime_rl.utils.logger import format_time, get_logger, setup_logger from prime_rl.utils.monitor import setup_monitor from prime_rl.utils.pathing import get_log_dir, get_rollout_dir, get_step_path -from prime_rl.utils.run_assets import configure_run_asset_env +from prime_rl.utils.run_assets import apply_run_asset_env from prime_rl.utils.usage_reporter import UsageReporter from prime_rl.utils.utils import ( clean_exit, @@ -841,7 +841,7 @@ async def run_orchestrator(config: OrchestratorConfig) -> None: """Top-level entrypoint. Wrapped in ``@clean_exit`` so wandb is flushed on exit (success or crash); keeps that out of the class. """ - configure_run_asset_env(config.output_dir, config.multimodal) + apply_run_asset_env(config.output_dir, config.multimodal) await Orchestrator(config).start() diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py index 9202ad47f1..a359e45ec7 100644 --- a/src/prime_rl/utils/run_assets.py +++ b/src/prime_rl/utils/run_assets.py @@ -7,6 +7,7 @@ from prime_rl.configs.shared import MultimodalConfig +# Contract: must match renderers.mm_store.IMAGE_OFFLOAD_DIR_ENV. IMAGE_OFFLOAD_DIR_ENV = "VF_RENDERER_IMAGE_OFFLOAD_DIR" RUN_ID_ENV = "RUN_ID" @@ -32,6 +33,7 @@ def resolve_image_offload_dir( multimodal: MultimodalConfig, env: Mapping[str, str], ) -> Path: + """Resolve image asset dir by precedence: offload_dir, RUN_ID hosted path, then output_dir/assets/images.""" explicit = multimodal.offload_dir if explicit is not None: return _expand_path(explicit, env) @@ -41,7 +43,7 @@ def resolve_image_offload_dir( return (output_dir.resolve() / IMAGE_ASSET_SUBDIR).resolve() -def run_asset_env( +def build_run_asset_env( output_dir: Path, multimodal: MultimodalConfig | None = None, base: Mapping[str, str] | None = None, @@ -60,7 +62,7 @@ def run_asset_env( return env -def configure_run_asset_env(output_dir: Path, multimodal: MultimodalConfig) -> None: +def apply_run_asset_env(output_dir: Path, multimodal: MultimodalConfig) -> None: if os.environ.get(IMAGE_OFFLOAD_DIR_ENV): return - os.environ.update(run_asset_env(output_dir, multimodal=multimodal)) + os.environ.update(build_run_asset_env(output_dir, multimodal=multimodal)) From 805486f6d35287f4c1a50f8dde042f1332ca01d5 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 29 Jun 2026 19:44:25 +0000 Subject: [PATCH 12/52] Use URI-complete raw multimodal refs --- deps/renderers | 2 +- deps/verifiers | 2 +- src/prime_rl/entrypoints/rl.py | 9 ++-- src/prime_rl/inference/server.py | 3 -- src/prime_rl/inference/vllm/serving_tokens.py | 15 +++---- src/prime_rl/multimodal/schema.py | 4 +- src/prime_rl/orchestrator/trajectories.py | 2 +- src/prime_rl/utils/mm.py | 44 +++++-------------- tests/unit/inference/test_serving_tokens.py | 5 +-- tests/unit/orchestrator/test_batch.py | 2 +- tests/unit/utils/test_mm.py | 22 +++------- 11 files changed, 38 insertions(+), 72 deletions(-) diff --git a/deps/renderers b/deps/renderers index af84c192df..2a19d75986 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit af84c192df5a8cbaa5fa51b1ee5c85fdd8d8e5e2 +Subproject commit 2a19d759860f46eb269a56b6f69f5efe0f38b7a1 diff --git a/deps/verifiers b/deps/verifiers index 0b1d73fa1d..0dc57a1493 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 0b1d73fa1d148b48e2f8165a99c9b10a8777ef1c +Subproject commit 0dc57a149390b11aa879fc3b7a0f69b4f8f59b21 diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index bc754876eb..54f1174959 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -126,7 +126,8 @@ def rl_local(config: RLConfig): "WANDB_SHARED_MODE": "1", "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } - shared_run_asset_env = build_run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal) + inherited_env = dict(os.environ) + writer_run_asset_env = build_run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.model.client.is_elastic: @@ -172,7 +173,7 @@ def sigterm_handler(signum, frame): inference_process = Popen( inference_cmd, env={ - **shared_run_asset_env, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_INFERENCE_ENV_VARS, **config.env_vars, @@ -227,7 +228,7 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ - **shared_run_asset_env, + **writer_run_asset_env, **DEFAULT_COMMON_ENV_VARS, "LOGURU_FORCE_COLORS": "1", "WANDB_PROGRAM": "uv run rl", @@ -276,7 +277,7 @@ def sigterm_handler(signum, frame): trainer_process = Popen( trainer_cmd, env={ - **shared_run_asset_env, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, "LOGURU_FORCE_COLORS": "1", diff --git a/src/prime_rl/inference/server.py b/src/prime_rl/inference/server.py index 0970b6c291..4429146c5f 100644 --- a/src/prime_rl/inference/server.py +++ b/src/prime_rl/inference/server.py @@ -3,14 +3,11 @@ from prime_rl.configs.inference import InferenceConfig from prime_rl.utils.config import cli -from prime_rl.utils.run_assets import apply_run_asset_env def setup_vllm_env(config: InferenceConfig): """Set vLLM environment variables based on config. Must be called before importing vLLM.""" - apply_run_asset_env(config.output_dir, config.multimodal) - # spawn is more robust in vLLM nightlies and Qwen3-VL (fork can deadlock with multithreaded processes) os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn") diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index e29a5c0998..2a95f1398a 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -61,6 +61,7 @@ from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture from prime_rl.multimodal.registry import get_multimodal_adapter from prime_rl.multimodal.schema import RawMMItem +from prime_rl.utils.mm import file_uri_to_path @dataclass @@ -194,12 +195,10 @@ def _parse_raw_image_ref(raw_ref: str, *, feature_modality: str, mm_hash: str): def _read_verified_raw_image(ref) -> bytes: - from renderers.mm_store import raw_image_path - try: - raw = raw_image_path(raw_image_id=ref.raw_image_id).read_bytes() + raw = file_uri_to_path(ref.raw_image_uri).read_bytes() except OSError as exc: - raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_id!r}: {exc}") from exc + raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_uri!r}: {exc}") from exc actual_hash = hashlib.sha256(raw).hexdigest()[:32] if actual_hash != ref.mm_hash: @@ -207,13 +206,13 @@ def _read_verified_raw_image(ref) -> bytes: return raw -def _decode_raw_image(raw: bytes, *, raw_image_id: str): +def _decode_raw_image(raw: bytes, *, raw_image_uri: str): from PIL import Image try: return Image.open(BytesIO(raw)).convert("RGB") except OSError as exc: - raise _MMImageRefError(f"Unable to decode raw image asset {raw_image_id!r}: {exc}") from exc + raise _MMImageRefError(f"Unable to decode raw image asset {raw_image_uri!r}: {exc}") from exc def _materialize_raw_image_ref_sync( @@ -227,13 +226,13 @@ def _materialize_raw_image_ref_sync( ): ref = _parse_raw_image_ref(raw_ref, feature_modality=feature_modality, mm_hash=mm_hash) raw = _read_verified_raw_image(ref) - image = _decode_raw_image(raw, raw_image_id=ref.raw_image_id) + image = _decode_raw_image(raw, raw_image_uri=ref.raw_image_uri) image_processor = _load_image_processor(processor_model_name, trust_remote_code) item = RawMMItem( modality=ref.modality, family=ref.family, layout_fingerprint=ref.fingerprint, - raw_image_id=ref.raw_image_id, + raw_image_uri=ref.raw_image_uri, payload=dict(ref.payload), raw_ref=raw_ref, ) diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index 4b3267acc0..50eda23731 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -12,7 +12,7 @@ class RawMMItem: modality: str family: str layout_fingerprint: str - raw_image_id: str + raw_image_uri: str payload: dict[str, Any] raw_ref: str | None = None vllm_modality: str | None = None @@ -59,7 +59,7 @@ def parse_raw_mm_item(value: Any) -> RawMMItem: modality=_required_str(descriptor, "modality"), family=_required_str(descriptor, "family"), layout_fingerprint=_required_str(descriptor, "layout_fingerprint"), - raw_image_id=_required_str(descriptor, "raw_image_id"), + raw_image_uri=_required_str(descriptor, "raw_image_uri"), payload=_payload(descriptor), raw_ref=_optional_str(descriptor, "raw_ref"), vllm_modality=_optional_str(descriptor, "vllm_modality"), diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 3e1d76f58a..9dc555ccd9 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -69,7 +69,7 @@ def trace_to_samples( mm_token_type_ids: list[int] | None = None mmd = branch.multi_modal_data if mmd is not None: - mm_refs = build_mm_refs(mmd, branch.messages) + mm_refs = build_mm_refs(mmd) mapping = mm_token_type_ids_mapping or {} mm_token_type_ids = [mapping.get(t, 0) for t in token_ids] samples.append( diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index c3c73af761..38010de684 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -42,22 +42,6 @@ def missing_file_uris(uris: Iterable[str]) -> list[str]: return missing -def image_uris_from_messages(messages: Iterable[Any]) -> list[str]: - uris: list[str] = [] - for message in messages: - content = _field(message, "content") - if not isinstance(content, list): - continue - for part in content: - if _field(part, "type") != "image_url": - continue - image_url = _field(part, "image_url") - url = image_url if isinstance(image_url, str) else _field(image_url, "url") - if isinstance(url, str): - uris.append(url) - return uris - - def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: unsupported = sorted( modality for modality, items in mm_items.items() if items and modality not in SUPPORTED_MODALITIES @@ -69,20 +53,25 @@ def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: ) -def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: - item_dicts = [dict(item) for item in items] - for item in item_dicts: - parse_raw_mm_item(item) - return item_dicts +def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> tuple[list[dict[str, Any]], list[str]]: + item_dicts: list[dict[str, Any]] = [] + uris: list[str] = [] + for item in items: + item_dict = dict(item) + parsed = parse_raw_mm_item(item_dict) + file_uri_to_path(parsed.raw_image_uri) + item_dicts.append(item_dict) + uris.append(parsed.raw_image_uri) + return item_dicts, uris -def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | None: +def build_mm_refs(multi_modal_data: Any) -> MMRefs | None: mm_items = _field(multi_modal_data, "mm_items", None) if not mm_items: return None _validate_modalities(mm_items) - image_items = _raw_item_dicts(mm_items.get(IMAGE_MODALITY, [])) + image_items, uris = _raw_item_dicts(mm_items.get(IMAGE_MODALITY, [])) if not image_items: return None @@ -94,15 +83,6 @@ def build_mm_refs(multi_modal_data: Any, messages: Iterable[Any]) -> MMRefs | No f"{len(image_items)} image descriptors but {len(image_hashes)} image hashes" ) - uris = image_uris_from_messages(messages) - if len(uris) != len(image_items): - raise ValueError( - "Raw image URI/descriptor mismatch: " - f"{len(uris)} image refs in messages but {len(image_items)} image descriptors" - ) - for uri in uris: - file_uri_to_path(uri) - return MMRefs( descriptor={ "mm_items": {IMAGE_MODALITY: image_items}, diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index fa8954fd67..f3e3ea9eba 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -280,7 +280,6 @@ def test_materialize_raw_image_ref_uses_generic_family_payload(tmp_path, monkeyp image_dir.mkdir(parents=True) image_path = image_dir / "image.png" Image.new("RGB", (8, 6), color=(32, 64, 128)).save(image_path) - monkeypatch.setenv("VF_RENDERER_IMAGE_OFFLOAD_DIR", str(image_dir)) mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] fingerprint = "f" * 32 @@ -289,7 +288,7 @@ def test_materialize_raw_image_ref_uses_generic_family_payload(tmp_path, monkeyp fingerprint=fingerprint, modality="image", mm_hash=mm_hash, - raw_image_id=image_path.name, + raw_image_uri=image_path.as_uri(), payload={"adapter_owned": [1, 2, 3]}, ) processor = object() @@ -327,5 +326,5 @@ def _get_adapter(family): item = captured["item"] assert item.family == "test_family" assert item.layout_fingerprint == fingerprint - assert item.raw_image_id == image_path.name + assert item.raw_image_uri == image_path.as_uri() assert item.payload == {"adapter_owned": [1, 2, 3]} diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index a5c89bcb11..6428571226 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -419,7 +419,7 @@ def _mm_refs() -> MMRefs: "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, - "raw_image_id": "image.png", + "raw_image_uri": "file:///tmp/image.png", "payload": {"image_grid_thw": [[1, 1, 1]]}, } ] diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index 3d4b812bc6..886a034eb4 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -23,14 +23,14 @@ def synthesize_placeholder(self, refs): ) -def _qwen_item(grid): +def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): return { "kind": RAW_MM_ITEM_KIND, "version": RAW_MM_ITEM_VERSION, "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, - "raw_image_id": "image.png", + "raw_image_uri": uri, "payload": {"image_grid_thw": grid}, } @@ -38,7 +38,7 @@ def _qwen_item(grid): def _refs(uri: str = "file:///tmp/missing-image.png") -> MMRefs: return MMRefs( descriptor={ - "mm_items": {"image": [_qwen_item([[1, 1, 1]])]}, + "mm_items": {"image": [_qwen_item([[1, 1, 1]], uri)]}, "mm_hashes": {"image": ["a" * 32]}, }, uris=[uri], @@ -76,21 +76,11 @@ def test_build_mm_refs_accepts_offloaded_file_uri(tmp_path): image_path = tmp_path / "image.png" image_path.write_bytes(b"image") multi_modal_data = SimpleNamespace( - mm_items={"image": [_qwen_item([[1, 1, 1]])]}, + mm_items={"image": [_qwen_item([[1, 1, 1]], image_path.as_uri())]}, mm_hashes={"image": ["a" * 32]}, ) - messages = [ - { - "content": [ - { - "type": "image_url", - "image_url": {"url": image_path.as_uri()}, - } - ] - } - ] - - refs = build_mm_refs(multi_modal_data, messages) + + refs = build_mm_refs(multi_modal_data) assert refs is not None assert refs.uris == [image_path.as_uri()] From dce2759cc5b0c1032a58b56e76c004ffedc8fc03 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 29 Jun 2026 20:50:44 +0000 Subject: [PATCH 13/52] Drop raw multimodal descriptor version --- deps/renderers | 2 +- src/prime_rl/multimodal/schema.py | 4 +--- tests/unit/orchestrator/test_batch.py | 1 - tests/unit/utils/test_mm.py | 3 +-- 4 files changed, 3 insertions(+), 7 deletions(-) diff --git a/deps/renderers b/deps/renderers index 2a19d75986..aa2d44d06a 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 2a19d759860f46eb269a56b6f69f5efe0f38b7a1 +Subproject commit aa2d44d06a40ee31aaa68acd2dae072a8dcca585 diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index 50eda23731..0685274b20 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -4,7 +4,7 @@ from dataclasses import dataclass from typing import Any -from renderers.mm_store import RAW_MM_ITEM_KIND, RAW_MM_ITEM_VERSION +from renderers.mm_store import RAW_MM_ITEM_KIND @dataclass(frozen=True) @@ -48,8 +48,6 @@ def _payload(value: Mapping[str, Any]) -> dict[str, Any]: def _validate_envelope(value: Mapping[str, Any]) -> None: if value.get("kind") != RAW_MM_ITEM_KIND: raise ValueError("raw multimodal descriptor is missing the common envelope kind") - if value.get("version") != RAW_MM_ITEM_VERSION: - raise ValueError(f"unsupported raw multimodal descriptor version: {value.get('version')!r}") def parse_raw_mm_item(value: Any) -> RawMMItem: diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 6428571226..a907de706e 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -415,7 +415,6 @@ def _mm_refs() -> MMRefs: "image": [ { "kind": "prime_raw_mm_item", - "version": 1, "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index 886a034eb4..eec7fce482 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -3,7 +3,7 @@ import torch from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM -from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND, RAW_MM_ITEM_VERSION +from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND from prime_rl.trainer.rl.data import DataLoader from prime_rl.transport.types import MicroBatch, MMRefs from prime_rl.utils.mm import build_mm_refs @@ -26,7 +26,6 @@ def synthesize_placeholder(self, refs): def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): return { "kind": RAW_MM_ITEM_KIND, - "version": RAW_MM_ITEM_VERSION, "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, From 0b0f1fbb1625c136a7c817968da3cc7622919206 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 29 Jun 2026 22:24:28 +0000 Subject: [PATCH 14/52] Bump verifiers raw multimodal client prep --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 0dc57a1493..18b0fbea5c 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 0dc57a149390b11aa879fc3b7a0f69b4f8f59b21 +Subproject commit 18b0fbea5cc8ec37d447305f22ce8ec5c4b765ef From d18ed1072032384db6932e1ac023264a6f5c232d Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 30 Jun 2026 18:16:06 +0000 Subject: [PATCH 15/52] Bump renderers for processed multimodal output --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index aa2d44d06a..ed5b404edf 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit aa2d44d06a40ee31aaa68acd2dae072a8dcca585 +Subproject commit ed5b404edf690440b9845a974bb097d9458c308a From 0da60a8b776ac06c8bb1ccb322cb8f16310c7cc6 Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 30 Jun 2026 18:21:05 +0000 Subject: [PATCH 16/52] Bump renderers lock cleanup --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index ed5b404edf..a7953b99a9 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit ed5b404edf690440b9845a974bb097d9458c308a +Subproject commit a7953b99a96488777c6f9b0c25783e47b4ec3afb From c7deb3b5fc9deb4510ba16c21e809f9038a72e57 Mon Sep 17 00:00:00 2001 From: eligotts Date: Wed, 1 Jul 2026 17:30:16 +0000 Subject: [PATCH 17/52] Bump verifiers for main sync --- deps/verifiers | 2 +- uv.lock | 4 ++++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 22c7cf4ce8..9b3e7eee07 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 22c7cf4ce8a87fe84069b7244cd27ef80a350980 +Subproject commit 9b3e7eee07c91bdcfb8da9ee49ec99c9e0855ed0 diff --git a/uv.lock b/uv.lock index f18b38307c..ff11e87706 100644 --- a/uv.lock +++ b/uv.lock @@ -4673,10 +4673,14 @@ requires-dist = [ { name = "numpy" }, { name = "openai", specifier = ">=1.108.1" }, { name = "openai-harmony", specifier = ">=0.0.4" }, + { name = "pillow", marker = "extra == 'vision'", specifier = ">=12.2.0" }, { name = "prime-pydantic-config", specifier = ">=0.3.0.dev83" }, { name = "tiktoken" }, + { name = "torch", marker = "extra == 'vision'", specifier = ">=2.11.0" }, + { name = "torchvision", marker = "extra == 'vision'", specifier = ">=0.26.0" }, { name = "transformers", specifier = ">=4.50.0" }, ] +provides-extras = ["vision"] [package.metadata.requires-dev] dev = [ From 36b5042a007dc6875e032b14923405ef4e84856d Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 3 Jul 2026 03:50:29 +0000 Subject: [PATCH 18/52] fix: truncate raw multimodal refs safely --- docs/advanced.md | 2 +- src/prime_rl/trainer/batch.py | 30 +++++++++--- src/prime_rl/transport/types.py | 31 ++++++++++--- src/prime_rl/utils/mm.py | 67 +++++++++++++++------------ tests/unit/orchestrator/test_batch.py | 63 ++++++++++++++----------- tests/unit/utils/test_mm.py | 12 ++--- 6 files changed, 124 insertions(+), 81 deletions(-) diff --git a/docs/advanced.md b/docs/advanced.md index ff342975f8..2bbe9aa865 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -88,7 +88,7 @@ To add a new model family permanently, append an entry to `VLM_REGISTRY` in `src ### Limitations - **Vision encoder frozen by default.** Set `freeze_vision_encoder = false` to fine-tune it; in that case it's FSDP-sharded per block. The combination `freeze_vision_encoder = false` + LoRA is rejected by a config validator — LoRA freezes everything non-adapter, so unfreezing the encoder under LoRA would be a silent no-op. -- **No multimodal truncation.** Raw multimodal samples that exceed `seq_len` are rejected before packing, because truncating image placeholder spans would desynchronize tokens from the image descriptors. Set `seq_len` to cover your longest sample. +- **Truncation cuts at image boundaries.** Raw multimodal samples that exceed `seq_len` are truncated before packing at the start of the first image that doesn't fit — an image placeholder span is never split, and refs for dropped images are discarded. A sample whose leading image alone exceeds `seq_len` is rejected. - **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/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index 85c46a3480..f93f6e0621 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -3,7 +3,7 @@ from dataclasses import dataclass from prime_rl.trainer.utils import balanced_partition -from prime_rl.transport.types import MicroBatch, RoutedExperts, TrainingSample +from prime_rl.transport.types import MicroBatch, MMRefs, RoutedExperts, TrainingSample ROUTED_EXPERTS_DTYPE_ITEMSIZE = { "uint8": 1, @@ -46,6 +46,25 @@ def _pad_routed_experts(micro_batch: MicroBatch, padding_size: int) -> None: routed_experts.shape[0] += padding_size +def _truncate_mm_refs(mm_refs: MMRefs, seq_len: int) -> tuple[int, MMRefs | None]: + """Return a token cut that never splits an image placeholder, plus the surviving refs.""" + cut, kept = seq_len, 0 + for image in mm_refs.images: # token order, non-overlapping — enforced by build_mm_refs + if image.offset + image.length <= seq_len: + kept += 1 + continue + if image.offset < seq_len: + cut = image.offset + break + if cut == 0: + raise ValueError(f"Cannot truncate multimodal sample: leading image does not fit in seq_len={seq_len}") + if kept == len(mm_refs.images): + return seq_len, mm_refs + if kept == 0: + return cut, None + return cut, MMRefs(images=mm_refs.images[:kept]) + + def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch: """ Prepare a problem for sequence packing training. @@ -74,7 +93,7 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch mm_token_type_ids = training_example.mm_token_type_ids if training_example.mm_kwargs is not None: raise ValueError("Processed multimodal mm_kwargs are unsupported in v1; use raw mm_refs") - mm_refs = copy.deepcopy(training_example.mm_refs) + mm_refs = training_example.mm_refs assert training_example.env_name != "all", "env_name='all' is reserved for aggregate metric keys" env_names = [training_example.env_name] * len(input_ids) @@ -89,12 +108,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ) if len(input_ids) > seq_len: - if mm_refs is not None: - raise ValueError( - "Multimodal samples cannot be truncated after raw image offload. " - f"Got {len(input_ids)} tokens with seq_len={seq_len}." - ) cut = seq_len + if mm_refs is not None: + cut, mm_refs = _truncate_mm_refs(mm_refs, seq_len) input_ids = input_ids[:cut] loss_mask = loss_mask[:cut] inference_logprobs = inference_logprobs[:cut] diff --git a/src/prime_rl/transport/types.py b/src/prime_rl/transport/types.py index 5b8b03358a..9169ea8529 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -18,17 +18,34 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru dtype: str +class MMImageRef(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): + """One raw image reference for trainer-side materialization. + + ``item`` is the JSON-safe renderer descriptor (adapter family + payload, + parseable via ``parse_raw_mm_item``), ``hash``/``uri`` identify the raw image + file, and ``offset``/``length`` are the image's placeholder range in token + space (needed to truncate at image boundaries). + """ + + item: dict + hash: str + uri: str + offset: int + length: int + + class MMRefs(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): - """Raw multimodal sidecar references for one sample. + """Raw multimodal sidecar references for one sample, in token order. - ``descriptor`` carries JSON-safe renderer metadata (hashes and adapter - payloads). ``uris`` carries the raw image files that the trainer materializes - with its own processor. Processed tensors are intentionally not part of this - transport. + The trainer materializes the referenced image files with its own processor. + Processed tensors are intentionally not part of this transport. """ - descriptor: dict - uris: list[str] + images: list[MMImageRef] + + @property + def uris(self) -> list[str]: + return [image.uri for image in self.images] # Orchestrator -> Packer diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index 38010de684..cbdcff38c8 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -10,7 +10,7 @@ from prime_rl.multimodal.adapters.base import MaterializedMM, MultimodalAdapter from prime_rl.multimodal.registry import get_multimodal_adapter from prime_rl.multimodal.schema import RawMMItem, parse_raw_mm_item -from prime_rl.transport.types import MMRefs +from prime_rl.transport.types import MMImageRef, MMRefs IMAGE_MODALITY = "image" SUPPORTED_MODALITIES = {IMAGE_MODALITY} @@ -65,6 +65,16 @@ def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> tuple[list[dict[str, return item_dicts, uris +def _placeholder_bounds(placeholder: Any) -> tuple[int, int]: + offset = _field(placeholder, "offset") + length = _field(placeholder, "length") + if not isinstance(offset, int) or not isinstance(length, int): + raise ValueError(f"Raw image placeholder must have integer offset/length, got {placeholder!r}") + if offset < 0 or length <= 0: + raise ValueError(f"Raw image placeholder must have offset >= 0 and length > 0, got {placeholder!r}") + return offset, length + + def build_mm_refs(multi_modal_data: Any) -> MMRefs | None: mm_items = _field(multi_modal_data, "mm_items", None) if not mm_items: @@ -77,36 +87,33 @@ def build_mm_refs(multi_modal_data: Any) -> MMRefs | None: mm_hashes = _field(multi_modal_data, "mm_hashes", {}) or {} image_hashes = list(mm_hashes.get(IMAGE_MODALITY, [])) - if len(image_hashes) != len(image_items): + mm_placeholders = _field(multi_modal_data, "mm_placeholders", {}) or {} + image_placeholders = list(mm_placeholders.get(IMAGE_MODALITY, [])) + if len(image_hashes) != len(image_items) or len(image_placeholders) != len(image_items): raise ValueError( - "Raw image descriptor/hash mismatch: " - f"{len(image_items)} image descriptors but {len(image_hashes)} image hashes" + "Raw image descriptor/hash/placeholder mismatch: " + f"descriptors={len(image_items)}, hashes={len(image_hashes)}, placeholders={len(image_placeholders)}" ) - return MMRefs( - descriptor={ - "mm_items": {IMAGE_MODALITY: image_items}, - "mm_hashes": {IMAGE_MODALITY: image_hashes}, - }, - uris=uris, - ) + images: list[MMImageRef] = [] + prev_end = 0 + for item, uri, image_hash, placeholder in zip(image_items, uris, image_hashes, image_placeholders, strict=True): + offset, length = _placeholder_bounds(placeholder) + # Truncation cuts at a prefix of ``images``, which is only sound if + # placeholders arrive in token order without overlap. + if offset < prev_end: + raise ValueError(f"Raw image placeholders must be sorted and non-overlapping, got {image_placeholders!r}") + prev_end = offset + length + images.append(MMImageRef(item=item, hash=image_hash, uri=uri, offset=offset, length=length)) + return MMRefs(images=images) def sha256_32(data: bytes) -> str: return hashlib.sha256(data).hexdigest()[:32] -def _parse_image_refs(refs: MMRefs) -> tuple[list[RawMMItem], list[str]]: - image_item_dicts = refs.descriptor.get("mm_items", {}).get(IMAGE_MODALITY, []) - image_hashes = list(refs.descriptor.get("mm_hashes", {}).get(IMAGE_MODALITY, [])) - if not image_item_dicts: - return [], [] - if len(refs.uris) != len(image_item_dicts) or len(image_hashes) != len(image_item_dicts): - raise ValueError( - "Raw image refs must have matching URI, descriptor, and hash counts " - f"(uris={len(refs.uris)}, descriptors={len(image_item_dicts)}, hashes={len(image_hashes)})" - ) - return [parse_raw_mm_item(item) for item in image_item_dicts], image_hashes +def _parse_image_refs(refs: MMRefs) -> list[RawMMItem]: + return [parse_raw_mm_item(image.item) for image in refs.images] def _single_family_adapter(items: list[RawMMItem]) -> MultimodalAdapter: @@ -125,15 +132,15 @@ def _validate_processor_layout(adapter: MultimodalAdapter, image_processor: Any, ) -def _load_verified_images(uris: list[str], expected_hashes: list[str]) -> list[Any]: +def _load_verified_images(image_refs: list[MMImageRef]) -> list[Any]: from PIL import Image images = [] - for uri, expected_hash in zip(uris, expected_hashes, strict=True): - raw = file_uri_to_path(uri).read_bytes() + for ref in image_refs: + raw = file_uri_to_path(ref.uri).read_bytes() actual_hash = sha256_32(raw) - if actual_hash != expected_hash: - raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") + if actual_hash != ref.hash: + raise ValueError(f"Raw image hash mismatch for {ref.uri}: expected {ref.hash}, got {actual_hash}") with Image.open(BytesIO(raw)) as image: images.append(image.convert("RGB")) return images @@ -160,19 +167,19 @@ def image_processor(self): return self._image_processor def materialize(self, refs: MMRefs) -> MaterializedMM | None: - image_items, image_hashes = _parse_image_refs(refs) + image_items = _parse_image_refs(refs) if not image_items: return None image_processor = self.image_processor adapter = _single_family_adapter(image_items) _validate_processor_layout(adapter, image_processor, image_items) - images = _load_verified_images(refs.uris, image_hashes) + images = _load_verified_images(refs.images) return adapter.materialize_for_trainer(image_processor, image_items, images) def synthesize_placeholder(self, refs: MMRefs) -> MaterializedMM | None: """Build zero-valued multimodal tensors via the owning adapter.""" - image_items, _ = _parse_image_refs(refs) + image_items = _parse_image_refs(refs) if not image_items: return None image_processor = self.image_processor diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index a907de706e..c45557b6b4 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -5,7 +5,7 @@ from prime_rl.trainer.batch import pad_micro_batch, prepare_batch, prepare_sample from prime_rl.trainer.utils import build_bin_cost -from prime_rl.transport.types import MicroBatch, MMRefs, RoutedExperts, TrainingSample +from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, TrainingSample def _routed_experts(data, dtype=np.uint8): @@ -408,41 +408,48 @@ def test_prepare_sample_truncates_routed_experts(): assert micro_batch.env_names == ["test-env"] * 3 -def _mm_refs() -> MMRefs: - return MMRefs( - descriptor={ - "mm_items": { - "image": [ - { - "kind": "prime_raw_mm_item", - "modality": "image", - "family": "qwen_vl", - "layout_fingerprint": "f" * 32, - "raw_image_uri": "file:///tmp/image.png", - "payload": {"image_grid_thw": [[1, 1, 1]]}, - } - ] - }, - "mm_hashes": {"image": ["a" * 32]}, +def _image_ref(uri: str, offset: int, length: int) -> MMImageRef: + return MMImageRef( + item={ + "kind": "prime_raw_mm_item", + "modality": "image", + "family": "qwen_vl", + "layout_fingerprint": "f" * 32, + "raw_image_uri": uri, + "payload": {"image_grid_thw": [[1, 1, 1]]}, }, - uris=["file:///tmp/image.png"], + hash="a" * 32, + uri=uri, + offset=offset, + length=length, ) -def test_prepare_sample_rejects_overlong_raw_mm_refs(): +def test_prepare_sample_truncates_raw_mm_refs_at_image_boundary(): + first_image = _image_ref("file:///tmp/image-0.png", offset=1, length=2) + second_image = _image_ref("file:///tmp/image-1.png", offset=4, length=2) sample = TrainingSample( - token_ids=[10, 11, 12, 13], - mask=[False, False, True, True], - logprobs=[0.0] * 4, - temperatures=[1.0] * 4, - advantages=[0.0, 0.0, 1.0, 1.0], + token_ids=[10, 11, 12, 13, 14, 15, 16], + mask=[False, False, True, True, False, True, True], + logprobs=[0.0] * 7, + temperatures=[1.0] * 7, + advantages=[0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 1.0], env_name="test-env", - mm_token_type_ids=[0, 1, 1, 0], - mm_refs=_mm_refs(), + mm_token_type_ids=[0, 1, 1, 0, 1, 1, 0], + mm_refs=MMRefs(images=[first_image, second_image]), ) - with pytest.raises(ValueError, match="Multimodal samples cannot be truncated"): - prepare_sample(sample, seq_len=3) + # seq_len=5 splits the second image: cut at its start, keep refs for the first. + micro_batch = prepare_sample(sample, seq_len=5) + assert micro_batch.input_ids == [10, 11, 12, 13] + assert micro_batch.mm_token_type_ids == [0, 1, 1, 0] + assert micro_batch.mm_refs == MMRefs(images=[first_image]) + + # seq_len=2 splits the first image: no image survives. + micro_batch = prepare_sample(sample, seq_len=2) + assert micro_batch.input_ids == [10] + assert micro_batch.mm_token_type_ids == [0] + assert micro_batch.mm_refs is None def test_prepare_sample_none_routed_experts(): diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index eec7fce482..59ab7092de 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -5,7 +5,7 @@ from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND from prime_rl.trainer.rl.data import DataLoader -from prime_rl.transport.types import MicroBatch, MMRefs +from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs from prime_rl.utils.mm import build_mm_refs @@ -36,11 +36,7 @@ def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): def _refs(uri: str = "file:///tmp/missing-image.png") -> MMRefs: return MMRefs( - descriptor={ - "mm_items": {"image": [_qwen_item([[1, 1, 1]], uri)]}, - "mm_hashes": {"image": ["a" * 32]}, - }, - uris=[uri], + images=[MMImageRef(item=_qwen_item([[1, 1, 1]], uri), hash="a" * 32, uri=uri, offset=1, length=1)], ) @@ -77,12 +73,12 @@ def test_build_mm_refs_accepts_offloaded_file_uri(tmp_path): multi_modal_data = SimpleNamespace( mm_items={"image": [_qwen_item([[1, 1, 1]], image_path.as_uri())]}, mm_hashes={"image": ["a" * 32]}, + mm_placeholders={"image": [SimpleNamespace(offset=1, length=1)]}, ) refs = build_mm_refs(multi_modal_data) - assert refs is not None - assert refs.uris == [image_path.as_uri()] + assert refs == _refs(image_path.as_uri()) def test_dataloader_uses_zero_loss_placeholder_for_missing_raw_image(): From 53d9ed8b0158712f5d5026afdfa27e2d333097d2 Mon Sep 17 00:00:00 2001 From: eligotts Date: Sat, 4 Jul 2026 19:53:01 +0000 Subject: [PATCH 19/52] Tighten raw multimodal validation and trainer mm plumbing - Kimi adapter reads the actual processor layout instead of hard-equating it to the renderer constant; the fingerprint comparison is the single drift detector, matching the Qwen adapter. - Orchestrator validates every image ref's placeholder span lands on image-typed tokens before a sample ships, so offset drift anywhere upstream fails loudly instead of silently truncating wrong. - FakeDataLoader carries the mm stat counters; trainer metrics read them directly. Drop the duplicate apply_run_asset_env in the entrypoint. Co-Authored-By: Claude Fable 5 --- src/prime_rl/entrypoints/orchestrator.py | 2 -- src/prime_rl/multimodal/adapters/kimi_k25.py | 12 ++++-------- src/prime_rl/orchestrator/trajectories.py | 19 +++++++++++++++++++ src/prime_rl/trainer/rl/data.py | 3 +++ src/prime_rl/trainer/rl/train.py | 6 +++--- 5 files changed, 29 insertions(+), 13 deletions(-) diff --git a/src/prime_rl/entrypoints/orchestrator.py b/src/prime_rl/entrypoints/orchestrator.py index 842ddb01f5..48a8a354f8 100644 --- a/src/prime_rl/entrypoints/orchestrator.py +++ b/src/prime_rl/entrypoints/orchestrator.py @@ -14,13 +14,11 @@ from prime_rl.configs.orchestrator import OrchestratorConfig from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title -from prime_rl.utils.run_assets import apply_run_asset_env def main(): set_proc_title("Orchestrator") config = cli(OrchestratorConfig) - apply_run_asset_env(config.output_dir, config.multimodal) from prime_rl.orchestrator.orchestrator import run_orchestrator asyncio.run(run_orchestrator(config)) diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index 39923a2c9f..4a2aba7aa9 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -4,7 +4,7 @@ from collections.abc import Mapping from typing import Any -from renderers.kimi_k25 import KIMI_K25_IMAGE_LAYOUT, KimiK25ImageLayoutSpec +from renderers.kimi_k25 import KimiK25ImageLayoutSpec from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem @@ -44,8 +44,10 @@ def _float_triple(value: Any, *, name: str) -> tuple[float, float, float]: def _processor_layout(image_processor: Any) -> KimiK25ImageLayoutSpec: + """Read the actual processor's layout; drift from the renderer's baked + layout surfaces as a fingerprint mismatch at materialization.""" cfg = _media_proc_cfg(image_processor) - layout = KimiK25ImageLayoutSpec( + return KimiK25ImageLayoutSpec( patch_size=int(_required_cfg(cfg, "patch_size")), merge_kernel_size=int(_required_cfg(cfg, "merge_kernel_size")), in_patch_limit=int(_required_cfg(cfg, "in_patch_limit")), @@ -54,12 +56,6 @@ def _processor_layout(image_processor: Any) -> KimiK25ImageLayoutSpec: image_mean=_float_triple(_required_cfg(cfg, "image_mean"), name="image_mean"), image_std=_float_triple(_required_cfg(cfg, "image_std"), name="image_std"), ) - if layout != KIMI_K25_IMAGE_LAYOUT: - raise ValueError( - "Kimi image processor layout does not match renderer layout contract: " - f"expected {KIMI_K25_IMAGE_LAYOUT}, got {layout}" - ) - return layout def _grid_payload(item: RawMMItem) -> list[int]: diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 9dc555ccd9..5240c68339 100644 --- a/src/prime_rl/orchestrator/trajectories.py +++ b/src/prime_rl/orchestrator/trajectories.py @@ -41,6 +41,23 @@ def _encode_routed_experts(arr: np.ndarray | None, num_tokens: int) -> RoutedExp return RoutedExperts(data=arr.tobytes(), shape=list(arr.shape), dtype=str(arr.dtype)) +def _validate_image_spans(mm_refs: MMRefs, mm_token_type_ids: list[int]) -> None: + """Every image ref's placeholder span must land on image-typed tokens. + + Placeholder offsets flow through the renderer, bridge extension, and node + attribution before arriving here — and the trainer truncates on them — so + any drift anywhere upstream must fail loudly before the sample ships. + """ + for image in mm_refs.images: + span = mm_token_type_ids[image.offset : image.offset + image.length] + if len(span) != image.length or any(t != 1 for t in span): + raise ValueError( + f"Raw image placeholder [{image.offset}, {image.offset + image.length}) does not " + f"cover image-typed tokens (branch length {len(mm_token_type_ids)}) — placeholder " + "offsets have drifted from the branch token stream" + ) + + def trace_to_samples( trace: vf.Trace, *, @@ -72,6 +89,8 @@ def trace_to_samples( mm_refs = build_mm_refs(mmd) mapping = mm_token_type_ids_mapping or {} mm_token_type_ids = [mapping.get(t, 0) for t in token_ids] + if mm_refs is not None and mapping: + _validate_image_spans(mm_refs, mm_token_type_ids) samples.append( TrainingSample( token_ids=token_ids, diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 2c9e9287e9..3fb0640d84 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -77,6 +77,9 @@ def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: in self.generate_samples = config.generate_samples self.batch_counter = 0 self.multi_run_manager = get_multi_run_manager() + self.last_mm_materialize_time = 0.0 + self.last_mm_images_materialized = 0 + self.last_mm_images_placeholdered = 0 def wait_for_batch(self) -> None: return diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 11d7ff3669..98f3a88ca4 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -688,12 +688,12 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: "time/step": step_time, "time/wait_for_batch": wait_for_batch_time, "time/load_data": load_data_time, - "time/mm_materialize": getattr(dataloader, "last_mm_materialize_time", 0.0), + "time/mm_materialize": dataloader.last_mm_materialize_time, "time/broadcast_weights": broadcast_weights_time, "time/save_ckpt": save_ckpt_time, "time/forward_backward": forward_backward_time, - "mm/images_materialized": getattr(dataloader, "last_mm_images_materialized", 0), - "mm/images_placeholdered": getattr(dataloader, "last_mm_images_placeholdered", 0), + "mm/images_materialized": dataloader.last_mm_images_materialized, + "mm/images_placeholdered": dataloader.last_mm_images_placeholdered, "step": progress.step, } monitor.log(time_metrics, step=progress.step) From a13a2141102cc307d10637991040e38f568149fa Mon Sep 17 00:00:00 2001 From: eligotts Date: Sat, 4 Jul 2026 19:53:01 +0000 Subject: [PATCH 20/52] Bump renderers and verifiers for raw multimodal hardening renderers: bridge sidecar aliasing fix, HF smart_resize import, full-hash asset filenames, layout parity tests. verifiers: ingress offload covers every image part shape. Co-Authored-By: Claude Fable 5 --- deps/renderers | 2 +- deps/verifiers | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deps/renderers b/deps/renderers index a7953b99a9..e3c12e96a5 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit a7953b99a96488777c6f9b0c25783e47b4ec3afb +Subproject commit e3c12e96a5c75ec1ee511126619d719228a22b16 diff --git a/deps/verifiers b/deps/verifiers index 9b3e7eee07..2c2824ae4e 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 9b3e7eee07c91bdcfb8da9ee49ec99c9e0855ed0 +Subproject commit 2c2824ae4ef494316dd23b0a0dc475a841ad09fa From 1f290dc32cde62d010d81cfae9e73a74ed8bbacb Mon Sep 17 00:00:00 2001 From: eligotts Date: Sun, 5 Jul 2026 20:15:48 +0000 Subject: [PATCH 21/52] Document multimodal offload monitoring in the monitor-run skill Co-Authored-By: Claude Fable 5 --- skills/training/monitor-run/SKILL.md | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 078f62188b..24a81b69c3 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -154,6 +154,32 @@ A few warnings are normal. Escalate when errors are persistent, growing, or hit - **Trainer**: NCCL/CUDA errors, OOM, NaN loss or gradients. - **Inference**: NCCL/CUDA errors, OOM, request timeouts. +### Multimodal image offload checks + +v1 multimodal RL offloads images exactly once, at verifiers ingress: every image +content part is rewritten to a `file://` asset under the run image directory +before rendering, and renderers/inference/trainer all work from those refs. A +`data:` URL reaching a renderer ("requires offloaded file:// image assets") +means ingress was bypassed. + +- The image directory comes from `[multimodal].offload_dir` in the resolved + config; unset, it defaults to a run-scoped path (`{output_dir}/assets/images` + or the hosted `RUN_ID` path). The launcher exports it to the orchestrator as + `VF_RENDERER_IMAGE_OFFLOAD_DIR` (protected — not settable via `env_vars`). +- While multimodal rollouts are in flight, image files should accumulate under + that directory. Zero files means the offload path isn't being exercised or + image preparation failed before request submission. +- Inference rejects bad refs with `invalid_mm_image_ref` 400s (hash mismatch, + fingerprint mismatch, unreadable asset) — grep the inference log. +- The orchestrator raises on placeholder/token drift ("does not cover + image-typed tokens") before a sample ships — treat any occurrence as a bug, + not noise. +- Trainer metrics: `mm/images_materialized`, `time/mm_materialize`, and + `mm/images_placeholdered`. A nonzero placeholder count means image files + disappeared before materialization (the batch trains with zero-loss + placeholders and logs "raw image materialization missing image(s)") — check + whether something cleaned the offload directory mid-run. + ### Process tree All processes use `setproctitle` so they're visible in `ps`/`htop`/`pstree`: From a3f52519ed50fe7e5e8204ad015ee0cdf7ac218e Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 23 Jul 2026 07:21:33 +0000 Subject: [PATCH 22/52] Bump renderers and verifiers for main reconciliation Re-pins to the commits where each companion branch merged its own origin/main (renderers e64cc58, verifiers 2b1627d03), ahead of merging main into this branch. --- deps/renderers | 2 +- deps/verifiers | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/deps/renderers b/deps/renderers index a49e0fc254..e64cc58803 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit a49e0fc25472d17a98cd1b49b67ce21e2718d27c +Subproject commit e64cc588037a75ff99d37c2968329805147d820d diff --git a/deps/verifiers b/deps/verifiers index 9bc3cc3616..2b1627d031 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 9bc3cc36161b46b5382822e4d72099351e95eee6 +Subproject commit 2b1627d031f13ec825a9d61dfa308412a1828ac0 From a2fdf6a88f582aefbb4d5f59487f77743aca65b5 Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 23 Jul 2026 17:06:02 +0000 Subject: [PATCH 23/52] Default SFT renderers to processed multimodal output MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The renderers library defaults multimodal_output to 'raw' (built for the RL offload path), but SFT's data path consumes processed pixel tensors straight from the renderer and has no raw-ref materializer — the library default would silently ship JSON descriptors into training. SFTConfig now defaults its renderer to 'processed' and rejects an explicit 'raw'. --- .../src/prime_rl/configs/sft.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) 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 941c531deb..8308b7130e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -248,6 +248,23 @@ def normalize_deployment(cls, data): ### Validate configs (e.g. raise for unsupported (combinations of) configs) + @model_validator(mode="after") + def renderer_emits_processed_multimodal(self): + """SFT consumes processed pixel tensors straight from the renderer — it has no + raw-ref materializer, so the renderers-library default of ``multimodal_output='raw'`` + (built for the RL offload path) would silently ship JSON descriptors into training. + Default SFT renderers to ``'processed'`` and reject an explicit ``'raw'``.""" + if "multimodal_output" in self.renderer.model_fields_set: + if self.renderer.multimodal_output == "raw": + raise ValueError( + "multimodal_output='raw' is unsupported for SFT: the SFT data path materializes " + "images from processed renderer output, not raw refs. Remove the override " + "(SFT defaults to 'processed')." + ) + else: + self.renderer = self.renderer.model_copy(update={"multimodal_output": "processed"}) + return self + @model_validator(mode="after") def deepep_disables_grad_clipping(self): if self.model.ep_comm_backend == "deepep" and self.optim.max_norm is not None: From 76761553529bb308f2b3c3520a7fbac8767ff7e0 Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 23 Jul 2026 17:11:08 +0000 Subject: [PATCH 24/52] Cache materialized raw image refs in the serving layer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Every request carries a raw ref for every image in its prompt (prior turns included), so a 20-turn rollout with 5 accumulated images paid ~100 materializations (read + sha256 + PIL decode + HF processor forward) where 5 would do — multiplied by group size hammering the same prompt images. vLLM's processor cache can't help: production happens in our handler before vLLM sees the request. Adds a byte-bounded, single-flight LRU in serving_tokens.py keyed by (raw_ref, expected_placeholder_length, processor_model_name) — content- addressed, so hits are sound and evicted entries can only go cold, never stale. _decode_raw_mm_kwargs now gathers all of a request's images concurrently instead of the sequential await loop, and single-flight collapses concurrent misses for one new image into one materialization. Failures are never cached; they propagate to all awaiters and the next request retries cleanly. One knob: PRIME_RL_MM_MATERIALIZE_CACHE_GB (default 2.0). 0 disables all bookkeeping and single-flight — byte-identical to the uncached path, and the kill switch. Logs hits/misses/bytes/evictions every 1000 lookups; monitor-run skill documents the signature. --- skills/training/monitor-run/SKILL.md | 6 + src/prime_rl/inference/vllm/serving_tokens.py | 162 ++++++++++++++++-- tests/unit/inference/test_serving_tokens.py | 143 ++++++++++++++++ 3 files changed, 296 insertions(+), 15 deletions(-) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 49bb312c37..fa62ef42ab 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -182,6 +182,12 @@ means ingress was bypassed. image preparation failed before request submission. - Inference rejects bad refs with `invalid_mm_image_ref` 400s (hash mismatch, fingerprint mismatch, unreadable asset) — grep the inference log. +- Inference caches materialized refs and logs + `mm materialize cache: hits=X misses=Y hit_rate=Z% bytes=A/B evictions=C` + every 1000 lookups. Hit rate should climb after turn 1 of multi-turn + multimodal rollouts; a stuck-at-zero hit rate with repeat images means the + cache is disabled or thrashing (sized by `PRIME_RL_MM_MATERIALIZE_CACHE_GB`, + default 2.0, `0` disables). - The orchestrator raises on placeholder/token drift ("does not cover image-typed tokens") before a sample ships — treat any occurrence as a bug, not noise. diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 2a95f1398a..efc6ea435a 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -34,7 +34,10 @@ import asyncio import hashlib -from collections.abc import AsyncGenerator, AsyncIterable +import logging +import os +from collections import OrderedDict +from collections.abc import AsyncGenerator, AsyncIterable, Callable from dataclasses import dataclass from functools import cached_property, lru_cache from http import HTTPStatus @@ -55,6 +58,7 @@ ) from vllm.entrypoints.serve.disagg.serving import ServingTokens from vllm.entrypoints.serve.utils.api_utils import get_max_tokens +from vllm.multimodal.cache import MultiModalCache from vllm.outputs import RequestOutput from vllm.sampling_params import RequestOutputKind, SamplingParams @@ -63,6 +67,8 @@ from prime_rl.multimodal.schema import RawMMItem from prime_rl.utils.mm import file_uri_to_path +logger = logging.getLogger(__name__) + @dataclass class _MMImageRefError(Exception): @@ -245,6 +251,101 @@ def _materialize_raw_image_ref_sync( ) +# (raw_ref, expected_placeholder_length, processor_model_name). The ref string is +# content-addressed (it embeds the image hash, layout fingerprint, family, and URI), +# so identical keys describe byte-identical materialization work. +_MaterializeKey = tuple[str, int, str] + +_MM_MATERIALIZE_CACHE_GB_ENV = "PRIME_RL_MM_MATERIALIZE_CACHE_GB" +_MM_MATERIALIZE_LOG_EVERY = 1000 + + +class _MaterializedRefCache: + """Byte-bounded LRU of materialized raw image refs, with single-flight misses. + + Every request carries a ref for every image in its prompt (prior turns included), + so multi-turn rollouts re-materialize the same images once per turn per rollout — + this cache turns those repeats into lookups. Keys are content-addressed, so + entries can only go cold, never stale: a post-eviction request simply misses and + re-materializes from the durable ``file://`` asset. + + All mutation happens on the event loop thread (materialization itself runs in a + worker thread via ``asyncio.to_thread``, but lookup/insertion/eviction happen in + the async caller) — so no lock is needed. Do not touch this cache from sync code. + + ``max_bytes == 0`` disables caching and single-flight entirely: every call runs + the materializer, byte-identical to the uncached path. + + Host-RAM note: this budget is additive to vLLM's own processor cache + (``mm_processor_cache_gb × (api_server_count + data_parallel_size)``). + """ + + def __init__(self, max_bytes: int) -> None: + self.max_bytes = max_bytes + self._items: OrderedDict[_MaterializeKey, tuple[Any, int]] = OrderedDict() + self._inflight: dict[_MaterializeKey, asyncio.Future] = {} + self._bytes = 0 + self.hits = 0 + self.misses = 0 + self.evictions = 0 + + async def get_or_materialize(self, key: _MaterializeKey, materialize_fn: Callable[[], Any]) -> Any: + if self.max_bytes == 0: + return await asyncio.to_thread(materialize_fn) + if key in self._items: + self._items.move_to_end(key) + self.hits += 1 + self._maybe_log() + return self._items[key][0] + if (inflight := self._inflight.get(key)) is not None: + self.hits += 1 + self._maybe_log() + return await inflight + self.misses += 1 + self._maybe_log() + future: asyncio.Future = asyncio.get_running_loop().create_future() + self._inflight[key] = future + try: + item = await asyncio.to_thread(materialize_fn) + except BaseException as exc: + # Never cache a failure: the exception propagates to every awaiter and + # the next request for this key retries cleanly. + future.set_exception(exc) + # A never-awaited errored future logs a spurious warning at GC. + future.exception() + del self._inflight[key] + raise + future.set_result(item) + del self._inflight[key] + self._insert(key, item) + return item + + def _insert(self, key: _MaterializeKey, item: Any) -> None: + nbytes = MultiModalCache.get_item_size(item) + if nbytes > self.max_bytes: + # Don't churn the whole cache for one oversized entry. + return + self._items[key] = (item, nbytes) + self._bytes += nbytes + while self._bytes > self.max_bytes: + _, (_, evicted_bytes) = self._items.popitem(last=False) + self._bytes -= evicted_bytes + self.evictions += 1 + + def _maybe_log(self) -> None: + total = self.hits + self.misses + if total % _MM_MATERIALIZE_LOG_EVERY == 0: + logger.info( + "mm materialize cache: hits=%d misses=%d hit_rate=%.1f%% bytes=%d/%d evictions=%d", + self.hits, + self.misses, + 100.0 * self.hits / total, + self._bytes, + self.max_bytes, + self.evictions, + ) + + def _raw_ref_payloads_for_feature( features: Any, feature_modality: str, hashes: list[str] ) -> tuple[list[str], list[Any]]: @@ -276,30 +377,60 @@ async def _decode_raw_mm_kwargs( *, processor_model_name: str, trust_remote_code: bool, + cache: _MaterializedRefCache, ) -> dict[str, list[Any]]: - mm_kwargs: dict[str, list[Any]] = {} + # Flatten across modalities so every image in the request materializes + # concurrently; single-flight in the cache dedupes identical refs across + # (and within) requests. Fresh lists per request — vLLM's cache-injection + # path replaces list elements in place, so never hand it cache-owned lists. + flat: list[tuple[str, str, str, Any]] = [] for feature_modality, hashes in features.mm_hashes.items(): raw_refs, placeholders = _raw_ref_payloads_for_feature(features, feature_modality, hashes) - decoded: list[Any] = [] - for raw_ref, mm_hash, placeholder in zip(raw_refs, hashes, placeholders, strict=True): - decoded.append( - await asyncio.to_thread( - _materialize_raw_image_ref_sync, - raw_ref, - feature_modality=feature_modality, - mm_hash=mm_hash, - expected_placeholder_length=placeholder.length, - processor_model_name=processor_model_name, - trust_remote_code=trust_remote_code, - ) + flat.extend(zip([feature_modality] * len(hashes), raw_refs, hashes, placeholders, strict=True)) + + def _materialize(feature_modality: str, raw_ref: str, mm_hash: str, placeholder: Any): + return _materialize_raw_image_ref_sync( + raw_ref, + feature_modality=feature_modality, + mm_hash=mm_hash, + expected_placeholder_length=placeholder.length, + processor_model_name=processor_model_name, + trust_remote_code=trust_remote_code, + ) + + decoded = await asyncio.gather( + *( + cache.get_or_materialize( + (raw_ref, placeholder.length, processor_model_name), + lambda fm=feature_modality, r=raw_ref, h=mm_hash, p=placeholder: _materialize(fm, r, h, p), ) - mm_kwargs[feature_modality] = decoded + for feature_modality, raw_ref, mm_hash, placeholder in flat + ) + ) + + mm_kwargs: dict[str, list[Any]] = {feature_modality: [] for feature_modality in features.mm_hashes} + for (feature_modality, _, _, _), item in zip(flat, decoded, strict=True): + mm_kwargs[feature_modality].append(item) return mm_kwargs class PrimeRlServingTokens(ServingTokens): """ServingTokens + DP-rank routing + routed experts + raw image refs + max_tokens defaulting.""" + @cached_property + def _mm_materialize_cache(self) -> _MaterializedRefCache: + """Materialized-ref cache, sized by ``PRIME_RL_MM_MATERIALIZE_CACHE_GB`` + (float GiB, default 2.0, ``0`` disables — the kill switch: disabled is + byte-identical to the uncached path). + + A ``cached_property`` because ``custom_init_app_state`` grafts this + subclass via ``object.__new__`` + ``__dict__.update``, so ``__init__`` + never runs (see ``_max_tokens_defaults``). One frontend process serves + all DP ranks of an engine, so the cache is shared across ranks. + """ + gb = float(os.environ.get(_MM_MATERIALIZE_CACHE_GB_ENV, "2.0")) + return _MaterializedRefCache(max_bytes=int(gb * (1 << 30))) + @cached_property def _max_tokens_defaults(self) -> tuple[dict, int | None]: """Server-side ``max_tokens`` defaulting inputs, mirroring upstream ``ServingTokens``. @@ -363,6 +494,7 @@ async def serve_tokens( features, processor_model_name=processor_model_name, trust_remote_code=trust_remote_code, + cache=self._mm_materialize_cache, ) except _MMImageRefError as exc: return self.create_error_response( diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index f3e3ea9eba..5d3602dff4 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -328,3 +328,146 @@ def _get_adapter(family): assert item.layout_fingerprint == fingerprint assert item.raw_image_uri == image_path.as_uri() assert item.payload == {"adapter_owned": [1, 2, 3]} + + +def _mm_features(tmp_path, *, image_name: str = "image.png", placeholder_length: int = 7): + """A minimal ``GenerateRequest.features``-shaped object carrying one real raw ref.""" + from types import SimpleNamespace + + from PIL import Image + from renderers.mm_store import raw_mm_ref + + image_dir = tmp_path / "assets" / "images" + image_dir.mkdir(parents=True, exist_ok=True) + image_path = image_dir / image_name + Image.new("RGB", (8, 6), color=(len(image_name) % 255, 64, 128)).save(image_path) + + mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] + raw_ref = raw_mm_ref( + family="test_family", + fingerprint="f" * 32, + modality="image", + mm_hash=mm_hash, + raw_image_uri=image_path.as_uri(), + payload={}, + ) + return SimpleNamespace( + mm_hashes={"image": [mm_hash]}, + kwargs_data={"image": [raw_ref]}, + mm_placeholders={"image": [SimpleNamespace(offset=0, length=placeholder_length)]}, + ) + + +def _patch_adapter(monkeypatch, calls: list, materialize=None): + from prime_rl.inference.vllm import serving_tokens + + class _Adapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + calls.append(item.raw_ref) + if materialize is not None: + return materialize(item) + return {"materialized": item.raw_image_uri} + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: object()) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _Adapter()) + + +def test_mm_materialize_cache_hit_skips_work(tmp_path, monkeypatch): + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + features = _mm_features(tmp_path) + + async def _run(): + first = await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + second = await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + return first, second + + first, second = asyncio.run(_run()) + assert len(calls) == 1 + assert first["image"][0] is second["image"][0] + assert (cache.hits, cache.misses) == (1, 1) + + +def test_mm_materialize_cache_byte_budget_evicts_oldest(tmp_path, monkeypatch): + from vllm.multimodal.cache import MultiModalCache + + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + monkeypatch.setattr(MultiModalCache, "get_item_size", classmethod(lambda _cls, _item: 60)) + cache = serving_tokens._MaterializedRefCache(max_bytes=100) + features_a = _mm_features(tmp_path, image_name="a.png") + features_b = _mm_features(tmp_path, image_name="b.png") + + async def _decode(features): + return await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + + asyncio.run(_decode(features_a)) + asyncio.run(_decode(features_b)) # 60 + 60 > 100: evicts a + assert cache.evictions == 1 + asyncio.run(_decode(features_a)) # miss again: re-materializes + assert len(calls) == 3 + assert cache.misses == 3 + + +def test_mm_materialize_cache_failure_not_cached(tmp_path, monkeypatch): + import pytest + + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + fail_first = {"remaining": 1} + + def _materialize(item): + if fail_first["remaining"]: + fail_first["remaining"] -= 1 + raise serving_tokens._MMImageRefError("transient failure") + return {"materialized": True} + + _patch_adapter(monkeypatch, calls, materialize=_materialize) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + features = _mm_features(tmp_path) + + async def _decode(): + return await serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ) + + with pytest.raises(serving_tokens._MMImageRefError): + asyncio.run(_decode()) + out = asyncio.run(_decode()) + assert out["image"][0] == {"materialized": True} + assert len(calls) == 2 + + +def test_mm_materialize_cache_single_flight(tmp_path, monkeypatch): + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + features = _mm_features(tmp_path) + + async def _run(): + return await asyncio.gather( + serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ), + serving_tokens._decode_raw_mm_kwargs( + features, processor_model_name="model", trust_remote_code=False, cache=cache + ), + ) + + first, second = asyncio.run(_run()) + assert len(calls) == 1 + assert first["image"][0] is second["image"][0] From d00ddacb7d0275f9ef6f11cd01037d7b1c714028 Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 23 Jul 2026 17:24:11 +0000 Subject: [PATCH 25/52] Bump verifiers: drop the orphaned prepare_messages hook --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index 2b1627d031..d9e79d6ca2 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 2b1627d031f13ec825a9d61dfa308412a1828ac0 +Subproject commit d9e79d6ca215e75a269b5ecc6773f2afcf22e870 From e0c22ee5eee01ebcf9f0c1539021ea0e94d4a74d Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 23 Jul 2026 17:48:27 +0000 Subject: [PATCH 26/52] Audit fixes: SFT mm position ids, cache cancellation, CI test adaptations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Post-merge audit findings (adversarial review of the main reconciliation): - forward(): the ForwardPolicy fallback inverted main's gate for callers that don't thread an adapter policy — SFT passes mm_kwargs without one, so Qwen-VL SFT got packed 1D position_ids and the model skipped its internal MRoPE construction. The no-policy default now reproduces the key-presence heuristic (image_grid_thw => model owns position ids); RL's explicit adapter policies are unaffected. - _MaterializedRefCache: the owner request's cancellation (client disconnect) set CancelledError on the shared single-flight future, poisoning every deduped awaiter — and a cancelled awaiter could cancel the future out from under the rest. Materialization now runs as a detached task and awaiters shield the shared future. - _is_multimodal_sample: also treat eager mm_kwargs samples as multimodal so main's packing compatibility machinery stays live-correct, not dead code guarded by prepare_sample's rejection alone. - multimodal_sample_error: raw mm_refs samples require mm_token_type_ids (the orchestrator always stamps them; trainer truncation and forward policies rely on them). - CI runs all of tests/unit on GPU runners (not just the CPU subset): adapted main's two mm_kwargs packer tests to the raw-ref contract (raw samples never pack, even within a run), added the missing seq_lens kwarg to the branch-only forward-policy test, extended the batch no-pack test to cover the mm+mm direction, and migrated main's new configs/ci/nightly-fft/wordle.toml to the verifiers textarena player seat. - Bumped verifiers for a docstring completeness fix. --- configs/ci/nightly-fft/wordle.toml | 2 +- deps/verifiers | 2 +- src/prime_rl/inference/vllm/serving_tokens.py | 37 ++++++++----- src/prime_rl/trainer/batch.py | 6 +- src/prime_rl/trainer/model.py | 5 +- tests/unit/orchestrator/test_batch.py | 40 ++++++++------ tests/unit/train/rl/test_packer.py | 55 ++++++++++--------- tests/unit/train/test_model_forward.py | 1 + 8 files changed, 84 insertions(+), 64 deletions(-) diff --git a/configs/ci/nightly-fft/wordle.toml b/configs/ci/nightly-fft/wordle.toml index 694986f2cc..8f97c36df4 100644 --- a/configs/ci/nightly-fft/wordle.toml +++ b/configs/ci/nightly-fft/wordle.toml @@ -21,7 +21,7 @@ group_size = 16 [[orchestrator.train.env]] name = "wordle" env.taskset = { id = "wordle-v1" } -env.agent.harness = { id = "null", runtime = { type = "subprocess" } } +env.player.harness = { id = "null", runtime = { type = "subprocess" } } [orchestrator.train.sampling] max_completion_tokens = 1024 diff --git a/deps/verifiers b/deps/verifiers index d9e79d6ca2..ae516c19b0 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit d9e79d6ca215e75a269b5ecc6773f2afcf22e870 +Subproject commit ae516c19b0b63ea164bb1c0113611e7811efacbc diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index efc6ea435a..e8fc154f39 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -300,25 +300,32 @@ async def get_or_materialize(self, key: _MaterializeKey, materialize_fn: Callabl if (inflight := self._inflight.get(key)) is not None: self.hits += 1 self._maybe_log() - return await inflight + return await asyncio.shield(inflight) self.misses += 1 self._maybe_log() future: asyncio.Future = asyncio.get_running_loop().create_future() self._inflight[key] = future - try: - item = await asyncio.to_thread(materialize_fn) - except BaseException as exc: - # Never cache a failure: the exception propagates to every awaiter and - # the next request for this key retries cleanly. - future.set_exception(exc) - # A never-awaited errored future logs a spurious warning at GC. - future.exception() - del self._inflight[key] - raise - future.set_result(item) - del self._inflight[key] - self._insert(key, item) - return item + + async def _materialize_and_resolve() -> None: + try: + item = await asyncio.to_thread(materialize_fn) + except BaseException as exc: + # Never cache a failure: the exception propagates to every awaiter + # and the next request for this key retries cleanly. + future.set_exception(exc) + # A never-awaited errored future logs a spurious warning at GC. + future.exception() + else: + future.set_result(item) + self._insert(key, item) + finally: + del self._inflight[key] + + # The work runs as its own task and awaiters shield the shared future: + # a cancelled request (client disconnect) must neither poison the future + # for deduped awaiters nor cancel it out from under them. + asyncio.get_running_loop().create_task(_materialize_and_resolve()) + return await asyncio.shield(future) def _insert(self, key: _MaterializeKey, item: Any) -> None: nbytes = MultiModalCache.get_item_size(item) diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index cb8e1f74c7..fe3eddf116 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -68,6 +68,10 @@ def multimodal_sample_error(sample: TrainingSample) -> str | None: "mm_token_type_ids length must match token_ids length " f"({len(mm_token_type_ids)} != {len(sample.token_ids)})" ) + if sample.mm_refs is not None and mm_token_type_ids is None: + # The orchestrator always stamps mm_token_type_ids alongside mm_refs; the + # trainer's image-boundary truncation and forward policies rely on them. + return "raw multimodal samples require mm_token_type_ids" if sample.mm_kwargs is not None and "image_grid_thw" in sample.mm_kwargs and mm_token_type_ids is None: return "image_grid_thw multimodal samples require mm_token_type_ids" return None @@ -191,7 +195,7 @@ 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 + return sample.mm_refs is not None or sample.mm_kwargs is not None @dataclass diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index f21c1b0f5d..97f8200414 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -1396,7 +1396,10 @@ def forward( kwargs.update(mm_kwargs) if mm_token_type_ids is not None: kwargs["mm_token_type_ids"] = mm_token_type_ids - policy = mm_forward_policy or ForwardPolicy() + # Callers without an adapter policy (SFT) fall back to key presence: + # models whose kwargs carry image_grid_thw (Qwen-VL family) build + # their own MRoPE position ids and must not receive packed 1D ones. + policy = mm_forward_policy or ForwardPolicy(pass_position_ids_with_mm="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_with_mm: diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 7f0cc1ecd3..320f6ad946 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -455,18 +455,21 @@ def test_prepare_sample_truncates_raw_mm_refs_at_image_boundary(): assert micro_batch.mm_refs is None -def test_prepare_batch_keeps_raw_mm_sample_unpacked(): +def test_prepare_batch_keeps_raw_mm_samples_unpacked(): """Raw-ref multimodal samples never pack — not with text, not with each other.""" - mm_sample = TrainingSample( - token_ids=[10, 11, 12], - mask=[False, True, True], - logprobs=[0.0, -0.1, -0.2], - temperatures=[1.0, 1.0, 1.0], - advantages=[0.0, 1.0, 1.0], - env_name="mm-env", - mm_token_type_ids=[0, 1, 0], - mm_refs=MMRefs(images=[_image_ref("file:///tmp/image-0.png", offset=1, length=1)]), - ) + + def mm_sample(uri: str) -> TrainingSample: + return TrainingSample( + token_ids=[10, 11, 12], + mask=[False, True, True], + logprobs=[0.0, -0.1, -0.2], + temperatures=[1.0, 1.0, 1.0], + advantages=[0.0, 1.0, 1.0], + env_name="mm-env", + mm_token_type_ids=[0, 1, 0], + mm_refs=MMRefs(images=[_image_ref(uri, offset=1, length=1)]), + ) + text_sample = TrainingSample( token_ids=[20, 21], mask=[False, True], @@ -477,20 +480,21 @@ def test_prepare_batch_keeps_raw_mm_sample_unpacked(): ) batches_per_gpu = prepare_batch( - rollouts=[mm_sample, text_sample], - seq_len=8, + rollouts=[mm_sample("file:///tmp/image-0.png"), mm_sample("file:///tmp/image-1.png"), text_sample], + seq_len=16, num_train_workers=2, - idxs=[0, 0], + idxs=[0, 0, 0], num_loras=1, bin_cost=build_bin_cost(None), ) real_batches = [batch for batch in _flatten_batches(batches_per_gpu) if _has_loss_tokens(batch)] - assert len(real_batches) == 2 + assert len(real_batches) == 3 mm_batches = [batch for batch in real_batches if batch.mm_refs is not None] - assert len(mm_batches) == 1 - assert mm_batches[0].env_names == ["mm-env"] * 3 - assert mm_batches[0].mm_refs == mm_sample.mm_refs + assert len(mm_batches) == 2 + for batch in mm_batches: + assert batch.env_names == ["mm-env"] * 3 + assert len(batch.mm_refs.images) == 1 def test_prepare_sample_none_routed_experts(): diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 890e290d20..8633ee9938 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -1,7 +1,6 @@ from pathlib import Path from typing import Generator -import numpy as np import pytest import tomli_w import torch @@ -13,7 +12,7 @@ from prime_rl.trainer.runs import setup_multi_run_manager from prime_rl.trainer.utils import build_bin_cost from prime_rl.trainer.world import reset_world -from prime_rl.transport.types import EncodedTensor, TrainingSample +from prime_rl.transport.types import MMImageRef, MMRefs, TrainingSample @pytest.fixture(autouse=True, scope="module") @@ -53,16 +52,8 @@ def make_training_sample() -> TrainingSample: ) -def _encoded_tensor(data, dtype) -> EncodedTensor: - arr = np.asarray(data, dtype=dtype) - return EncodedTensor(dtype=str(arr.dtype), shape=list(arr.shape), data=arr.tobytes()) - - -def _decode_encoded_tensor(encoded: EncodedTensor): - return np.frombuffer(encoded.data, dtype=np.dtype(encoded.dtype)).reshape(encoded.shape).tolist() - - def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: + uri = f"file:///tmp/image-{value}.png" return TrainingSample( token_ids=[1, 250, 2], mask=[False, True, True], @@ -71,10 +62,24 @@ def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: env_name=env_name, advantages=[0.0, 1.0, 1.0], mm_token_type_ids=[0, 1, 0], - mm_kwargs={ - "pixel_values": _encoded_tensor([[value, value + 1]], np.float32), - "image_grid_thw": _encoded_tensor([[1, 2, 2]], np.int64), - }, + mm_refs=MMRefs( + images=[ + MMImageRef( + item={ + "kind": "prime_raw_mm_item", + "modality": "image", + "family": "qwen_vl", + "layout_fingerprint": "f" * 32, + "raw_image_uri": uri, + "payload": {"image_grid_thw": [[1, 1, 1]]}, + }, + hash="a" * 32, + uri=uri, + offset=1, + length=1, + ) + ] + ), ) @@ -177,7 +182,7 @@ def fake_sender(_output_dir, _data_world_size, _current_step, _config): assert micro_batch.run_step == 1 -def test_multipacker_pack_preserves_mm_kwargs_modality_and_run_tagging(tmp_path, monkeypatch): +def test_multipacker_pack_preserves_mm_modality_alignment_and_run_tagging(tmp_path, monkeypatch): """MultiPacker keeps multimodal and text microbatches aligned across ranks.""" from prime_rl.trainer.batch import _is_multimodal_sample @@ -205,7 +210,7 @@ def test_multipacker_pack_preserves_mm_kwargs_modality_and_run_tagging(tmp_path, assert mm_mbs, "no MM microbatches produced" real_run_idxs = set() for mb in mm_mbs: - assert mb.mm_kwargs is not None + assert mb.mm_refs is not None if any(mb.loss_mask): tagged = [i for i, n in enumerate(mb.lora_num_tokens) if n > 0] assert len(tagged) == 1 and mb.lora_num_tokens[tagged[0]] == len(mb.input_ids) @@ -213,8 +218,8 @@ def test_multipacker_pack_preserves_mm_kwargs_modality_and_run_tagging(tmp_path, assert real_run_idxs == {a, b}, f"both runs' MM should be tagged; got {real_run_idxs}" -def test_multipacker_packs_mm_kwargs_within_each_run(tmp_path, monkeypatch): - """Compatible eager multimodal samples pack within a run but never across runs.""" +def test_multipacker_keeps_raw_mm_samples_unpacked_within_a_run(tmp_path, monkeypatch): + """Raw-ref multimodal samples never pack, even within one run — one micro batch each.""" from prime_rl.trainer.batch import _is_multimodal_sample manager, packer, sent = _packer_with_two_runs(tmp_path, monkeypatch, dp_world_size=1, seq_len=12) @@ -228,14 +233,10 @@ def test_multipacker_packs_mm_kwargs_within_each_run(tmp_path, monkeypatch): grid = sent[-1] real_mm_mbs = [mb for rank in grid for mb in rank if _is_multimodal_sample(mb) and any(mb.loss_mask)] - assert len(real_mm_mbs) == 2 + assert len(real_mm_mbs) == 4 for mb in real_mm_mbs: - assert len(mb.input_ids) == 6 - assert mb.position_ids == [0, 1, 2, 0, 1, 2] - assert mb.seq_lens == [3, 3] - assert mb.mm_kwargs is not None - assert mb.mm_kwargs["pixel_values"].shape == [2, 2] - assert mb.mm_kwargs["image_grid_thw"].shape == [2, 3] - assert len(_decode_encoded_tensor(mb.mm_kwargs["pixel_values"])) == 2 + assert len(mb.input_ids) == 3 + assert mb.seq_lens == [3] + assert mb.mm_refs is not None and len(mb.mm_refs.images) == 1 tagged = [i for i, n in enumerate(mb.lora_num_tokens) if n > 0] assert len(tagged) == 1 diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py index 763c97954f..8071809655 100644 --- a/tests/unit/train/test_model_forward.py +++ b/tests/unit/train/test_model_forward.py @@ -99,6 +99,7 @@ def test_forward_policy_can_require_mm_token_type_ids(): model, input_ids, position_ids, + seq_lens=torch.tensor([input_ids.shape[1]]), mm_kwargs={"pixel_values": torch.ones(2, 3)}, mm_forward_policy=ForwardPolicy(requires_mm_token_type_ids=True), ) From 6a1a54bcf0baf553d369cb9d55a05e31a3991ca4 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 24 Jul 2026 07:29:54 +0000 Subject: [PATCH 27/52] Pack raw-ref multimodal samples like main packs eager ones MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The main merge resolved the mm packing conflict too conservatively: raw-ref samples never packed, regressing main's feature where multimodal samples pack with text spans and compatible mm samples from the same run/LoRA. Restore it for the mm_refs representation: - can_add: raw-ref samples are pack-compatible when they share an adapter family (the materializer enforces one family per micro batch; grids and image sizes vary freely within a family) — the simpler analogue of the eager path's dtype/trailing-dim check. The run/LoRA gate applies as on main. - _materialize_bin: merges refs across the bin's samples, rebasing each sample's placeholder offsets by its start position in the packed token stream. mm_token_type_ids already concatenate with zero-fill and seq_lens already preserve boundaries; the trainer-side adapter already materializes a list of refs into concatenated tensors, so the packed batch reaches forward() shaped exactly like main's packed eager batches. - Restored main-shaped packing tests (pack within run with rebased offsets, family mismatch splits, never across runs). Also map adapter validation errors (fingerprint/grid/placeholder-length mismatches) to invalid_mm_image_ref 400s in serving_tokens instead of letting them escape as 500s — ported from the inline bundle's review. --- src/prime_rl/inference/vllm/serving_tokens.py | 17 ++++---- src/prime_rl/trainer/batch.py | 42 +++++++++++++------ tests/unit/inference/test_serving_tokens.py | 29 +++++++++++++ tests/unit/orchestrator/test_batch.py | 38 ++++++++++++----- tests/unit/train/rl/test_packer.py | 15 +++---- 5 files changed, 103 insertions(+), 38 deletions(-) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index e8fc154f39..5cea362616 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -242,13 +242,16 @@ def _materialize_raw_image_ref_sync( payload=dict(ref.payload), raw_ref=raw_ref, ) - adapter = get_multimodal_adapter(ref.family) - return adapter.materialize_for_vllm( - image_processor, - item, - image, - expected_placeholder_length, - ) + try: + adapter = get_multimodal_adapter(ref.family) + return adapter.materialize_for_vllm( + image_processor, + item, + image, + expected_placeholder_length, + ) + except (TypeError, ValueError) as exc: + raise _MMImageRefError(str(exc)) from exc # (raw_ref, expected_placeholder_length, processor_model_name). The ref string is diff --git a/src/prime_rl/trainer/batch.py b/src/prime_rl/trainer/batch.py index fe3eddf116..bb3eb257fd 100644 --- a/src/prime_rl/trainer/batch.py +++ b/src/prime_rl/trainer/batch.py @@ -2,10 +2,11 @@ from collections.abc import Callable, Sequence from dataclasses import dataclass +import msgspec import numpy as np from prime_rl.trainer.utils import balanced_partition -from prime_rl.transport.types import MicroBatch, MMRefs, RoutedExperts, TrainingSample +from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs, RoutedExperts, TrainingSample # Backfill value per component weight stream when a packed sample doesn't # carry it: absent rl means weight 1.0 on the loss mask, absent ce/ref_kl @@ -198,6 +199,15 @@ def _is_multimodal_sample(sample: MicroBatch) -> bool: return sample.mm_refs is not None or sample.mm_kwargs is not None +def _mm_refs_family(mm_refs: MMRefs) -> str | None: + """The adapter family of a sample's raw image refs. + + ``build_mm_refs`` validates every descriptor and the materializer enforces + exactly one family per micro batch, so the first image's family stands for + the sample.""" + return mm_refs.images[0].item.get("family") if mm_refs.images else None + + @dataclass class _MicroBatchBin: samples: list[tuple[int, MicroBatch]] @@ -234,17 +244,17 @@ def can_add(self, sample: MicroBatch, max_seq_len: int, lora_idx: int) -> bool: sample_is_mm = _is_multimodal_sample(sample) existing_mm_sample = self.first_multimodal_sample - # Raw image refs (this codebase's multimodal representation) never pack — with - # text or with each other: their placeholder offsets are relative to the sample's - # own token stream and ``_materialize_bin`` carries only the first sample's refs, - # so any packing would silently drop or misalign images. - if sample.mm_refs is not None or (existing_mm_sample is not None and existing_mm_sample.mm_refs is not None): - return False if existing_mm_sample is None and not sample_is_mm: return True if self.first_lora_idx != lora_idx: return False if existing_mm_sample is not None and sample_is_mm: + if (existing_mm_sample.mm_refs is None) != (sample.mm_refs is None): + return False + if sample.mm_refs is not None: + # Raw refs materialize downstream through exactly one adapter family + # per micro batch; within a family, grids and image sizes vary freely. + return _mm_refs_family(existing_mm_sample.mm_refs) == _mm_refs_family(sample.mm_refs) dst = existing_mm_sample.mm_kwargs src = sample.mm_kwargs assert dst is not None and src is not None, "multimodal samples must carry mm_kwargs" @@ -304,10 +314,12 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: streams: dict[str, list[float] | None] = {name: ([] if has_stream[name] else None) for name in STREAM_FILL} seq_lens: list[int] = [] routed_experts: RoutedExperts | None = None + mm_ref_images: list[MMImageRef] = [] lora_num_tokens = [0] * num_loras for lora_idx, sample in bin_content.samples: sample_len = len(sample.input_ids) + sample_start = len(input_ids) input_ids.extend(sample.input_ids) loss_mask.extend(sample.loss_mask) advantages.extend(sample.advantages) @@ -341,6 +353,11 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: for key in mm_kwargs: mm_kwargs[key].data += sample.mm_kwargs[key].data mm_kwargs[key].shape[0] += sample.mm_kwargs[key].shape[0] + if sample.mm_refs is not None: + # Placeholder offsets are sample-relative; rebase them to the packed stream. + mm_ref_images.extend( + msgspec.structs.replace(image, offset=image.offset + sample_start) for image in sample.mm_refs.images + ) seq_lens.extend(sample.seq_lens) lora_num_tokens[lora_idx] += sample_len @@ -361,7 +378,7 @@ def _materialize_bin(bin_content: _MicroBatchBin, num_loras: int) -> MicroBatch: routed_experts=routed_experts, mm_token_type_ids=mm_token_type_ids, env_names=env_names, - mm_refs=bin_content.first_sample.mm_refs if _is_multimodal_sample(bin_content.first_sample) else None, + mm_refs=MMRefs(images=mm_ref_images) if mm_ref_images else None, mm_kwargs=mm_kwargs, rl_weights=streams["rl_weights"], ce_weights=streams["ce_weights"], @@ -399,11 +416,10 @@ def packed_samples_into_micro_bs( We follow the First Fit Decreasing algorithm to pack the samples into bins and minimize potential padding while never truncating. With per-token temperatures, samples can be packed together regardless of their temperature values. - NOTE: Multimodal samples never pack with anything (including each other) in this - codebase's raw-ref representation (``mm_refs``) — each becomes its own micro batch. - Eager ``mm_kwargs`` sidecars (a separate, tensor-based multimodal path) do support - packing with compatible same-run/LoRA samples; packed batches preserve sample - boundaries in ``seq_lens`` either way. + Multimodal samples pack with text spans from the same run/LoRA and with + same-family raw-ref samples (image ref offsets are rebased to the packed + stream in ``_materialize_bin``). Packed batches preserve sample boundaries + in ``seq_lens``. """ # Sort by (lora_idx, -length) for packing efficiency samples.sort(key=lambda x: (x[0], -len(x[1].input_ids))) diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 5d3602dff4..e15a2148f2 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -330,6 +330,35 @@ def _get_adapter(family): assert item.payload == {"adapter_owned": [1, 2, 3]} +def test_materialize_raw_image_ref_maps_adapter_validation_error(tmp_path, monkeypatch): + import pytest + + from prime_rl.inference.vllm import serving_tokens + + features = _mm_features(tmp_path) + raw_ref = features.kwargs_data["image"][0] + mm_hash = features.mm_hashes["image"][0] + + class _InvalidAdapter: + def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): + raise ValueError("image layout fingerprint mismatch") + + monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: object()) + monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _InvalidAdapter()) + + with pytest.raises(serving_tokens._MMImageRefError, match="image layout fingerprint mismatch") as exc_info: + serving_tokens._materialize_raw_image_ref_sync( + raw_ref, + feature_modality="image", + mm_hash=mm_hash, + expected_placeholder_length=7, + processor_model_name="model", + trust_remote_code=False, + ) + + assert exc_info.value.status_code == 400 + + def _mm_features(tmp_path, *, image_name: str = "image.png", placeholder_length: int = 7): """A minimal ``GenerateRequest.features``-shaped object carrying one real raw ref.""" from types import SimpleNamespace diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 320f6ad946..5b0236596d 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -455,10 +455,14 @@ def test_prepare_sample_truncates_raw_mm_refs_at_image_boundary(): assert micro_batch.mm_refs is None -def test_prepare_batch_keeps_raw_mm_samples_unpacked(): - """Raw-ref multimodal samples never pack — not with text, not with each other.""" - - def mm_sample(uri: str) -> TrainingSample: +def test_prepare_batch_packs_raw_mm_samples_with_rebased_offsets(): + """Same-family raw-ref samples pack with each other and with text from the same + run/LoRA; merged image ref offsets are rebased to the packed token stream. A + different adapter family never joins the bin.""" + + def mm_sample(uri: str, family: str = "qwen_vl") -> TrainingSample: + ref = _image_ref(uri, offset=1, length=1) + ref.item["family"] = family return TrainingSample( token_ids=[10, 11, 12], mask=[False, True, True], @@ -467,7 +471,7 @@ def mm_sample(uri: str) -> TrainingSample: advantages=[0.0, 1.0, 1.0], env_name="mm-env", mm_token_type_ids=[0, 1, 0], - mm_refs=MMRefs(images=[_image_ref(uri, offset=1, length=1)]), + mm_refs=MMRefs(images=[ref]), ) text_sample = TrainingSample( @@ -480,21 +484,33 @@ def mm_sample(uri: str) -> TrainingSample: ) batches_per_gpu = prepare_batch( - rollouts=[mm_sample("file:///tmp/image-0.png"), mm_sample("file:///tmp/image-1.png"), text_sample], + rollouts=[ + mm_sample("file:///tmp/image-0.png"), + mm_sample("file:///tmp/image-1.png"), + mm_sample("file:///tmp/image-2.png", family="kimi_k25"), + text_sample, + ], seq_len=16, num_train_workers=2, - idxs=[0, 0, 0], + idxs=[0, 0, 0, 0], num_loras=1, bin_cost=build_bin_cost(None), ) real_batches = [batch for batch in _flatten_batches(batches_per_gpu) if _has_loss_tokens(batch)] - assert len(real_batches) == 3 mm_batches = [batch for batch in real_batches if batch.mm_refs is not None] assert len(mm_batches) == 2 - for batch in mm_batches: - assert batch.env_names == ["mm-env"] * 3 - assert len(batch.mm_refs.images) == 1 + packed = max(mm_batches, key=lambda batch: len(batch.mm_refs.images)) + assert len(packed.mm_refs.images) == 2 + assert packed.seq_lens[:2] == [3, 3] + # Sample-relative offset 1, rebased by each sample's start in the packed stream. + assert [image.offset for image in packed.mm_refs.images] == [1, 4] + assert packed.mm_token_type_ids[:6] == [0, 1, 0, 0, 1, 0] + assert packed.env_names[:6] == ["mm-env"] * 6 + # The kimi-family sample stayed out of the qwen bin. + (other,) = [batch for batch in mm_batches if batch is not packed] + assert len(other.mm_refs.images) == 1 + assert other.mm_refs.images[0].item["family"] == "kimi_k25" def test_prepare_sample_none_routed_experts(): diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 8633ee9938..f9adf44fd7 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -218,8 +218,8 @@ def test_multipacker_pack_preserves_mm_modality_alignment_and_run_tagging(tmp_pa assert real_run_idxs == {a, b}, f"both runs' MM should be tagged; got {real_run_idxs}" -def test_multipacker_keeps_raw_mm_samples_unpacked_within_a_run(tmp_path, monkeypatch): - """Raw-ref multimodal samples never pack, even within one run — one micro batch each.""" +def test_multipacker_packs_raw_mm_samples_within_each_run(tmp_path, monkeypatch): + """Same-family raw-ref samples pack within a run (offsets rebased) but never across runs.""" from prime_rl.trainer.batch import _is_multimodal_sample manager, packer, sent = _packer_with_two_runs(tmp_path, monkeypatch, dp_world_size=1, seq_len=12) @@ -233,10 +233,11 @@ def test_multipacker_keeps_raw_mm_samples_unpacked_within_a_run(tmp_path, monkey grid = sent[-1] real_mm_mbs = [mb for rank in grid for mb in rank if _is_multimodal_sample(mb) and any(mb.loss_mask)] - assert len(real_mm_mbs) == 4 + assert len(real_mm_mbs) == 2 for mb in real_mm_mbs: - assert len(mb.input_ids) == 3 - assert mb.seq_lens == [3] - assert mb.mm_refs is not None and len(mb.mm_refs.images) == 1 + assert len(mb.input_ids) == 6 + assert mb.seq_lens == [3, 3] + assert mb.mm_refs is not None and len(mb.mm_refs.images) == 2 + assert [image.offset for image in mb.mm_refs.images] == [1, 4] tagged = [i for i, n in enumerate(mb.lora_num_tokens) if n > 0] - assert len(tagged) == 1 + assert len(tagged) == 1 and mb.lora_num_tokens[tagged[0]] == 6 From 3bfaf11da1a7ed9e9e8790ec78862598406ac7bb Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 24 Jul 2026 19:22:37 +0000 Subject: [PATCH 28/52] Bump renderers for the unwrapped ref payload; harden the cache key MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Renderers no longer base64-wrap the ref payload (parity with the inline bundle, where the wrapper cost ~32 KiB per image slot). Adopts the inline bundle's materialize-cache key so both bundles key identically: digest the whole ref rather than trusting the caller's outer metadata, and include feature_modality and trust_remote_code. The digest matters less here (an offload ref is a short URI, so keying on the ref string was already sound) but trust_remote_code is a real gap — it selects a different image processor, so two values could previously alias one cache entry. Ports the inline anti-aliasing test. The test fixture now seeds image content explicitly: it previously colored images by filename length, so two same-length names produced byte-identical images and a forged-hash case couldn't be expressed. --- deps/renderers | 2 +- src/prime_rl/inference/vllm/serving_tokens.py | 18 +++++++--- tests/unit/inference/test_serving_tokens.py | 36 ++++++++++++++++--- 3 files changed, 45 insertions(+), 11 deletions(-) diff --git a/deps/renderers b/deps/renderers index e64cc58803..7c58fd62c1 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit e64cc588037a75ff99d37c2968329805147d820d +Subproject commit 7c58fd62c188c8eea61e05792adcd0662b5b85f7 diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 5cea362616..719ef9a8b6 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -254,10 +254,11 @@ def _materialize_raw_image_ref_sync( raise _MMImageRefError(str(exc)) from exc -# (raw_ref, expected_placeholder_length, processor_model_name). The ref string is -# content-addressed (it embeds the image hash, layout fingerprint, family, and URI), -# so identical keys describe byte-identical materialization work. -_MaterializeKey = tuple[str, int, str] +# (raw_ref_digest, feature_modality, mm_hash, expected_placeholder_length, +# processor_model_name, trust_remote_code). The digest keeps keys small while +# ensuring a distinct descriptor can never alias an already-validated cache +# entry, even when its outer hash/placeholder metadata is forged or stale. +_MaterializeKey = tuple[str, str, str, int, str, bool] _MM_MATERIALIZE_CACHE_GB_ENV = "PRIME_RL_MM_MATERIALIZE_CACHE_GB" _MM_MATERIALIZE_LOG_EVERY = 1000 @@ -411,7 +412,14 @@ def _materialize(feature_modality: str, raw_ref: str, mm_hash: str, placeholder: decoded = await asyncio.gather( *( cache.get_or_materialize( - (raw_ref, placeholder.length, processor_model_name), + ( + hashlib.sha256(raw_ref.encode("utf-8")).hexdigest(), + feature_modality, + mm_hash, + placeholder.length, + processor_model_name, + trust_remote_code, + ), lambda fm=feature_modality, r=raw_ref, h=mm_hash, p=placeholder: _materialize(fm, r, h, p), ) for feature_modality, raw_ref, mm_hash, placeholder in flat diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index e15a2148f2..1200807066 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -359,7 +359,7 @@ def materialize_for_vllm(self, image_processor, item, image, expected_placeholde assert exc_info.value.status_code == 400 -def _mm_features(tmp_path, *, image_name: str = "image.png", placeholder_length: int = 7): +def _mm_features(tmp_path, *, image_seed: int = 0, placeholder_length: int = 7): """A minimal ``GenerateRequest.features``-shaped object carrying one real raw ref.""" from types import SimpleNamespace @@ -368,8 +368,8 @@ def _mm_features(tmp_path, *, image_name: str = "image.png", placeholder_length: image_dir = tmp_path / "assets" / "images" image_dir.mkdir(parents=True, exist_ok=True) - image_path = image_dir / image_name - Image.new("RGB", (8, 6), color=(len(image_name) % 255, 64, 128)).save(image_path) + image_path = image_dir / f"image-{image_seed}.png" + Image.new("RGB", (8, 6), color=(image_seed % 255, 64, 128)).save(image_path) mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] raw_ref = raw_mm_ref( @@ -424,6 +424,32 @@ async def _run(): assert (cache.hits, cache.misses) == (1, 1) +def test_mm_materialize_cache_does_not_alias_distinct_raw_refs(tmp_path, monkeypatch): + import pytest + + from prime_rl.inference.vllm import serving_tokens + + calls: list = [] + _patch_adapter(monkeypatch, calls) + cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) + valid = _mm_features(tmp_path, image_seed=1) + forged = _mm_features(tmp_path, image_seed=2) + forged.mm_hashes["image"][0] = valid.mm_hashes["image"][0] + + async def _run(): + await serving_tokens._decode_raw_mm_kwargs( + valid, processor_model_name="model", trust_remote_code=False, cache=cache + ) + await serving_tokens._decode_raw_mm_kwargs( + forged, processor_model_name="model", trust_remote_code=False, cache=cache + ) + + with pytest.raises(serving_tokens._MMImageRefError, match="Expected image hash"): + asyncio.run(_run()) + assert len(calls) == 1 + assert (cache.hits, cache.misses) == (0, 2) + + def test_mm_materialize_cache_byte_budget_evicts_oldest(tmp_path, monkeypatch): from vllm.multimodal.cache import MultiModalCache @@ -433,8 +459,8 @@ def test_mm_materialize_cache_byte_budget_evicts_oldest(tmp_path, monkeypatch): _patch_adapter(monkeypatch, calls) monkeypatch.setattr(MultiModalCache, "get_item_size", classmethod(lambda _cls, _item: 60)) cache = serving_tokens._MaterializedRefCache(max_bytes=100) - features_a = _mm_features(tmp_path, image_name="a.png") - features_b = _mm_features(tmp_path, image_name="b.png") + features_a = _mm_features(tmp_path, image_seed=1) + features_b = _mm_features(tmp_path, image_seed=2) async def _decode(features): return await serving_tokens._decode_raw_mm_kwargs( From fd3f379d50ab41464a317acb33add3cbf552b461 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 24 Jul 2026 19:25:02 +0000 Subject: [PATCH 29/52] Inherit the parent env explicitly at the orchestrator launch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The orchestrator subprocess inherited os.environ only as a side effect of build_run_asset_env returning dict(os.environ) plus the image dir — the same implicit coupling that broke inheritance outright in the inline bundle when that helper was removed. Spread inherited_env first, as the trainer and inference launches already do, so the image-dir helper only has to supply the var it owns. --- src/prime_rl/entrypoints/rl.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 1546fda14b..f9ffb7718e 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -227,6 +227,7 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ + **inherited_env, **writer_run_asset_env, **DEFAULT_COMMON_ENV_VARS, "LOGURU_FORCE_COLORS": "1", From 3b2dffd441f6ef2287d7c5209a74c35a0ba9bc57 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 31 Jul 2026 01:03:45 +0000 Subject: [PATCH 30/52] Relock against the merged research-environments pin; bump verifiers MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The merge brought a new deps/research-environments pin, but this worktree still had the previous commit checked out, so the regenerated lockfile resolved against a stale tree — dropping charxiv-v1 / mmmu-pro-v1 and downgrading bfcl-v3-v1, which failed CI's `uv sync --locked`. Also bumps deps/verifiers for the ruff 0.16 import sort in its multimodal client type test. --- deps/verifiers | 2 +- uv.lock | 64 +++++++++++++++++++++++++++++++++++++++----------- 2 files changed, 51 insertions(+), 15 deletions(-) diff --git a/deps/verifiers b/deps/verifiers index 0e41666080..88c759b0f6 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit 0e416660801d67c26b0686ed8fbcf880a25ee771 +Subproject commit 88c759b0f68a1018b9c5015e4640979184b34cf9 diff --git a/uv.lock b/uv.lock index 00b59560a5..0e8fc133b1 100644 --- a/uv.lock +++ b/uv.lock @@ -45,6 +45,7 @@ members = [ "bfcl-v3-v1", "browsecomp-plus-v1", "browsecomp-v1", + "charxiv-v1", "clbench-v1", "code-golf-v1", "color-codeword-v1", @@ -78,6 +79,7 @@ members = [ "math500-v1", "mcp-atlas-v1", "mmlu-pro-v1", + "mmmu-pro-v1", "mmmu-v1", "mrcr-v2-v1", "multiswe-v1", @@ -594,7 +596,7 @@ dependencies = [ [[package]] name = "bfcl-v3-v1" -version = "0.2.0" +version = "0.3.0" source = { editable = "deps/research-environments/environments/tool_use/bfcl_v3_v1" } dependencies = [ { name = "bfcl-eval" }, @@ -606,7 +608,7 @@ dependencies = [ requires-dist = [ { name = "bfcl-eval", git = "https://github.com/mikasenghaas/gorilla.git?subdirectory=berkeley-function-call-leaderboard&rev=898763a" }, { name = "soundfile", specifier = ">=0.13.0" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev21" }, ] [[package]] @@ -872,6 +874,23 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/db/8f/61959034484a4a7c527811f4721e75d02d653a35afb0b6054474d8185d4c/charset_normalizer-3.4.7-py3-none-any.whl", hash = "sha256:3dce51d0f5e7951f8bb4900c257dad282f49190fdbebecd4ba99bcc41fef404d", size = 61958, upload-time = "2026-04-02T09:28:37.794Z" }, ] +[[package]] +name = "charxiv-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/multimodal/charxiv_v1" } +dependencies = [ + { name = "datasets" }, + { name = "pillow" }, + { name = "verifiers" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets" }, + { name = "pillow" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "chess" version = "1.11.2" @@ -3736,6 +3755,23 @@ requires-dist = [ { name = "verifiers", specifier = ">=0.2.0" }, ] +[[package]] +name = "mmmu-pro-v1" +version = "0.1.0" +source = { editable = "deps/research-environments/environments/multimodal/mmmu_pro_v1" } +dependencies = [ + { name = "datasets" }, + { name = "pillow" }, + { name = "verifiers" }, +] + +[package.metadata] +requires-dist = [ + { name = "datasets" }, + { name = "pillow" }, + { name = "verifiers", specifier = ">=0.2.0" }, +] + [[package]] name = "mmmu-v1" version = "0.1.0" @@ -3947,7 +3983,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "multi-swe-bench" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -4574,7 +4610,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "huggingface-hub" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -5959,7 +5995,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -6278,7 +6314,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -6397,7 +6433,7 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.1" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev27" }] [[package]] name = "sentence-transformers" @@ -6815,7 +6851,7 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.1" }] +requires-dist = [{ name = "verifiers", extras = ["harbor"], specifier = ">=0.2.2.dev27" }] [[package]] name = "swebench-pro-v1" @@ -6826,7 +6862,7 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.1" }] +requires-dist = [{ name = "verifiers", specifier = ">=0.2.2.dev27" }] [[package]] name = "swebench-verified-v1" @@ -6837,7 +6873,7 @@ dependencies = [ ] [package.metadata] -requires-dist = [{ name = "verifiers", specifier = ">=0.2.1" }] +requires-dist = [{ name = "verifiers", specifier = ">=0.2.2.dev27" }] [[package]] name = "swelego-v1" @@ -6851,7 +6887,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -6866,7 +6902,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "datasets" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -6890,7 +6926,7 @@ requires-dist = [ { name = "datasets" }, { name = "swebench" }, { name = "swesmith", git = "https://github.com/SWE-bench/SWE-smith.git?rev=9b74ac0" }, - { name = "verifiers", specifier = ">=0.2.2.dev5" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] @@ -8478,7 +8514,7 @@ dependencies = [ requires-dist = [ { name = "datasets" }, { name = "openai" }, - { name = "verifiers", specifier = ">=0.2.0" }, + { name = "verifiers", specifier = ">=0.2.2.dev27" }, ] [[package]] From 95afd1fce788c79bd9d9af9bc1fde9d8a99d7910 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 31 Jul 2026 01:31:14 +0000 Subject: [PATCH 31/52] Adapt orchestrator timing metrics to the verifiers agent-span rename MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit This branch's verifiers pin is ahead of the one prime-rl main tracks, and includes the runtime-on-agent work that renamed Trace.timing.generation to .agent (identical shape: duration + model/harness splits). The orchestrator's timing metrics still read .generation, which surfaced as 'Timing' object has no attribute 'generation' in the integration tests. Reads the renamed field but keeps the emitted metric names (time/generation, generation/model, generation/harness) unchanged so W&B dashboards and the monitor-run skill stay valid. PHASES did double duty as both metric names and trace field names — those have now diverged, so the trace-side lookup goes through an explicit TRACE_PHASES mapping. --- src/prime_rl/orchestrator/metrics.py | 13 +++++++++---- tests/unit/orchestrator/test_metrics.py | 2 +- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/src/prime_rl/orchestrator/metrics.py b/src/prime_rl/orchestrator/metrics.py index ce860cb284..aacd5b13bd 100644 --- a/src/prime_rl/orchestrator/metrics.py +++ b/src/prime_rl/orchestrator/metrics.py @@ -91,6 +91,9 @@ class TimingMetrics(StatGroup): ``total`` is the per-rollout sum across all phases.""" PHASES = ("setup", "generation", "finalize", "scoring") + # Verifiers renamed its generation span to `agent`; the metric keeps the + # `generation` name so dashboards and the monitor-run skill stay stable. + TRACE_PHASES = {"generation": "agent"} @property def setup(self) -> Stat: @@ -98,18 +101,18 @@ def setup(self) -> Stat: @property def generation(self) -> Stat: - return Stat([r.timing.generation.duration for r in self.rollouts]) + return Stat([r.timing.agent.duration for r in self.rollouts]) @property def generation_model(self) -> Stat: """The share of the generation phase spent inside model calls (inference).""" - return Stat([r.timing.generation.model.duration for r in self.rollouts]) + return Stat([r.timing.agent.model.duration for r in self.rollouts]) @property def generation_harness(self) -> Stat: """The share of the generation phase spent outside model calls (harness, tools, user simulation).""" - return Stat([r.timing.generation.harness.duration for r in self.rollouts]) + return Stat([r.timing.agent.harness.duration for r in self.rollouts]) @property def finalize(self) -> Stat: @@ -121,7 +124,9 @@ def scoring(self) -> Stat: @property def total(self) -> Stat: - return Stat([sum(getattr(r.timing, p).duration for p in self.PHASES) for r in self.rollouts]) + return Stat( + [sum(getattr(r.timing, self.TRACE_PHASES.get(p, p)).duration for p in self.PHASES) for r in self.rollouts] + ) def stats(self) -> dict[str, Stat]: return { diff --git a/tests/unit/orchestrator/test_metrics.py b/tests/unit/orchestrator/test_metrics.py index 5362037833..2e2cc26d8d 100644 --- a/tests/unit/orchestrator/test_metrics.py +++ b/tests/unit/orchestrator/test_metrics.py @@ -59,7 +59,7 @@ def mk( filter_results=filter_results or {}, timing=SimpleNamespace( setup=SimpleNamespace(duration=setup), - generation=SimpleNamespace( + agent=SimpleNamespace( duration=generation, model=SimpleNamespace(duration=generation_model), harness=SimpleNamespace(duration=generation_harness), From 37ff96a75a654ed355f51e99147547eff784dd25 Mon Sep 17 00:00:00 2001 From: eligotts Date: Sat, 1 Aug 2026 21:55:17 +0000 Subject: [PATCH 32/52] Bump renderers: raw mode accepts file:// URLs only --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index 1af6589aab..66505d2806 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 1af6589aabc5c37b20859ca59e7174010c321dd1 +Subproject commit 66505d280683654a98d97a5eb7d627f41650c923 From 743a02290c3a2d75ef414ff79cd815d31bae2940 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 3 Aug 2026 20:13:17 +0000 Subject: [PATCH 33/52] refactor: derive processor fingerprints from renderer layout specs The renderer-side layout dataclasses are now the single canonical knob list per family, so the adapters stop re-listing the fields: fingerprints come from qwen_layout_from(image_processor).fingerprint() / kimi_layout_from(...).fingerprint(), and the kimi materialize path reads patch shapes through the same extractor. Deletes the hand-duplicated _processor_value / _processor_layout helper stacks. Co-Authored-By: Claude Fable 5 --- deps/renderers | 2 +- src/prime_rl/multimodal/adapters/kimi_k25.py | 61 ++------------------ src/prime_rl/multimodal/adapters/qwen_vl.py | 29 ++-------- 3 files changed, 12 insertions(+), 80 deletions(-) diff --git a/deps/renderers b/deps/renderers index 66505d2806..f0d1dd945e 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 66505d280683654a98d97a5eb7d627f41650c923 +Subproject commit f0d1dd945e13935665633c6fdc04de8462e0bf40 diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index 4a2aba7aa9..319492b361 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -1,10 +1,9 @@ from __future__ import annotations import math -from collections.abc import Mapping from typing import Any -from renderers.kimi_k25 import KimiK25ImageLayoutSpec +from renderers.kimi_k25 import kimi_layout_from from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem @@ -18,46 +17,6 @@ def _tensorize(value: Any): return torch.as_tensor(value).contiguous() -def _media_proc_cfg(image_processor: Any) -> Mapping[str, Any]: - cfg = getattr(image_processor, "media_proc_cfg", None) - if not isinstance(cfg, Mapping): - raise ValueError("Kimi image processor must expose media_proc_cfg") - return cfg - - -def _required_cfg(cfg: Mapping[str, Any], name: str) -> Any: - if name not in cfg: - raise ValueError(f"Kimi image processor media_proc_cfg is missing {name!r}") - return cfg[name] - - -def _optional_int(value: Any) -> int | None: - if value is None: - return None - return int(value) - - -def _float_triple(value: Any, *, name: str) -> tuple[float, float, float]: - if not isinstance(value, list | tuple) or len(value) != 3: - raise ValueError(f"Kimi image processor media_proc_cfg[{name!r}] must be a length-3 sequence") - return (float(value[0]), float(value[1]), float(value[2])) - - -def _processor_layout(image_processor: Any) -> KimiK25ImageLayoutSpec: - """Read the actual processor's layout; drift from the renderer's baked - layout surfaces as a fingerprint mismatch at materialization.""" - cfg = _media_proc_cfg(image_processor) - return KimiK25ImageLayoutSpec( - patch_size=int(_required_cfg(cfg, "patch_size")), - merge_kernel_size=int(_required_cfg(cfg, "merge_kernel_size")), - in_patch_limit=int(_required_cfg(cfg, "in_patch_limit")), - patch_limit_on_one_side=int(_required_cfg(cfg, "patch_limit_on_one_side")), - fixed_output_tokens=_optional_int(_required_cfg(cfg, "fixed_output_tokens")), - image_mean=_float_triple(_required_cfg(cfg, "image_mean"), name="image_mean"), - image_std=_float_triple(_required_cfg(cfg, "image_std"), name="image_std"), - ) - - def _grid_payload(item: RawMMItem) -> list[int]: grid = item.payload.get("grid_thws") if grid is None: @@ -92,19 +51,9 @@ def validate_item(self, item: RawMMItem) -> None: _grid_payload(item) def processor_fingerprint(self, image_processor: Any) -> str: - from renderers.mm_store import image_layout_fingerprint - - layout = _processor_layout(image_processor) - return image_layout_fingerprint( - family=self.family, - patch_size=layout.patch_size, - merge_kernel_size=layout.merge_kernel_size, - in_patch_limit=layout.in_patch_limit, - patch_limit_on_one_side=layout.patch_limit_on_one_side, - fixed_output_tokens=layout.fixed_output_tokens, - image_mean=list(layout.image_mean), - image_std=list(layout.image_std), - ) + # Same canonical knob list and hash as the renderer used at layout time + # (the spec dataclass in renderers.kimi_k25 is the single field list). + return kimi_layout_from(image_processor).fingerprint() def materialize_for_trainer( self, @@ -164,7 +113,7 @@ def synthesize_placeholder( return None import torch - layout = _processor_layout(image_processor) + layout = kimi_layout_from(image_processor) grids: list[list[int]] = [] pixel_values: list[torch.Tensor] = [] for item in items: diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py index 6adf9c0bfa..36735a1ab1 100644 --- a/src/prime_rl/multimodal/adapters/qwen_vl.py +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -3,22 +3,12 @@ import math from typing import Any +from renderers.qwen3_vl import qwen_layout_from + from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem -def _processor_value(processor: Any, name: str, *, size_key: str | None = None) -> int: - value = getattr(processor, name, None) - if value is None and size_key is not None: - size = getattr(processor, "size", None) - get_size_value = getattr(size, "get", None) - if callable(get_size_value): - value = get_size_value(size_key) - if value is None: - raise ValueError(f"Image processor is missing {name}") - return int(value) - - def _tensorize(value: Any): import torch @@ -69,16 +59,9 @@ def validate_item(self, item: RawMMItem) -> None: _grid_payload(item) def processor_fingerprint(self, image_processor: Any) -> str: - from renderers.mm_store import image_layout_fingerprint - - return image_layout_fingerprint( - family=self.family, - patch_size=_processor_value(image_processor, "patch_size"), - merge_size=_processor_value(image_processor, "merge_size"), - temporal_patch_size=_processor_value(image_processor, "temporal_patch_size"), - min_pixels=_processor_value(image_processor, "min_pixels", size_key="shortest_edge"), - max_pixels=_processor_value(image_processor, "max_pixels", size_key="longest_edge"), - ) + # Same canonical knob list and hash as the renderer used at layout time + # (the spec dataclass in renderers.qwen3_vl is the single field list). + return qwen_layout_from(image_processor).fingerprint() def materialize_for_trainer( self, @@ -116,7 +99,7 @@ def materialize_for_vllm( f"Image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" ) hf_inputs = image_processor(images=[image], return_tensors="pt") - merge_size = _processor_value(image_processor, "merge_size") + merge_size = qwen_layout_from(image_processor).merge_size config_by_key = _create_qwen2vl_field_factory(merge_size)(hf_inputs) mm_item = MultiModalKwargsItems.from_hf_inputs(hf_inputs, config_by_key)["image"][0] expected_grid = _grid_payload(item) From 4b4f9e3a6ff6a5d88b96132d2d644216800dcb98 Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 19:15:20 +0000 Subject: [PATCH 34/52] refactor: trust the train-time checkpoint contract at materialize MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Render and materialize are bound to the same model.name, so materialize stops re-litigating layout policy: processor_fingerprint and the fingerprint comparisons are deleted from both materialize paths (vLLM front + trainer) along with the RawMMItem.layout_fingerprint field. The output-level grid and placeholder-length asserts stay — they compare processor output against the ref payload with values already in hand, and turn checkpoint skew into a clean per-request error instead of an engine crash or silent training corruption. Adapters read the few knobs they still need (merge_size, patch_size) straight off the live processor; the renderers-side layout extractors are config-JSON-only now. Co-Authored-By: Claude Fable 5 --- deps/renderers | 2 +- src/prime_rl/inference/vllm/serving_tokens.py | 1 - src/prime_rl/multimodal/adapters/base.py | 2 -- src/prime_rl/multimodal/adapters/kimi_k25.py | 18 ++---------------- src/prime_rl/multimodal/adapters/qwen_vl.py | 14 +------------- src/prime_rl/multimodal/schema.py | 2 -- src/prime_rl/utils/mm.py | 10 ---------- tests/unit/inference/test_serving_tokens.py | 8 ++------ tests/unit/orchestrator/test_batch.py | 1 - tests/unit/train/rl/test_packer.py | 1 - tests/unit/utils/test_mm.py | 1 - 11 files changed, 6 insertions(+), 54 deletions(-) diff --git a/deps/renderers b/deps/renderers index f0d1dd945e..29d443c357 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit f0d1dd945e13935665633c6fdc04de8462e0bf40 +Subproject commit 29d443c3570745661d5d0595c6a5e6db9b564366 diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 67bd910ba6..5ba5736431 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -237,7 +237,6 @@ def _materialize_raw_image_ref_sync( item = RawMMItem( modality=ref.modality, family=ref.family, - layout_fingerprint=ref.fingerprint, raw_image_uri=ref.raw_image_uri, payload=dict(ref.payload), raw_ref=raw_ref, diff --git a/src/prime_rl/multimodal/adapters/base.py b/src/prime_rl/multimodal/adapters/base.py index b746b9f90a..ce5b9e2628 100644 --- a/src/prime_rl/multimodal/adapters/base.py +++ b/src/prime_rl/multimodal/adapters/base.py @@ -28,8 +28,6 @@ class MultimodalAdapter(Protocol): def validate_item(self, item: "RawMMItem") -> None: ... - def processor_fingerprint(self, image_processor: Any) -> str: ... - def materialize_for_trainer( self, image_processor: Any, diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index 319492b361..89653bb7a1 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -3,8 +3,6 @@ import math from typing import Any -from renderers.kimi_k25 import kimi_layout_from - from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem @@ -50,11 +48,6 @@ def validate_item(self, item: RawMMItem) -> None: raise ValueError(f"Kimi adapter cannot handle family {item.family!r}") _grid_payload(item) - def processor_fingerprint(self, image_processor: Any) -> str: - # Same canonical knob list and hash as the renderer used at layout time - # (the spec dataclass in renderers.kimi_k25 is the single field list). - return kimi_layout_from(image_processor).fingerprint() - def materialize_for_trainer( self, image_processor: Any, @@ -84,11 +77,6 @@ def materialize_for_vllm( from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems self.validate_item(item) - actual_fingerprint = self.processor_fingerprint(image_processor) - if actual_fingerprint != item.layout_fingerprint: - raise ValueError( - f"Image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" - ) hf_inputs = _process_images(image_processor, [image], return_tensors="pt") tensors = {str(k): _tensorize(v) for k, v in dict(hf_inputs).items()} expected_grid = _grid_payload(item) @@ -113,16 +101,14 @@ def synthesize_placeholder( return None import torch - layout = kimi_layout_from(image_processor) + patch_size = int(image_processor.media_proc_cfg["patch_size"]) grids: list[list[int]] = [] pixel_values: list[torch.Tensor] = [] for item in items: self.validate_item(item) grid = _grid_payload(item) grids.append(grid) - pixel_values.append( - torch.zeros((math.prod(grid), 3, layout.patch_size, layout.patch_size), dtype=torch.float32) - ) + pixel_values.append(torch.zeros((math.prod(grid), 3, patch_size, patch_size), dtype=torch.float32)) return MaterializedMM( kwargs={ "pixel_values": torch.cat(pixel_values, dim=0).contiguous(), diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py index 36735a1ab1..f6e7105bdf 100644 --- a/src/prime_rl/multimodal/adapters/qwen_vl.py +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -3,8 +3,6 @@ import math from typing import Any -from renderers.qwen3_vl import qwen_layout_from - from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RawMMItem @@ -58,11 +56,6 @@ def validate_item(self, item: RawMMItem) -> None: raise ValueError(f"Qwen adapter cannot handle family {item.family!r}") _grid_payload(item) - def processor_fingerprint(self, image_processor: Any) -> str: - # Same canonical knob list and hash as the renderer used at layout time - # (the spec dataclass in renderers.qwen3_vl is the single field list). - return qwen_layout_from(image_processor).fingerprint() - def materialize_for_trainer( self, image_processor: Any, @@ -93,13 +86,8 @@ def materialize_for_vllm( from vllm.multimodal.inputs import MultiModalKwargsItems self.validate_item(item) - actual_fingerprint = self.processor_fingerprint(image_processor) - if actual_fingerprint != item.layout_fingerprint: - raise ValueError( - f"Image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" - ) hf_inputs = image_processor(images=[image], return_tensors="pt") - merge_size = qwen_layout_from(image_processor).merge_size + merge_size = int(image_processor.merge_size) config_by_key = _create_qwen2vl_field_factory(merge_size)(hf_inputs) mm_item = MultiModalKwargsItems.from_hf_inputs(hf_inputs, config_by_key)["image"][0] expected_grid = _grid_payload(item) diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index 0685274b20..ca8f9453be 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -11,7 +11,6 @@ class RawMMItem: modality: str family: str - layout_fingerprint: str raw_image_uri: str payload: dict[str, Any] raw_ref: str | None = None @@ -56,7 +55,6 @@ def parse_raw_mm_item(value: Any) -> RawMMItem: return RawMMItem( modality=_required_str(descriptor, "modality"), family=_required_str(descriptor, "family"), - layout_fingerprint=_required_str(descriptor, "layout_fingerprint"), raw_image_uri=_required_str(descriptor, "raw_image_uri"), payload=_payload(descriptor), raw_ref=_optional_str(descriptor, "raw_ref"), diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index cbdcff38c8..5cf9d80a53 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -123,15 +123,6 @@ def _single_family_adapter(items: list[RawMMItem]) -> MultimodalAdapter: return get_multimodal_adapter(next(iter(families))) -def _validate_processor_layout(adapter: MultimodalAdapter, image_processor: Any, image_items: list[RawMMItem]) -> None: - actual_fingerprint = adapter.processor_fingerprint(image_processor) - for item in image_items: - if item.layout_fingerprint != actual_fingerprint: - raise ValueError( - f"Raw image layout fingerprint mismatch: expected {item.layout_fingerprint}, got {actual_fingerprint}" - ) - - def _load_verified_images(image_refs: list[MMImageRef]) -> list[Any]: from PIL import Image @@ -173,7 +164,6 @@ def materialize(self, refs: MMRefs) -> MaterializedMM | None: image_processor = self.image_processor adapter = _single_family_adapter(image_items) - _validate_processor_layout(adapter, image_processor, image_items) images = _load_verified_images(refs.images) return adapter.materialize_for_trainer(image_processor, image_items, images) diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 80be7c3099..723f526a23 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -282,10 +282,8 @@ def test_materialize_raw_image_ref_uses_generic_family_payload(tmp_path, monkeyp Image.new("RGB", (8, 6), color=(32, 64, 128)).save(image_path) mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] - fingerprint = "f" * 32 raw_ref = raw_mm_ref( family="test_family", - fingerprint=fingerprint, modality="image", mm_hash=mm_hash, raw_image_uri=image_path.as_uri(), @@ -325,7 +323,6 @@ def _get_adapter(family): assert captured["expected_placeholder_length"] == 7 item = captured["item"] assert item.family == "test_family" - assert item.layout_fingerprint == fingerprint assert item.raw_image_uri == image_path.as_uri() assert item.payload == {"adapter_owned": [1, 2, 3]} @@ -341,12 +338,12 @@ def test_materialize_raw_image_ref_maps_adapter_validation_error(tmp_path, monke class _InvalidAdapter: def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): - raise ValueError("image layout fingerprint mismatch") + raise ValueError("image grid mismatch") monkeypatch.setattr(serving_tokens, "_load_image_processor", lambda _model, _trust: object()) monkeypatch.setattr(serving_tokens, "get_multimodal_adapter", lambda _family: _InvalidAdapter()) - with pytest.raises(serving_tokens._MMImageRefError, match="image layout fingerprint mismatch") as exc_info: + with pytest.raises(serving_tokens._MMImageRefError, match="image grid mismatch") as exc_info: serving_tokens._materialize_raw_image_ref_sync( raw_ref, feature_modality="image", @@ -374,7 +371,6 @@ def _mm_features(tmp_path, *, image_seed: int = 0, placeholder_length: int = 7): mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] raw_ref = raw_mm_ref( family="test_family", - fingerprint="f" * 32, modality="image", mm_hash=mm_hash, raw_image_uri=image_path.as_uri(), diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 5b0236596d..f842b52fd2 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -417,7 +417,6 @@ def _image_ref(uri: str, offset: int, length: int) -> MMImageRef: "kind": "prime_raw_mm_item", "modality": "image", "family": "qwen_vl", - "layout_fingerprint": "f" * 32, "raw_image_uri": uri, "payload": {"image_grid_thw": [[1, 1, 1]]}, }, diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index eab4ac31ab..1c1e83724d 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -68,7 +68,6 @@ def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: "kind": "prime_raw_mm_item", "modality": "image", "family": "qwen_vl", - "layout_fingerprint": "f" * 32, "raw_image_uri": uri, "payload": {"image_grid_thw": [[1, 1, 1]]}, }, diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index abb80d42a1..1be2f28812 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -28,7 +28,6 @@ def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): "kind": RAW_MM_ITEM_KIND, "modality": "image", "family": "qwen_vl", - "layout_fingerprint": "f" * 32, "raw_image_uri": uri, "payload": {"image_grid_thw": grid}, } From 52adc5919e63d672cd45ec7a3eb8a5f1cc66fc25 Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 20:05:58 +0000 Subject: [PATCH 35/52] refactor: prune RawMMItem to the fields materialize reads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RawMMItem mirrors only what its consumers use now: family (adapter routing), the image source, and the adapter-owned payload. The modality, raw_ref, and vllm_modality mirror fields had zero readers — modality lives on the container key and the mmraw: ref, and vllm_modality is consumed at the renderers client. Drops the now-unused _optional_str helper and the dead envelope keys from test fixtures. deps/renderers advances to the matching envelope change. Co-Authored-By: Claude Fable 5 --- deps/renderers | 2 +- src/prime_rl/inference/vllm/serving_tokens.py | 2 -- src/prime_rl/multimodal/schema.py | 13 ------------- tests/unit/inference/test_serving_tokens.py | 2 +- tests/unit/orchestrator/test_batch.py | 1 - tests/unit/train/rl/test_packer.py | 1 - tests/unit/utils/test_mm.py | 1 - 7 files changed, 2 insertions(+), 20 deletions(-) diff --git a/deps/renderers b/deps/renderers index 29d443c357..aa4b0c66f9 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 29d443c3570745661d5d0595c6a5e6db9b564366 +Subproject commit aa4b0c66f992b4ca0fef0ed6ba4da773bdf7f619 diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 5ba5736431..86c5b9ea8f 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -235,11 +235,9 @@ def _materialize_raw_image_ref_sync( image = _decode_raw_image(raw, raw_image_uri=ref.raw_image_uri) image_processor = _load_image_processor(processor_model_name, trust_remote_code) item = RawMMItem( - modality=ref.modality, family=ref.family, raw_image_uri=ref.raw_image_uri, payload=dict(ref.payload), - raw_ref=raw_ref, ) try: adapter = get_multimodal_adapter(ref.family) diff --git a/src/prime_rl/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index ca8f9453be..f26ec1673b 100644 --- a/src/prime_rl/multimodal/schema.py +++ b/src/prime_rl/multimodal/schema.py @@ -9,12 +9,9 @@ @dataclass(frozen=True) class RawMMItem: - modality: str family: str raw_image_uri: str payload: dict[str, Any] - raw_ref: str | None = None - vllm_modality: str | None = None def _descriptor_mapping(value: Any) -> Mapping[str, Any]: @@ -30,13 +27,6 @@ def _required_str(value: Mapping[str, Any], field: str) -> str: raise ValueError(f"raw multimodal descriptor is missing {field}") -def _optional_str(value: Mapping[str, Any], field: str) -> str | None: - item = value.get(field) - if item is None or isinstance(item, str): - return item - raise ValueError(f"raw multimodal descriptor {field} must be a string when present") - - def _payload(value: Mapping[str, Any]) -> dict[str, Any]: payload = value.get("payload") if not isinstance(payload, Mapping): @@ -53,10 +43,7 @@ def parse_raw_mm_item(value: Any) -> RawMMItem: descriptor = _descriptor_mapping(value) _validate_envelope(descriptor) return RawMMItem( - modality=_required_str(descriptor, "modality"), family=_required_str(descriptor, "family"), raw_image_uri=_required_str(descriptor, "raw_image_uri"), payload=_payload(descriptor), - raw_ref=_optional_str(descriptor, "raw_ref"), - vllm_modality=_optional_str(descriptor, "vllm_modality"), ) diff --git a/tests/unit/inference/test_serving_tokens.py b/tests/unit/inference/test_serving_tokens.py index 723f526a23..9b01b84a26 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -388,7 +388,7 @@ def _patch_adapter(monkeypatch, calls: list, materialize=None): class _Adapter: def materialize_for_vllm(self, image_processor, item, image, expected_placeholder_length): - calls.append(item.raw_ref) + calls.append(item.raw_image_uri) if materialize is not None: return materialize(item) return {"materialized": item.raw_image_uri} diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index f842b52fd2..c85115d0eb 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -415,7 +415,6 @@ def _image_ref(uri: str, offset: int, length: int) -> MMImageRef: return MMImageRef( item={ "kind": "prime_raw_mm_item", - "modality": "image", "family": "qwen_vl", "raw_image_uri": uri, "payload": {"image_grid_thw": [[1, 1, 1]]}, diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 1c1e83724d..09a59a598a 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -66,7 +66,6 @@ def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: MMImageRef( item={ "kind": "prime_raw_mm_item", - "modality": "image", "family": "qwen_vl", "raw_image_uri": uri, "payload": {"image_grid_thw": [[1, 1, 1]]}, diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index 1be2f28812..b67f544b11 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -26,7 +26,6 @@ def synthesize_placeholder(self, refs): def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): return { "kind": RAW_MM_ITEM_KIND, - "modality": "image", "family": "qwen_vl", "raw_image_uri": uri, "payload": {"image_grid_thw": grid}, From f753800d24c361675f665301d025a5ac1d429ebb Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 21:27:32 +0000 Subject: [PATCH 36/52] chore: bump renderers for shared mm_image helpers Points deps/renderers at the commit that lifts cross-family image I/O helpers out of qwen3_vl into renderers.mm_image. Co-authored-by: Cursor --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index aa4b0c66f9..6b2d20822a 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit aa4b0c66f992b4ca0fef0ed6ba4da773bdf7f619 +Subproject commit 6b2d20822ad863c34f88a4f31ee04412ca789662 From 21e0430ce0d575d17c8b8302bf4540964d9c218c Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 21:57:11 +0000 Subject: [PATCH 37/52] chore: bump renderers after dropping image cache Points deps/renderers at the commit that removes the processed-mode per-renderer image cache and stale image_cache_max docs. Co-authored-by: Cursor --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index 6b2d20822a..acd28d21a7 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 6b2d20822ad863c34f88a4f31ee04412ca789662 +Subproject commit acd28d21a74e953ea1b18934f3f70691b3ba7fc4 From 8e002917e4050eb973dcd8b0866f6f4e620b41e8 Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 22:11:46 +0000 Subject: [PATCH 38/52] chore: bump renderers for local-only processed image sources Co-authored-by: Cursor --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index acd28d21a7..2fbf8bdb1a 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit acd28d21a74e953ea1b18934f3f70691b3ba7fc4 +Subproject commit 2fbf8bdb1aa3cb1a0016e8d7dc2aab57b140664c From 0e5403243cadee811cb38fa109177120421e2da7 Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 6 Aug 2026 03:18:39 +0000 Subject: [PATCH 39/52] chore: bump renderers for the Pillow-only vision extra Co-Authored-By: Claude Fable 5 --- deps/renderers | 2 +- uv.lock | 2 -- 2 files changed, 1 insertion(+), 3 deletions(-) diff --git a/deps/renderers b/deps/renderers index 2fbf8bdb1a..3aeb6eef36 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 2fbf8bdb1aa3cb1a0016e8d7dc2aab57b140664c +Subproject commit 3aeb6eef363d8c6a1a44863f16a5d6fc9ef021bb diff --git a/uv.lock b/uv.lock index b49ef38c89..ca12338e13 100644 --- a/uv.lock +++ b/uv.lock @@ -6190,8 +6190,6 @@ requires-dist = [ { name = "pillow", marker = "extra == 'vision'", specifier = ">=12.2.0" }, { name = "prime-pydantic-config", specifier = ">=0.3.0.dev83" }, { name = "tiktoken" }, - { name = "torch", marker = "extra == 'vision'", specifier = ">=2.11.0" }, - { name = "torchvision", marker = "extra == 'vision'", specifier = ">=0.26.0" }, { name = "transformers", specifier = ">=4.50.0" }, ] provides-extras = ["vision"] From 98480ba6a45e7ef99a69f6f91ccf3cf34ca00590 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 00:10:46 +0000 Subject: [PATCH 40/52] chore: bump verifiers to latest raw-image-offload head Co-authored-by: Cursor --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index c9d1f134c5..d754774ae2 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit c9d1f134c540a327bfe52d7e53560263a756efa7 +Subproject commit d754774ae2d11806792ac854c8e547fd23c4eac0 From 8bcd8117e19bcd4d27e40d4169c4b519b641d124 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 02:17:08 +0000 Subject: [PATCH 41/52] refactor: env servers resolve their own image-asset dir MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follows the review pattern on #2836: the launcher stops painting VF_RENDERER_IMAGE_OFFLOAD_DIR onto env-server processes. EnvServerConfig carries multimodal + the run output_dir (the launcher already writes one toml per server), and the env-server entrypoint resolves the dir once at startup via os.environ.setdefault — so an operator-set env var wins, config-derived defaults otherwise. Standalone env servers now get the offload dir from config instead of requiring a hand-exported var. Deletes the orphaned build_run_asset_env. Co-Authored-By: Claude Fable 5 --- deps/renderers | 2 +- .../src/prime_rl/configs/env_server.py | 5 ++++- src/prime_rl/entrypoints/env_server.py | 7 +++++++ src/prime_rl/entrypoints/rl.py | 8 ++------ src/prime_rl/utils/run_assets.py | 19 ------------------- 5 files changed, 14 insertions(+), 27 deletions(-) diff --git a/deps/renderers b/deps/renderers index 3aeb6eef36..ab603543ec 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 3aeb6eef363d8c6a1a44863f16a5d6fc9ef021bb +Subproject commit ab603543ecd45f14d684a512f5b1af1853aaa287 diff --git a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py index 8e5ed3f3a0..935227652b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/env_server.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/env_server.py @@ -3,7 +3,7 @@ import verifiers.v1 as vf from pydantic import SerializeAsAny, model_validator -from prime_rl.configs.shared import LogConfig +from prime_rl.configs.shared import LogConfig, MultimodalConfig from prime_rl.utils.config import BaseConfig @@ -26,6 +26,9 @@ class EnvServerConfig(BaseConfig): output_dir: Path = Path("outputs") """Directory to write outputs to — logs and any generated artifacts are written as subdirectories.""" + multimodal: MultimodalConfig = MultimodalConfig() + """Raw multimodal image offload settings; with ``output_dir``, resolves the run image-asset dir this server offloads into.""" + @model_validator(mode="before") @classmethod def _resolve_env(cls, data): diff --git a/src/prime_rl/entrypoints/env_server.py b/src/prime_rl/entrypoints/env_server.py index 05951ba031..af245fc087 100644 --- a/src/prime_rl/entrypoints/env_server.py +++ b/src/prime_rl/entrypoints/env_server.py @@ -1,3 +1,4 @@ +import os from functools import partial from verifiers.v1 import pool_serve_kwargs @@ -7,11 +8,17 @@ from prime_rl.orchestrator.utils import setup_env_server_logging from prime_rl.utils.config import cli from prime_rl.utils.process import set_proc_title +from prime_rl.utils.run_assets import IMAGE_OFFLOAD_DIR_ENV, resolve_image_offload_dir from prime_rl.utils.utils import clean_exit @clean_exit def run_server(config: EnvServerConfig): + # Renderers offload images to the dir named by this env var; resolve it from + # this server's own config, letting an operator-set env win (multi-node). + os.environ.setdefault( + IMAGE_OFFLOAD_DIR_ENV, str(resolve_image_offload_dir(config.output_dir, config.multimodal, os.environ)) + ) # ``serve.pool`` (static or elastic) sizes the server; a v0/legacy env runs through # the bridge, a v1 env is a native env block — both speak the same serve protocol, # so the orchestrator is agnostic. serve_env applies the logging setup in this process diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index c35dc317f9..2aef461cd0 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -37,7 +37,6 @@ monitor_process, set_proc_title, ) -from prime_rl.utils.run_assets import build_run_asset_env RL_TOML = "rl.toml" RL_SBATCH = "rl.sbatch" @@ -106,6 +105,8 @@ def write_subconfigs(config: RLConfig, output_dir: Path) -> None: "serve": {**source_dict.get("serve", {}), "address": address}, "legacy": source_dict.get("legacy", {}), "log": {"level": config.orchestrator.log.vf_level, "json_logging": config.orchestrator.log.json_logging}, + "output_dir": str(config.orchestrator.output_dir), + "multimodal": to_toml_dict(config.multimodal), } with open(env_dir / f"{source.resolved_name}.toml", "wb") as f: tomli_w.dump(env_server_dict, f) @@ -159,7 +160,6 @@ def rl_local(config: RLConfig): "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } inherited_env = dict(os.environ) - writer_run_asset_env = build_run_asset_env(config.orchestrator.output_dir, multimodal=config.multimodal) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.model.client.is_elastic: @@ -264,10 +264,6 @@ def sigterm_handler(signum, frame): env_server_cmd, env={ **os.environ, - # Env servers render prompts and offload raw images, so they - # need the run's image-asset dir (standalone servers set it - # themselves). - **writer_run_asset_env, **DEFAULT_COMMON_ENV_VARS, **config.env_vars, **config.orchestrator.env_vars, diff --git a/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py index 45ec1fd449..d15ede0cf1 100644 --- a/src/prime_rl/utils/run_assets.py +++ b/src/prime_rl/utils/run_assets.py @@ -41,22 +41,3 @@ def resolve_image_offload_dir( if hosted_run_dir is not None: return (hosted_run_dir / IMAGE_ASSET_SUBDIR).resolve() return (output_dir.resolve() / IMAGE_ASSET_SUBDIR).resolve() - - -def build_run_asset_env( - output_dir: Path, - multimodal: MultimodalConfig | None = None, - base: Mapping[str, str] | None = None, -) -> dict[str, str]: - """Resolve the environment used by subprocesses that share run image assets. - - Prime-RL config owns the multimodal image offload path. Env vars are only the - transport used by verifiers/renderers running in subprocesses. - """ - - env = dict(os.environ if base is None else base) - config = multimodal or MultimodalConfig() - - env[IMAGE_OFFLOAD_DIR_ENV] = str(resolve_image_offload_dir(output_dir, config, env)) - - return env From 3888738a50d8ca9dbd4da81520f014f79bdc424a Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 02:19:23 +0000 Subject: [PATCH 42/52] chore: relock for the verifiers pin; exclude nemo-gym-weather-v1 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit nemo-gym-weather-v1 (verifiers #2212) needs verifiers[nemo-gym], whose nemo-gym==0.4.0 pins openai<=2.7.2 against verifiers' own openai>=2.9.0 — same shared-lock conflict class as the excluded tau2 forks. Co-Authored-By: Claude Fable 5 --- pyproject.toml | 3 +++ uv.lock | 5 +++-- 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2274aa3aa0..7a468e3a45 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -133,9 +133,12 @@ members = [ ] # tau2-synth-v1 and tau3-bench-v1 pin `tau2` forks/revs that conflict with the # canonical tau2-bench-v1; a workspace shares one lock, so these can't coexist. +# nemo-gym-weather-v1 needs verifiers[nemo-gym], whose nemo-gym==0.4.0 pins +# openai<=2.7.2 against verifiers' own openai>=2.9.0. exclude = [ "deps/research-environments/environments/tool_use/tau2_synth_v1", "deps/research-environments/environments/tool_use/tau3_bench_v1", + "deps/verifiers/environments/nemo_gym_weather_v1", ] # ModelExpress 0.3.0 publishes protobuf<6 metadata, but its generated proto is diff --git a/uv.lock b/uv.lock index ca12338e13..7403c9927e 100644 --- a/uv.lock +++ b/uv.lock @@ -6200,7 +6200,7 @@ dev = [ { name = "pre-commit" }, { name = "pytest", specifier = ">=7.0.0" }, { name = "pytest-asyncio", specifier = ">=0.21.0" }, - { name = "ruff" }, + { name = "ruff", specifier = "==0.15.12" }, { name = "torch", specifier = ">=2.11.0" }, { name = "torchvision", specifier = ">=0.26.0" }, { name = "ty", specifier = ">=0.0.1a29,<0.0.22" }, @@ -7920,6 +7920,7 @@ requires-dist = [ { name = "mcp", specifier = ">=1.24.0,<2" }, { name = "modal", marker = "extra == 'modal'", specifier = ">=1.4.0" }, { name = "msgpack", specifier = ">=1.1.2" }, + { name = "nemo-gym", marker = "python_full_version >= '3.12' and extra == 'nemo-gym'", specifier = "==0.4.0" }, { name = "nest-asyncio", marker = "extra == 'notebook'", specifier = ">=1.6.0" }, { name = "nltk", marker = "extra == 'ta'", specifier = ">=3.9.2" }, { name = "numpy", specifier = ">=2.1.0" }, @@ -7944,7 +7945,7 @@ requires-dist = [ { name = "typing-extensions", specifier = ">=4.12.2" }, { name = "uvloop", marker = "platform_python_implementation != 'PyPy' and sys_platform != 'cygwin' and sys_platform != 'win32'", specifier = ">=0.21.0" }, ] -provides-extras = ["browser", "harbor", "modal", "notebook", "openenv", "rg", "ta"] +provides-extras = ["browser", "harbor", "modal", "nemo-gym", "notebook", "openenv", "rg", "ta"] [package.metadata.requires-dev] dev = [ From b92dbd59b30980a7562fecc5a1e83e1dc6fe4832 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 02:20:36 +0000 Subject: [PATCH 43/52] chore: bump verifiers for the nemo-gym relock Co-Authored-By: Claude Fable 5 --- deps/verifiers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/verifiers b/deps/verifiers index d754774ae2..333052c0bd 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit d754774ae2d11806792ac854c8e547fd23c4eac0 +Subproject commit 333052c0bd6e443b1485481ae9c5d48a3f5a8352 From ab09483fb6916b66c9baf8b411baa43b0a2416fb Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 02:27:47 +0000 Subject: [PATCH 44/52] refactor(trainer): always fail-fast on missing MM images Drop missing_mm_image_policy and the zero-loss placeholder soft path so a missing offloaded image kills the run instead of training through zeros. Co-authored-by: Cursor --- .../src/prime_rl/configs/trainer.py | 4 -- skills/training/monitor-run/SKILL.md | 8 ++- src/prime_rl/multimodal/adapters/base.py | 6 --- src/prime_rl/multimodal/adapters/kimi_k25.py | 26 ---------- src/prime_rl/multimodal/adapters/qwen_vl.py | 50 ------------------- src/prime_rl/trainer/rl/data.py | 35 +++---------- src/prime_rl/trainer/rl/train.py | 2 - src/prime_rl/utils/mm.py | 20 -------- tests/unit/utils/test_mm.py | 26 ++-------- 9 files changed, 13 insertions(+), 164 deletions(-) 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 bfc43e5a32..eb6edf9c70 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -22,7 +22,6 @@ AttnImplementation: TypeAlias = Literal["flash_attention_2", "flash_attention_3", "flash_attention_4", "auto"] EPCommBackend: TypeAlias = Literal["torch", "deepep"] -MissingMMImagePolicy: TypeAlias = Literal["error", "placeholder_zero_loss"] class GCConfig(BaseConfig): @@ -630,9 +629,6 @@ class TrainerConfig(BaseConfig): max_concurrent_runs: int = Field(1, ge=1) """Maximum number of concurrent runs to allow. If 1, only one run may run at a time.""" - missing_mm_image_policy: MissingMMImagePolicy = "placeholder_zero_loss" - """Policy when raw multimodal image files disappear before trainer materialization. ``placeholder_zero_loss`` warns, synthesizes zero-valued image tensors with the original descriptor geometry, and masks out the affected microbatch loss; ``error`` preserves fail-fast behavior.""" - multimodal: MultimodalConfig = MultimodalConfig() """Raw multimodal image offload settings shared with orchestrator and inference.""" diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 8e67d91106..bb056ea1ff 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -199,11 +199,9 @@ means ingress was bypassed. - The orchestrator raises on placeholder/token drift ("does not cover image-typed tokens") before a sample ships — treat any occurrence as a bug, not noise. -- Trainer metrics: `mm/images_materialized`, `time/mm_materialize`, and - `mm/images_placeholdered`. A nonzero placeholder count means image files - disappeared before materialization (the batch trains with zero-loss - placeholders and logs "raw image materialization missing image(s)") — check - whether something cleaned the offload directory mid-run. +- Trainer metrics: `mm/images_materialized` and `time/mm_materialize`. Missing + image files fail the trainer hard ("raw image materialization failed") — + check whether something cleaned the offload directory mid-run. ### Process tree diff --git a/src/prime_rl/multimodal/adapters/base.py b/src/prime_rl/multimodal/adapters/base.py index ce5b9e2628..c6217e13dc 100644 --- a/src/prime_rl/multimodal/adapters/base.py +++ b/src/prime_rl/multimodal/adapters/base.py @@ -42,9 +42,3 @@ def materialize_for_vllm( image: "Image", expected_placeholder_length: int, ) -> Any: ... - - def synthesize_placeholder( - self, - image_processor: Any, - items: list["RawMMItem"], - ) -> MaterializedMM | None: ... diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index 89653bb7a1..cf3410b16f 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -1,6 +1,5 @@ from __future__ import annotations -import math from typing import Any from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM @@ -91,28 +90,3 @@ def materialize_for_vllm( "grid_thws": MultiModalFieldConfig.batched("vision_chunk"), } return MultiModalKwargsItems.from_hf_inputs(tensors, config_by_key)["vision_chunk"][0] - - def synthesize_placeholder( - self, - image_processor: Any, - items: list[RawMMItem], - ) -> MaterializedMM | None: - if not items: - return None - import torch - - patch_size = int(image_processor.media_proc_cfg["patch_size"]) - grids: list[list[int]] = [] - pixel_values: list[torch.Tensor] = [] - for item in items: - self.validate_item(item) - grid = _grid_payload(item) - grids.append(grid) - pixel_values.append(torch.zeros((math.prod(grid), 3, patch_size, patch_size), dtype=torch.float32)) - return MaterializedMM( - kwargs={ - "pixel_values": torch.cat(pixel_values, dim=0).contiguous(), - "grid_thws": torch.tensor(grids, dtype=torch.long), - }, - forward_policy=self.forward_policy, - ) diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py index f6e7105bdf..d793419d2b 100644 --- a/src/prime_rl/multimodal/adapters/qwen_vl.py +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -1,6 +1,5 @@ from __future__ import annotations -import math from typing import Any from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM @@ -31,19 +30,6 @@ def _grid_payload(item: RawMMItem) -> list[int]: return out -def _patch_area(patch_size: Any) -> int: - if isinstance(patch_size, list | tuple): - return math.prod(int(dim) for dim in patch_size) - size = int(patch_size) - return size * size - - -def _temporal_patch_extent(temporal_patch_size: Any) -> int: - if isinstance(temporal_patch_size, list | tuple): - return math.prod(int(dim) for dim in temporal_patch_size) - return int(temporal_patch_size) - - class QwenVLAdapter: family = "qwen_vl" forward_policy = ForwardPolicy( @@ -100,39 +86,3 @@ def materialize_for_vllm( f"Image placeholder length mismatch: expected {expected_placeholder_length}, got {num_image_tokens}" ) return mm_item - - def placeholder_feature_dim(self, image_processor: Any) -> int: - patch_size = getattr(image_processor, "patch_size", None) - temporal_patch_size = getattr(image_processor, "temporal_patch_size", None) - image_mean = getattr(image_processor, "image_mean", None) - channels = len(image_mean) if image_mean is not None else getattr(image_processor, "num_channels", 3) - if patch_size is None or temporal_patch_size is None: - raise ValueError( - "Cannot synthesize raw image placeholders without image processor patch_size and temporal_patch_size" - ) - return int(channels) * _temporal_patch_extent(temporal_patch_size) * _patch_area(patch_size) - - def synthesize_placeholder( - self, - image_processor: Any, - items: list[RawMMItem], - ) -> MaterializedMM | None: - if not items: - return None - import torch - - feature_dim = self.placeholder_feature_dim(image_processor) - pixel_values: list[torch.Tensor] = [] - image_grid_thw: list[list[int]] = [] - for item in items: - self.validate_item(item) - grid = _grid_payload(item) - pixel_values.append(torch.zeros((math.prod(grid), feature_dim), dtype=torch.float32)) - image_grid_thw.append(grid) - return MaterializedMM( - kwargs={ - "pixel_values": torch.cat(pixel_values, dim=0).contiguous(), - "image_grid_thw": torch.tensor(image_grid_thw, dtype=torch.long), - }, - forward_policy=self.forward_policy, - ) diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 5e98d9636a..0f09e00fea 100644 --- a/src/prime_rl/trainer/rl/data.py +++ b/src/prime_rl/trainer/rl/data.py @@ -7,7 +7,7 @@ from jaxtyping import Bool, Float, Int from torch import Tensor -from prime_rl.configs.trainer import FakeDataLoaderConfig, MissingMMImagePolicy +from prime_rl.configs.trainer import FakeDataLoaderConfig from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.trainer.rl.packer import BasePacker, setup_packer from prime_rl.trainer.runs import get_multi_run_manager @@ -20,7 +20,7 @@ ) from prime_rl.transport.types import MMRefs from prime_rl.utils.logger import get_logger -from prime_rl.utils.mm import RawImageMaterializer, missing_file_uris +from prime_rl.utils.mm import RawImageMaterializer class TensorMicroBatch(TypedDict): @@ -79,7 +79,6 @@ def __init__(self, config: FakeDataLoaderConfig, seq_len: int, dp_world_size: in self.multi_run_manager = get_multi_run_manager() self.last_mm_materialize_time = 0.0 self.last_mm_images_materialized = 0 - self.last_mm_images_placeholdered = 0 def wait_for_batch(self) -> None: return @@ -199,7 +198,6 @@ def __init__( config: TransportConfig, model_name: str, model_trust_remote_code: bool, - missing_mm_image_policy: MissingMMImagePolicy = "placeholder_zero_loss", ): self.world = get_world() @@ -219,7 +217,6 @@ def __init__( self.receiver: MicroBatchReceiver = setup_micro_batch_receiver(output_dir, dp_rank, start_step, config) self.mm_materializer = RawImageMaterializer(model_name, trust_remote_code=model_trust_remote_code) - self.missing_mm_image_policy = missing_mm_image_policy self._reset_mm_stats() def wait_for_batch(self) -> None: @@ -240,7 +237,6 @@ def get_batch(self) -> list[TensorMicroBatch]: def _reset_mm_stats(self) -> None: self.last_mm_materialize_time = 0.0 self.last_mm_images_materialized = 0 - self.last_mm_images_placeholdered = 0 @staticmethod def _materialized_mm_parts(materialized: MaterializedMM | None) -> MaterializedMMParts: @@ -259,34 +255,15 @@ def _materialize_mm_refs(self, micro_batch: MicroBatch, refs: MMRefs) -> Materia materialized = self.mm_materializer.materialize(refs) except FileNotFoundError as exc: self.last_mm_materialize_time += time.perf_counter() - materialize_start - if self.missing_mm_image_policy == "error": - get_logger().error( - f"raw image materialization failed ({self._mm_run_context(micro_batch)}, uris={refs.uris}): {exc!r}" - ) - raise - return self._synthesize_missing_mm_placeholder(micro_batch, refs) + get_logger().error( + f"raw image materialization failed ({self._mm_run_context(micro_batch)}, uris={refs.uris}): {exc!r}" + ) + raise self.last_mm_materialize_time += time.perf_counter() - materialize_start self.last_mm_images_materialized += len(refs.uris) return self._materialized_mm_parts(materialized) - def _synthesize_missing_mm_placeholder(self, micro_batch: MicroBatch, refs: MMRefs) -> MaterializedMMParts: - placeholder_start = time.perf_counter() - materialized = self.mm_materializer.synthesize_placeholder(refs) - self.last_mm_materialize_time += time.perf_counter() - placeholder_start - self.last_mm_images_placeholdered += len(refs.uris) - micro_batch.loss_mask = [False] * len(micro_batch.loss_mask) - micro_batch.advantages = [0.0] * len(micro_batch.advantages) - - missing_uris = missing_file_uris(refs.uris) - get_logger().warning( - "raw image materialization missing image(s); using zero-loss placeholder " - f"({self._mm_run_context(micro_batch)}, " - f"missing_uris={missing_uris or ['']}, " - f"uris={refs.uris})" - ) - return self._materialized_mm_parts(materialized) - def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: """Convert a MicroBatch (msgspec struct with lists) to a TensorMicroBatch (dict with tensors).""" if micro_batch.lora_num_tokens is None: diff --git a/src/prime_rl/trainer/rl/train.py b/src/prime_rl/trainer/rl/train.py index 4e3497999e..0b5149a9a2 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -253,7 +253,6 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: config.rollout_transport, config.model.name, config.model.trust_remote_code, - missing_mm_image_policy=config.missing_mm_image_policy, ) token_exporter = setup_token_exporter(config, parallel_dims, world, logger) @@ -724,7 +723,6 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: "time/save_ckpt": save_ckpt_time, "time/forward_backward": forward_backward_time, "mm/images_materialized": dataloader.last_mm_images_materialized, - "mm/images_placeholdered": dataloader.last_mm_images_placeholdered, "step": progress.step, } monitor.log(time_metrics, step=progress.step) diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index 5cf9d80a53..089dbf2c02 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -31,17 +31,6 @@ def file_uri_to_path(uri: str) -> Path: return Path(unquote(parsed.path)) -def missing_file_uris(uris: Iterable[str]) -> list[str]: - """Return missing local ``file://`` image refs; non-file refs are ignored.""" - missing: list[str] = [] - for uri in uris: - if urlparse(uri).scheme != "file": - continue - if not file_uri_to_path(uri).exists(): - missing.append(uri) - return missing - - def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: unsupported = sorted( modality for modality, items in mm_items.items() if items and modality not in SUPPORTED_MODALITIES @@ -166,12 +155,3 @@ def materialize(self, refs: MMRefs) -> MaterializedMM | None: adapter = _single_family_adapter(image_items) images = _load_verified_images(refs.images) return adapter.materialize_for_trainer(image_processor, image_items, images) - - def synthesize_placeholder(self, refs: MMRefs) -> MaterializedMM | None: - """Build zero-valued multimodal tensors via the owning adapter.""" - image_items = _parse_image_refs(refs) - if not image_items: - return None - image_processor = self.image_processor - adapter = _single_family_adapter(image_items) - return adapter.synthesize_placeholder(image_processor, image_items) diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index b67f544b11..e891435693 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -1,8 +1,7 @@ from types import SimpleNamespace -import torch +import pytest -from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND from prime_rl.trainer.rl.data import DataLoader from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs @@ -13,15 +12,6 @@ class _MissingMaterializer: def materialize(self, refs): raise FileNotFoundError("missing image") - def synthesize_placeholder(self, refs): - return MaterializedMM( - kwargs={ - "pixel_values": torch.zeros((1, 24), dtype=torch.float32), - "image_grid_thw": torch.tensor([[1, 1, 1]], dtype=torch.long), - }, - forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), - ) - def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): return { @@ -42,10 +32,8 @@ def _loader() -> DataLoader: loader = object.__new__(DataLoader) loader.multi_run_manager = SimpleNamespace(max_runs=1) loader.mm_materializer = _MissingMaterializer() - loader.missing_mm_image_policy = "placeholder_zero_loss" loader.last_mm_materialize_time = 0.0 loader.last_mm_images_materialized = 0 - loader.last_mm_images_placeholdered = 0 return loader @@ -80,12 +68,6 @@ def test_build_mm_refs_accepts_offloaded_file_uri(tmp_path): assert refs == _refs(image_path.as_uri()) -def test_dataloader_uses_zero_loss_placeholder_for_missing_raw_image(): - tensor_batch = _loader()._micro_batch_to_tensor(_micro_batch()) - - assert tensor_batch["mm_kwargs"] is not None - assert tensor_batch["mm_kwargs"]["pixel_values"].shape == (1, 24) - assert tensor_batch["mm_forward_policy"] == ForwardPolicy(pass_position_ids_with_mm=False) - assert tensor_batch["loss_mask"].tolist() == [[False, False, False]] - assert tensor_batch["advantages"].tolist() == [[0.0, 0.0, 0.0]] - assert tensor_batch["mm_token_type_ids"].tolist() == [[0, 1, 0]] +def test_dataloader_raises_on_missing_raw_image(): + with pytest.raises(FileNotFoundError, match="missing image"): + _loader()._micro_batch_to_tensor(_micro_batch()) From bbb9dce69404f1419c46f892c491f431e6821d01 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 02:46:03 +0000 Subject: [PATCH 45/52] fix(inference): retain strong refs to mm materialize tasks asyncio only weakly references tasks, so create_task without a holder can drop in-flight materialization before the shared future resolves. Co-authored-by: Cursor --- src/prime_rl/inference/vllm/serving_tokens.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 86c5b9ea8f..c906d809d7 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -285,6 +285,8 @@ def __init__(self, max_bytes: int) -> None: self.max_bytes = max_bytes self._items: OrderedDict[_MaterializeKey, tuple[Any, int]] = OrderedDict() self._inflight: dict[_MaterializeKey, asyncio.Future] = {} + # asyncio only weakly refs tasks; keep strong refs until they finish. + self._tasks: set[asyncio.Task] = set() self._bytes = 0 self.hits = 0 self.misses = 0 @@ -325,7 +327,9 @@ async def _materialize_and_resolve() -> None: # The work runs as its own task and awaiters shield the shared future: # a cancelled request (client disconnect) must neither poison the future # for deduped awaiters nor cancel it out from under them. - asyncio.get_running_loop().create_task(_materialize_and_resolve()) + task = asyncio.get_running_loop().create_task(_materialize_and_resolve()) + self._tasks.add(task) + task.add_done_callback(self._tasks.discard) return await asyncio.shield(future) def _insert(self, key: _MaterializeKey, item: Any) -> None: From 49fa7091fd2fca5a17c7bd27ed1fe6145252de1a Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 02:58:15 +0000 Subject: [PATCH 46/52] docs: drop advanced.md VLM limitation churn from this PR Co-authored-by: Cursor --- docs/advanced.md | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/docs/advanced.md b/docs/advanced.md index e9d6e5f41f..c4c1ba9b9c 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -98,8 +98,7 @@ 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 it; in that case it's FSDP-sharded per block. The combination `freeze_vision_encoder = false` + LoRA is rejected by a config validator — LoRA freezes everything non-adapter, so unfreezing the encoder under LoRA would be a silent no-op. -- **Truncation cuts at image boundaries.** Raw multimodal samples that exceed `seq_len` are truncated before packing at the start of the first image that doesn't fit — an image placeholder span is never split, and refs for dropped images are discarded. A sample whose leading image alone exceeds `seq_len` is rejected. +- **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. From 0ea481a5d66c19f8ff296cb00626243cdba0fd21 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 03:01:23 +0000 Subject: [PATCH 47/52] docs: correct VF_RENDERER_IMAGE_OFFLOAD_DIR wiring notes Env-servers resolve the path from [multimodal].offload_dir; the launcher no longer paints that env var last on subprocesses. Co-authored-by: Cursor --- docs/configuration.md | 4 +++- .../prime-rl-configs/src/prime_rl/configs/shared.py | 12 ++++++++---- skills/training/monitor-run/SKILL.md | 5 +++-- 3 files changed, 14 insertions(+), 7 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index 7f51824665..a25ed02ffc 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -226,7 +226,9 @@ The `rl` launcher applies these the same way in both single-node and multi-node 1. The launcher's own defaults — **your `env_vars` override these**. 2. Your top-level `[env_vars]`. 3. Your `[component.env_vars]`. -4. Orchestration-critical vars the launcher always sets last — `CUDA_VISIBLE_DEVICES` (GPU partitioning), `WANDB_SHARED_*` (the single shared W&B run), and `VF_RENDERER_IMAGE_OFFLOAD_DIR` (raw multimodal image asset path) — **these cannot be overridden** from `env_vars`. +4. Orchestration-critical vars the launcher always sets last — `CUDA_VISIBLE_DEVICES` (GPU partitioning) and `WANDB_SHARED_*` (the single shared W&B run) — **these cannot be overridden** from `env_vars`. + +`VF_RENDERER_IMAGE_OFFLOAD_DIR` is also rejected in `env_vars`. For v1 multimodal RL, set the shared image-asset directory with top-level `[multimodal].offload_dir` (propagated into trainer / orchestrator / inference / env-server configs). Each env-server resolves that into `VF_RENDERER_IMAGE_OFFLOAD_DIR` for verifiers/renderers (`setdefault`, so an already-exported process env wins — e.g. the multi-node SLURM template). When unset, the default is `{orchestrator.output_dir}/assets/images`, or `/data/outputs/run_${RUN_ID}/assets/images` when `RUN_ID` is set. Inference and the trainer read the `file://` URIs on the refs and do not need the env var. For standalone `sft` and `inference` configs, `[env_vars]` applies to that entrypoint's process(es). For disaggregated P/D inference, the role-specific [`deployment.{prefill,decode}_env_vars`](inference.md) layer on top of any shared inference env vars. 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 a8ce6f22f1..24d1e6e50b 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -6,9 +6,11 @@ from prime_rl.utils.config import BaseConfig -# Launcher-managed env vars that a component's `env_vars` must not set: GPU partitioning -# and the single shared W&B run. The launcher always sets these last, so allowing them in -# `env_vars` would be a silent no-op (or, on multi-node, a footgun) — reject them instead. +# Env vars that a component's `env_vars` must not set: +# - CUDA_VISIBLE_DEVICES / WANDB_SHARED_*: the launcher always sets these last. +# - VF_RENDERER_IMAGE_OFFLOAD_DIR: each env-server resolves this from +# ``[multimodal].offload_dir`` (operator process env wins via setdefault). +# Allowing them in ``env_vars`` would be a silent no-op or a footgun — reject instead. PROTECTED_ENV_VARS = frozenset( { "CUDA_VISIBLE_DEVICES", @@ -24,7 +26,9 @@ def reject_protected_env_vars(env_vars: dict[str, str]) -> dict[str, str]: clobbered = sorted(PROTECTED_ENV_VARS & env_vars.keys()) if clobbered: raise ValueError( - f"env_vars cannot set launcher-managed vars {clobbered} — set by the launcher, not overridable" + f"env_vars cannot set protected vars {clobbered}; " + "use [multimodal].offload_dir for VF_RENDERER_IMAGE_OFFLOAD_DIR, " + "and leave CUDA_VISIBLE_DEVICES / WANDB_SHARED_* to the launcher" ) return env_vars diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index bb056ea1ff..a2b2892bd0 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -183,8 +183,9 @@ means ingress was bypassed. - The image directory comes from `[multimodal].offload_dir` in the resolved config; unset, it defaults to a run-scoped path (`{output_dir}/assets/images` - or the hosted `RUN_ID` path). The launcher exports it to the orchestrator as - `VF_RENDERER_IMAGE_OFFLOAD_DIR` (protected — not settable via `env_vars`). + or the hosted `RUN_ID` path). Each env-server resolves that into + `VF_RENDERER_IMAGE_OFFLOAD_DIR` for verifiers/renderers (protected — set via + `[multimodal].offload_dir`, not `env_vars`). - While multimodal rollouts are in flight, image files should accumulate under that directory. Zero files means the offload path isn't being exercised or image preparation failed before request submission. From 454b4385d0f8972ffc39f4a22ea5441845165e03 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 03:05:50 +0000 Subject: [PATCH 48/52] docs: drop configuration.md multimodal offload note Co-authored-by: Cursor --- docs/configuration.md | 2 -- 1 file changed, 2 deletions(-) diff --git a/docs/configuration.md b/docs/configuration.md index a25ed02ffc..fbfdc5b67c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -228,8 +228,6 @@ The `rl` launcher applies these the same way in both single-node and multi-node 3. Your `[component.env_vars]`. 4. Orchestration-critical vars the launcher always sets last — `CUDA_VISIBLE_DEVICES` (GPU partitioning) and `WANDB_SHARED_*` (the single shared W&B run) — **these cannot be overridden** from `env_vars`. -`VF_RENDERER_IMAGE_OFFLOAD_DIR` is also rejected in `env_vars`. For v1 multimodal RL, set the shared image-asset directory with top-level `[multimodal].offload_dir` (propagated into trainer / orchestrator / inference / env-server configs). Each env-server resolves that into `VF_RENDERER_IMAGE_OFFLOAD_DIR` for verifiers/renderers (`setdefault`, so an already-exported process env wins — e.g. the multi-node SLURM template). When unset, the default is `{orchestrator.output_dir}/assets/images`, or `/data/outputs/run_${RUN_ID}/assets/images` when `RUN_ID` is set. Inference and the trainer read the `file://` URIs on the refs and do not need the env var. - For standalone `sft` and `inference` configs, `[env_vars]` applies to that entrypoint's process(es). For disaggregated P/D inference, the role-specific [`deployment.{prefill,decode}_env_vars`](inference.md) layer on top of any shared inference env vars. ## Examples From 17b1735c796374a73113215369ec535c0b12c7d2 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 04:02:22 +0000 Subject: [PATCH 49/52] docs: shorten SFT multimodal_output validator messaging Co-authored-by: Cursor --- packages/prime-rl-configs/src/prime_rl/configs/sft.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) 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 8308b7130e..e9506efb7c 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/sft.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/sft.py @@ -250,17 +250,10 @@ def normalize_deployment(cls, data): @model_validator(mode="after") def renderer_emits_processed_multimodal(self): - """SFT consumes processed pixel tensors straight from the renderer — it has no - raw-ref materializer, so the renderers-library default of ``multimodal_output='raw'`` - (built for the RL offload path) would silently ship JSON descriptors into training. - Default SFT renderers to ``'processed'`` and reject an explicit ``'raw'``.""" + """SFT needs processed MM tensors; default ``multimodal_output`` to that and reject ``raw``.""" if "multimodal_output" in self.renderer.model_fields_set: if self.renderer.multimodal_output == "raw": - raise ValueError( - "multimodal_output='raw' is unsupported for SFT: the SFT data path materializes " - "images from processed renderer output, not raw refs. Remove the override " - "(SFT defaults to 'processed')." - ) + raise ValueError("multimodal_output='raw' is unsupported for SFT (defaults to 'processed')") else: self.renderer = self.renderer.model_copy(update={"multimodal_output": "processed"}) return self From 4cde3d87368bb86bf764fca5b6c69cd74d477d1c Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 04:09:12 +0000 Subject: [PATCH 50/52] refactor: lean PROTECTED_ENV_VARS change for image offload dir Co-authored-by: Cursor --- .../prime-rl-configs/src/prime_rl/configs/shared.py | 12 ++++-------- 1 file changed, 4 insertions(+), 8 deletions(-) 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 24d1e6e50b..a8ce6f22f1 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -6,11 +6,9 @@ from prime_rl.utils.config import BaseConfig -# Env vars that a component's `env_vars` must not set: -# - CUDA_VISIBLE_DEVICES / WANDB_SHARED_*: the launcher always sets these last. -# - VF_RENDERER_IMAGE_OFFLOAD_DIR: each env-server resolves this from -# ``[multimodal].offload_dir`` (operator process env wins via setdefault). -# Allowing them in ``env_vars`` would be a silent no-op or a footgun — reject instead. +# Launcher-managed env vars that a component's `env_vars` must not set: GPU partitioning +# and the single shared W&B run. The launcher always sets these last, so allowing them in +# `env_vars` would be a silent no-op (or, on multi-node, a footgun) — reject them instead. PROTECTED_ENV_VARS = frozenset( { "CUDA_VISIBLE_DEVICES", @@ -26,9 +24,7 @@ def reject_protected_env_vars(env_vars: dict[str, str]) -> dict[str, str]: clobbered = sorted(PROTECTED_ENV_VARS & env_vars.keys()) if clobbered: raise ValueError( - f"env_vars cannot set protected vars {clobbered}; " - "use [multimodal].offload_dir for VF_RENDERER_IMAGE_OFFLOAD_DIR, " - "and leave CUDA_VISIBLE_DEVICES / WANDB_SHARED_* to the launcher" + f"env_vars cannot set launcher-managed vars {clobbered} — set by the launcher, not overridable" ) return env_vars From f92bb052daa56b91bd70a60db8627e9685e8e3dd Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 04:11:23 +0000 Subject: [PATCH 51/52] docs: drop monitor-run multimodal offload skill section Co-authored-by: Cursor --- skills/training/monitor-run/SKILL.md | 31 ---------------------------- 1 file changed, 31 deletions(-) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index a2b2892bd0..4643c3cd2b 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -173,37 +173,6 @@ A few warnings are normal. Escalate when errors are persistent, growing, or hit - **Trainer**: NCCL/CUDA errors, OOM, NaN loss or gradients. - **Inference**: NCCL/CUDA errors, OOM, request timeouts. -### Multimodal image offload checks - -v1 multimodal RL offloads images exactly once, at verifiers ingress: every image -content part is rewritten to a `file://` asset under the run image directory -before rendering, and renderers/inference/trainer all work from those refs. A -`data:` URL reaching a renderer ("requires offloaded file:// image assets") -means ingress was bypassed. - -- The image directory comes from `[multimodal].offload_dir` in the resolved - config; unset, it defaults to a run-scoped path (`{output_dir}/assets/images` - or the hosted `RUN_ID` path). Each env-server resolves that into - `VF_RENDERER_IMAGE_OFFLOAD_DIR` for verifiers/renderers (protected — set via - `[multimodal].offload_dir`, not `env_vars`). -- While multimodal rollouts are in flight, image files should accumulate under - that directory. Zero files means the offload path isn't being exercised or - image preparation failed before request submission. -- Inference rejects bad refs with `invalid_mm_image_ref` 400s (hash mismatch, - fingerprint mismatch, unreadable asset) — grep the inference log. -- Inference caches materialized refs and logs - `mm materialize cache: hits=X misses=Y hit_rate=Z% bytes=A/B evictions=C` - every 1000 lookups. Hit rate should climb after turn 1 of multi-turn - multimodal rollouts; a stuck-at-zero hit rate with repeat images means the - cache is disabled or thrashing (sized by `PRIME_RL_MM_MATERIALIZE_CACHE_GB`, - default 2.0, `0` disables). -- The orchestrator raises on placeholder/token drift ("does not cover - image-typed tokens") before a sample ships — treat any occurrence as a bug, - not noise. -- Trainer metrics: `mm/images_materialized` and `time/mm_materialize`. Missing - image files fail the trainer hard ("raw image materialization failed") — - check whether something cleaned the offload directory mid-run. - ### Process tree All processes use `setproctitle` so they're visible in `ps`/`htop`/`pstree`: From b57c4ac69f6ac2639f37e0b1e81e0fc51c183753 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 7 Aug 2026 05:40:16 +0000 Subject: [PATCH 52/52] refactor(mm): share image load/verify; trim serving_tokens helpers Centralize sha256_32 verify and HF image_processor load in utils/mm.py so inference and trainer cannot drift. Fold serving micro-helpers into the sync materializer and keep only the load-bearing cache docs. Co-authored-by: Cursor --- src/prime_rl/inference/vllm/serving_tokens.py | 80 +++++-------------- src/prime_rl/utils/mm.py | 33 +++++--- 2 files changed, 44 insertions(+), 69 deletions(-) diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index c906d809d7..ad75460eda 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -65,7 +65,7 @@ from prime_rl.inference.vllm.routed_experts import RoutedExpertsCapture from prime_rl.multimodal.registry import get_multimodal_adapter from prime_rl.multimodal.schema import RawMMItem -from prime_rl.utils.mm import file_uri_to_path +from prime_rl.utils.mm import load_image_processor, read_verified_image_bytes logger = logging.getLogger(__name__) @@ -176,63 +176,43 @@ async def _client_set_max_tokens(raw_request: Request | None) -> bool: @lru_cache(maxsize=8) def _load_image_processor(model_name: str, trust_remote_code: bool): - from transformers import AutoProcessor + # Named def so tests can monkeypatch ``serving_tokens._load_image_processor``. + return load_image_processor(model_name, trust_remote_code) - processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=trust_remote_code) - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - raise ValueError(f"{model_name!r} does not expose an image_processor") - return image_processor - -def _parse_raw_image_ref(raw_ref: str, *, feature_modality: str, mm_hash: str): +def _materialize_raw_image_ref_sync( + raw_ref: str, + *, + feature_modality: str, + mm_hash: str, + expected_placeholder_length: int, + processor_model_name: str, + trust_remote_code: bool, +): + from PIL import Image from renderers.mm_store import split_raw_mm_ref try: ref = split_raw_mm_ref(raw_ref) except ValueError as exc: raise _MMImageRefError(str(exc)) from exc - if ref.modality != feature_modality: raise _MMImageRefError(f"Expected feature modality {feature_modality!r}, got {ref.modality!r}") if ref.mm_hash != mm_hash: raise _MMImageRefError(f"Expected image hash {mm_hash}, got {ref.mm_hash}") - return ref - -def _read_verified_raw_image(ref) -> bytes: try: - raw = file_uri_to_path(ref.raw_image_uri).read_bytes() + raw = read_verified_image_bytes(ref.raw_image_uri, ref.mm_hash) except OSError as exc: raise _MMImageRefError(f"Unable to read raw image asset {ref.raw_image_uri!r}: {exc}") from exc - - actual_hash = hashlib.sha256(raw).hexdigest()[:32] - if actual_hash != ref.mm_hash: - raise _MMImageRefError(f"Raw image hash mismatch: expected {ref.mm_hash}, got {actual_hash}") - return raw - - -def _decode_raw_image(raw: bytes, *, raw_image_uri: str): - from PIL import Image + except ValueError as exc: + raise _MMImageRefError(str(exc)) from exc try: - return Image.open(BytesIO(raw)).convert("RGB") + image = Image.open(BytesIO(raw)).convert("RGB") except OSError as exc: - raise _MMImageRefError(f"Unable to decode raw image asset {raw_image_uri!r}: {exc}") from exc - + raise _MMImageRefError(f"Unable to decode raw image asset {ref.raw_image_uri!r}: {exc}") from exc -def _materialize_raw_image_ref_sync( - raw_ref: str, - *, - feature_modality: str, - mm_hash: str, - expected_placeholder_length: int, - processor_model_name: str, - trust_remote_code: bool, -): - ref = _parse_raw_image_ref(raw_ref, feature_modality=feature_modality, mm_hash=mm_hash) - raw = _read_verified_raw_image(ref) - image = _decode_raw_image(raw, raw_image_uri=ref.raw_image_uri) image_processor = _load_image_processor(processor_model_name, trust_remote_code) item = RawMMItem( family=ref.family, @@ -252,9 +232,8 @@ def _materialize_raw_image_ref_sync( # (raw_ref_digest, feature_modality, mm_hash, expected_placeholder_length, -# processor_model_name, trust_remote_code). The digest keeps keys small while -# ensuring a distinct descriptor can never alias an already-validated cache -# entry, even when its outer hash/placeholder metadata is forged or stale. +# processor_model_name, trust_remote_code). Digest the ref so forged outer +# hash/placeholder metadata cannot alias a cached entry. _MaterializeKey = tuple[str, str, str, int, str, bool] _MM_MATERIALIZE_CACHE_GB_ENV = "PRIME_RL_MM_MATERIALIZE_CACHE_GB" @@ -264,21 +243,9 @@ def _materialize_raw_image_ref_sync( class _MaterializedRefCache: """Byte-bounded LRU of materialized raw image refs, with single-flight misses. - Every request carries a ref for every image in its prompt (prior turns included), - so multi-turn rollouts re-materialize the same images once per turn per rollout — - this cache turns those repeats into lookups. Keys are content-addressed, so - entries can only go cold, never stale: a post-eviction request simply misses and - re-materializes from the durable ``file://`` asset. - All mutation happens on the event loop thread (materialization itself runs in a worker thread via ``asyncio.to_thread``, but lookup/insertion/eviction happen in the async caller) — so no lock is needed. Do not touch this cache from sync code. - - ``max_bytes == 0`` disables caching and single-flight entirely: every call runs - the materializer, byte-identical to the uncached path. - - Host-RAM note: this budget is additive to vLLM's own processor cache - (``mm_processor_cache_gb × (api_server_count + data_parallel_size)``). """ def __init__(self, max_bytes: int) -> None: @@ -438,14 +405,11 @@ class PrimeRlServingTokens(ServingTokens): @cached_property def _mm_materialize_cache(self) -> _MaterializedRefCache: - """Materialized-ref cache, sized by ``PRIME_RL_MM_MATERIALIZE_CACHE_GB`` - (float GiB, default 2.0, ``0`` disables — the kill switch: disabled is - byte-identical to the uncached path). + """Sized by ``PRIME_RL_MM_MATERIALIZE_CACHE_GB`` (GiB, default 2.0, ``0`` disables). A ``cached_property`` because ``custom_init_app_state`` grafts this subclass via ``object.__new__`` + ``__dict__.update``, so ``__init__`` - never runs (see ``_max_tokens_defaults``). One frontend process serves - all DP ranks of an engine, so the cache is shared across ranks. + never runs (see ``_max_tokens_defaults``). """ gb = float(os.environ.get(_MM_MATERIALIZE_CACHE_GB_ENV, "2.0")) return _MaterializedRefCache(max_bytes=int(gb * (1 << 30))) diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index 089dbf2c02..d46cf93c75 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -101,6 +101,26 @@ def sha256_32(data: bytes) -> str: return hashlib.sha256(data).hexdigest()[:32] +def read_verified_image_bytes(uri: str, expected_hash: str) -> bytes: + """Read a local ``file://`` image and verify its content hash (``sha256_32``).""" + raw = file_uri_to_path(uri).read_bytes() + actual_hash = sha256_32(raw) + if actual_hash != expected_hash: + raise ValueError(f"Raw image hash mismatch for {uri}: expected {expected_hash}, got {actual_hash}") + return raw + + +def load_image_processor(model_name: str, trust_remote_code: bool = False): + """Load the HF ``image_processor`` for ``model_name``.""" + from transformers import AutoProcessor + + processor = AutoProcessor.from_pretrained(model_name, trust_remote_code=trust_remote_code) + image_processor = getattr(processor, "image_processor", None) + if image_processor is None: + raise ValueError(f"{model_name!r} does not expose an image_processor") + return image_processor + + def _parse_image_refs(refs: MMRefs) -> list[RawMMItem]: return [parse_raw_mm_item(image.item) for image in refs.images] @@ -117,10 +137,7 @@ def _load_verified_images(image_refs: list[MMImageRef]) -> list[Any]: images = [] for ref in image_refs: - raw = file_uri_to_path(ref.uri).read_bytes() - actual_hash = sha256_32(raw) - if actual_hash != ref.hash: - raise ValueError(f"Raw image hash mismatch for {ref.uri}: expected {ref.hash}, got {actual_hash}") + raw = read_verified_image_bytes(ref.uri, ref.hash) with Image.open(BytesIO(raw)) as image: images.append(image.convert("RGB")) return images @@ -137,13 +154,7 @@ def __init__(self, model_name: str, *, trust_remote_code: bool): @property def image_processor(self): if self._image_processor is None: - from transformers import AutoProcessor - - processor = AutoProcessor.from_pretrained(self.model_name, trust_remote_code=self.trust_remote_code) - image_processor = getattr(processor, "image_processor", None) - if image_processor is None: - raise ValueError(f"{self.model_name!r} does not expose an image_processor") - self._image_processor = image_processor + self._image_processor = load_image_processor(self.model_name, self.trust_remote_code) return self._image_processor def materialize(self, refs: MMRefs) -> MaterializedMM | None: