Skip to content

Support v1 inline raw multimodal images - #3128

Draft
eligotts wants to merge 39 commits into
mainfrom
feat/v1-inline-mm
Draft

Support v1 inline raw multimodal images#3128
eligotts wants to merge 39 commits into
mainfrom
feat/v1-inline-mm

Conversation

@eligotts

@eligotts eligotts commented Jul 23, 2026

Copy link
Copy Markdown
Contributor

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, PrimeRlServingTokens materializes 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:

  • Serving: raw 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—return invalid_mm_image_ref 400s. 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 by PRIME_RL_MM_MATERIALIZE_CACHE_GB.
  • Trainer: RawImageMaterializer decodes inline bytes from the descriptor. MMImageRef drops the uri field — the image lives once, in the descriptor.
  • Removed entirely (nothing to configure, nothing to go missing): [multimodal].offload_dir / MultimodalConfig, utils/run_assets.py, VF_RENDERER_IMAGE_OFFLOAD_DIR launcher + SLURM plumbing, and missing_mm_image_policy + the adapters' synthesize_placeholder machinery — an inline image cannot disappear, so the zero-loss placeholder path has nothing to guard.
  • Kept: SFT renderers default to multimodal_output='processed' with a hard guard against 'raw' (storage-independent).
  • Rollout records/traces keep the inline base64 (deliberate; noted in the monitor-run skill — expect payloads and results.jsonl to grow with turn count).
  • The local RL launcher preserves the inherited process environment when spawning the orchestrator, including the virtualenv PATH and 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-addressed file:// storage on top: ingress offload walker, offload-dir plumbing, missing-image policy, and disk-backed refs.

Validation

  • Original branch validation: tests/unit minus GPU-kernel model tests — 482 passed (398 non-train + 84 train non-models); uv lock --check clean; ruff check + format clean.
  • Materialization hardening: uv run pytest -q tests/unit/inference/test_serving_tokens.py21 passed; touched-file ruff check and format check clean.

Update: hardening + restored main's multimodal packing (7d85b170, fb6ef306)

  • Review hardening (7d85b170): subprocess env inheritance restored in rl_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 to invalid_mm_image_ref 400s.
  • Packing parity with main (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_bin merges refs with placeholder offsets rebased to the packed stream. Mirrors the offload branch's fix; tests restored to main's semantics.

Validation: tests/unit minus GPU-kernel model tests: 484 passed.

Update: unwrapped the ref payload (603118b / e40dded5)

raw_mm_ref no 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.

eligotts and others added 30 commits June 25, 2026 06:42
/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
- 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'.
eligotts added 9 commits July 23, 2026 17:11
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).
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants