Support v1 inline raw multimodal images - #3128
Draft
eligotts wants to merge 39 commits into
Draft
Conversation
/inference/v1/generate materializes every raw ref (no cache-only None branch); an unresolved ref is a hard error, not a silent None. Removes the cache-miss 409 path + helpers (_cache_only_mm_hashes/_is_missing_mm_cache_error/_missing_mm_cache_message). Bumps renderers/verifiers submodules to the matching cleanup commits. Deployment-agnostic; mm_hash encoder cache still skips re-encode. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
renderers.mm_store dropped the split_mmraw_ref backcompat alias; use split_raw_mm_ref. Bump renderers submodule pin to the cleanup commit. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
# Conflicts: # deps/verifiers # src/prime_rl/orchestrator/orchestrator.py # src/prime_rl/trainer/batch.py # tests/unit/orchestrator/test_batch.py
# Conflicts: # packages/prime-rl-configs/src/prime_rl/configs/inference.py # packages/prime-rl-configs/src/prime_rl/configs/rl.py # src/prime_rl/entrypoints/rl.py
# Conflicts: # deps/verifiers
- Kimi adapter reads the actual processor layout instead of hard-equating it to the renderer constant; the fingerprint comparison is the single drift detector, matching the Qwen adapter. - Orchestrator validates every image ref's placeholder span lands on image-typed tokens before a sample ships, so offset drift anywhere upstream fails loudly instead of silently truncating wrong. - FakeDataLoader carries the mm stat counters; trainer metrics read them directly. Drop the duplicate apply_run_asset_env in the entrypoint. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
renderers: bridge sidecar aliasing fix, HF smart_resize import, full-hash asset filenames, layout parity tests. verifiers: ingress offload covers every image part shape. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
# Conflicts: # deps/renderers # deps/verifiers # packages/prime-rl-configs/src/prime_rl/configs/inference.py # packages/prime-rl-configs/src/prime_rl/configs/orchestrator.py # src/prime_rl/orchestrator/trajectories.py
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Re-pins to the commits where each companion branch merged its own origin/main (renderers e64cc58, verifiers 2b1627d03), ahead of merging main into this branch.
Reconciles the raw multimodal offload work with main's iter_trainable_branches dedup, mm_kwargs packing, MXFP8, and seq_lens contract: - trajectories.py: main's iter_trainable_branches() loop is authoritative; this branch's _validate_image_spans + mm_refs build spliced into it. - trainer/model.py: kept the ForwardPolicy generalization over main's image_grid_thw string check in forward() — Kimi K2.5 and Qwen VL need opposite position_ids behavior, which the key-presence check can't express. Restored the numpy import main's _routed_experts_row_size needs (auto-merge dropped it; caught by tests, not conflicts). - trainer/batch.py: took main's mm_kwargs bin-packing machinery, with a guard that raw-ref (mm_refs) samples never pack — with text or each other: their placeholder offsets are sample-relative and _materialize_bin carries only the first sample's refs, so packing would drop or misalign images. Rewrote main's packing test to assert this branch's contract. - trainer/rl/train.py: kept both mm_forward_policy threading (ours) and the [model.vlm] guard for multimodal samples (main's). - Bumped deps/renderers (e64cc58) and deps/verifiers (2b1627d03) to their own main-reconciled companion-PR heads and relocked. The verifiers bump crosses the interleaving-agents refactor (#2049), which renamed textarena's seat: migrated env.agent -> env.player in the two wordle rl.toml configs. Validation: uv lock --check clean; tests/unit minus GPU-only tests/unit/train: 393 passed; tests/unit/train collects (160 tests) without import errors.
The renderers library defaults multimodal_output to 'raw' (built for the RL offload path), but SFT's data path consumes processed pixel tensors straight from the renderer and has no raw-ref materializer — the library default would silently ship JSON descriptors into training. SFTConfig now defaults its renderer to 'processed' and rejects an explicit 'raw'.
Every request carries a raw ref for every image in its prompt (prior turns included), so a 20-turn rollout with 5 accumulated images paid ~100 materializations (read + sha256 + PIL decode + HF processor forward) where 5 would do — multiplied by group size hammering the same prompt images. vLLM's processor cache can't help: production happens in our handler before vLLM sees the request. Adds a byte-bounded, single-flight LRU in serving_tokens.py keyed by (raw_ref, expected_placeholder_length, processor_model_name) — content- addressed, so hits are sound and evicted entries can only go cold, never stale. _decode_raw_mm_kwargs now gathers all of a request's images concurrently instead of the sequential await loop, and single-flight collapses concurrent misses for one new image into one materialization. Failures are never cached; they propagate to all awaiters and the next request retries cleanly. One knob: PRIME_RL_MM_MATERIALIZE_CACHE_GB (default 2.0). 0 disables all bookkeeping and single-flight — byte-identical to the uncached path, and the kill switch. Logs hits/misses/bytes/evictions every 1000 lookups; monitor-run skill documents the signature.
…ions Post-merge audit findings (adversarial review of the main reconciliation): - forward(): the ForwardPolicy fallback inverted main's gate for callers that don't thread an adapter policy — SFT passes mm_kwargs without one, so Qwen-VL SFT got packed 1D position_ids and the model skipped its internal MRoPE construction. The no-policy default now reproduces the key-presence heuristic (image_grid_thw => model owns position ids); RL's explicit adapter policies are unaffected. - _MaterializedRefCache: the owner request's cancellation (client disconnect) set CancelledError on the shared single-flight future, poisoning every deduped awaiter — and a cancelled awaiter could cancel the future out from under the rest. Materialization now runs as a detached task and awaiters shield the shared future. - _is_multimodal_sample: also treat eager mm_kwargs samples as multimodal so main's packing compatibility machinery stays live-correct, not dead code guarded by prepare_sample's rejection alone. - multimodal_sample_error: raw mm_refs samples require mm_token_type_ids (the orchestrator always stamps them; trainer truncation and forward policies rely on them). - CI runs all of tests/unit on GPU runners (not just the CPU subset): adapted main's two mm_kwargs packer tests to the raw-ref contract (raw samples never pack, even within a run), added the missing seq_lens kwarg to the branch-only forward-policy test, extended the batch no-pack test to cover the mm+mm direction, and migrated main's new configs/ci/nightly-fft/wordle.toml to the verifiers textarena player seat. - Bumped verifiers for a docstring completeness fix.
Intermediate storage mode built off feat/v1-raw-mm-offload: the image processor stays out of the env worker (vLLM-front materialization + cache, trainer adapter re-materialization, renderer geometry math all kept), but images travel inline as base64 data URLs instead of offloaded file:// run assets. - Serving: materializes from the ref's inline payload (decode + hash verify); the materialize cache re-keys to (mm_hash, placeholder_len, model) so keys stay small now that the ref embeds the full image. - Trainer: RawImageMaterializer decodes inline bytes from the descriptor; MMImageRef drops the uri field (the data lives once, in the descriptor). - Removed the entire offload plumbing layer: [multimodal].offload_dir, MultimodalConfig, run_assets.py, VF_RENDERER_IMAGE_OFFLOAD_DIR launcher/SLURM exports, and missing_mm_image_policy + the adapters' synthesize_placeholder machinery — an inline image cannot go missing, so the zero-loss placeholder path has nothing to guard. - Rollout records and traces keep the inline base64 (documented in the monitor-run skill). - Pins deps/renderers and deps/verifiers to the inline companion branches. tests/unit minus GPU-kernel model tests: 482 passed.
Mirrors the offload branch's restoration of main's mm packing feature, which the merge had regressed to never-pack for raw-ref samples: can_add treats same-adapter-family raw-ref samples as pack-compatible (one family per micro batch is what the materializer enforces), and _materialize_bin merges refs across the bin's samples with placeholder offsets rebased to the packed token stream. Same run/LoRA gating, seq_lens boundaries, and modality alignment as main; the trainer-side adapter already materializes a list of refs into concatenated tensors. Packing tests restored to main's semantics (pack within run with rebased offsets, family mismatch splits, never across runs).
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Intermediate multimodal storage mode — the inline bundle. Keeps the load-bearing architecture of the raw multimodal redesign — the image processor runs nowhere near the env worker: renderers compute token layout via pure geometry math,
PrimeRlServingTokensmaterializes pixels in front of vLLM (with the byte-bounded single-flight materialize cache), and the trainer re-materializes from the same source via family adapters — but images travel inline as base64 data URLs through messages, traces, transport, and refs. No shared image directory, no offload.Companion inline PRs:
Design:
mmraw:refs embed the inline image source; materialization decodes and sha256-verifies the inline payload. Expected descriptor failures—including hash, fingerprint, grid, and placeholder-length mismatches—returninvalid_mm_image_ref400s. The byte-bounded single-flight cache keys on a digest of the complete ref plus its modality, outer hash, placeholder length, processor model, and trust setting, so a distinct or stale descriptor cannot alias a validated entry. Cache size is controlled byPRIME_RL_MM_MATERIALIZE_CACHE_GB.RawImageMaterializerdecodes inline bytes from the descriptor.MMImageRefdrops theurifield — the image lives once, in the descriptor.[multimodal].offload_dir/MultimodalConfig,utils/run_assets.py,VF_RENDERER_IMAGE_OFFLOAD_DIRlauncher + SLURM plumbing, andmissing_mm_image_policy+ the adapters'synthesize_placeholdermachinery — an inline image cannot disappear, so the zero-loss placeholder path has nothing to guard.multimodal_output='processed'with a hard guard against'raw'(storage-independent).results.jsonlto grow with turn count).PATHand credentials.Relationship to the offload bundle
This branch is cut from
feat/v1-raw-mm-offload(#2836) and targets main independently. The offload bundle (renderers #89 / verifiers #1746 / prime-rl #2836) is the follow-on that layers content-addressedfile://storage on top: ingress offload walker, offload-dir plumbing, missing-image policy, and disk-backed refs.Validation
tests/unitminus GPU-kernel model tests — 482 passed (398non-train +84train non-models);uv lock --checkclean; ruff check + format clean.uv run pytest -q tests/unit/inference/test_serving_tokens.py— 21 passed; touched-file ruff check and format check clean.Update: hardening + restored main's multimodal packing (
7d85b170,fb6ef306)7d85b170): subprocess env inheritance restored inrl_local's orchestrator launch; the materialize-cache key now digests the full ref ((sha256(raw_ref), modality, mm_hash, placeholder_len, model, trust_remote_code)) so a distinct descriptor can never alias an already-validated entry; adapter validation errors map toinvalid_mm_image_ref400s.fb6ef306): the merge had left raw-ref multimodal samples never packing, regressing main's feature. Same-family raw-ref samples now pack with each other and with text from the same run/LoRA;_materialize_binmerges refs with placeholder offsets rebased to the packed stream. Mirrors the offload branch's fix; tests restored to main's semantics.Validation:
tests/unitminus GPU-kernel model tests: 484 passed.Update: unwrapped the ref payload (
603118b/e40dded5)raw_mm_refno longer base64-wraps its payload — it serializes compact JSON directly. The ref travels as a string inside a JSON request body, so the only encoding cost is escaping the payload's own quotes (~200 bytes); the base64 wrapper instead inflated the whole payload by a third, and in inline mode the payload contains the image.Measured on a 75 KiB JPEG: 130.3 KiB → 97.9 KiB per image slot on the wire (the 97.7 KiB data URL plus ~200 bytes of metadata/escapes) — a 25% cut, applied to every image slot of every request. Refs now parse with a single
partition(":")since the payload contains colons; the base64-alphabet guard retired with the wrapper.