Skip to content

Carry image inputs from rollout traces to the VERL training batch - #560

Merged
Zhiyuan He (hzy46) merged 6 commits into
microsoft:mainfrom
shuming-dev:multimodal-vision-training
Aug 26, 2026
Merged

Carry image inputs from rollout traces to the VERL training batch#560
Zhiyuan He (hzy46) merged 6 commits into
microsoft:mainfrom
shuming-dev:multimodal-vision-training

Conversation

@shuming-dev

Copy link
Copy Markdown
Contributor

Summary

Fixes #559.

In v1.0.0, VLM rollouts go through the proxy with images embedded in the OpenAI
messages, and vLLM rolls out with the images — but the trace → VERL training
data conversion drops all visual information. Training rows only carry
prompt_token_ids / response_token_ids with 2D cumsum position_ids; no
pixel_values / multi_modal_inputs is ever produced, so the FSDP engine's
extract_multi_modal_inputs(micro_batch.get("multi_modal_inputs", [])) sees
nothing and the training forward runs without the image features. Training
still "works" (loss decreases, no error), which makes this a silent correctness
bug: GRPO gradients are computed for a policy that cannot see the image.

This PR restores the vision signal end-to-end with minimal, opt-in changes:
text-only training is byte-identical to the current behavior.

Root cause

  1. The server stores the full request body (including base64 images) as raw
    model_request events, but the triplet view keeps only token ids, and the
    rollout manager builds Triplets from the trimmed view — image URLs never
    reach the training side.
  2. RolloutAdapter.get_train_data_batch assembles rows purely from token ids;
    it has no access to an HF processor and never builds multi_modal_inputs.
  3. For mrope models (Qwen-VL family), position ids are computed as a plain
    cumsum, ignoring image_grid_thw, so even the rotary position geometry is
    wrong for image-bearing rows.
  4. entrypoint relies on verl's hf_processor(), which only recognizes a fixed
    set of processor class names and silently returns None for others (e.g.
    Qwen3.5-VL), so the processor may never reach the trainer.

Changes

  • agentlightning/verl/agl_rollout_manager.py
    • Triplet gains an optional image_urls: list[str] | None field.
    • New helpers recover image URLs from the raw model_request payloads
      (OpenAI content parts, including JSON-serialized content), replicate the
      server-side triplet view (dedupe by prompt_token_ids, keep last) and the
      manager-side filtering (error / http_status >= 400 / empty response), and
      align the result one-to-one with the kept triplets. If alignment cannot be
      guaranteed, a warning is printed and image_urls is left unset (safe
      fallback). Text-only rollouts return early — zero behavior change, no new
      log output.
  • agentlightning/verl/rollout_adapter.py
    • RolloutAdapter accepts an optional processor (default None = exact
      original behavior).
    • Rows whose triplet carries image_urls are decoded (data: base64,
      file://, http(s)://) and processed into per-row multi_modal_inputs
      (pixel_values, image_grid_thw, …) in non_tensor_batch, matching the
      format produced by verl 0.8.0's own agent loop
      (AgentLoopWorker._compute_multi_modal_inputs), which the FSDP engine cats
      across rows and merges into model_inputs.
    • mrope processors (Qwen-VL family) get (batch, 4, seq_len) mrope position
      ids (1 text row + 3 vision rows, mirroring verl 0.8.0 _compute_position_ids);
      pure-text rows use the get_rope_index(image_grid_thw=None) variant.
    • Robustness guards:
      • trajectory-level aggregation with image-bearing traces raises a clear
        error pointing to trace_aggregator.level: transition (prompt merging
        breaks image-to-token alignment; previously the images were silently
        dropped).
      • Rows whose prompt was truncated through the image tokens (is_drop) fall
        back to text-only with a warning; these rows are filtered by
        is_drop_mask before the training forward and never contribute gradients.
      • Non-mrope VLM processors that return pixel_values without
        image_grid_thw still attach multi_modal_inputs and keep the 2D cumsum
        position ids.
      • Any per-row failure (image decode, processor error, mrope failure) degrades
        that row to text-only with a warning instead of aborting the step.
  • agentlightning/verl/trainer.py — forward self.processor (already stored
    by RayPPOTrainer.__init__) to RolloutAdapter.
  • agentlightning/verl/entrypoint.py — when verl's hf_processor() returns
    None, fall back to a plain AutoProcessor.from_pretrained load (text models
    are unaffected: tokenizer-backend results are still treated as None).
  • tests/verl/test_rollout_adapter.py / test_agl_rollout_manager.py — new
    unit tests (see Validation).
  • examples/multimodal_qa/ — new minimal multimodal example (see Validation).

All changed lines in the four library files are marked with
# [multimodal-patch] comments for easy review.

Validation

Real GRPO training (Qwen3.5-VL-2B, custom 8-node vision workflow, 33 steps)

Verified on a real multimodal GRPO run (single-image per node, transition-level
aggregation, verl 0.8.0):

Metric Step 0 Step 33 Delta
train reward (mean) -0.33 4.21 (peak 4.91) +4.5
val reward (55 samples) -0.40 4.69 (step 32) +5.1
rollout_corr/kl (train/inference engine consistency) ~4e-4 (0.0004–0.002) healthy

Without the patch the same pipeline runs silently with empty
multi_modal_inputs and 2D position ids (vision signal absent from the training
forward).

Unit tests (this PR)

  • test_agl_rollout_manager.py: image URL extraction (content parts and
    JSON-serialized content), raw→triplet alignment including dedupe / error /
    http≥400 / empty-response filtering, count-mismatch fallback with warning,
    and text-only early return with no warnings.
  • test_rollout_adapter.py: per-row multi_modal_inputs keys/structure with
    stub processors (no GPU / transformers / real model needed), (batch, 4, seq_len) mrope position ids, trajectory-level error for image-bearing traces,
    truncated-prompt fallback, non-mrope processor path, processors returning no
    vision tensors, data: URL decoding, and a pure-text regression asserting a
    processor-configured adapter produces byte-identical batches to an
    unconfigured one.

Example

examples/multimodal_qa/ is a self-contained synthetic single-image QA task
(PIL draws 1–5 red circles; the agent answers "How many red circles are in the
image?"; rule-based 0/1 reward; default model Qwen/Qwen2.5-VL-3B-Instruct,
no dataset download). Its README documents how to launch it and what to observe
(multi_modal_inputs present, position_ids.dim() == 3).

Limitations / non-goals

  • Transition-level aggregation only. trajectory level merges multi-turn
    prompts and cannot preserve the image-to-token correspondence; the adapter now
    raises a clear error for image-bearing traces instead of silently dropping
    images. (Consistent with Support of trajectory aggregation for mrope multimodal model, and add multimodal prefix checks for trajectory merge #469, which disabled trajectory aggregation for mrope
    models in v0.3.0.)
  • Truncated (is_drop) image rows degrade to text-only with a warning instead
    of attempting partial-image position reconstruction (cf. Fix mRoPE position ID crash on Qwen2-VL prompt truncation #482).
  • Critic / reward models are not covered: a text-only critic receiving rows with
    multi_modal_inputs may reject the unexpected keys. GRPO without a critic
    (the verl default for agentic RL here) is unaffected.
  • file:// image URLs are decoded in the trainer process, so the path must be
    readable there; data: (base64) URLs — the common case through the proxy —
    have no such constraint.
  • Each image row is re-processed by the HF processor in the trainer process
    (ViT preprocessing), adding tens to hundreds of ms of CPU per step.

Checklist

hushuming added 3 commits August 24, 2026 13:54
VLM rollouts went through the proxy with images in the OpenAI messages, but
the trace-to-training-data conversion dropped all visual information: training
rows only carried prompt/response token ids with 2D cumsum position ids, so
the FSDP engine never received pixel_values and the training forward could not
see the images (silent no-op for the vision signal; mrope position ids were
also wrong without image_grid_thw).

- agl_rollout_manager: add Triplet.image_urls and recover per-triplet image
  URLs from the raw model_request events, aligned with the server triplet
  view (dedupe by prompt_token_ids, keep last) and the manager-side filtering
  (error / http>=400 / empty response); on alignment mismatch leave unset
  with a warning. Text-only rollouts return early and keep byte-identical
  behavior.
- rollout_adapter: RolloutAdapter accepts an optional HF processor (None =
  original behavior). Rows with images get per-row multi_modal_inputs
  (pixel_values + image_grid_thw, matching verl 0.8.0's agent loop format)
  and, for mrope processors, (batch, 4, seq_len) position ids. Rows whose
  prompt was truncated through the image tokens (is_drop) fall back to
  text-only with a warning; they are filtered by is_drop_mask downstream.
  Non-mrope VLM processors (pixel_values without image_grid_thw) attach
  multi_modal_inputs and keep 2D cumsum position ids. Image-bearing traces
  at trajectory aggregation level raise a clear error pointing to the
  transition level, since prompt merging breaks image-to-token alignment.
- trainer: forward self.processor to RolloutAdapter.
- entrypoint: fall back to AutoProcessor when verl's hf_processor does not
  recognize the processor class (e.g. Qwen3.5-VL) and would silently
  return None.

Refs microsoft#559
…mbly

- agl_rollout_manager: image_url extraction from OpenAI-style messages
  (content parts and JSON-serialized content), raw-to-triplet alignment
  (dedupe / error / http>=400 / empty response filtering), count-mismatch
  fallback with warning, text-only early return without warnings.
- rollout_adapter: per-row multi_modal_inputs keys and structure with stub
  processors (no GPU / transformers needed), (batch, 4, seq_len) mrope
  position ids, trajectory-level error for image-bearing traces, truncated
  (is_drop) image prompts falling back to text-only, non-mrope processors
  keeping 2D position ids, processors without vision outputs, data: URL
  decoding, and a pure-text regression asserting processor-configured
  adapters produce byte-identical batches.
Single-image QA RL on fully synthetic data (no downloads): PIL renders
1-5 non-overlapping red circles at random positions, the agent asks the
proxied VLM how many red circles are in the image, and a rule-based reward
scores the numeric answer. Defaults to Qwen2.5-VL-3B-Instruct with
transition-level trace aggregation. Demonstrates that pixel_values and
mrope position ids reach the training forward pass.
Copilot AI lite review requested due to automatic review settings August 24, 2026 06:17

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This PR fixes a silent multimodal-correctness bug in the Agent Lightning → VERL training bridge by preserving image inputs from rollout traces and carrying them into the VERL training batch, including the mRoPE-specific position-id path needed for Qwen-VL family models.

Changes:

  • Extend rollout triplets to optionally carry image_urls, recovered from raw model_request events and aligned with the filtered/deduped triplet view.
  • Teach RolloutAdapter.get_train_data_batch to (optionally) build per-row multi_modal_inputs (e.g., pixel_values, image_grid_thw) and compute (batch, 4, seq_len) mRoPE position_ids for image-bearing rows.
  • Plumb an HF processor through entrypoint → trainer → adapter (with an AutoProcessor fallback) and add unit tests + a minimal multimodal example.

Reviewed changes

Copilot reviewed 10 out of 10 changed files in this pull request and generated 3 comments.

Show a summary per file
File Description
agentlightning/verl/agl_rollout_manager.py Adds Triplet.image_urls and helpers to extract/align image URLs from raw rollout events.
agentlightning/verl/rollout_adapter.py Adds optional processor-driven multimodal batch assembly and mRoPE position id computation; guards against unsupported aggregation.
agentlightning/verl/trainer.py Forwards the stored processor into RolloutAdapter construction.
agentlightning/verl/entrypoint.py Adds AutoProcessor fallback when verl’s hf_processor() returns None.
tests/verl/test_agl_rollout_manager.py Adds unit tests for image URL extraction and raw↔triplet alignment behavior.
tests/verl/test_rollout_adapter.py Adds unit tests for multimodal multi_modal_inputs, mRoPE (batch,4,seq) position ids, and fallbacks.
examples/multimodal_qa/train_multimodal_qa.py New synthetic image QA training script exercising the multimodal training path end-to-end.
examples/multimodal_qa/multimodal_qa_agent.py New local-mode agent that sends image + question and posts a rule-based reward.
examples/multimodal_qa/run_local.sh New helper script to launch server + controller + training locally.
examples/multimodal_qa/README.md New documentation describing the example and what to observe for multimodal correctness.

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

Comment on lines +207 to +212
# Server-side _dedupe_model_requests_by_prompt_token_ids: keep last per prompt key.
last_index_by_prompt: dict[tuple[Any, ...], int] = {}
for index, event in enumerate(raw_requests):
prompt_token_ids, _ = _raw_prompt_response_token_ids(event.data)
last_index_by_prompt[tuple(prompt_token_ids)] = index
kept_indexes = set(last_index_by_prompt.values())

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Good catch — fixed at the source in fac005e: _raw_prompt_response_token_ids now normalizes null/malformed token ids to empty lists, so this dedupe key and all other callers are safe.

Comment on lines +223 to +228
if url.startswith(("http://", "https://")):
import httpx

response = httpx.get(url, timeout=60.0, follow_redirects=True)
response.raise_for_status()
return Image.open(io.BytesIO(response.content)).convert("RGB")

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in fac005e: remote http(s) image fetching is now opt-in via AGENTLIGHTNING_ALLOW_REMOTE_IMAGE_FETCH=1, with image/* content-type and 50 MB size validation. data: and file:// URLs are unaffected.

Comment on lines +127 to +132
try:
fallback_processor = AutoProcessor.from_pretrained(
local_path, trust_remote_code=trust_remote_code, use_fast=True
)
except Exception:
fallback_processor = None

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Fixed in fac005e: the AutoProcessor fallback now logs a warning with the underlying exception instead of leaving processor=None silently.

@shuming-dev

Copy link
Copy Markdown
Contributor Author

@microsoft-github-policy-service agree

- normalize null prompt/response token ids to empty lists so malformed
  raw events cannot raise in tuple() during dedupe alignment
- gate http(s) image fetching behind AGENTLIGHTNING_ALLOW_REMOTE_IMAGE_FETCH
  with content-type and size validation (SSRF/egress hardening)
- log a warning when the AutoProcessor fallback fails instead of
  silently leaving processor=None
@hzy46

Copy link
Copy Markdown
Contributor

Thank you for submitting this PR. I will check it.

@hzy46

Copy link
Copy Markdown
Contributor

Ming (@shuming-dev) I am generally okay with this PR. Please address the following issues, and then I’ll merge it:

  1. The three functions _extract_image_urls_from_messages, _raw_prompt_response_token_ids, and _aligned_image_urls feel overly fragmented for simple logic. Please simplify them appropriately.

  2. Please avoid using environment variables to control program behavior, as this can make the behavior harder for users to understand. For example, I suggest removing AGENTLIGHTNING_ALLOW_REMOTE_IMAGE_FETCH and allowing remote image fetching by default.

  3. Please move the example README to docs instead of placing it under examples/multimodal_qa. Please follow the documentation format used by the other examples. For instance, the document should start with a table summarizing the required resources and supported run modes, and the main body should explain the dataset and how to run the example.

…efault, move example doc to docs

- merge token-id extraction into _aligned_image_urls (single pass)
- drop AGENTLIGHTNING_ALLOW_REMOTE_IMAGE_FETCH env gate; remote http(s)
  image fetching is now allowed by default (content-type/size checks kept)
- move multimodal_qa example doc to docs/80-example-multimodal-qa.md
  following the example doc format; register in mkdocs nav and docs index
- example default model: Qwen/Qwen3.5-2B
@shuming-dev

Copy link
Copy Markdown
Contributor Author

Zhiyuan He (@hzy46) All three points are addressed in 106dc7a:

  1. Simplified: _raw_prompt_response_token_ids is merged into _aligned_image_urls (single pass over the events); _extract_image_urls_from_messages stays as the small message-parsing helper.
  2. AGENTLIGHTNING_ALLOW_REMOTE_IMAGE_FETCH removed — remote http(s) image fetching is now allowed by default, with the content-type and 50 MB size checks kept.
  3. Example doc moved to docs/80-example-multimodal-qa.md following the Calc-X/GSM8K format (resource table on top, dataset + how-to-run sections), registered in the mkdocs nav and the docs index. The examples dir keeps a short pointer README.

Tests (18 multimodal unit tests) and ruff are green.

@hzy46

Copy link
Copy Markdown
Contributor

Ming (@shuming-dev) There're some lint and type check problem, please fix, thx

@shuming-dev

Copy link
Copy Markdown
Contributor Author

Zhiyuan He (@hzy46) Fixed the two CI failures in 29986b2:

  • Lint: examples/multimodal_qa/run_local.sh now has the executable bit set (shebang check).
  • Type-check: the two pyright Sized errors in tests/verl/test_rollout_adapter.py are fixed (images: list | None).

Local pytest (18 multimodal tests) and ruff are green.

@hzy46
Zhiyuan He (hzy46) merged commit e43cbf2 into microsoft:main Aug 26, 2026
6 checks passed
@hzy46

Copy link
Copy Markdown
Contributor

merged, thx for this contribution!

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.

Multimodal RL (v1.0.0): training forward pass never receives pixel_values — vision features silently dropped between rollout and training

3 participants