Skip to content

Multi-turn VLM generator corrupts trajectories for all thinking models: history-editing chat templates break the deferred obs-token offsets #2075

Description

@dzorlu

Summary

SkyRLVLMGymGenerator.agent_loop collects observation tokens by slicing the next turn's render at a predicted offset: pending_obs_offset = len(previous_render_ids) + len(generated_ids), flushed at the top of the next loop iteration. This assumes every re-render is a token-prefix extension of the previous one. The chat templates of all current mainstream thinking models violate that assumption: they remove reasoning content from historical assistant turns on every render. Each completed think block therefore makes the next render shorter than predicted by len(reasoning) + 3 tokens, the flush slice starts that many tokens inside the following observation, and the head of the observation is silently dropped from response_ids. When the observation contains an image, the dropped tokens include <|image_pad|> tokens while the multimodal features (built from the full render) remain complete, and training fails at the first policy update inside the HF forward:

ValueError: Image features and image tokens do not match, tokens: 1261, features: 1680

The code documents this exact limitation in a NOTE ("This only works for standard tokenizers which preserve sequence extension... will not hold for thinking models, e.g., Qwen3-Thinking"), present since the file was created in #1486. There is no guard, so the failure surfaces ~30 minutes later in torch._check rather than at trajectory creation.

Observed

Two independent GRPO runs, Qwen3.5-9B, multi-turn computer-use tasks (screenshots in observations), max_turns=64. Both completed generation and reward computation, then crashed in the reference-model forward at step 1:

  • run 1: tokens: 1261, features: 1680 (deficit 419)
  • run 2: tokens: 1280, features: 1344 (deficit 64)

Arithmetic consistency: at 768px, each screenshot is exactly 336 image tokens for this model family (patch_size=16, spatial_merge_size=2). Both feature counts are exact multiples of 336 (5x336, 4x336); neither deficit is a multiple, and 419 cannot be any whole image's token count. The deficits are partial images, i.e. mid-span loss at the front of an observation, matching the offset-overshoot mechanism and ruling out whole-observation loss. The deficits back-solve to reasoning blocks of roughly 80 and 350+100 tokens.

Why this affects every current thinking model

We read the shipped chat templates:

  • Qwen3/Qwen3.5: historical assistant content is content.split('</think>')[-1]; history gets no <think> prefix.
  • DeepSeek-R1 (tokenizer_config.json): same split('</think>')[-1] idiom.
  • Kimi-K2-Thinking (chat_template.jinja): history turns rendered with an empty <think></think>; only the final turn keeps reasoning.
  • GLM-4.6 (chat_template.jinja): reasoning kept only when loop.index0 > ns.last_user_index.

None preserve prefix extension across renders. enable_thinking=false is not a sufficient workaround for Qwen: the generation prompt and history still differ by 4-5 tokens per turn.

The silent failure mode is worse than the crash

The image/feature pairing in the model forward is the only integrity check in the pipeline. With text-only observations the same offset drift silently deletes the front of observations from the training sequence (e.g. an observation Error: file not found can enter training as found). Loss, rewards, and metrics look normal. Nothing in validate_generator_output compares image-token counts to image_grid_thw.

Fixes (implemented and validated; PR incoming)

  1. Forward chat_template / chat_template_kwargs in _render_conversation so a thinking-preserving template can be selected (the VLM analogue of the existing qwen3_with_thinking text-path template). With such a template, renders are append-only again, the existing offset arithmetic is exact, and training is fully on-policy (rendered history is byte-identical to the generated stream). Note the server needs trust_request_chat_template for per-request templates.
  2. A one-time startup contract check: render a synthetic thinking conversation and its extension through the real endpoint and hard-fail unless renders are token-prefix extensions. Catches both history-editing templates and the endpoint ignoring the template field, at launch instead of silently.
  3. A trajectory-end integrity guard using the render response's own mm_placeholders: image tokens present in the assembled sequence vs declared placeholder lengths; fail with an actionable per-trajectory message instead of the opaque torch._check.

Validation: with (1)-(3) plus a Qwen3.5 template whose only change is rendering assistant history verbatim inside the <think> scaffold, a full multi-turn VLM GRPO run (Qwen3.5-9B, 64-turn computer-use episodes with screenshots, 6 GRPO steps, one epoch) completed end to end; property tests against the real tokenizer confirm token-prefix extension holds and the offset arithmetic lands exactly on observations. The template-agnostic default fix (measuring the observation boundary instead of predicting it, or independent observation tokenization as in the text path) is a natural follow-up; happy to discuss in the PR.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions