diff --git a/deps/renderers b/deps/renderers index d4707862ac..20b3f2dbff 160000 --- a/deps/renderers +++ b/deps/renderers @@ -1 +1 @@ -Subproject commit d4707862ac83aa3773c21f4096aec72bd17b91e4 +Subproject commit 20b3f2dbff84a39260191b799e283bdbcb91886e diff --git a/deps/verifiers b/deps/verifiers index d30a3f48e5..10da86af08 160000 --- a/deps/verifiers +++ b/deps/verifiers @@ -1 +1 @@ -Subproject commit d30a3f48e5f14b06b3081b2102ec32cc3149b849 +Subproject commit 10da86af0850867d970155dbde5d28c7df1d2a3a diff --git a/docs/advanced.md b/docs/advanced.md index c4c1ba9b9c..e9d6e5f41f 100644 --- a/docs/advanced.md +++ b/docs/advanced.md @@ -98,7 +98,8 @@ VLM training requires a registered custom PrimeRL implementation. ### Limitations -- **Vision encoder frozen by default.** The default LoRA targets do not match Qwen3.5 vision modules. Set `freeze_vision_encoder = false` to fine-tune the encoder; this is incompatible with LoRA because LoRA freezes all non-adapter parameters. +- **Vision encoder frozen by default.** The default LoRA targets do not match Qwen3.5 vision modules. Set `freeze_vision_encoder = false` to fine-tune it; in that case it's FSDP-sharded per block. The combination `freeze_vision_encoder = false` + LoRA is rejected by a config validator — LoRA freezes everything non-adapter, so unfreezing the encoder under LoRA would be a silent no-op. +- **Truncation cuts at image boundaries.** Raw multimodal samples that exceed `seq_len` are truncated before packing at the start of the first image that doesn't fit — an image placeholder span is never split, and refs for dropped images are discarded. A sample whose leading image alone exceeds `seq_len` is rejected. - **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/docs/configuration.md b/docs/configuration.md index fbfdc5b67c..0e7c73588e 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -226,7 +226,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) and `WANDB_SHARED_*` (the single shared W&B run) — **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/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: 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 74f2a859b8..02dcf5e57e 100644 --- a/packages/prime-rl-configs/src/prime_rl/configs/shared.py +++ b/packages/prime-rl-configs/src/prime_rl/configs/shared.py @@ -10,7 +10,12 @@ # and the single shared W&B run. The launcher always sets these last, so allowing them in # `env_vars` would be a silent no-op (or, on multi-node, a footgun) — reject them instead. PROTECTED_ENV_VARS = frozenset( - {"CUDA_VISIBLE_DEVICES", "WANDB_SHARED_MODE", "WANDB_SHARED_RUN_ID", "WANDB_SHARED_LABEL"} + { + "CUDA_VISIBLE_DEVICES", + "WANDB_SHARED_MODE", + "WANDB_SHARED_RUN_ID", + "WANDB_SHARED_LABEL", + } ) diff --git a/skills/training/monitor-run/SKILL.md b/skills/training/monitor-run/SKILL.md index 70e95165b6..c087e260c0 100644 --- a/skills/training/monitor-run/SKILL.md +++ b/skills/training/monitor-run/SKILL.md @@ -173,6 +173,30 @@ 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 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, 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 + cache is disabled or thrashing (sized by `PRIME_RL_MM_MATERIALIZE_CACHE_GB`, + default 2.0, `0` disables). +- The orchestrator raises on placeholder/token drift ("does not cover + image-typed tokens") before a sample ships — treat any occurrence as a bug, + not noise. +- Trainer metrics: `mm/images_materialized` and `time/mm_materialize`. +- 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 All processes use `setproctitle` so they're visible in `ps`/`htop`/`pstree`: diff --git a/src/prime_rl/entrypoints/rl.py b/src/prime_rl/entrypoints/rl.py index 1c00ee4a74..020264ad74 100644 --- a/src/prime_rl/entrypoints/rl.py +++ b/src/prime_rl/entrypoints/rl.py @@ -128,6 +128,7 @@ def rl_local(config: RLConfig): "WANDB_SHARED_MODE": "1", "WANDB_SHARED_RUN_ID": os.environ.get("WANDB_SHARED_RUN_ID", uuid.uuid4().hex), } + inherited_env = dict(os.environ) # Validate client port matches inference server port if config.inference is not None and not config.orchestrator.model.client.is_elastic: @@ -173,7 +174,7 @@ def sigterm_handler(signum, frame): inference_process = Popen( inference_cmd, env={ - **os.environ, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_INFERENCE_ENV_VARS, **config.env_vars, @@ -228,7 +229,7 @@ def sigterm_handler(signum, frame): stdout=log_file, stderr=log_file, env={ - **os.environ, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, "LOGURU_FORCE_COLORS": "1", "WANDB_PROGRAM": "uv run rl", @@ -277,7 +278,7 @@ def sigterm_handler(signum, frame): trainer_process = Popen( trainer_cmd, env={ - **os.environ, + **inherited_env, **DEFAULT_COMMON_ENV_VARS, **DEFAULT_TRAINER_ENV_VARS, "LOGURU_FORCE_COLORS": "1", @@ -375,7 +376,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 "", ) - # 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 = { diff --git a/src/prime_rl/inference/vllm/serving_tokens.py b/src/prime_rl/inference/vllm/serving_tokens.py index e14a5ac83e..60fa3c7460 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.scale_out.token_in_token_out.serving.ServingTokens`` that 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 a lightweight raw + descriptor ref at every image slot (current and prior turns alike). This + 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`` + ``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,11 +32,20 @@ from __future__ import annotations -from collections.abc import AsyncGenerator, AsyncIterable -from functools import cached_property +import asyncio +import hashlib +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 +from io import BytesIO 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, @@ -44,10 +59,22 @@ ) from vllm.entrypoints.scale_out.token_in_token_out.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 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 + +logger = logging.getLogger(__name__) + + +@dataclass +class _MMImageRefError(Exception): + message: str + err_type: str = "invalid_mm_image_ref" + status_code: HTTPStatus = HTTPStatus.BAD_REQUEST class PrimeRlGenerateResponseChoice(GenerateResponseChoice): @@ -147,8 +174,277 @@ async def _client_set_max_tokens(raw_request: Request | None) -> bool: return isinstance(sp, dict) and "max_tokens" in sp +@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 _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: + try: + 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: + raise _MMImageRefError(f"Raw image hash mismatch: expected {ref.mm_hash}, got {actual_hash}") + return raw + + +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 inline raw image bytes: {exc}") from exc + + +def _materialize_raw_image_ref_sync( + raw_ref: str, + *, + feature_modality: str, + mm_hash: str, + expected_placeholder_length: int, + processor_model_name: str, + trust_remote_code: bool, +): + ref = _parse_raw_image_ref(raw_ref, feature_modality=feature_modality, mm_hash=mm_hash) + raw = _read_verified_raw_image(ref) + image = _decode_raw_image(raw) + image_processor = _load_image_processor(processor_model_name, trust_remote_code) + item = RawMMItem( + family=ref.family, + raw_image_data=ref.raw_image_data, + payload=dict(ref.payload), + ) + try: + adapter = get_multimodal_adapter(ref.family) + return adapter.materialize_for_vllm( + image_processor, + item, + image, + expected_placeholder_length, + ) + except (TypeError, ValueError) as exc: + raise _MMImageRefError(str(exc)) from exc + + +# (raw_ref_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 + + +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 asyncio.shield(inflight) + self.misses += 1 + self._maybe_log() + future: asyncio.Future = asyncio.get_running_loop().create_future() + self._inflight[key] = future + + 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) + 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]]: + 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)}" + ) + if len(placeholders) != len(hashes): + raise _MMImageRefError( + f"Multimodal placeholder/hash length mismatch for {feature_modality}: {len(placeholders)} != {len(hashes)}" + ) + 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( + features: Any, + *, + processor_model_name: str, + trust_remote_code: bool, + cache: _MaterializedRefCache, +) -> 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) + 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( + ( + 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 + ) + ) + + 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 + compact routed experts + max_tokens defaulting.""" + """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]: @@ -198,27 +494,29 @@ 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.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item from vllm.inputs import mm_input - from vllm.multimodal.inputs import ( - MultiModalKwargsItem, - MultiModalKwargsItems, - PlaceholderRange, - ) + from vllm.multimodal.inputs import MultiModalKwargsItems, 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, + cache=self._mm_materialize_cache, + ) + 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=MultiModalKwargsItems(mm_kwargs), 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..c6217e13dc --- /dev/null +++ b/src/prime_rl/multimodal/adapters/base.py @@ -0,0 +1,44 @@ +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 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, + ) -> Any: ... 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..cf3410b16f --- /dev/null +++ b/src/prime_rl/multimodal/adapters/kimi_k25.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from typing import Any + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RawMMItem + + +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("grid_thws") + if grid is None: + raise ValueError("Kimi raw descriptor payload is missing grid_thws") + 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}") + 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 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, actual_grid) in enumerate(zip(items, actual_grids, strict=True)): + expected = _grid_payload(item) + 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( + self, + image_processor: Any, + item: RawMMItem, + image: Any, + expected_placeholder_length: int, + ) -> Any: + from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems + + self.validate_item(item) + 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 != 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] 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..ad96b0faa7 --- /dev/null +++ b/src/prime_rl/multimodal/adapters/qwen_vl.py @@ -0,0 +1,102 @@ +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 _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 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}") + 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 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, actual_grid) in enumerate(zip(items, actual_grids, strict=True)): + expected = _grid_payload(item) + 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( + self, + image_processor: Any, + item: RawMMItem, + image: Any, + expected_placeholder_length: int, + ) -> Any: + from vllm.model_executor.models.qwen2_vl import _create_qwen2vl_field_factory + from vllm.multimodal.inputs import MultiModalKwargsItems + + self.validate_item(item) + hf_inputs = image_processor(images=[image], return_tensors="pt") + 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) + 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 != num_image_tokens: + raise ValueError( + f"Image placeholder length mismatch: expected {expected_placeholder_length}, got {num_image_tokens}" + ) + return mm_item 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..be3ad9b0e6 --- /dev/null +++ b/src/prime_rl/multimodal/schema.py @@ -0,0 +1,49 @@ +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +from renderers.mm_store import RAW_MM_ITEM_KIND + + +@dataclass(frozen=True) +class RawMMItem: + family: str + raw_image_data: str + payload: dict[str, Any] + + +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 _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 value.get("kind") != RAW_MM_ITEM_KIND: + raise ValueError("raw multimodal descriptor is missing the common envelope kind") + + +def parse_raw_mm_item(value: Any) -> RawMMItem: + descriptor = _descriptor_mapping(value) + _validate_envelope(descriptor) + return RawMMItem( + family=_required_str(descriptor, "family"), + raw_image_data=_required_str(descriptor, "raw_image_data"), + payload=_payload(descriptor), + ) diff --git a/src/prime_rl/orchestrator/trajectories.py b/src/prime_rl/orchestrator/trajectories.py index 8053453ac8..e62b47e13d 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 @@ -21,32 +21,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: @@ -66,6 +43,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 iter_trainable_branches(trace: vf.Trace) -> Iterator[tuple[vf.Branch, list[bool]]]: """Yield each branch that yields a training sample, with its trainable-token mask. @@ -101,21 +95,23 @@ 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. Errored rollouts are dropped upstream (`TrainSink.process_rollout`), so no error - handling happens here. 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. + handling happens here. A branch carrying images also 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. """ samples: list[TrainingSample] = [] for branch, mask in iter_trainable_branches(trace): 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) 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, @@ -123,7 +119,7 @@ def trace_to_samples( logprobs=branch.logprobs, temperatures=[], # filled by TrainSink.process_group 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/trainer/batch.py b/src/prime_rl/trainer/batch.py index 4eb8ff4baa..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 EncodedTensor, MicroBatch, 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 @@ -42,56 +43,23 @@ 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: +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 - 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 + return cut, MMRefs(images=mm_refs.images[:kept]) def multimodal_sample_error(sample: TrainingSample) -> str | None: @@ -101,6 +69,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 @@ -134,7 +106,9 @@ def prepare_sample(training_example: TrainingSample, seq_len: int) -> MicroBatch ref_kl_weights = list(training_example.ref_kl_weights) if training_example.ref_kl_weights is not None else None 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 = 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) @@ -149,11 +123,9 @@ 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. 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) + 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] @@ -214,7 +186,7 @@ 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, rl_weights=rl_weights, ce_weights=ce_weights, ref_kl_weights=ref_kl_weights, @@ -224,7 +196,16 @@ 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_kwargs is not None + 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 @@ -268,6 +249,12 @@ def can_add(self, sample: MicroBatch, max_seq_len: int, lora_idx: int) -> bool: 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" @@ -327,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) @@ -364,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 @@ -384,6 +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=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"], @@ -422,8 +417,9 @@ def packed_samples_into_micro_bs( With per-token temperatures, samples can be packed together regardless of their temperature values. Multimodal samples pack with text spans from the same run/LoRA and with - compatible eager ``mm_kwargs`` samples. Packed batches preserve sample - boundaries in ``seq_lens``. + 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/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index d2d6703819..34e03dd9a8 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -34,6 +34,7 @@ MXFP8Config, TokenizerConfig, ) +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.distributed import DeepEPExpertParallel, MXFP8AllToAllExpertParallel 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 ( @@ -1373,12 +1374,14 @@ 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 comes from the renderer rather than the processor. + # Generic multimodal kwargs materialized by the trainer's model processor + # (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 + # renderer token ids, not part of the processor's own output. mm_kwargs: dict[str, Tensor] | None = None, + mm_forward_policy: ForwardPolicy | None = None, mm_token_type_ids: Int[Tensor, "batch seq"] | None = None, # True when seq_lens holds the full pre-CP-shard document boundaries # (kept global because documents can straddle the shard cut). @@ -1390,14 +1393,19 @@ def forward( "temperature": temperature, } - if mm_kwargs: + if mm_kwargs is not None: # 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 - if "image_grid_thw" not in mm_kwargs: + # 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: kwargs["position_ids"] = position_ids else: kwargs["position_ids"] = position_ids diff --git a/src/prime_rl/trainer/rl/data.py b/src/prime_rl/trainer/rl/data.py index 1ad250cf72..f3863c7c9b 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,6 +8,7 @@ from torch import Tensor 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 from prime_rl.trainer.world import get_world @@ -16,6 +18,8 @@ TransportConfig, setup_micro_batch_receiver, ) +from prime_rl.transport.types import MMRefs +from prime_rl.utils.mm import RawImageMaterializer class TensorMicroBatch(TypedDict): @@ -39,12 +43,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 @@ -59,6 +60,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() @@ -72,6 +76,8 @@ 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 def wait_for_batch(self) -> None: return @@ -133,6 +139,7 @@ def _get_sample_micro_batch(self, generator: torch.Generator) -> TensorMicroBatc "seq_lens": torch.tensor(sequence_lengths, dtype=torch.long), "routed_experts": None, "mm_kwargs": None, + "mm_forward_policy": None, "mm_token_type_ids": None, "rl_weights": None, "ce_weights": None, @@ -166,6 +173,7 @@ def _get_micro_batch(self, generator: torch.Generator) -> TensorMicroBatch: "seq_lens": torch.tensor([self.seq_len], dtype=torch.long), "routed_experts": None, "mm_kwargs": None, + "mm_forward_policy": None, "mm_token_type_ids": None, "rl_weights": None, "ce_weights": None, @@ -187,6 +195,8 @@ def __init__( pad_to_multiple_of: int, bin_cost: Callable[[Sequence[int]], int], config: TransportConfig, + model_name: str, + model_trust_remote_code: bool, ): self.world = get_world() @@ -205,6 +215,8 @@ 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._reset_mm_stats() def wait_for_batch(self) -> None: if self.world.is_master: @@ -218,22 +230,37 @@ 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 + + @staticmethod + def _materialized_mm_parts(materialized: MaterializedMM | None) -> MaterializedMMParts: + if materialized is None: + return None, None + return materialized.kwargs, materialized.forward_policy + + def _materialize_mm_refs(self, micro_batch: MicroBatch, refs: MMRefs) -> MaterializedMMParts: + materialize_start = time.perf_counter() + materialized = self.mm_materializer.materialize(refs) + self.last_mm_materialize_time += time.perf_counter() - materialize_start + self.last_mm_images_materialized += len(refs.images) + return self._materialized_mm_parts(materialized) + def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: """Convert a MicroBatch (msgspec struct with lists) to a TensorMicroBatch (dict with tensors).""" if micro_batch.lora_num_tokens is None: 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 - 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() - } + mm_forward_policy: ForwardPolicy | None = None + 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) routed_experts = None packed_routed_experts = micro_batch.routed_experts if packed_routed_experts is not None: @@ -261,6 +288,7 @@ def _micro_batch_to_tensor(self, micro_batch: MicroBatch) -> TensorMicroBatch: lora_num_tokens=torch.tensor(micro_batch.lora_num_tokens, dtype=torch.int32), seq_lens=torch.tensor(micro_batch.seq_lens, dtype=torch.long), 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 295d3ad15e..0b5149a9a2 100644 --- a/src/prime_rl/trainer/rl/train.py +++ b/src/prime_rl/trainer/rl/train.py @@ -251,6 +251,8 @@ def load_run_checkpoint(_optimizer, idx: int) -> None: config.model.cp, build_bin_cost(model.config), config.rollout_transport, + config.model.name, + config.model.trust_remote_code, ) token_exporter = setup_token_exporter(config, parallel_dims, world, logger) @@ -365,12 +367,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") if mm_kwargs is not None and config.model.vlm is None: raise ValueError( "Received multimodal samples but [model.vlm] is not set. " @@ -433,6 +433,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, seq_lens=seq_lens, seq_lens_are_pre_shard=seq_lens_are_pre_shard, @@ -717,9 +718,11 @@ 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": 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": dataloader.last_mm_images_materialized, "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 121da68d92..64011c71c3 100644 --- a/src/prime_rl/transport/types.py +++ b/src/prime_rl/transport/types.py @@ -18,6 +18,32 @@ 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 + + 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 + offset: int + length: int + + +class MMRefs(msgspec.Struct, array_like=True, gc=False, omit_defaults=True): + """Raw multimodal sidecar references for one sample, in token order. + + The trainer materializes the referenced image files with its own processor. + Processed tensors are intentionally not part of this transport. + """ + + images: list[MMImageRef] + + # 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. @@ -34,16 +60,11 @@ 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) - # 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. + # Processed multimodal payloads are rejected by the v1 raw-image-ref path. 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) @@ -94,8 +115,9 @@ 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. + # 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) 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..cdc9843943 --- /dev/null +++ b/src/prime_rl/utils/mm.py @@ -0,0 +1,145 @@ +from __future__ import annotations + +import hashlib +from collections.abc import Iterable, Mapping +from io import BytesIO +from typing import Any + +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 +from prime_rl.multimodal.schema import RawMMItem, parse_raw_mm_item +from prime_rl.transport.types import MMImageRef, MMRefs + +IMAGE_MODALITY = "image" +SUPPORTED_MODALITIES = {IMAGE_MODALITY} + + +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 _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 _raw_item_dicts(items: Iterable[Mapping[str, Any]]) -> list[dict[str, Any]]: + item_dicts: list[dict[str, Any]] = [] + for item in items: + item_dict = dict(item) + parse_raw_mm_item(item_dict) + item_dicts.append(item_dict) + return item_dicts + + +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: + return None + _validate_modalities(mm_items) + + image_items = _raw_item_dicts(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, [])) + 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/placeholder mismatch: " + f"descriptors={len(image_items)}, hashes={len(image_hashes)}, placeholders={len(image_placeholders)}" + ) + + images: list[MMImageRef] = [] + prev_end = 0 + 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, 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) -> list[RawMMItem]: + return [parse_raw_mm_item(image.item) for image in refs.images] + + +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 _load_verified_images(image_refs: list[MMImageRef], image_items: list[RawMMItem]) -> list[Any]: + from PIL import Image + + images = [] + 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: expected {ref.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.""" + + 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) -> MaterializedMM | None: + image_items = _parse_image_refs(refs) + if not image_items: + return None + + image_processor = self.image_processor + adapter = _single_family_adapter(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 951ef5c9d8..38e87978e1 100644 --- a/tests/unit/inference/test_serving_tokens.py +++ b/tests/unit/inference/test_serving_tokens.py @@ -11,9 +11,11 @@ from __future__ import annotations import asyncio +import hashlib import numpy as np import pybase64 +import pytest from vllm.entrypoints.openai.engine.protocol import UsageInfo from vllm.entrypoints.scale_out.token_in_token_out.protocol import GenerateResponse, GenerateResponseChoice @@ -267,3 +269,258 @@ 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(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 + + 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(raw_bytes).hexdigest()[:32] + raw_ref = raw_mm_ref( + family="test_family", + modality="image", + mm_hash=mm_hash, + raw_image_data=image_data, + 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, + 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["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.raw_image_data == image_data + 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 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 grid 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 + from io import BytesIO + from types import SimpleNamespace + + from PIL import Image + from renderers.mm_store import raw_mm_ref + + 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(raw_bytes).hexdigest()[:32] + raw_ref = raw_mm_ref( + family="test_family", + modality="image", + mm_hash=mm_hash, + raw_image_data=image_data, + 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_image_data) + if materialize is not None: + return materialize(item) + 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(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() + + 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_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 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(image_seed=1) + features_b = _mm_features(image_seed=2) + + 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(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() + + 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(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() + + 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] diff --git a/tests/unit/orchestrator/test_batch.py b/tests/unit/orchestrator/test_batch.py index 8b4b31f7ed..66ed08991c 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, MMImageRef, MMRefs, RoutedExperts, TrainingSample def _routed_experts(data, dtype=np.uint8): @@ -17,11 +17,6 @@ def _routed_experts(data, dtype=np.uint8): ) -def _encoded(arr) -> EncodedTensor: - a = np.asarray(arr) - return EncodedTensor(data=a.tobytes(), shape=list(a.shape), dtype=str(a.dtype)) - - @pytest.fixture def make_training_example(): def _make_training_example( @@ -416,52 +411,66 @@ def test_prepare_sample_truncates_routed_experts(): assert micro_batch.env_names == ["test-env"] * 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 _image_ref(image_data: str, offset: int, length: int) -> MMImageRef: + return MMImageRef( + item={ + "kind": "prime_raw_mm_item", + "family": "qwen_vl", + "raw_image_data": image_data, + "payload": {"image_grid_thw": [[1, 1, 1]]}, + }, + hash="a" * 32, + offset=offset, + length=length, + ) + + +def test_prepare_sample_truncates_raw_mm_refs_at_image_boundary(): + 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, False, False, False, True, True], + mask=[False, False, True, True, False, True, True], logprobs=[0.0] * 7, temperatures=[1.0] * 7, - advantages=[0.0] * 6 + [1.0], + advantages=[0.0, 0.0, 1.0, 1.0, 0.0, 1.0, 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_token_type_ids=[0, 1, 1, 0, 1, 1, 0], + mm_refs=MMRefs(images=[first_image, second_image]), ) - # 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 - - -def test_prepare_batch_packs_multimodal_with_text(): - 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_kwargs={ - "pixel_values": _encoded(np.array([[1.0, 2.0]], dtype=np.float32)), - "image_grid_thw": _encoded(np.array([[1, 2, 2]], dtype=np.int64)), - }, - ) + # 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_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], + 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=[ref]), + ) + text_sample = TrainingSample( token_ids=[20, 21], mask=[False, True], @@ -472,25 +481,33 @@ def test_prepare_batch_packs_multimodal_with_text(): ) batches_per_gpu = prepare_batch( - rollouts=[mm_sample, text_sample], - seq_len=8, + 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], + 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) == 1 - batch = real_batches[0] - assert batch.seq_lens == [3, 2] - assert batch.sequence_lengths == [3, 2] - assert batch.position_ids == [0, 1, 2, 0, 1] - assert batch.mm_token_type_ids == [0, 1, 0, 0, 0] - assert batch.mm_kwargs is not None - assert batch.mm_kwargs["pixel_values"].shape == [1, 2] - assert batch.mm_kwargs["image_grid_thw"].shape == [1, 3] - assert batch.env_names == ["mm-env"] * 3 + ["text-env"] * 2 + mm_batches = [batch for batch in real_batches if batch.mm_refs is not None] + assert len(mm_batches) == 2 + 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/orchestrator/test_qwen3_vl_e2e.py b/tests/unit/orchestrator/test_qwen3_vl_e2e.py deleted file mode 100644 index d125e8a469..0000000000 --- a/tests/unit/orchestrator/test_qwen3_vl_e2e.py +++ /dev/null @@ -1,190 +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(): - """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, - 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.qwen3_vl import Qwen3VLRenderer - from transformers import AutoProcessor - from verifiers.clients.renderer_client import RendererClient - from verifiers.types import ( - ClientConfig, - UserMessage, - ) - from vllm.entrypoints.scale_out.token_in_token_out.mm_serde import decode_mm_kwargs_item - from vllm.entrypoints.scale_out.token_in_token_out.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. ──────────── - img = Image.new("RGB", (224, 224), color=(64, 128, 255)) - # 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}, - ] - ) - - # 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 decodes to MultiModalKwargsItem with the right keys. ─ - 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 - - # ── 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/train/rl/test_packer.py b/tests/unit/train/rl/test_packer.py index 93a1f04355..2abd747c30 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") @@ -52,16 +51,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: + image_data = f"data:image/png;base64,{value}" return TrainingSample( token_ids=[1, 250, 2], mask=[False, True, True], @@ -70,10 +61,21 @@ 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", + "family": "qwen_vl", + "raw_image_data": image_data, + "payload": {"image_grid_thw": [[1, 1, 1]]}, + }, + hash="a" * 32, + offset=1, + length=1, + ) + ] + ), ) @@ -176,7 +178,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 @@ -204,7 +206,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) @@ -212,8 +214,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_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) @@ -230,11 +232,8 @@ def test_multipacker_packs_mm_kwargs_within_each_run(tmp_path, monkeypatch): assert len(real_mm_mbs) == 2 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 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 diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py index 14fd521d5e..8071809655 100644 --- a/tests/unit/train/test_model_forward.py +++ b/tests/unit/train/test_model_forward.py @@ -1,8 +1,10 @@ from types import SimpleNamespace +import pytest import torch import torch.nn as nn +from prime_rl.multimodal.adapters.base import ForwardPolicy from prime_rl.trainer.model import forward @@ -35,11 +37,11 @@ def test_forward_passes_renderer_mm_token_type_ids_through(): position_ids, seq_lens=torch.tensor([input_ids.shape[1]]), mm_kwargs={"pixel_values": pixel_values, "image_grid_thw": image_grid_thw}, + mm_forward_policy=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) @@ -60,6 +62,7 @@ def test_forward_omits_mm_token_type_ids_when_renderer_does_not_supply(): position_ids, seq_lens=torch.tensor([input_ids.shape[1]]), mm_kwargs={"pixel_values": torch.ones(2, 3), "image_grid_thw": torch.tensor([[1, 1, 2]])}, + mm_forward_policy=ForwardPolicy(pass_position_ids_with_mm=False), ) assert model.kwargs is not None @@ -68,8 +71,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) @@ -80,7 +82,24 @@ def test_forward_keeps_position_ids_for_non_mrope_vlm(): position_ids, seq_lens=torch.tensor([input_ids.shape[1]]), 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) + + with pytest.raises(ValueError, match="mm_token_type_ids"): + forward( + 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), + ) diff --git a/tests/unit/utils/test_mm.py b/tests/unit/utils/test_mm.py new file mode 100644 index 0000000000..6af0ae9c39 --- /dev/null +++ b/tests/unit/utils/test_mm.py @@ -0,0 +1,86 @@ +import base64 +from types import SimpleNamespace + +import torch + +from prime_rl.multimodal.adapters.base import ForwardPolicy, MaterializedMM +from prime_rl.multimodal.schema import RAW_MM_ITEM_KIND +from prime_rl.trainer.rl.data import DataLoader +from prime_rl.transport.types import MicroBatch, MMImageRef, MMRefs +from prime_rl.utils.mm import build_mm_refs + +_IMAGE_DATA = "data:image/png;base64," + base64.b64encode(b"png-ish bytes").decode("ascii") + + +class _FakeMaterializer: + def materialize(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, image_data: str = _IMAGE_DATA): + return { + "kind": RAW_MM_ITEM_KIND, + "family": "qwen_vl", + "raw_image_data": image_data, + "payload": {"image_grid_thw": grid}, + } + + +def _refs() -> MMRefs: + return MMRefs( + 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 = _FakeMaterializer() + loader.last_mm_materialize_time = 0.0 + loader.last_mm_images_materialized = 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], + seq_lens=[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_build_mm_refs_accepts_inline_descriptors(): + multi_modal_data = SimpleNamespace( + 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() + + +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, True, True]] + assert tensor_batch["mm_token_type_ids"].tolist() == [[0, 1, 0]] diff --git a/uv.lock b/uv.lock index e59c12c400..d578003188 100644 --- a/uv.lock +++ b/uv.lock @@ -6187,10 +6187,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 = [