[Bugfix][Rollout] Fix Geo3K VLM multi-turn rollout by FeathBow · Pull Request #341 · vllm-project/vime

FeathBow · 2026-07-11T21:53:19Z

Summary

Fixes #331.

Remove the stale _apply_vllm_routed_experts and _inference_generate_tokens_and_logprobs imports left behind by refactor(rollout): align vllm_rollout.py with slime sglang_rollout.py #178.
Render the initial multimodal prompt once and reuse the returned vLLM features for later turns.
Preserve generated token IDs as the canonical conversation history and append only the chat-template observation suffix between turns.
Add a deterministic single-GPU Geo3K custom-rollout E2E regression to the manual Buildkite short suite.

#178 removed or inlined two private helpers from vime.rollout.vllm_rollout, but the Geo3K example continued importing them, causing the module-level ImportError reported in #331.

Restoring those helpers alone was insufficient. The legacy implementation decoded the generated assistant response and re-rendered the complete conversation before every turn. Chat-template rendering can tokenize that decoded response differently from the exact token stream returned by vLLM, which caused the prefix-stability failure reported in #331.

The fixed path renders only the initial prompt. Later turns preserve the exact generated token IDs and append only the EOS-aware observation suffix.

Validation

The E2E uses Qwen/Qwen3-VL-2B-Instruct, VeraIsHere/geo3k_imgurl_processed, one dataset row, two deterministic turns, and one GPU, against a real vLLM worker and router with one multimodal render request and two generation requests.

It verifies:

public dotted-path loading of the Geo3K generate entry point
exactly render -> generate -> generate
first-turn generated token IDs are preserved verbatim in the second request
both generation requests carry the render-derived multimodal features
response tokens, loss masks, and rollout log probabilities remain aligned

Manual A/B validation reproduced the original ImportError and demonstrated the legacy re-render divergence: with the deleted helpers restored, the second-turn re-rendered request contained a duplicated EOS token that is absent from the vLLM-generated token stream, while the fixed path issues the canonical request. The rollout patch was applied to a fresh worktree at b929921 (origin/main) and passed the single-GPU E2E; the vLLM worker, router, and engine processes were gone afterward. A later test-only cleanup removed redundant assertions and inlined test configuration without changing the rollout implementation.

Scope

This test exercises the Geo3K model/data custom-rollout path. It is not a full training-step E2E and uses a deterministic two-turn test environment rather than the real env_geo3k behavior. Its assertions assume the pinned model's greedy trajectory completes two turns within the token budget on the CI image.

gemini-code-assist

Code Review

This pull request refactors the multi-turn rollout logic for the Geo3K example package by encapsulating the execution flow into a structured _Geo3kRollout class, improving validation, and refining feature decoding. It also adds a single-GPU end-to-end regression test (test_geo3k_vlm_multi_turn_e2e.py) and updates Buildkite configurations to support 1-GPU test suites. The review feedback highlights two key areas for improvement: first, the logic for identifying the end-of-turn boundary using tokenizer.eos_token_id is fragile and should instead rely on the last token of the formatted prefix; second, the consecutive user messages in DUMMY_MESSAGES may violate strict alternating role constraints in certain chat templates and should include an intermediate assistant message.

gemini-code-assist · 2026-07-11T21:55:21Z

+    eos_token_id = getattr(tokenizer, "eos_token_id", None)
+    if not isinstance(eos_token_id, int):
+        raise ValueError("tokenizer.eos_token_id must be an integer")
    try:
-        env.reset()
-        latest_features = None
-        pending_obs_offset: int | None = None
-        rendered_body = await render()
-        prompt_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))
-        if not sample.tokens:
-            sample.tokens = list(prompt_ids)
-        if args.rollout_max_context_len is not None:
-            max_response_budget = max(0, args.rollout_max_context_len - len(sample.tokens))
-
-        for turn_idx in range(args.max_turns):
-            input_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))
-            latest_features = rendered_body.get("features")
-
-            if pending_obs_offset is not None:
-                obs_tokens = input_ids[pending_obs_offset:]
-                remaining = remaining_budget()
-                if remaining is not None and len(obs_tokens) > remaining:
-                    append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0))
-                    sample.status = Sample.Status.TRUNCATED
-                    break
-                append_response_window(obs_tokens, [0] * len(obs_tokens))
-                pending_obs_offset = None
-
-            current_sampling_params = sampling_params_for_turn()
-            if current_sampling_params is None:
-                sample.status = Sample.Status.TRUNCATED
-                break
+        eos_index = len(prefix_ids) - 1 - prefix_ids[::-1].index(eos_token_id)
+    except ValueError as exc:
+        raise ValueError(f"Dummy observation prefix lacks eos_token_id={eos_token_id}") from exc
+    boundary = prefix_ids[eos_index:] + rendered_ids[len(prefix_ids) :]
+    return boundary[1:] if canonical_ids and canonical_ids[-1] == eos_token_id else boundary


Searching for tokenizer.eos_token_id in the tokenized prefix is fragile and can fail with a ValueError if the chat template uses a different end-of-turn token (e.g., <|im_end|> or <|eot_id|>) than the model's default eos_token_id (e.g., <|endoftext|>).

Since prefix is formatted with add_generation_prompt=False, the very last token of prefix_ids (prefix_ids[-1]) is guaranteed to be the template's end-of-turn token. We can use this token directly to simplify the logic and make it robust across different models and templates.

Suggested change

eos_token_id = getattr(tokenizer, "eos_token_id", None)

if not isinstance(eos_token_id, int):

raise ValueError("tokenizer.eos_token_id must be an integer")

try:

env.reset()

latest_features = None

pending_obs_offset: int | None = None

rendered_body = await render()

prompt_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))

if not sample.tokens:

sample.tokens = list(prompt_ids)

if args.rollout_max_context_len is not None:

max_response_budget = max(0, args.rollout_max_context_len - len(sample.tokens))

for turn_idx in range(args.max_turns):

input_ids = _coerce_flat_int_token_ids(rendered_body.get("token_ids"))

latest_features = rendered_body.get("features")

if pending_obs_offset is not None:

obs_tokens = input_ids[pending_obs_offset:]

remaining = remaining_budget()

if remaining is not None and len(obs_tokens) > remaining:

append_response_window(obs_tokens[: max(remaining, 0)], [0] * max(remaining, 0))

sample.status = Sample.Status.TRUNCATED

break

append_response_window(obs_tokens, [0] * len(obs_tokens))

pending_obs_offset = None

current_sampling_params = sampling_params_for_turn()

if current_sampling_params is None:

sample.status = Sample.Status.TRUNCATED

break

eos_index = len(prefix_ids) - 1 - prefix_ids[::-1].index(eos_token_id)

except ValueError as exc:

raise ValueError(f"Dummy observation prefix lacks eos_token_id={eos_token_id}") from exc

boundary = prefix_ids[eos_index:] + rendered_ids[len(prefix_ids) :]

return boundary[1:] if canonical_ids and canonical_ids[-1] == eos_token_id else boundary

end_token_id = prefix_ids[-1]

boundary = [end_token_id] + rendered_ids[len(prefix_ids) :]

return boundary[1:] if canonical_ids and canonical_ids[-1] == end_token_id else boundary

Thanks for flagging this. The suggestion rests on the premise that prefix_ids[-1] is the turn terminator — but that doesn't hold for the Qwen templates this example targets. Each Qwen turn renders as ...<|im_end|>\n, so the last prefix token is the newline separator (id 198), not <|im_end|> (id 151645). Verified with the pinned Qwen3-VL-2B tokenizer.

This matters because the boundary is also responsible for closing the assistant turn when the generated stream doesn't already end with EOS (e.g. a stop_token_ids finish). In that case the current code emits <|im_end|>\n<|im_start|>user…, while the suggested code would emit \n<|im_start|>user… — the assistant turn is never terminated, silently corrupting the training stream.

For templates whose turn terminator differs from eos_token_id, the current behavior is an explicit ValueError rather than silent mis-splicing, which is intentional fail-fast; this example is scoped to Qwen-family templates (the E2E pins Qwen3-VL).

FeathBow · 2026-07-11T22:17:58Z

