From 034e83e6147b3f4e17be9883fec2e0ed60742823 Mon Sep 17 00:00:00 2001 From: eligotts Date: Wed, 24 Jun 2026 01:11:43 +0000 Subject: [PATCH 01/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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/33] 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 e9bea1f0e8d0c0aaafb0db5d40fd950dfe6649fb Mon Sep 17 00:00:00 2001 From: eligotts Date: Thu, 23 Jul 2026 23:42:42 +0000 Subject: [PATCH 27/33] Inline raw multimodal images MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Intermediate storage mode built off feat/v1-raw-mm-offload: the image processor stays out of the env worker (vLLM-front materialization + cache, trainer adapter re-materialization, renderer geometry math all kept), but images travel inline as base64 data URLs instead of offloaded file:// run assets. - Serving: materializes from the ref's inline payload (decode + hash verify); the materialize cache re-keys to (mm_hash, placeholder_len, model) so keys stay small now that the ref embeds the full image. - Trainer: RawImageMaterializer decodes inline bytes from the descriptor; MMImageRef drops the uri field (the data lives once, in the descriptor). - Removed the entire offload plumbing layer: [multimodal].offload_dir, MultimodalConfig, run_assets.py, VF_RENDERER_IMAGE_OFFLOAD_DIR launcher/SLURM exports, and missing_mm_image_policy + the adapters' synthesize_placeholder machinery — an inline image cannot go missing, so the zero-loss placeholder path has nothing to guard. - Rollout records and traces keep the inline base64 (documented in the monitor-run skill). - Pins deps/renderers and deps/verifiers to the inline companion branches. tests/unit minus GPU-kernel model tests: 482 passed. --- deps/renderers | 2 +- deps/verifiers | 2 +- docs/configuration.md | 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 | 6 -- .../src/prime_rl/configs/trainer.py | 8 --- .../src/prime_rl/utils/validation.py | 1 - skills/training/monitor-run/SKILL.md | 36 ++++------ src/prime_rl/entrypoints/rl.py | 10 --- src/prime_rl/inference/vllm/serving_tokens.py | 25 +++---- src/prime_rl/multimodal/adapters/base.py | 6 -- src/prime_rl/multimodal/adapters/kimi_k25.py | 28 -------- src/prime_rl/multimodal/adapters/qwen_vl.py | 36 ---------- src/prime_rl/multimodal/schema.py | 4 +- src/prime_rl/orchestrator/orchestrator.py | 2 - .../templates/multi_node_rl.sbatch.j2 | 12 ---- src/prime_rl/trainer/rl/data.py | 45 ++---------- src/prime_rl/trainer/rl/train.py | 2 - src/prime_rl/transport/types.py | 14 ++-- src/prime_rl/utils/mm.py | 58 ++++------------ src/prime_rl/utils/run_assets.py | 68 ------------------- tests/unit/inference/test_serving_tokens.py | 55 ++++++++------- tests/unit/orchestrator/test_batch.py | 19 +++--- tests/unit/train/rl/test_packer.py | 5 +- tests/unit/utils/test_mm.py | 33 ++++----- 27 files changed, 112 insertions(+), 380 deletions(-) delete mode 100644 src/prime_rl/utils/run_assets.py diff --git a/deps/renderers b/deps/renderers index e64cc58803..d54e7fc0f9 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit e64cc588037a75ff99d37c2968329805147d820d +Subproject commit d54e7fc0f9455c05e45e6d6fd423179325f6a563 diff --git a/deps/verifiers b/deps/verifiers index ae516c19b0..6e095591e5 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit ae516c19b0b63ea164bb1c0113611e7811efacbc +Subproject commit 6e095591e549ce746e60b583d8f8321977dbcfa8 diff --git a/docs/configuration.md b/docs/configuration.md index 2aae50704e..3e6b76151b 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -200,7 +200,7 @@ 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), `WANDB_SHARED_*` (the single shared W&B run) — **these cannot be overridden** from `env_vars`. 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/inference.py b/packages/prime-rl-configs/src/prime_rl/configs/inference.py index 2f699394d8..04dbebeca5 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, EnvVars, LogConfig, MultimodalConfig, SlurmConfig +from prime_rl.configs.shared import BaseModelConfig, EnvVars, LogConfig, 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 @@ -429,9 +429,6 @@ class InferenceConfig(BaseConfig): dry_run: bool = False """Only validate and dump resolved configs, then exit early.""" - multimodal: MultimodalConfig = MultimodalConfig() - """Raw multimodal image offload settings 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 f5c2ea1b0f..fe6b696203 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py @@ -18,7 +18,6 @@ FileSystemTransportConfig, HeartbeatConfig, LogConfig, - MultimodalConfig, PrimeMonitorConfig, TransportConfig, WandbWithExtrasConfig, @@ -553,9 +552,6 @@ class OrchestratorConfig(BaseConfig): heartbeat: HeartbeatConfig | None = None """BetterStack heartbeat configuration for monitoring training progress.""" - multimodal: MultimodalConfig = MultimodalConfig() - """Raw multimodal image offload settings 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 77bd2e9380..48b9e6a543 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/rl.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/rl.py @@ -21,7 +21,6 @@ from prime_rl.configs.shared import ( EnvVars, FileMonitorConfig, - MultimodalConfig, SlurmConfig, VLMConfig, ) @@ -261,9 +260,6 @@ class RLConfig(BaseConfig): weight_broadcast: SharedWeightBroadcastConfig | None = None - multimodal: MultimodalConfig = MultimodalConfig() - """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 159f0898bd..5774bdcafe 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -12,7 +12,6 @@ PROTECTED_ENV_VARS = frozenset( { "CUDA_VISIBLE_DEVICES", - "VF_RENDERER_IMAGE_OFFLOAD_DIR", "WANDB_SHARED_MODE", "WANDB_SHARED_RUN_ID", "WANDB_SHARED_LABEL", @@ -92,11 +91,6 @@ def resolve_project_dir(self): ServerType = Literal["vllm", "openai"] -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 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 07c706806c..d18ce525b4 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/trainer.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/trainer.py @@ -11,7 +11,6 @@ FileSystemTransportConfig, HeartbeatConfig, MetricsServerConfig, - MultimodalConfig, TrainerLogConfig, TransportConfig, WandbConfig, @@ -22,7 +21,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,12 +628,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.""" - 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 92619f20f5..6280228238 100644 --- a/packages/prime-rl-configs/src/prime_rl/utils/validation.py +++ b/packages/prime-rl-configs/src/prime_rl/utils/validation.py @@ -129,7 +129,6 @@ 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/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index fa62ef42ab..26e01b91a9 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -165,24 +165,17 @@ 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. +### Multimodal image checks + +v1 multimodal RL keeps images inline: message content carries +`data:image/...;base64` URLs end to end, and renderers emit raw descriptors +embedding that inline source. Inference and trainer each materialize pixels +from the inline data with their own image processor; nothing is written to a +shared image directory. + - 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 + fingerprint mismatch, undecodable inline data) — grep the inference log. +- Inference caches materialized images by content hash 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 @@ -191,11 +184,10 @@ 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`. +- Rollout records and traces carry the inline base64 images, so multimodal + runs produce large `results.jsonl` files and wire payloads that grow with + turn count — expected, not a leak. ### Process tree diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 1546fda14b..0bb1348130 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -35,7 +35,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" @@ -126,7 +125,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: @@ -227,7 +225,6 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ - **writer_run_asset_env, **DEFAULT_COMMON_ENV_VARS, "LOGURU_FORCE_COLORS": "1", "WANDB_PROGRAM": "uv run rl", @@ -374,10 +371,6 @@ 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 "" - ) - # Per-component env vars: launcher defaults (shared + multi-node-specific) with the # user's config merged on top. Runtime wiring stays in the template. trainer_env_vars = { @@ -398,7 +391,6 @@ def write_slurm_script(config: RLConfig, config_dir: Path, script_path: Path) -> **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": @@ -410,7 +402,6 @@ 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, @@ -448,7 +439,6 @@ 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.infer_nodes_per_replica, diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index e8fc154f39..57e1ee4381 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -45,6 +45,7 @@ from typing import Any from fastapi import Request +from renderers.mm_store import decode_data_image_url from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, PromptTokenUsageInfo, @@ -65,7 +66,6 @@ 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 logger = logging.getLogger(__name__) @@ -202,9 +202,9 @@ def _parse_raw_image_ref(raw_ref: str, *, feature_modality: str, mm_hash: str): def _read_verified_raw_image(ref) -> bytes: try: - 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_uri!r}: {exc}") from exc + raw = decode_data_image_url(ref.raw_image_data) + except ValueError as exc: + raise _MMImageRefError(f"Unable to decode inline raw image data: {exc}") from exc actual_hash = hashlib.sha256(raw).hexdigest()[:32] if actual_hash != ref.mm_hash: @@ -212,13 +212,13 @@ def _read_verified_raw_image(ref) -> bytes: return raw -def _decode_raw_image(raw: bytes, *, raw_image_uri: str): +def _decode_raw_image(raw: bytes): 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_uri!r}: {exc}") from exc + raise _MMImageRefError(f"Unable to decode inline raw image bytes: {exc}") from exc def _materialize_raw_image_ref_sync( @@ -232,13 +232,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_uri=ref.raw_image_uri) + image = _decode_raw_image(raw) 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_uri=ref.raw_image_uri, + raw_image_data=ref.raw_image_data, payload=dict(ref.payload), raw_ref=raw_ref, ) @@ -251,9 +251,10 @@ 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. +# (mm_hash, expected_placeholder_length, processor_model_name). The hash is the +# content identity of the raw image bytes (verified on every miss), so identical +# keys describe byte-identical materialization work — and the key stays small +# even though the ref string itself embeds the full inline image payload. _MaterializeKey = tuple[str, int, str] _MM_MATERIALIZE_CACHE_GB_ENV = "PRIME_RL_MM_MATERIALIZE_CACHE_GB" @@ -408,7 +409,7 @@ 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), + (mm_hash, placeholder.length, processor_model_name), 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/src/prime_rl/multimodal/adapters/base.py b/src/prime_rl/multimodal/adapters/base.py index b746b9f90a..4db78f988e 100644 --- a/src/prime_rl/multimodal/adapters/base.py +++ b/src/prime_rl/multimodal/adapters/base.py @@ -44,9 +44,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 4a2aba7aa9..7d8b9d08c9 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 collections.abc import Mapping from typing import Any @@ -154,30 +153,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 - - 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, layout.patch_size, layout.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 6adf9c0bfa..b555f6eb61 100644 --- a/src/prime_rl/multimodal/adapters/qwen_vl.py +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -129,39 +129,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/multimodal/schema.py b/src/prime_rl/multimodal/schema.py index 0685274b20..9a40dadddc 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_uri: str + raw_image_data: str payload: dict[str, Any] raw_ref: str | None = None vllm_modality: str | None = None @@ -57,7 +57,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_uri=_required_str(descriptor, "raw_image_uri"), + raw_image_data=_required_str(descriptor, "raw_image_data"), 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 6fc2cf2ff8..400d078f79 100644 --- a/src/prime_rl/orchestrator/orchestrator.py +++ b/src/prime_rl/orchestrator/orchestrator.py @@ -81,7 +81,6 @@ 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_trace_path -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, @@ -1000,7 +999,6 @@ 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. """ - apply_run_asset_env(config.output_dir, config.multimodal) await Orchestrator(config).start() diff --git a/src/prime_rl/templates/multi_node_rl.sbatch.j2 b/src/prime_rl/templates/multi_node_rl.sbatch.j2 index c4e0cec88c..f401f46a62 100755 --- a/src/prime_rl/templates/multi_node_rl.sbatch.j2 +++ b/src/prime_rl/templates/multi_node_rl.sbatch.j2 @@ -59,18 +59,6 @@ export PROJECT_DIR={{ project_dir }} export CONFIG_DIR={{ config_dir }} export OUTPUT_DIR={{ output_dir }} export ORCHESTRATOR_OUTPUT_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/rl/data.py b/src/prime_rl/trainer/rl/data.py index 5e98d9636a..f3863c7c9b 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 @@ -19,8 +19,7 @@ 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 +from prime_rl.utils.mm import RawImageMaterializer class TensorMicroBatch(TypedDict): @@ -79,7 +78,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 +197,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 +216,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 +236,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: @@ -248,43 +243,11 @@ def _materialized_mm_parts(materialized: MaterializedMM | None) -> MaterializedM 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) - + materialized = self.mm_materializer.materialize(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})" - ) + self.last_mm_images_materialized += len(refs.images) return self._materialized_mm_parts(materialized) def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: 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/transport/types.py b/src/prime_rl/transport/types.py index c6d8ed490a..64011c71c3 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -21,15 +21,15 @@ class RoutedExperts(msgspec.Struct, array_like=True, gc=False, omit_defaults=Tru 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`` is the JSON-safe renderer descriptor (adapter family + payload + + the inline base64 image source, parseable via ``parse_raw_mm_item``), + ``hash`` is the content identity of the raw bytes, 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 @@ -43,10 +43,6 @@ class MMRefs(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): images: list[MMImageRef] - @property - def uris(self) -> list[str]: - return [image.uri for image in self.images] - # Orchestrator -> Packer class TrainingSample(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): diff --git a/src/prime_rl/utils/mm.py b/src/prime_rl/utils/mm.py index cbdcff38c8..2889239caf 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -3,9 +3,9 @@ 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 renderers.mm_store import decode_data_image_url from prime_rl.multimodal.adapters.base import MaterializedMM, MultimodalAdapter from prime_rl.multimodal.registry import get_multimodal_adapter @@ -22,26 +22,6 @@ def _field(value: Any, name: str, default: Any = None) -> Any: 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: - 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 @@ -53,16 +33,13 @@ def _validate_modalities(mm_items: Mapping[str, list[Any]]) -> None: ) -def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> tuple[list[dict[str, Any]], list[str]]: +def _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: 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) + parse_raw_mm_item(item_dict) item_dicts.append(item_dict) - uris.append(parsed.raw_image_uri) - return item_dicts, uris + return item_dicts def _placeholder_bounds(placeholder: Any) -> tuple[int, int]: @@ -81,7 +58,7 @@ def build_mm_refs(multi_modal_data: Any) -> MMRefs | None: return None _validate_modalities(mm_items) - image_items, uris = _raw_item_dicts(mm_items.get(IMAGE_MODALITY, [])) + image_items = _raw_item_dicts(mm_items.get(IMAGE_MODALITY, [])) if not image_items: return None @@ -97,14 +74,14 @@ def build_mm_refs(multi_modal_data: Any) -> MMRefs | None: images: list[MMImageRef] = [] prev_end = 0 - for item, uri, image_hash, placeholder in zip(image_items, uris, image_hashes, image_placeholders, strict=True): + for item, image_hash, placeholder in zip(image_items, 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)) + images.append(MMImageRef(item=item, hash=image_hash, offset=offset, length=length)) return MMRefs(images=images) @@ -132,15 +109,15 @@ def _validate_processor_layout(adapter: MultimodalAdapter, image_processor: Any, ) -def _load_verified_images(image_refs: list[MMImageRef]) -> list[Any]: +def _load_verified_images(image_refs: list[MMImageRef], image_items: list[RawMMItem]) -> list[Any]: from PIL import Image images = [] - for ref in image_refs: - raw = file_uri_to_path(ref.uri).read_bytes() + for ref, item in zip(image_refs, image_items, strict=True): + raw = decode_data_image_url(item.raw_image_data) 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}") + raise ValueError(f"Raw image hash mismatch: expected {ref.hash}, got {actual_hash}") with Image.open(BytesIO(raw)) as image: images.append(image.convert("RGB")) return images @@ -174,14 +151,5 @@ 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) + images = _load_verified_images(refs.images, image_items) 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/src/prime_rl/utils/run_assets.py b/src/prime_rl/utils/run_assets.py deleted file mode 100644 index a359e45ec7..0000000000 --- a/src/prime_rl/utils/run_assets.py +++ /dev/null @@ -1,68 +0,0 @@ -from __future__ import annotations - -import os -from collections.abc import Mapping -from pathlib import Path -from string import Template - -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" - -RUN_OUTPUT_ROOT = Path("/data/outputs") -IMAGE_ASSET_SUBDIR = Path("assets/images") - - -def _expand_path(path: Path, env: Mapping[str, str]) -> Path: - expanded = Template(os.path.expanduser(str(path))).safe_substitute(env) - 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: - """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) - 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 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 - - -def apply_run_asset_env(output_dir: Path, multimodal: MultimodalConfig) -> None: - if os.environ.get(IMAGE_OFFLOAD_DIR_ENV): - return - os.environ.update(build_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 5d3602dff4..8b0f72aeba 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -270,25 +270,28 @@ def test_client_set_max_tokens_assumes_set_when_body_unreadable(): 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): +def test_materialize_raw_image_ref_uses_generic_family_payload(monkeypatch): + import base64 + from io import BytesIO + 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) + buf = BytesIO() + Image.new("RGB", (8, 6), color=(32, 64, 128)).save(buf, format="PNG") + raw_bytes = buf.getvalue() + image_data = "data:image/png;base64," + base64.b64encode(raw_bytes).decode("ascii") - mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] + mm_hash = hashlib.sha256(raw_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(), + raw_image_data=image_data, payload={"adapter_owned": [1, 2, 3]}, ) processor = object() @@ -326,29 +329,31 @@ def _get_adapter(family): item = captured["item"] assert item.family == "test_family" assert item.layout_fingerprint == fingerprint - assert item.raw_image_uri == image_path.as_uri() + assert item.raw_image_data == image_data assert item.payload == {"adapter_owned": [1, 2, 3]} -def _mm_features(tmp_path, *, image_name: str = "image.png", placeholder_length: int = 7): +def _mm_features(*, image_seed: int = 0, placeholder_length: int = 7): """A minimal ``GenerateRequest.features``-shaped object carrying one real raw ref.""" + import base64 + from io import BytesIO 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) + buf = BytesIO() + Image.new("RGB", (8, 6), color=(image_seed % 255, 64, 128)).save(buf, format="PNG") + raw_bytes = buf.getvalue() + image_data = "data:image/png;base64," + base64.b64encode(raw_bytes).decode("ascii") - mm_hash = hashlib.sha256(image_path.read_bytes()).hexdigest()[:32] + mm_hash = hashlib.sha256(raw_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(), + raw_image_data=image_data, payload={}, ) return SimpleNamespace( @@ -366,19 +371,19 @@ def materialize_for_vllm(self, image_processor, item, image, expected_placeholde calls.append(item.raw_ref) if materialize is not None: return materialize(item) - return {"materialized": item.raw_image_uri} + return {"materialized": item.raw_image_data} 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): +def test_mm_materialize_cache_hit_skips_work(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) + features = _mm_features() async def _run(): first = await serving_tokens._decode_raw_mm_kwargs( @@ -395,7 +400,7 @@ async def _run(): assert (cache.hits, cache.misses) == (1, 1) -def test_mm_materialize_cache_byte_budget_evicts_oldest(tmp_path, monkeypatch): +def test_mm_materialize_cache_byte_budget_evicts_oldest(monkeypatch): from vllm.multimodal.cache import MultiModalCache from prime_rl.inference.vllm import serving_tokens @@ -404,8 +409,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(image_seed=1) + features_b = _mm_features(image_seed=2) async def _decode(features): return await serving_tokens._decode_raw_mm_kwargs( @@ -420,7 +425,7 @@ async def _decode(features): assert cache.misses == 3 -def test_mm_materialize_cache_failure_not_cached(tmp_path, monkeypatch): +def test_mm_materialize_cache_failure_not_cached(monkeypatch): import pytest from prime_rl.inference.vllm import serving_tokens @@ -436,7 +441,7 @@ def _materialize(item): _patch_adapter(monkeypatch, calls, materialize=_materialize) cache = serving_tokens._MaterializedRefCache(max_bytes=1 << 20) - features = _mm_features(tmp_path) + features = _mm_features() async def _decode(): return await serving_tokens._decode_raw_mm_kwargs( @@ -450,13 +455,13 @@ async def _decode(): assert len(calls) == 2 -def test_mm_materialize_cache_single_flight(tmp_path, monkeypatch): +def test_mm_materialize_cache_single_flight(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) + features = _mm_features() async def _run(): return await asyncio.gather( diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 320f6ad946..b0c1e320e4 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -411,26 +411,25 @@ def test_prepare_sample_truncates_routed_experts(): assert micro_batch.env_names == ["test-env"] * 3 -def _image_ref(uri: str, offset: int, length: int) -> MMImageRef: +def _image_ref(image_data: 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, + "raw_image_data": image_data, "payload": {"image_grid_thw": [[1, 1, 1]]}, }, hash="a" * 32, - uri=uri, offset=offset, length=length, ) 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) + first_image = _image_ref("data:image/png;base64,aW1nMA==", offset=1, length=2) + second_image = _image_ref("data:image/png;base64,aW1nMQ==", offset=4, length=2) sample = TrainingSample( token_ids=[10, 11, 12, 13, 14, 15, 16], mask=[False, False, True, True, False, True, True], @@ -458,7 +457,7 @@ def test_prepare_sample_truncates_raw_mm_refs_at_image_boundary(): 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 mm_sample(image_data: str) -> TrainingSample: return TrainingSample( token_ids=[10, 11, 12], mask=[False, True, True], @@ -467,7 +466,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=[_image_ref(image_data, offset=1, length=1)]), ) text_sample = TrainingSample( @@ -480,7 +479,11 @@ 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("data:image/png;base64,aW1nMA=="), + mm_sample("data:image/png;base64,aW1nMQ=="), + text_sample, + ], seq_len=16, num_train_workers=2, idxs=[0, 0, 0], diff --git a/tests/unit/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 8633ee9938..60b2dd003b 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -53,7 +53,7 @@ def make_training_sample() -> TrainingSample: def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: - uri = f"file:///tmp/image-{value}.png" + image_data = f"data:image/png;base64,{value}" return TrainingSample( token_ids=[1, 250, 2], mask=[False, True, True], @@ -70,11 +70,10 @@ def _mm_sample(value: float, env_name: str = "test-env") -> TrainingSample: "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, - "raw_image_uri": uri, + "raw_image_data": image_data, "payload": {"image_grid_thw": [[1, 1, 1]]}, }, hash="a" * 32, - uri=uri, offset=1, length=1, ) diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index abb80d42a1..f6d3e8dd92 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -1,3 +1,4 @@ +import base64 from types import SimpleNamespace import torch @@ -8,12 +9,11 @@ from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs from prime_rl.utils.mm import build_mm_refs +_IMAGE_DATA = "data:image/png;base64," + base64.b64encode(b"png-ish bytes").decode("ascii") -class _MissingMaterializer: - def materialize(self, refs): - raise FileNotFoundError("missing image") - def synthesize_placeholder(self, refs): +class _FakeMaterializer: + def materialize(self, refs): return MaterializedMM( kwargs={ "pixel_values": torch.zeros((1, 24), dtype=torch.float32), @@ -23,31 +23,29 @@ def synthesize_placeholder(self, refs): ) -def _qwen_item(grid, uri: str = "file:///tmp/missing-image.png"): +def _qwen_item(grid, image_data: str = _IMAGE_DATA): return { "kind": RAW_MM_ITEM_KIND, "modality": "image", "family": "qwen_vl", "layout_fingerprint": "f" * 32, - "raw_image_uri": uri, + "raw_image_data": image_data, "payload": {"image_grid_thw": grid}, } -def _refs(uri: str = "file:///tmp/missing-image.png") -> MMRefs: +def _refs() -> MMRefs: return MMRefs( - images=[MMImageRef(item=_qwen_item([[1, 1, 1]], uri), hash="a" * 32, uri=uri, offset=1, length=1)], + images=[MMImageRef(item=_qwen_item([[1, 1, 1]]), hash="a" * 32, offset=1, length=1)], ) 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.mm_materializer = _FakeMaterializer() loader.last_mm_materialize_time = 0.0 loader.last_mm_images_materialized = 0 - loader.last_mm_images_placeholdered = 0 return loader @@ -68,26 +66,23 @@ def _micro_batch() -> MicroBatch: ) -def test_build_mm_refs_accepts_offloaded_file_uri(tmp_path): - image_path = tmp_path / "image.png" - image_path.write_bytes(b"image") +def test_build_mm_refs_accepts_inline_descriptors(): multi_modal_data = SimpleNamespace( - mm_items={"image": [_qwen_item([[1, 1, 1]], image_path.as_uri())]}, + mm_items={"image": [_qwen_item([[1, 1, 1]])]}, mm_hashes={"image": ["a" * 32]}, mm_placeholders={"image": [SimpleNamespace(offset=1, length=1)]}, ) refs = build_mm_refs(multi_modal_data) - assert refs == _refs(image_path.as_uri()) + assert refs == _refs() -def test_dataloader_uses_zero_loss_placeholder_for_missing_raw_image(): +def test_dataloader_materializes_raw_mm_refs(): 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["loss_mask"].tolist() == [[False, True, True]] assert tensor_batch["mm_token_type_ids"].tolist() == [[0, 1, 0]] From 7d85b170782162bda40d9f556d81e1eef63ddcd6 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 24 Jul 2026 00:25:16 +0000 Subject: [PATCH 28/33] Harden inline multimodal materialization --- src/prime_rl/entrypoints/rl.py | 1 + src/prime_rl/inference/vllm/serving_tokens.py | 36 ++++++++----- tests/unit/inference/test_serving_tokens.py | 52 +++++++++++++++++++ 3 files changed, 76 insertions(+), 13 deletions(-) diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 0bb1348130..dca9c600da 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -225,6 +225,7 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ + **inherited_env, **DEFAULT_COMMON_ENV_VARS, "LOGURU_FORCE_COLORS": "1", "WANDB_PROGRAM": "uv run rl", diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 57e1ee4381..53c8c98ace 100644 --- a/src/prime_rl/inference/vllm/serving_tokens.py +++ b/src/prime_rl/inference/vllm/serving_tokens.py @@ -242,20 +242,23 @@ 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 -# (mm_hash, expected_placeholder_length, processor_model_name). The hash is the -# content identity of the raw image bytes (verified on every miss), so identical -# keys describe byte-identical materialization work — and the key stays small -# even though the ref string itself embeds the full inline image payload. -_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 @@ -409,7 +412,14 @@ def _materialize(feature_modality: str, raw_ref: str, mm_hash: str, placeholder: decoded = await asyncio.gather( *( cache.get_or_materialize( - (mm_hash, 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 8b0f72aeba..6946a29c38 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -15,6 +15,7 @@ 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 @@ -333,6 +334,33 @@ def _get_adapter(family): assert item.payload == {"adapter_owned": [1, 2, 3]} +def test_materialize_raw_image_ref_maps_adapter_validation_error(monkeypatch): + from prime_rl.inference.vllm import serving_tokens + + features = _mm_features() + 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(*, image_seed: int = 0, placeholder_length: int = 7): """A minimal ``GenerateRequest.features``-shaped object carrying one real raw ref.""" import base64 @@ -400,6 +428,30 @@ async def _run(): assert (cache.hits, cache.misses) == (1, 1) +def test_mm_materialize_cache_does_not_alias_distinct_raw_refs(monkeypatch): + 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(image_seed=1) + forged = _mm_features(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(monkeypatch): from vllm.multimodal.cache import MultiModalCache From fb6ef3063e2e0175a0748cc3bba1cc29d6feb9cd Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 24 Jul 2026 07:32:43 +0000 Subject: [PATCH 29/33] Pack raw-ref multimodal samples like main packs eager ones Mirrors the offload branch's restoration of main's mm packing feature, which the merge had regressed to never-pack for raw-ref samples: can_add treats same-adapter-family raw-ref samples as pack-compatible (one family per micro batch is what the materializer enforces), and _materialize_bin merges refs across the bin's samples with placeholder offsets rebased to the packed token stream. Same run/LoRA gating, seq_lens boundaries, and modality alignment as main; the trainer-side adapter already materializes a list of refs into concatenated tensors. Packing tests restored to main's semantics (pack within run with rebased offsets, family mismatch splits, never across runs). --- src/prime_rl/trainer/batch.py | 42 ++++++++++++++++++--------- tests/unit/orchestrator/test_batch.py | 32 +++++++++++++------- tests/unit/train/rl/test_packer.py | 15 +++++----- 3 files changed, 59 insertions(+), 30 deletions(-) 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/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index b0c1e320e4..41751dbc42 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -454,10 +454,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(image_data: 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(image_data: str, family: str = "qwen_vl") -> TrainingSample: + ref = _image_ref(image_data, offset=1, length=1) + ref.item["family"] = family return TrainingSample( token_ids=[10, 11, 12], mask=[False, True, True], @@ -466,7 +470,7 @@ def mm_sample(image_data: 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(image_data, offset=1, length=1)]), + mm_refs=MMRefs(images=[ref]), ) text_sample = TrainingSample( @@ -482,22 +486,30 @@ def mm_sample(image_data: str) -> TrainingSample: rollouts=[ mm_sample("data:image/png;base64,aW1nMA=="), mm_sample("data:image/png;base64,aW1nMQ=="), + mm_sample("data:image/png;base64,aW1nMg==", 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 60b2dd003b..2486b460e6 100644 --- a/tests/unit/train/rl/test_packer.py +++ b/tests/unit/train/rl/test_packer.py @@ -217,8 +217,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) @@ -232,10 +232,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 e40dded5886726190532967f98ee5e8fe1909f84 Mon Sep 17 00:00:00 2001 From: eligotts Date: Fri, 24 Jul 2026 19:23:07 +0000 Subject: [PATCH 30/33] Bump renderers for the unwrapped ref payload (~32 KiB saved per image slot) --- deps/renderers | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deps/renderers b/deps/renderers index d54e7fc0f9..603118bda6 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit d54e7fc0f9455c05e45e6d6fd423179325f6a563 +Subproject commit 603118bda6a8300eab62b7dac71e49ac5ddfc6c5 From d453dfa8202157da18d712a4b6b6ed8acbe08eb4 Mon Sep 17 00:00:00 2001 From: eligotts Date: Mon, 3 Aug 2026 20:13:17 +0000 Subject: [PATCH 31/33] 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 | 59 ++------------------ src/prime_rl/multimodal/adapters/qwen_vl.py | 29 ++-------- 3 files changed, 11 insertions(+), 79 deletions(-) diff --git a/deps/renderers b/deps/renderers index 50ce3c3e6d..9208b40c14 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 50ce3c3e6d58fe8b1ba722bafee374db5189615b +Subproject commit 9208b40c14b9f2a06ea9c8432f462a38661fc257 diff --git a/src/prime_rl/multimodal/adapters/kimi_k25.py b/src/prime_rl/multimodal/adapters/kimi_k25.py index 7d8b9d08c9..02f6f0b722 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -1,9 +1,8 @@ from __future__ import annotations -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 @@ -17,46 +16,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: @@ -91,19 +50,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, diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py index b555f6eb61..46c8a7964c 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 8e9c8c75ab435114ea9f03450533c5eeb44b4da8 Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 19:22:56 +0000 Subject: [PATCH 32/33] 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 knobs they still need 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 | 12 ------------ 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, 4 insertions(+), 50 deletions(-) diff --git a/deps/renderers b/deps/renderers index 27fe3c6b4d..ce7078aa5c 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit 27fe3c6b4dcbd2d25cfa2497eb8814842e4d4931 +Subproject commit ce7078aa5cce0c434507b7c159643e8914db076d diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index dce5c5b746..8792e35032 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_data=ref.raw_image_data, 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 4db78f988e..c6217e13dc 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 02f6f0b722..cf3410b16f 100644 --- a/src/prime_rl/multimodal/adapters/kimi_k25.py +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -2,8 +2,6 @@ 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 @@ -49,11 +47,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, @@ -83,11 +76,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) diff --git a/src/prime_rl/multimodal/adapters/qwen_vl.py b/src/prime_rl/multimodal/adapters/qwen_vl.py index 46c8a7964c..ad96b0faa7 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 9a40dadddc..6d3bc70d4a 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_data: 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_data=_required_str(descriptor, "raw_image_data"), 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 2889239caf..cdc9843943 100644 --- a/src/prime_rl/utils/mm.py +++ b/src/prime_rl/utils/mm.py @@ -100,15 +100,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], image_items: list[RawMMItem]) -> list[Any]: from PIL import Image @@ -150,6 +141,5 @@ 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, image_items) 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 58c1a2b12f..6e6cfab0fd 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -286,10 +286,8 @@ def test_materialize_raw_image_ref_uses_generic_family_payload(monkeypatch): image_data = "data:image/png;base64," + base64.b64encode(raw_bytes).decode("ascii") mm_hash = hashlib.sha256(raw_bytes).hexdigest()[:32] - fingerprint = "f" * 32 raw_ref = raw_mm_ref( family="test_family", - fingerprint=fingerprint, modality="image", mm_hash=mm_hash, raw_image_data=image_data, @@ -329,7 +327,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_data == image_data assert item.payload == {"adapter_owned": [1, 2, 3]} @@ -343,12 +340,12 @@ def test_materialize_raw_image_ref_maps_adapter_validation_error(monkeypatch): 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", @@ -378,7 +375,6 @@ def _mm_features(*, image_seed: int = 0, placeholder_length: int = 7): mm_hash = hashlib.sha256(raw_bytes).hexdigest()[:32] raw_ref = raw_mm_ref( family="test_family", - fingerprint="f" * 32, modality="image", mm_hash=mm_hash, raw_image_data=image_data, diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 41751dbc42..9a78cccaf7 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -417,7 +417,6 @@ def _image_ref(image_data: str, offset: int, length: int) -> MMImageRef: "kind": "prime_raw_mm_item", "modality": "image", "family": "qwen_vl", - "layout_fingerprint": "f" * 32, "raw_image_data": image_data, "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 4f182db2e7..93090c5239 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_data": image_data, "payload": {"image_grid_thw": [[1, 1, 1]]}, }, diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index f6d3e8dd92..4c841e14b6 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -28,7 +28,6 @@ def _qwen_item(grid, image_data: str = _IMAGE_DATA): "kind": RAW_MM_ITEM_KIND, "modality": "image", "family": "qwen_vl", - "layout_fingerprint": "f" * 32, "raw_image_data": image_data, "payload": {"image_grid_thw": grid}, } From 56087658bea4f2465d68303c4ba585ce08026c8d Mon Sep 17 00:00:00 2001 From: eligotts Date: Tue, 4 Aug 2026 20:05:56 +0000 Subject: [PATCH 33/33] 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 ce7078aa5c..20b3f2dbff 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit ce7078aa5cce0c434507b7c159643e8914db076d +Subproject commit 20b3f2dbff84a39260191b799e283bdbcb91886e diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index 8792e35032..60fa3c7460 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) image_processor = _load_image_processor(processor_model_name, trust_remote_code) item = RawMMItem( - modality=ref.modality, family=ref.family, raw_image_data=ref.raw_image_data, 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 6d3bc70d4a..be3ad9b0e6 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_data: 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_data=_required_str(descriptor, "raw_image_data"), 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 6e6cfab0fd..38e87978e1 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -392,7 +392,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_data) if materialize is not None: return materialize(item) return {"materialized": item.raw_image_data} diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 9a78cccaf7..66ed08991c 100644 --- a/tests/unit/orchestrator/test_batch.py +++ b/tests/unit/orchestrator/test_batch.py @@ -415,7 +415,6 @@ def _image_ref(image_data: str, offset: int, length: int) -> MMImageRef: return MMImageRef( item={ "kind": "prime_raw_mm_item", - "modality": "image", "family": "qwen_vl", "raw_image_data": image_data, "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 93090c5239..2abd747c30 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_data": image_data, "payload": {"image_grid_thw": [[1, 1, 1]]}, diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py index 4c841e14b6..6af0ae9c39 100644 --- a/tests/unit/utils/test_mm.py +++ b/tests/unit/utils/test_mm.py @@ -26,7 +26,6 @@ def materialize(self, refs): def _qwen_item(grid, image_data: str = _IMAGE_DATA): return { "kind": RAW_MM_ITEM_KIND, - "modality": "image", "family": "qwen_vl", "raw_image_data": image_data, "payload": {"image_grid_thw": grid},