Skip to content
Markdown

Token-In, Token-Out: rollout token correctness in agentic RL

Scope: the token-level correctness invariant of a multi-turn / agentic RL rollout loop. Token-In, Token-Out (TITO) is not a model, a library, or a service: it is a rule about how the rollout buffer is built. The rule is never re-encode tokens you have decoded. This page explains why decode then encode silently corrupts the policy gradient, what the fix costs, and the two properties a chat template must have for the fix to work. It is the token-level counterpart of the trajectory-level masking argument in agentic RL, and it argues a position that partly contradicts chat rendering and loss masking; both pages flag the disagreement.

The three Python blocks below were executed with system python3 (numpy 2.4.6, jinja2 3.1.2) and every assert passes; the pasted output is byte-identical to the run. The prefix-preservation block fetches live chat templates from the Hugging Face Hub, so its result is a snapshot of the templates as of 2026-07-13, not a durable fact. TRL line references are pinned to TRL v1.8.0 (released 2026-07-09).

What it is

TITO names a single invariant on the rollout loop: the token IDs the sampler emitted are the token IDs the trainer scores, element for element. Sampled IDs go into a running buffer that is the source of truth. The messages list survives only as bookkeeping for tool dispatch and logging. Anything the policy generated enters the buffer once, as IDs, and is never converted to text and back.

The rule exists because the natural multi-turn rollout breaks it. That loop keeps a messages list, re-renders it with apply_chat_template on every turn, generates, parses the completion into a message, appends it, and repeats; at the end it tokenizes the whole conversation and backpropagates on the result. Every turn of that loop performs a decode then encode round trip on tokens the policy already produced.

Failure 1: token drift

decode then encode is not injective. A tokenizer's encode returns the canonical BPE segmentation of a string, but a sampler is free to emit any sequence of vocabulary IDs, and nothing constrains it to the canonical one. BPE merges are unstable across token boundaries, so a token pair that the sampler emitted separately may merge when the joined string is re-encoded. Structured tool-call output widens the gap further: JSON whitespace, key order, and false versus False all give the re-render freedom the sampler did not have.

The consequence is not a crash. The re-tokenized sequence differs from the sampled one, so the gradient lands on a sequence the policy never produced, and the PPO/GRPO importance ratio exp(new_logp - old_logp) compares the logprob of one token against the recorded logprob of a different token. The ratio is meaningless. The loss still computes, the run still proceeds, and the numbers are silently wrong. This is a correctness bug, not an efficiency one.

Failure 2: lost turn boundaries

Tokenizing the full conversation at the end destroys the per-turn spans. The loss mask that zeroes tool-output tokens (agentic RL) then has to be reconstructed by walking role markers through the re-rendered string, which is a per-model-family parser and a second source of drift. When the tokenization drifts, the reconstructed mask drifts with it.

Why use it

  • The importance ratio is only defined on the sampled sequence. PPO and GRPO weight each token by the ratio of the current policy's logprob to the logprob recorded at sampling time. Score a different token at that position and the ratio is arbitrary. The executed block below shows that with an unchanged policy, where every ratio must be exactly 1.0, the re-encoded path returns ratios from 0.198 to 0.585.
  • The bug does not announce itself. When the re-encoded sequence happens to have the same length as the sampled one, nothing raises. Training continues on a corrupted gradient.
  • Masking gets easier, not harder. Building the mask as the buffer grows is a += [0] * tool_len + [1] * post_tool_len. Reconstructing it from a re-render is a parser per model family.
  • It removes the renderer from the inner loop. Exactly one template operation survives (the tool-response delta, below), and it runs on a dummy conversation, not on the trajectory.

When to use it (and when not)

  • Use TITO whenever you control the sampler and can read raw token IDs out of it (vLLM, SGLang, TRL's own generation path) and you are computing a policy gradient over multi-turn or tool-calling trajectories. This is the exact-token path for on-policy agentic RL.
  • Use it for single-turn RL too, where it is nearly free: there is only one generation, so holding its IDs costs nothing.
  • You cannot use it when the endpoint is a messages-only vendor API that returns text and never token IDs. There you have no sampled IDs to hold, so the exact on-policy invariant is unavailable. Use a token-returning sampler for PPO/GRPO; a per-model renderer can provide a deterministic approximation for text-only evaluation or offline data, but it cannot recover the vendor's sampled token path. See the counterpoint in chat rendering and loss masking.
  • It is irrelevant to SFT, where there is no sampled sequence: the tokens come from a dataset and are encoded exactly once.
  • The precondition can fail. TITO's one template operation requires prefix preservation for tool messages (below). A template that lacks it must be patched before the loop is correct.

Architecture

The buffer is the source of truth. Sampled IDs are appended verbatim with mask 1. Tool results enter as a delta of token IDs computed from a dummy render, appended by plain ID concatenation with mask 0. The parsed tool-call dict routes the dispatch and is then discarded: it never re-enters the prompt.

flowchart TB
  P["Prompt IDs"] --> BUF["Rollout buffer: token IDs<br/>(the source of truth)"]
  BUF --> GEN["Sampler (vLLM / SGLang)<br/>returns token IDs + logprobs"]
  GEN -->|"append IDs verbatim, mask = 1"| BUF
  GEN --> PARSE["Parse sampled tokens<br/>-> tool name + arguments"]
  PARSE -->|"routes dispatch only"| EXEC["Tool executor"]
  PARSE -.->|"NEVER re-enters the prompt"| BUF
  EXEC --> DELTA["Tool-response delta:<br/>render(dummy + tool) minus render(dummy)"]
  DELTA -->|"concatenate IDs, mask = 0"| BUF
  BUF --> LOSS["Policy update:<br/>ratio on the sampled IDs, loss on mask = 1"]

Two mechanisms make this work.

The tool-response delta. A tool result is text, so it has to be rendered into tokens somehow, and that is the one place a template is still needed. Render a dummy conversation twice, once ending at the assistant tool call and once with the tool message appended, tokenize both, and subtract: the suffix is the tool-response formatting plus content. Append it to the buffer by ID concatenation. The trajectory itself is never rendered.

Prefix preservation for tool messages. The subtraction is only valid if the shorter render is a genuine prefix of the longer one: render([user, asst_tool_call, tool_result]) must begin with render([user, asst_tool_call]), token for token. This is required only for tool messages, which is a much weaker demand than full prefix preservation across every role. Templates fail it when they make an assistant turn's rendering depend on what follows it.

How to use it

TITO is already shipped, production code in TRL; you do not have to build it. In TRL v1.8.0, GRPOTrainer runs a TITO tool loop, and every element of the blog's argument is present in the source:

Blog concept TRL v1.8.0 source
the property test trl/chat_template_utils.py, is_chat_template_prefix_preserving() (line 686), using the same dummy messages as the blog
auto-swap a patched template trl/trainer/grpo_trainer.py line 701: if self.tools and not is_chat_template_prefix_preserving(processing_class): then self.chat_template = get_training_chat_template(processing_class)
never re-tokenize grpo_trainer.py line 1941, comment: # Generate new completions after tool execution (using concatenated IDs, no re-tokenization)
the tool-response delta grpo_trainer.py, _get_tool_suffix_ids() (line 1737): renders a dummy conversation with and without the tool message, EOS-trims, and returns full_ids[len(prefix_ids):]
the mask built as you go grpo_trainer.py line 1794 initialises tool_mask = [[1] * len(ids) for ids in completion_ids], and line 1968 extends it inline: tool_mask[idx_with_tool] += [0] * tool_length + [1] * post_tool_length
patched templates trl/chat_templates/ ships 23 *_training.jinja files

Using it is therefore a matter of passing tools to GRPOTrainer and letting it swap the template. The snippet below is an unexecuted reference template (TRL and transformers are not installed in this KB's validation environment); pin trl==1.8.0 and verify against the installed release.

# Reference template, unexecuted. Pinned to trl==1.8.0.
# GRPOTrainer runs the TITO tool loop: sampled IDs are concatenated, never re-encoded,
# and tool_mask is built as the buffer grows.
from transformers import AutoTokenizer
from trl import GRPOConfig, GRPOTrainer
from trl.chat_template_utils import is_chat_template_prefix_preserving

def get_weather(city: str) -> str:
    """Return the current weather for a city."""
    return f"sunny in {city}"

model_id = "Qwen/Qwen3-8B"
revision = "b968826d9c46dd6066d109eabc6255188de91218"  # snapshot checked 2026-07-13
tok = AutoTokenizer.from_pretrained(model_id, revision=revision)
# This pinned Qwen3 template is not prefix-preserving. The trainer checks at init and swaps
# in trl/chat_templates/qwen3_training.jinja for the run.
assert not is_chat_template_prefix_preserving(tok)

trainer = GRPOTrainer(
    model=model_id,
    reward_funcs=my_reward_fn,
    args=GRPOConfig(max_tool_calling_iterations=8,
                    model_init_kwargs={"revision": revision}),
    train_dataset=ds,
    tools=[get_weather],          # triggers the prefix-preservation check + auto-swap
)
trainer.train()

How to develop with it

The failure rests on decode then encode not being injective, so that is the first property to prove. The block below builds a real BPE merge table, constructs a valid but non-canonical sampled sequence, and shows that the re-encoded sequence decodes to the same string while carrying different token IDs. It then computes the importance ratio on both paths with the policy held fixed, where the correct ratio is exactly 1.0 at every position.

# tito_roundtrip.py -- decode-then-encode is NOT injective, and the PPO/GRPO
# importance ratio computed on the re-encoded sequence is silently wrong.
# numpy only. A real BPE merge table, small enough to read.
import numpy as np

# Merge rules in rank order (best first). '_' stands in for a space, as BPE's marker does.
MERGES = [('"', 'a'), ('"a', '"'), ('"a"', ':'), (':', '_'), ('a', 'a'), ('aa', 'aa')]
RANK = {pair: i for i, pair in enumerate(MERGES)}
VOCAB = sorted({'"', 'a', ':', '_', '1'} | {"".join(p) for p in MERGES})
ID = {tok: i for i, tok in enumerate(VOCAB)}


def encode(text):
    """Canonical greedy BPE: repeatedly merge the lowest-rank adjacent pair."""
    parts = list(text)
    while True:
        pairs = {(a, b): RANK[(a, b)] for a, b in zip(parts, parts[1:]) if (a, b) in RANK}
        if not pairs:
            return [ID[p] for p in parts]
        best = min(pairs, key=pairs.get)
        merged, i = [], 0
        while i < len(parts):
            if i < len(parts) - 1 and (parts[i], parts[i + 1]) == best:
                merged.append(parts[i] + parts[i + 1])
                i += 2
            else:
                merged.append(parts[i])
                i += 1
        parts = merged


def decode(ids):
    return "".join(VOCAB[i] for i in ids)


# A bigram policy: logp(tok | prev). Deterministic, a genuine log-softmax.
rng = np.random.default_rng(0)
LOGITS = rng.normal(size=(len(VOCAB), len(VOCAB)))
LOGP = LOGITS - np.log(np.exp(LOGITS).sum(axis=1, keepdims=True))


def seq_logprobs(ids):
    """Per-token logprob of the sequence under the policy (token 0 conditioned on '"')."""
    prev = [ID['"']] + list(ids[:-1])
    return np.array([LOGP[p, t] for p, t in zip(prev, ids)])


# ---- 1. Non-injectivity: two DIFFERENT valid token sequences, one identical string.
sampled = [ID['"a"'], ID[':_'], ID['1']]       # a valid, NON-canonical sampled sequence
text = decode(sampled)
reencoded = encode(text)

assert decode(reencoded) == text, "both sequences decode to the same string"
assert reencoded != sampled, "but the token IDs differ: decode-then-encode is not injective"
print("sampled   :", sampled, "->", repr(text))
print("re-encoded:", reencoded, "->", repr(decode(reencoded)))
print("same text, different tokens:", decode(reencoded) == text and reencoded != sampled)

# ---- 2. The gradient lands on a sequence the policy never produced.
# The policy has NOT changed, so the true importance ratio must be exactly 1.0 everywhere.
old_logps = seq_logprobs(sampled)              # recorded by the sampler

tito_logps = seq_logprobs(sampled)             # TITO: score the buffer verbatim
tito_ratio = np.exp(tito_logps - old_logps)
assert np.allclose(tito_ratio, 1.0), tito_ratio

# The buggy path scores the RE-ENCODED sequence, then aligns positionally.
bug_logps = seq_logprobs(reencoded)
n = min(len(sampled), len(reencoded))
bug_ratio = np.exp(bug_logps[:n] - old_logps[:n])

misaligned = sum(s != r for s, r in zip(sampled[:n], reencoded[:n]))
assert not np.allclose(bug_ratio, 1.0), "the bug must show up as a ratio != 1"
print("\nTITO   ratio (policy unchanged):", np.round(tito_ratio, 6).tolist())
print("re-enc ratio (policy unchanged):", np.round(bug_ratio, 6).tolist())
print(f"misaligned positions: {misaligned}/{n}; ratio should be 1.0 at every one")
print("max ratio error:", round(float(np.abs(bug_ratio - 1.0).max()), 3))

# ---- 3. Length can change too: the crash the blog reports, but only sometimes.
sampled2 = [ID["aa"], ID["aa"]]                 # valid: sampler emitted two 'aa'
reencoded2 = encode(decode(sampled2))           # canonical: one 'aaaa'
assert decode(sampled2) == decode(reencoded2) == "aaaa"
assert len(reencoded2) == 1 and len(sampled2) == 2
print(f"\nshape mismatch: sampled len={len(sampled2)} vs re-encoded len={len(reencoded2)}")

# ---- 4. The adversarial point: EQUAL LENGTHS give no crash at all.
assert len(sampled) == len(reencoded) == 3      # same length, different IDs
print("equal-length case: no exception raised, ratio silently wrong ->",
      f"{misaligned} of {n} action IDs changed; every ratio uses a changed token or prefix")

# ---- 5. The TITO buffer preserves the sampled IDs verbatim.
buffer = []
for tok in sampled:
    buffer.append(tok)                          # IDs only; never a decoded string
assert buffer == sampled
print("\nTITO buffer == sampled IDs:", buffer == sampled)

Output:

sampled   : [2, 6, 4] -> '"a":_1'
re-encoded: [3, 7, 4] -> '"a":_1'
same text, different tokens: True

TITO   ratio (policy unchanged): [1.0, 1.0, 1.0]
re-enc ratio (policy unchanged): [0.585363, 0.198213, 0.563496]
misaligned positions: 2/3; ratio should be 1.0 at every one
max ratio error: 0.802

shape mismatch: sampled len=2 vs re-encoded len=1
equal-length case: no exception raised, ratio silently wrong -> 2 of 3 action IDs changed; every ratio uses a changed token or prefix

TITO buffer == sampled IDs: True

Read the two ratio lines together. The policy did not change, so every ratio must be 1.0. TITO returns 1.0 three times. The re-encoding path returns 0.585, 0.198, 0.563. Two action IDs changed, but all three ratios are wrong because the unchanged final ID is scored under a changed prefix. Since both sequences have length three, no shape check fires. The length-changing case in step 3 fails loudly; the equal-length case does not.

The mask must be built as the buffer grows

The second executed block shows why the mask cannot be recovered after the fact. It builds a trajectory whose tool-response text ends in a and whose next sampled token begins with a, with ('a','a') -> 'aa' in the merge table. Re-encoding the joined string merges across the tool/model boundary, producing a single token whose characters come from both regions. No mask value for that token is correct.

# tito_mask.py -- build the loss mask as the buffer grows, never from a re-render.
# The BPE merge that crosses the tool/model boundary makes the reconstructed mask
# not merely wrong but UNDEFINED: one token straddles both regions.
import numpy as np

MERGES = [("a", "a")]                       # the one merge that will cross the boundary
RANK = {p: i for i, p in enumerate(MERGES)}
VOCAB = sorted({"b", "T", "a", "x"} | {"".join(p) for p in MERGES})
ID = {t: i for i, t in enumerate(VOCAB)}


def encode(text):
    parts = list(text)
    while True:
        pairs = {(a, b): RANK[(a, b)] for a, b in zip(parts, parts[1:]) if (a, b) in RANK}
        if not pairs:
            return [ID[p] for p in parts]
        best = min(pairs, key=pairs.get)
        out, i = [], 0
        while i < len(parts):
            if i < len(parts) - 1 and (parts[i], parts[i + 1]) == best:
                out.append(parts[i] + parts[i + 1])
                i += 2
            else:
                out.append(parts[i])
                i += 1
        parts = out


def decode(ids):
    return "".join(VOCAB[i] for i in ids)


def masked_mean(logps, mask):
    """Mean over model-generated tokens only. Raises on an all-tool trajectory."""
    mask = np.asarray(mask, dtype=float)
    if mask.sum() == 0:
        raise ValueError("no model-generated tokens to train on")
    return float((np.asarray(logps) * mask).sum() / mask.sum())


# ---- TITO: the buffer is the source of truth; the mask is appended alongside it.
buffer, mask = [], []


def extend(ids, is_model):
    buffer.extend(ids)
    mask.extend([1 if is_model else 0] * len(ids))


extend([ID["b"]], is_model=True)                     # assistant samples 'b'
extend([ID["T"], ID["a"]], is_model=False)           # tool-response delta, by ID concat
extend([ID["a"], ID["x"]], is_model=True)            # assistant resumes, samples 'a','x'

assert buffer == [ID["b"], ID["T"], ID["a"], ID["a"], ID["x"]]
assert mask == [1, 0, 0, 1, 1]
assert len(mask) == len(buffer)
print("TITO buffer:", buffer, "->", repr(decode(buffer)))
print("TITO mask  :", mask, "(built as the buffer grew)")

# ---- The other path: re-render the conversation, re-tokenize, walk role markers.
text = decode(buffer)
reencoded = encode(text)
print("\nre-encoded :", reencoded, "->", repr(decode(reencoded)))
assert decode(reencoded) == text                     # same string
assert len(reencoded) != len(buffer)                 # but a different token count
print(f"length drift: buffer={len(buffer)} tokens, re-encoded={len(reencoded)} tokens")

# Per-character provenance: 1 where the model produced the character, 0 where the tool did.
char_src = []
for tok, m in zip(buffer, mask):
    char_src.extend([m] * len(VOCAB[tok]))
assert len(char_src) == len(text)

# Assign each re-encoded token the provenance of the characters it covers.
ambiguous, pos = [], 0
for tok in reencoded:
    span = char_src[pos : pos + len(VOCAB[tok])]
    if len(set(span)) > 1:                           # token straddles tool AND model text
        ambiguous.append((VOCAB[tok], span))
    pos += len(VOCAB[tok])

assert ambiguous, "expected at least one straddling token"
for tok, span in ambiguous:
    print(f"straddling token {tok!r}: char provenance {span} -> no correct mask value exists")

# ---- The consequence: the incremental mask is exact, the reconstructed one is not applicable.
logps = np.array([-0.2, -9.0, -8.0, -0.3, -0.4])     # tool tokens are wildly improbable
assert np.isclose(masked_mean(logps, mask), (-0.2 + -0.3 + -0.4) / 3)
print("\nmasked mean over model tokens only:", round(masked_mean(logps, mask), 4))

# A mask of the wrong length cannot even be applied to the re-encoded sequence.
try:
    np.asarray(logps[: len(reencoded)]) * np.asarray(mask, dtype=float)
    raise AssertionError("expected a shape mismatch")
except ValueError:
    print("applying the 5-token mask to the 4-token re-encode: ValueError (shape mismatch)")

# ---- Boundary: an all-tool-output trajectory must raise, not divide by zero.
try:
    masked_mean(logps, [0, 0, 0, 0, 0])
    raise AssertionError("expected ValueError")
except ValueError:
    print("all-tool trajectory: ValueError, not a divide-by-zero nan")

Output:

TITO buffer: [3, 0, 1, 1, 4] -> 'bTaax'
TITO mask  : [1, 0, 0, 1, 1] (built as the buffer grew)

re-encoded : [3, 0, 2, 4] -> 'bTaax'
length drift: buffer=5 tokens, re-encoded=4 tokens
straddling token 'aa': char provenance [0, 1] -> no correct mask value exists

masked mean over model tokens only: -0.3
applying the 5-token mask to the 4-token re-encode: ValueError (shape mismatch)
all-tool trajectory: ValueError, not a divide-by-zero nan

The straddling token is the point. It is not that the reconstructed mask is hard to compute; it is that it does not exist. One token carries a tool character and a model character, and any role-marker walk must assign it a single value. TITO's incremental mask never faces the question, because the boundary is a buffer append. Note also the boundary case at the end: an all-tool trajectory has no model tokens, so the masked mean must raise rather than divide by zero, exactly as in the masked-loss block on agentic RL.

Use the real tool name, not a dummy

The blog says of the delta's dummy prefix: "The prefix doesn't even have to be a real conversation. Any dummy that ends in an assistant tool call works." That claim is false for GPT-OSS, and TRL's own code says so. _get_tool_suffix_ids() opens with:

# Use the real tool name instead of a dummy: some templates (e.g. GPT-OSS) derive the tool response
# header from the assistant's tool call name.
dummy_tool_calls = [{"type": "function", "function": {"name": tool_messages[0]["name"], "arguments": {}}}]

The dummy conversation may have dummy content, but the tool name must be the real one, because some templates interpolate it into the tool-response header. Substitute a placeholder name and the delta you subtract carries the wrong header. Use the real name.

How to maintain it

Ship the property test, never the table. The blog's only table is a snapshot: 18 of 19 model families prefix-preserving, only Qwen3 failing, as of 2026-05-29. Chat templates drift with model revisions, so that table is a point-in-time observation, not a durable fact. The durable artefact is the test. Re-run it on every model bump and every transformers bump, in CI, and treat a failure as a blocker on the run.

The block below is TRL's is_chat_template_prefix_preserving() reimplemented on raw Jinja, so it runs without transformers. It pulls each template live from the Hub and asserts the expected result. Executed 2026-07-13.

# prefix_preservation.py -- TRL's is_chat_template_prefix_preserving(), reimplemented
# on raw Jinja against chat templates pulled live from the Hugging Face Hub.
# Requires jinja2 only (transformers is not needed to render a chat template).
import json, urllib.request
from jinja2 import Environment
from jinja2.exceptions import TemplateError

TOOL_CALLS = [{"type": "function", "function": {"name": "dummy", "arguments": {}}}]
NO_TOOL = [{"role": "user", "content": "dummy"},
           {"role": "assistant", "content": "", "tool_calls": TOOL_CALLS}]
WITH_TOOL = NO_TOOL + [{"role": "tool", "name": "dummy", "content": "dummy"}]

def fetch_template(repo):
    url = f"https://huggingface.co/{repo}/resolve/main/tokenizer_config.json"
    with urllib.request.urlopen(url) as r:
        return json.load(r)["chat_template"]

def render(tpl, messages, add_generation_prompt):
    """Mirror transformers' apply_chat_template Jinja environment."""
    env = Environment(trim_blocks=True, lstrip_blocks=True)
    env.globals["raise_exception"] = lambda m: (_ for _ in ()).throw(TemplateError(m))
    env.filters["tojson"] = lambda o, **kw: json.dumps(o, **kw)
    return env.from_string(tpl).render(
        messages=messages, tools=None, add_generation_prompt=add_generation_prompt)

def is_prefix_preserving(tpl):
    """TRL's property: render([user, asst_tool_call]) must be a verbatim prefix of
    render([user, asst_tool_call, tool_result], add_generation_prompt=True)."""
    a = render(tpl, NO_TOOL, add_generation_prompt=False)
    b = render(tpl, WITH_TOOL, add_generation_prompt=True)
    return b.startswith(a), a, b

def divergence(a, b):
    """First index where the two renders differ, i.e. where the prefix breaks."""
    n = 0
    while n < min(len(a), len(b)) and a[n] == b[n]:
        n += 1
    return n

MODELS = ["Qwen/Qwen3-8B", "Qwen/Qwen3.5-4B", "Qwen/Qwen3.6-35B-A3B",
          "Qwen/Qwen3-4B-Instruct-2507", "Qwen/Qwen2.5-7B-Instruct"]
EXPECTED = {"Qwen/Qwen3-8B": False}   # every other family below is expected to pass

for repo in MODELS:
    tpl = fetch_template(repo)
    ok, a, b = is_prefix_preserving(tpl)
    assert ok == EXPECTED.get(repo, True), f"{repo}: expected {EXPECTED.get(repo, True)}, got {ok}"
    print(f"{repo:<28} prefix-preserving={ok}")
    if not ok:
        n = divergence(a, b)
        print(f"    renders diverge at char {n}, inside the assistant header:")
        print(f"      no-tool  : ...{a[n - 22 : n + 30]!r}")
        print(f"      with-tool: ...{b[n - 22 : n + 30]!r}")

print("\nasserts passed: 1 failure (Qwen3-8B), 4 passes")

Output:

Qwen/Qwen3-8B                prefix-preserving=False
    renders diverge at char 57, inside the assistant header:
      no-tool  : ...'im_start|>assistant\n<think>\n\n</think>\n\n<tool_call>\n{'
      with-tool: ...'im_start|>assistant\n<tool_call>\n{"name": "dummy", "a'
Qwen/Qwen3.5-4B              prefix-preserving=True
Qwen/Qwen3.6-35B-A3B         prefix-preserving=True
Qwen/Qwen3-4B-Instruct-2507  prefix-preserving=True
Qwen/Qwen2.5-7B-Instruct     prefix-preserving=True

asserts passed: 1 failure (Qwen3-8B), 4 passes

The divergence shows the prefix violation directly. When the assistant tool-call message is last, Qwen3-8B injects an empty reasoning block (<think>\n\n</think>\n\n) before the tool call; when a tool message follows it, that block disappears. The assistant turn renders differently depending on what comes after it, which is exactly the prefix violation. The culprit is verbatim in the live Qwen/Qwen3-8B tokenizer_config.json:

{%- if loop.last or (not loop.last and reasoning_content) %}

One caveat on the method: TRL's check compares token IDs, this reimplementation compares rendered strings. String-prefix preservation is a necessary condition for token-prefix preservation, so the Qwen3-8B failure is real either way. For the four passes, a string prefix that ends on a special-token boundary (as these do, on <|im_end|>) forces a token boundary too, but if you have transformers installed, call TRL's function directly rather than this stand-in.

"The fix is one line" oversells it

The blog's fix for Qwen3 is {%- if true %} in place of the conditional above. TRL's actual qwen3_training.jinja does more, and documents it in the template's own header comment:

  • it deletes the conditional block entirely ("Removed the loop.index0 > ns.last_query_index conditional; always include thinking block"),
  • it requires both <think> and </think> before parsing a reasoning block ({%- if '</think>' in content %} becomes {%- if '<think>' in content and '</think>' in content %}), and
  • it adds {% generation %} / {% endgeneration %} markers for assistant-only loss masking.

Budget for a template patch, not a one-character edit.

How to run it in production

  • Read token IDs out of the sampler, not text. vLLM and SGLang expose token IDs and per-token logprobs through native or engine-specific paths. The standard OpenAI-compatible /v1/chat/completions response is text-oriented and does not guarantee token IDs. If the chosen endpoint exposes only text, exact TITO is unavailable; switch APIs or treat the rollout as approximate rather than claiming exact on-policy ratios. This is the single wiring decision the invariant forces (async RL systems).
  • Gate the run on the property test at init. TRL does this: GRPOTrainer checks prefix preservation the moment tools is passed and swaps in a patched template before step 0. Fail loudly at startup rather than silently at step 400.
  • Pin the template with the checkpoint. The template a run trained under belongs in run metadata, alongside the renderer name, for the same reason as in chat rendering and loss masking: a checkpoint must be re-served with what it was trained with.
  • Instrument the invariant. Assert len(buffer) == len(mask) == len(logprobs) on every trajectory before the update, and assert that a re-encode of the decoded buffer either matches the buffer or is logged as drift. The assertion is cheap and it is the only thing standing between you and a silently wrong run.
  • Watch step-0 KL. Non-zero sampler-versus-trainer KL before the first policy update is a diagnostic for token, template, checkpoint, dropout, or numerical skew; it is not unique proof of re-tokenization (GRPO).

Failure modes

  • Re-tokenizing the conversation between turns. The headline bug. decode then encode is not injective; the gradient lands on a sequence the policy never produced; the importance ratio is meaningless; nothing crashes. Hold the sampled IDs.
  • Feeding the parsed tool-call dict back into the prompt. Parsing the sampled tokens to route the tool dispatch is correct. Re-serialising that dict into messages and re-rendering is the same round trip by another name: JSON key order and whitespace will not survive it.
  • Reconstructing the loss mask from a re-render. A BPE merge across the tool/model boundary produces a token with no correct mask value (executed above). Build the mask as the buffer grows.
  • A dummy tool name in the response delta. GPT-OSS derives the tool-response header from the assistant's tool-call name, so a placeholder name yields the wrong delta. Use the real name (TRL does).
  • Trusting the "18 of 19" table. It is a 2026-05-29 snapshot. Templates drift. Run the property test.
  • Assuming a crash will tell you. The shape mismatch only fires when the re-encode changes the token count. Same-length drift is silent, and it is the common case.
  • All-tool trajectories. A trajectory whose completion is entirely tool output has zero unmasked tokens; the masked mean must raise rather than divide by zero.

Honest limits of the source

  • There are zero benchmark numbers in the blog post. No evals, no baselines, no model sizes, no GPU-hours, no released artefacts, no licence. It is a mechanism claim, not a measured one. The correctness argument is sound and reproduces (above), but nobody has published the delta in final reward from fixing it.
  • The motivating failure is asserted, not demonstrated. The post's hook is that "Loss occasionally spikes for no obvious reason. And eventually it fails with a shape mismatch error." There is no repro, no loss curve, and no logged run behind that sentence. The mechanism is real and this page proves it in isolation; the field evidence is not in the source. The nearest measured anchor for the adjacent claim (masking matters) is Search-R1's ablation, cited on agentic RL.
  • "No renderer needed" is softer than it reads. TRL, written by the same authors, still maintains 23 per-model *_training.jinja templates. Most are patched for {% generation %} SFT markers, which is a different requirement, so this is not a contradiction. But the implied "per-model maintenance burden vanishes" does not hold: the burden moves from the rollout loop to a template directory.

References

  • Agentic RL: Token-In, Token-Out Done Right (Quentin Gallouedec, Kashif Rasul, 2026-05-29): https://huggingface.co/blog/huggingface/tito
  • TRL v1.8.0, is_chat_template_prefix_preserving(): https://github.com/huggingface/trl/blob/v1.8.0/trl/chat_template_utils.py
  • TRL v1.8.0, GRPOTrainer._get_tool_suffix_ids() and the tool loop: https://github.com/huggingface/trl/blob/v1.8.0/trl/trainer/grpo_trainer.py
  • TRL v1.8.0, the 23 patched training templates: https://github.com/huggingface/trl/tree/v1.8.0/trl/chat_templates
  • Search-R1 (measured ablation: retrieved-token loss masking, Avg 0.431 with mask vs 0.343 without): https://arxiv.org/abs/2503.09516
  • verl AgentLoop ("Response mask, 1 for LLM generated token, 0 for tool response token"): https://verl.readthedocs.io/en/latest/advance/agent_loop.html
  • Qwen/Qwen3-8B chat template (the non-prefix-preserving conditional): https://huggingface.co/Qwen/Qwen3-8B/blob/main/tokenizer_config.json
  • HuggingFace chat templating: https://huggingface.co/docs/transformers/main/en/chat_templating

Related: Agentic RL · Chat rendering and loss masking · GRPO · PPO · TRL · verl · Async RL systems · RL libraries · Tools and function calling · Reward design · Glossary