I've kept this PR scoped to the Geo3K regression reported in #331. While working on it I noticed the same pattern elsewhere, which may be worth a separate repository-wide follow-up: examples/multi_agent/agent_system.py still imports the removed _inference_generate_tokens_and_logprobs, and examples/mem_agent/rollout_client.py imports both that helper and the removed _align_engine_tokens_and_logprobs, so loading either example currently fails with the same kind of ImportError.

There are also several independent implementations of the vLLM choices[0].token_ids / logprobs.content parsing contract in the tree (at least four non-streaming plus a streaming-specific variant), and their malformed/missing-logprob behavior differs slightly, so consolidation seems better handled as a separate API-design discussion than as an incidental refactor in this bugfix. Happy to open a follow-up issue for either point if maintainers think it would be useful.

aoshen02 · 2026-07-13T00:50:12Z

+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any


Could you add a specific test suite vime-customized and add the test to it?

Thanks! makes sense. Done.

read-the-docs-community · 2026-07-13T07:39:21Z

Documentation build overview

📚 vime | 🛠️ Build #33611922 | 📁 Comparing a7565de against latest (b929921)

🔍 Preview build

44 files changed · ± 44 modified

± Modified

aoshen02 · 2026-07-14T06:45:44Z

could you resolve the conflict?

FeathBow · 2026-07-14T07:40:51Z

Done, rebased with main

andakai

Thanks for your effort. The fix is correct. The refactor looks also generally good to me. The previous implementation basically refer to slime's design. Now the multi-turn's design is clear in the _Geo3kRollout.run(). These two features also look great.

Render the initial multimodal prompt once and reuse the returned vLLM features for later turns.
Preserve generated token IDs as the canonical conversation history and append only the chat-template observation suffix between turns.

Here are some suggestions. If you have time, could you also run an e2e test and paste a reward curve to check? I can also do it later.

andakai · 2026-07-14T08:11:41Z

+        )
+
+    async def _generate_turn(self, sampling_params: dict[str, Any]) -> _Turn:
+        body = copy.deepcopy(self.render_body)


deepcopy may increase the memory use. I think .copy is enough.

andakai · 2026-07-14T08:17:56Z

+def _decode_feature_data(features: Any) -> dict[str, Any] | None:
    if not features:
        return None
-    if isinstance(features, str):
-        try:
-            features = json.loads(features)
-        except json.JSONDecodeError:
-            return None
-    if not isinstance(features, dict):
-        return None
+    decoded = json.loads(features) if isinstance(features, str) else features
+    if not isinstance(decoded, dict):
+        raise TypeError(f"vLLM features must decode to a dictionary, got {type(decoded).__name__}")
+    return decoded
+
+
+def _encoded_image_features(features: dict[str, Any]) -> list[Any] | None:
    kwargs_data = features.get("kwargs_data")
-    if not isinstance(kwargs_data, dict):
+    if not isinstance(kwargs_data, dict) or "image" not in kwargs_data:
+        return None
+    encoded_images = kwargs_data["image"]
+    if not isinstance(encoded_images, list):
+        raise TypeError("vLLM features.kwargs_data.image must be a list")
+    return list(encoded_images)
+
+
+def _merge_feature_tensors(parts_by_key: dict[str, list[torch.Tensor]]) -> dict[str, torch.Tensor]:
+    merged: dict[str, torch.Tensor] = {}
+    for key, values in parts_by_key.items():
+        merged[key] = values[0] if len(values) == 1 else torch.cat(values, dim=0)
+    return merged
+
+
+def _multimodal_train_inputs_from_features(features: Any) -> dict[str, torch.Tensor] | None:


Why replace the previous codes with these three seperate helper functions? I think the previous code can be more clear. I think there is no need to be defensive here and can be merged to _multimodal_train_inputs_from_features()

andakai · 2026-07-14T08:20:22Z

+    if prompt_length < 0:
+        raise ValueError(f"prompt_length must be non-negative, got {prompt_length}")


no need to do this. the input is already len()

andakai · 2026-07-14T08:24:03Z

+        if self.sample.status == Sample.Status.PENDING:
+            self.sample.status = Sample.Status.COMPLETED


When reaching the finalize(), its status will not be PENDING.

andakai · 2026-07-14T08:26:59Z

+    reserved = {"tokenize", "add_generation_prompt"}.intersection(kwargs)
+    if reserved:
+        raise ValueError(f"template kwargs cannot override reserved keys: {sorted(reserved)}")
+    prefix = tokenizer.apply_chat_template(
+        list(DUMMY_MESSAGES),
+        tools=tools,
+        tokenize=True,
+        add_generation_prompt=False,
+        **kwargs,
+    )


no need to do reserved here. The override keys will be checked tokenizer.apply_chat_template here.

FeathBow · 2026-07-14T09:17:05Z

Thanks for the detailed review! I've adopted all five inline suggestions and will push them shortly — the single-GPU E2E regression still passes after the changes. My cluster is tied up with other tasks at the moment, so the reward-curve run may take a while on my side.

FeathBow · 2026-07-14T09:28:23Z

Btw it may be worth a separate follow-up for the remaining private-helper imports and duplicated token/logprob parsing. I’ve kept it out of scope for #341, but I’d be happy to track this if you think that would be useful :)

I've kept this PR scoped to the Geo3K regression reported in #331. While working on it I noticed the same pattern elsewhere, which may be worth a separate repository-wide follow-up: examples/multi_agent/agent_system.py still imports the removed _inference_generate_tokens_and_logprobs, and examples/mem_agent/rollout_client.py imports both that helper and the removed _align_engine_tokens_and_logprobs, so loading either example currently fails with the same kind of ImportError.

There are also several independent implementations of the vLLM choices[0].token_ids / logprobs.content parsing contract in the tree (at least four non-streaming plus a streaming-specific variant), and their malformed/missing-logprob behavior differs slightly, so consolidation seems better handled as a separate API-design discussion than as an incidental refactor in this bugfix. Happy to open a follow-up issue for either point if maintainers think it would be useful.

FeathBow · 2026-07-15T11:40:24Z

Update: I ran the example's training loop end-to-end with the fixed rollout on a single GPU (scaled config: 60 rollout steps, rollout batch 8 × 8 samples, --vllm-enforce-eager). Training is stable throughout: the truncation ratio drops from ~0.5–0.6 early on to ~0.2, and the smoothed raw reward ends the run at its high (~0.45–0.5). Curves below.

aoshen02 · 2026-07-15T11:48:29Z

Update: I ran the example's training loop end-to-end with the fixed rollout on a single GPU (scaled config: 60 rollout steps, rollout batch 8 × 8 samples, --vllm-enforce-eager). Training is stable throughout: the truncation ratio drops from ~0.5–0.6 early on to ~0.2, and the smoothed raw reward ends the run at its high (~0.45–0.5). Curves below.

Hi, I wonder why use --vllm-enforce-eager, any blocker?

FeathBow · 2026-07-15T11:53:48Z

Oh no new blocker from my side — I kept the example's stock flags for the run (only scaled num-rollout / batch sizes).

FeathBow · 2026-07-15T11:58:04Z

--vllm-enforce-eager comes from the example itself run_geo3k_vlm_multi_turn.py note that "vLLM 0.22.0 needs eager mode for Qwen3-VL logprob parity", to be dropped once vllm#43617 lands. Since 43617 shipped in vLLM v0.23.0 , the flag could be removable now.

aoshen02 · 2026-07-15T11:58:29Z

--vllm-enforce-eager comes from the example itself run_geo3k_vlm_multi_turn.py note that "vLLM 0.22.0 needs eager mode for Qwen3-VL logprob parity", to be dropped once vllm#43617 lands. Since 43617 shipped in vLLM v0.23.0 , the flag could be removable now.

Please remove it cc @andakai

FeathBow · 2026-07-15T12:03:12Z

Done. If useful, I can also run a short training without the flag on my side.

andakai · 2026-07-15T12:03:50Z

Yes, can remove --vllm-enforce-eager now.

By the way, I will also run an experiment tonight. Below is the previous curve. Compared with yours, I think we need to run a more obvious increase curve, which should basically align with the previous one. Because the current pr is a big refactor, I think we need to ensure its accuracy is correct.

FeathBow · 2026-07-15T12:08:23Z

Sounds good — I'll also arrange a run tonight if my cluster permits.

andakai · 2026-07-16T03:55:12Z

I have run a test. The reward curve seems good. Thanks for your effort. @FeathBow cc @aoshen02

aoshen02 · 2026-07-16T06:53:15Z

The pre-commit failed.

FeathBow · 2026-07-16T07:01:49Z

I'll check it.

Signed-off-by: Feathbow <feathbow@gmail.com>

FeathBow · 2026-07-16T07:07:11Z

There was only one code formatting issue, fixed.

Signed-off-by: Meihan-chen <zr010426ztt@outlook.com>

* vlm Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> * vlm Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> * fix nomask Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> * fix vl thd Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> * fix weight sync Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> * align #341 rollout Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> * fix Signed-off-by: Meihan-chen <zr010426ztt@outlook.com> --------- Signed-off-by: Meihan-chen <zr010426ztt@outlook.com>

gemini-code-assist Bot reviewed Jul 11, 2026

View reviewed changes

FeathBow marked this pull request as draft July 11, 2026 22:00

FeathBow force-pushed the fix/issue-331-geo3k-rollout branch from 4dc9f6d to 1845620 Compare July 11, 2026 22:11

FeathBow marked this pull request as ready for review July 11, 2026 22:18

aoshen02 reviewed Jul 13, 2026

View reviewed changes

FeathBow force-pushed the fix/issue-331-geo3k-rollout branch from 1845620 to 0191b59 Compare July 13, 2026 07:37

FeathBow force-pushed the fix/issue-331-geo3k-rollout branch from 0191b59 to 5647c4b Compare July 14, 2026 07:40

andakai reviewed Jul 14, 2026

View reviewed changes

FeathBow force-pushed the fix/issue-331-geo3k-rollout branch from 5647c4b to 5655415 Compare July 14, 2026 09:17

FeathBow force-pushed the fix/issue-331-geo3k-rollout branch from 5655415 to 121c686 Compare July 15, 2026 12:02

floatlibai mentioned this pull request Jul 16, 2026

[NPU] Enable Qwen3-VL GRPO (single- & multi-turn) on Ascend #314

Merged

[Bugfix][Rollout] Fix Geo3K VLM multi-turn rollout

a7565de

Signed-off-by: Feathbow <feathbow@gmail.com>

FeathBow force-pushed the fix/issue-331-geo3k-rollout branch from 121c686 to a7565de Compare July 16, 2026 07:06

aoshen02 merged commit db6c87d into vllm-project:main Jul 16, 2026
3 checks passed

Meihan-chen added a commit to Meihan-chen/vime that referenced this pull request Jul 16, 2026

align vllm-project#341 rollout

80b0c3d

Signed-off-by: Meihan-chen <zr010426ztt@outlook.com>

Meihan-chen added a commit to Meihan-chen/vime that referenced this pull request Jul 16, 2026

align vllm-project#341 rollout

10baf84

Signed-off-by: Meihan-chen <zr010426ztt@outlook.com>

		if prompt_length < 0:
		raise ValueError(f"prompt_length must be non-negative, got {prompt_length}")

		if self.sample.status == Sample.Status.PENDING:
		self.sample.status = Sample.Status.COMPLETED

Uh oh!

Conversation

FeathBow commented Jul 11, 2026

Summary

Validation

Scope

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 11, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

FeathBow commented Jul 11, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

read-the-docs-community Bot commented Jul 13, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Documentation build overview

Uh oh!

aoshen02 commented Jul 14, 2026

Uh oh!

FeathBow commented Jul 14, 2026

Uh oh!

andakai 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!

FeathBow commented Jul 14, 2026

Uh oh!

FeathBow commented Jul 14, 2026

Uh oh!

FeathBow commented Jul 15, 2026

Uh oh!

aoshen02 commented Jul 15, 2026

Uh oh!

FeathBow commented Jul 15, 2026

Uh oh!

FeathBow commented Jul 15, 2026

Uh oh!

aoshen02 commented Jul 15, 2026

Uh oh!

FeathBow commented Jul 15, 2026

Uh oh!

andakai commented Jul 15, 2026

Uh oh!

FeathBow commented Jul 15, 2026

Uh oh!

andakai commented Jul 16, 2026

Uh oh!

aoshen02 commented Jul 16, 2026

Uh oh!

FeathBow commented Jul 16, 2026

Uh oh!

FeathBow commented Jul 16, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

read-the-docs-community Bot commented Jul 13, 2026 •

edited

Loading