feat(openai): avoid re-rendering full chat history on multi-turn rollout (#1658) - #1663
Open
hsusul wants to merge 2 commits into
Open
feat(openai): avoid re-rendering full chat history on multi-turn rollout (#1658)#1663hsusul wants to merge 2 commits into
hsusul wants to merge 2 commits into
Conversation
- Use lazy evaluation in testing model path dictionaries to avoid eager downloads on test collection - Safely handle PackageNotFoundError in version check helper functions
…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
hsusul
requested review from
fishcrap,
garrett4wade and
sitabulaixizawaluduo
as code owners
September 2, 2026 01:12
| @@ -15,7 +15,6 @@ | |||
| from transformers import AutoConfig | |||
Collaborator
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.
| return -1 | ||
|
|
||
|
|
||
| class IncrementalPromptRenderer: |
Collaborator
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.
| delta_messages = tokenizer_messages[len(parent.messages) :] | ||
| parent_base = parent.prompt_base_token_ids | ||
| if parent_base is None: | ||
| parent_base = apply_chat_template( |
Collaborator
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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.).