Add VLM Geo3k multiturn experiment by andakai · Pull Request #120 · vllm-project/vime

andakai · 2026-06-02T04:48:10Z

Summary

This PR fixes the Geo3K VLM multi-turn rollout path so that Geo3k multiturn with vLLM reward score can increase.

What Changed

Geo3K VLM multi-turn rollout

Reworked examples/geo3k_vlm_multi_turn/rollout.py to use the full-conversation vLLM render route as the source of truth for prompt tokens and multimodal features.
Assistant tokens are appended from vLLM generation output.
Tool/observation tokens are sliced from the next full render using a pending offset and masked out from loss.
Added strict prefix-stability metadata/check support for detecting render-token drift.
Appends a masked EOS token after stop-string generation when needed, and uses that appended EOS in the next observation offset.
Propagates vLLM routed expert metadata through the custom rollout path.
Supports local image paths and normalized data:image/None URLs for render requests.

Geo3K environment behavior

examples/geo3k_vlm_multi_turn/env_geo3k.py now ends an episode as soon as calc_score returns a correct answer.
Wrong answers continue receiving feedback until max_turns.
If the model does not produce a valid tool call, the env now extracts a final boxed/text answer, scores it, and ends the episode. This preserves reward for correct final answers that are not emitted in tool-call format.

Run script defaults

examples/geo3k_vlm_multi_turn/run_geo3k_vlm_multi_turn.py now uses the validated Qwen3-VL-2B colocated vLLM/Megatron setup:
- Qwen3-VL-2B-Instruct
- 4 GPUs
- rollout TP = 1
- Megatron TP = 1
- --vllm-router-policy round_robin
- --vllm-max-model-len 32768
- --vllm-gpu-memory-utilization 0.9
- --vllm-generation-config vllm
- --vllm-enforce-eager
- --vllm-logprobs-mode processed_logprobs
- GRPO, lr 1e-6, global batch size 512

Validation

W&B run:

https://wandb.ai/danab2-hong-kong-university-of-science-and-technology/vime-dev/runs/pak7pj5q

Why vLLM eager mode

Without --vllm-enforce-eager, Qwen3-VL multimodal rollout showed a large silent parity drift between vLLM rollout logprobs and Megatron replay logprobs. The VLM two-turn diagnostic reached about 0.1399 train/rollout logprob absolute diff.

With eager mode, the same path dropped to text-only-level parity, around 0.020 in smoke diagnostics and around 0.014-0.016 in the long training run. This matches SkyRL's stability-oriented default of using eager vLLM inference for its training engine.

cc @CalvinXKY

Signed-off-by: Dakai An <dakaian108@gmail.com>

gemini-code-assist

Code Review

This pull request refactors the Geo3K multi-turn environment and rollout generation logic to support multi-turn interactions, boxed/text answer extraction when tool calls are missing, and prefix-stable rendering. It also updates the execution script configuration, including hardcoding model parameters and adjusting vLLM arguments. The code review feedback suggests making _image_to_render_url more robust by supporting pathlib.Path objects directly and adding a defensive check to ensure kwargs_data['image'] is a list before iterating over it.

gemini-code-assist · 2026-06-02T04:49:57Z

+    if isinstance(image, str):
+        if image.startswith("data:"):
+            # Some prepared Geo3K rows use data:image/None; vLLM expects a concrete media type.
+            return image.replace("data:image/None;", "data:image/png;", 1)
+        if image.startswith(("http://", "https://")):
+            return image
+        image_path = Path(image).expanduser()
+        if image_path.exists():
+            with Image.open(image_path) as loaded_image:
+                return encode_image_for_rollout_engine(loaded_image)
+        raise ValueError(f"Unsupported image string for vLLM render: {image!r}")


Support pathlib.Path objects directly in _image_to_render_url. Currently, if image is a Path object, isinstance(image, str) evaluates to False, and it falls back to encode_image_for_rollout_engine(image) which expects a PIL Image and will fail. Checking for (str, Path) makes the utility more robust.

