feat(openai): avoid re-rendering full chat history on multi-turn rollout (#1658) - #1663
feat(openai): avoid re-rendering full chat history on multi-turn rollout (#1658)#1663hsusul wants to merge 2 commits into
Conversation
|
|
||
| import torch | ||
| from huggingface_hub import snapshot_download | ||
| from transformers import AutoConfig |
There was a problem hiding this comment.
The modification of this file should have nothing to do with this PR, please remove it or submit it separately.
There was a problem hiding this comment.
You're right — that was an unrelated local fix that got carried along. Dropped from this branch; opened separately as #1689.
| return -1 | ||
|
|
||
|
|
||
| class IncrementalPromptRenderer: |
There was a problem hiding this comment.
Once reasoning tokens are saved into parent_base_token_ids, they cannot be removed, and for some models like Qwen3, when building context with apply_chat_template, it may remove the previous reasoning content, which can lead to misalignment between the incremental path and the normal path.
There was a problem hiding this comment.
Confirmed, thanks for catching this. Reproduced on Qwen3-0.6B: a cached base of [user, assistant "<think>r1</think>a1"] plus a [user] delta renders 31 tokens incrementally against 24 for the full render, since the template drops reasoning from every turn preceding the last user message. The original probe missed it because the synthetic 2-turn sequence contained no reasoning.
Fixed in the follow-up commit by adding a second capability probe for exactly that shape (_probe_reasoning_history) and refusing incremental rendering when the history carries reasoning that a later user turn supersedes (has_superseded_reasoning). The guard lives inside render_incremental rather than at the call site, so a poisoned prefix can neither be produced nor consumed even if a future caller forgets to check.
One case I deliberately kept on the incremental path: a tool-only delta (assistant <think> + tool_call -> tool result, with no new user turn). Qwen3 keeps the current turn's reasoning there, so it is token-identical to the full render — I verified this and it is covered by test_reasoning_tool_loop_stays_incremental. Blocking it would have removed the optimization for the main multi-turn tool-calling use case.
| delta_messages = tokenizer_messages[len(parent.messages) :] | ||
| parent_base = parent.prompt_base_token_ids | ||
| if parent_base is None: | ||
| parent_base = apply_chat_template( |
There was a problem hiding this comment.
If the tools included in the second round of requests are inconsistent with those in the first round, applying only to the incremental part here will lead to a discrepancy in content between the two.
There was a problem hiding this comment.
Confirmed. render_incremental accepted a tools argument but never passed it through to apply_chat_template, so the tool block was whatever the first turn baked into the prefix. Reproduced on Qwen3-0.6B: 1 tool on turn 1 and 2 tools on turn 2 gives 142 tokens incrementally against 178 for the full render, with the second tool's signature missing entirely.
Fixed by recording the tool set the prefix was rendered with (tools_signature(), stored as InteractionWithTokenLogpReward.prompt_tools_signature) and falling back to full-history rendering when it differs from the current turn's. Covered by a parametrized test over tools added, removed, cleared, and newly introduced between rounds.
|
Hi @hsusul, thanks for the effort on this PR Checking in: are you still planning to address the review feedback above (the reasoning-content misalignment for models like Qwen3, the tools-inconsistency across turns, and the unrelated testing_utils.py change)? There are also merge conflicts against main at this point. If you don't have bandwidth to continue, I'd be happy to take this over . Thanks! |
…out (areal-project#1658) - Introduce IncrementalPromptRenderer for O(1) prompt preparation per turn in multi-turn rollouts - Cache base prompt prefix tokens on InteractionWithTokenLogpReward.prompt_base_token_ids - Probe tokenizer template capability on first use to ensure 100% token-for-token mathematical parity with full-history rendering - Add incremental child token rendering in concat chat template mode - Fall back safely to canonical full-history rendering for dynamic or unsupported templates and multimodal processor inputs - Add comprehensive unit and regression test suite covering single/multi-tool calling, parallel tools, conversational turns, reasoning blocks, and fallback paths
Review found two cases where the cached prompt prefix diverges from a canonical full-history render: - Templates such as Qwen3 drop reasoning blocks from every turn preceding the last user message. An append-only cache cannot un-render tokens it already holds, so a prefix carrying superseded reasoning is stale. Probe the template for this behaviour and fall back to full rendering when the history contains reasoning a later user turn supersedes. Tool-only deltas, which keep the current turn's reasoning, stay on the incremental path. - Tool definitions live in the prompt head, which the delta render never revisits. A turn declaring a different tool set than the one baked into the prefix was silently served the stale tool block. Record the tool set the prefix was built with and fall back when it changes between turns.
f551f76 to
affdb29
Compare
|
Pushed a revision addressing all three comments (force-push, since dropping the unrelated commit required a rebase).
Both fixes fall back to the existing full-history path rather than attempting a partial-invalidation scheme, so parity is unconditional. Details in the thread replies; 10 new tests, each failing without the corresponding guard. |
3 participants
Description
Resolves #1658.
In multi-turn Agentic RL rollouts (such as SWE, search, reasoning, and tool-using agent episodes), trajectories frequently reach 50–200 turns. The OpenAI-compatible client (
ArealOpenAI) and Data Proxy previously re-rendered Jinja2 chat templates and tokenized the entire message history on every turn from scratch.For an episode of length$N$ , cumulative messages processed scaled as $\sum_{i=1}^N (2i-1) = N^2$ ($O(N^2)$). At $N=200$ , 40,000 messages were formatted, consuming over 1.1s of CPU time per episode and creating client-side rollout stalls that starve GPU inference backends (SGLang/vLLM).
This PR introduces an incremental prompt preparation mechanism (
IncrementalPromptRenderer):apply_chat_templateacross standard architectures (Qwen, ChatML, Llama-3, etc.).<think>modifications), missing parent tokens, or multimodal processor workflows.render_concat_child_tokensfor concat chat template mode, eliminating full-history re-tokenization during multi-turn concat rollouts.ArealOpenAIsignatures, CLI arguments, or trajectory export structures.Benchmarks & Scaling
Tested across episode lengths using
Qwen/Qwen3-0.6B:Verification Plan
tests/experimental/openai/test_prompt_renderer.py:apply_chat_template.<think>...</think>), custom system prompts._concat_prompt_token_ids_with_parent.tests/experimental/openai/(63 passed, 0 failed).pre-commithooks (ruff check,ruff format,check-yaml,spdx,check-json, etc.).