Suggested change

if isinstance(image, str):

if image.startswith("data:"):

# Some prepared Geo3K rows use data:image/None; vLLM expects a concrete media type.

return image.replace("data:image/None;", "data:image/png;", 1)

if image.startswith(("http://", "https://")):

return image

image_path = Path(image).expanduser()

if image_path.exists():

with Image.open(image_path) as loaded_image:

return encode_image_for_rollout_engine(loaded_image)

raise ValueError(f"Unsupported image string for vLLM render: {image!r}")

if isinstance(image, (str, Path)):

image_str = str(image)

if image_str.startswith("data:"):

# Some prepared Geo3K rows use data:image/None; vLLM expects a concrete media type.

return image_str.replace("data:image/None;", "data:image/png;", 1)

if image_str.startswith(("http://", "https://")):

return image_str

image_path = Path(image_str).expanduser()

if image_path.exists():

with Image.open(image_path) as loaded_image:

return encode_image_for_rollout_engine(loaded_image)

raise ValueError(f"Unsupported image path or string for vLLM render: {image!r}")

gemini-code-assist · 2026-06-02T04:49:57Z

+    if "image" not in kwargs_data:
        return None


Add a defensive check to ensure kwargs_data["image"] is a list before iterating over it. If the rendered features contain a malformed or non-list value for "image", it could cause a runtime TypeError during rollout.

Suggested change

if "image" not in kwargs_data:

return None

if "image" not in kwargs_data or not isinstance(kwargs_data["image"], list):

return None

andakai · 2026-06-02T07:06:31Z

The enforce eager problem is fixed by Fix Qwen3-VL and Qwen3-omni-thinker accuracy degradation from deepstack inputs under torch.compile vllm#43617. vllm 0.22.0 doesn't have this pr included. I have tested locally: if using this pr, the logprob will become low. Will update this in the later vllm version.
The router policy can be "consistent hash" instead of "round-robin". This will have higher prefix cache hit rate.

Prefix Cache Hit Rate Comparison

I compared the vLLM engine-side Prefix cache hit rate logs for two 5-step no-enforce-eager Geo3K VLM multi-turn runs:

Router policy	Log samples	Mean hit rate	Min	Max
`consistent_hash`	103	82.6%	60.9%	90.7%
`round_robin`	113	63.4%	46.5%	77.4%

consistent_hash improves the average engine-reported prefix cache hit rate by 19.2 percentage points over round_robin, which is about a 30% relative improvement. The rollout time is similar due to the long-tail prompt.

Both runs completed successfully:

consistent_hash + no enforce eager: https://wandb.ai/danab2-hong-kong-university-of-science-and-technology/vime-dev/runs/96i0vy4r
round_robin + no enforce eager: https://wandb.ai/danab2-hong-kong-university-of-science-and-technology/vime-dev/runs/93mq3y9f

Signed-off-by: Dakai An <dakaian108@gmail.com>

CalvinXKY · 2026-06-02T07:14:22Z

 import vime.utils.misc as U
 from vime.utils.external_utils.command_utils import execute_train

-MODEL_NAME = os.environ.get("SLIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct")


Remove "SLIME_SCRIPT_MODEL_NAME " from Readme. Consider refresh it too.

Will fix in this PR: https://github.com/vllm-project/vime/pull/121/changes#diff-c538ffca60ebb211024d65a44fbca17242c074e6e02504542ec09fa20cade232. pr 105 only rename the path.

CalvinXKY · 2026-06-02T07:25:51Z

+                    f"{sample.metadata['multiturn_render']}"
+                )
+
+        sample.multimodal_train_inputs = _multimodal_train_inputs_from_features(latest_features)


sample.tokens is stitched across turns, but multimodal features come from a single latest render. To assert that the count of image_token_id in sample.tokens matches image_grid_thw.shape[0] (or equivalent) before assigning multimodal_train_inputs, and fail loudly on mismatch rather than silently returning None. Is that better?

I tested and generally these two are the same. But there may be some mismatches in other scenarios, so I have added a validation here.

Signed-off-by: Dakai An <dakaian108@gmail.com>

CalvinXKY

LGTM

aoshen02

Code-review notes (Codex strict pass + manual). Inline comments below; 2 are correctness (multi-turn token alignment), the rest are simplifications. Pre-existing items not in this PR's diff (dead _extract_balanced_json, optional orjson, bare env.close() except, commented-out eval_args) are noted separately, not inlined since they predate this PR.

aoshen02 · 2026-06-03T13:25:35Z

+                            "Routing replay is not supported when appending an artificial EOS after a stop string, "
+                            "because vLLM does not return routed experts for that extra token."
+                        )
+                    train_tokens.append(int(eos_token_id))


#1 (correctness, multi-turn off-by-one) — part 1/2. This synthetic EOS is appended to train_tokens only; it is not added to messages (L298 stores just response_text). So the next render's assistant segment has 1 fewer token than len(train_tokens). See the offset note on L318.

aoshen02 · 2026-06-03T13:25:35Z

-        sample.multimodal_train_inputs = _merge_multimodal_train_inputs(multimodal_train_inputs_buffer)
+            next_user_message = env.format_observation(observation)
+            messages.append(next_user_message)
+            pending_obs_offset = len(input_ids) + len(train_tokens)


#1 (correctness) — part 2/2. len(train_tokens) includes the synthetic EOS appended at L290, but the re-rendered conversation (messages has only response_text, no EOS) has 1 fewer token before the observation. So next turn obs_tokens = input_ids[pending_obs_offset:] (L244) drops the first real observation token, desyncing sample.tokens from the context the model actually generated against.

Fix: compute the render-prefix length from len(new_tokens) (exclude the synthetic EOS), keep train_tokens separate:

render_prefix_len = len(input_ids) + len(new_tokens) pending_obs_offset = render_prefix_len

Only triggers when append_stop_eos fires.

These two parts are fixed by separating render-vs-training lengths:

Keep appending synthetic EOS to train_tokens as before.

Compute the render prefix with original generation tokens only: render_prefix_len = len(input_ids) + len(new_tokens) and use that for pending_obs_offset.

aoshen02 · 2026-06-03T13:25:36Z

+            messages.append(next_user_message)
+            pending_obs_offset = len(input_ids) + len(train_tokens)
+            rendered_body = await render()
+            latest_features = rendered_body.get("features")


Cleanup (redundant). This latest_features = rendered_body.get("features") is dead: the next loop iteration re-assigns it at the top (L241) from the same rendered_body, and the only consumer (L335) runs after a break that already saw L241's value. Safe to delete.

True, this is redundant. Deleted

aoshen02 · 2026-06-03T13:25:36Z

+                "sample_len": len(sample.tokens),
+                "rendered_len": len(rendered_ids),
+            }
+            if getattr(args, "strict_multiturn_render_token_match", False) and not is_prefix_stable:


#2 (correctness) — the alignment guard is off by default. The prefix-stability check at L322 is exactly what catches the #1 desync (and any render-vs-accumulation mismatch), but it only raises when strict_multiturn_render_token_match is True — which defaults to False. By default a mismatch is silently recorded in metadata and training proceeds on misaligned tokens.

Suggest defaulting strict to True (raise on mismatch), or making rendered_ids canonical and rebuilding sample.tokens/loss_mask/rollout_log_probs from the rendered conversation.

Yes this is false by default. Now I change it to fail fast whenever mismatch happens, without relying on "strict_multiturn_render_token_match" now.

aoshen02 · 2026-06-03T13:25:36Z

+        raise ValueError("Environment module must expose a callable `build_env(sample, args)`.")
+    try:
+        env = build_env(sample=sample, args=args)
+    except TypeError:


Cleanup. This except TypeError retry can swallow a real TypeError raised inside build_env (a constructor bug), making debugging hard. build_env here accepts keywords — just call env = build_env(sample=sample, args=args).

aoshen02 · 2026-06-03T13:25:36Z


        score = self._score_answer(parsed_answer)
        self.last_tool_score = score
+        self.correct = score == 1.0


Cleanup (redundant state). self.correct is only read immediately below at L249; it doesn't need to live on the env. A local correct = score == 1.0 plus return obs, correct or is_final_turn, info removes the field and its init/reset at L42/L49.

Fix. This is truly redundant.

aoshen02 · 2026-06-03T13:25:36Z


-MODEL_NAME = os.environ.get("SLIME_SCRIPT_MODEL_NAME", "Qwen3-VL-2B-Instruct")
+MODEL_NAME = "Qwen3-VL-2B-Instruct"
 assert MODEL_NAME in {


Cleanup (dead assert). MODEL_NAME is a hardcoded constant immediately asserted to be in a set — the assert is always true. Either restore an env override (MODEL_NAME = os.environ.get(...)) or drop the assert and keep just the constant.

Fixed. Use the "VIME_SCIRPT_MODEL_NAME" now.

Signed-off-by: Dakai An <dakaian108@gmail.com>

Add vlm multiturn experiment.

c30df68

Signed-off-by: Dakai An <dakaian108@gmail.com>

andakai requested a review from aoshen02 June 2, 2026 04:48

gemini-code-assist Bot reviewed Jun 2, 2026

View reviewed changes

update

372d6e0

Signed-off-by: Dakai An <dakaian108@gmail.com>

CalvinXKY reviewed Jun 2, 2026

View reviewed changes

update

95e89de

Signed-off-by: Dakai An <dakaian108@gmail.com>

CalvinXKY approved these changes Jun 2, 2026

View reviewed changes

CalvinXKY mentioned this pull request Jun 2, 2026

[Example][Bugfix]: align custom rollout paths with vLLM router API #75

Closed

3 tasks

aoshen02 reviewed Jun 3, 2026

View reviewed changes

CalvinXKY mentioned this pull request Jun 3, 2026

[Example] Add tau-bench multi-turn tool-use example #142

Merged

3 tasks

andakai added 6 commits June 4, 2026 05:16

Fix multi-turn EOS offset handling in render prefix

de7ae5d

Signed-off-by: Dakai An <dakaian108@gmail.com>

chore(geo3k): remove redundant latest_features assignment

2ca2d27

Signed-off-by: Dakai An <dakaian108@gmail.com>

enforce strict render-prefix alignment in geo3k rollout

7dfd9e8

Signed-off-by: Dakai An <dakaian108@gmail.com>

Simplify env construction call in geo3k multi-turn rollout

175f2be

Signed-off-by: Dakai An <dakaian108@gmail.com>

Keep tool call score history fields while removing env-correct state

ca939e2

Signed-off-by: Dakai An <dakaian108@gmail.com>

Merge origin/main into adk/vlm-multiturn

0696c9b

aoshen02 merged commit 5b499b5 into main Jun 4, 2026
11 of 13 checks passed

aoshen02 deleted the adk/vlm-multiturn branch June 8, 2026 14:17

Uh oh!

Conversation

andakai commented Jun 2, 2026

Summary

What Changed

Geo3K VLM multi-turn rollout

Geo3K environment behavior

Run script defaults

Validation

Why vLLM eager mode

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jun 2, 2026

Choose a reason for hiding this comment

Uh oh!

andakai commented Jun 2, 2026

Prefix Cache Hit Rate Comparison

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

CalvinXKY left a comment

Choose a reason for hiding this comment

Uh oh!

aoshen02 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants