From 329faaceb08efb3ec81014a0af0348761f71bc7d Mon Sep 17 00:00:00 2001 From: Daniel Peng Date: Sat, 17 Oct 2026 05:12:39 +0800 Subject: [PATCH 1/2] fix(verl): classify trajectory prefix mismatch diagnostics --- agentlightning/verl/rollout_adapter.py | 44 +++++++++- docs/20-trainer-configuration.md | 8 ++ tests/verl/test_rollout_adapter.py | 109 +++++++++++++++++++++++++ 3 files changed, 157 insertions(+), 4 deletions(-) diff --git a/agentlightning/verl/rollout_adapter.py b/agentlightning/verl/rollout_adapter.py index a898176b0..68649ae33 100644 --- a/agentlightning/verl/rollout_adapter.py +++ b/agentlightning/verl/rollout_adapter.py @@ -7,6 +7,7 @@ import base64 # [multimodal-patch] import io import json +import re import zipfile from typing import Any, cast @@ -60,6 +61,39 @@ def _decode_token_ids(tokenizer: Any | None, ids: list[int]) -> str: return text +def _classify_trace_merge_mismatch( + tokenizer: Any | None, previous_ids: list[int], prompt_ids: list[int] +) -> tuple[bool, bool, bool]: + """Return diagnostic (template, retokenization, other) flags, never merge decisions. + + Only decoded-prefix equality, whitespace and the explicit think/tool-call + delimiters are recognized. These are heuristics, not proof of the cause. + """ + if tokenizer is None: + return False, False, True + try: + try: + previous_text = tokenizer.decode(previous_ids, skip_special_tokens=False) + prompt_text = tokenizer.decode(prompt_ids, skip_special_tokens=False) + except TypeError: + previous_text = tokenizer.decode(previous_ids) + prompt_text = tokenizer.decode(prompt_ids) + except Exception: + return False, False, True + if not isinstance(previous_text, str) or not isinstance(prompt_text, str): + return False, False, True + if previous_text and prompt_text.startswith(previous_text): + return False, True, False + + # Keep all content, including reasoning/tool-call bodies and unknown tags. + # An empty normalized context cannot establish a meaningful prefix match. + previous_normalized = " ".join(re.sub(r"", "", previous_text).split()) + prompt_normalized = " ".join(re.sub(r"", "", prompt_text).split()) + if previous_normalized and prompt_normalized.startswith(previous_normalized): + return True, False, False + return False, False, True + + def _decode_trace_text(tokenizer: Any | None, ids: list[int]) -> str: text = _decode_token_ids(tokenizer, ids) if len(text) > _TRACE_MERGE_MISMATCH_TEXT_LIMIT: @@ -593,16 +627,18 @@ def append_training_row( continue if len(merge_mismatch_rows) < _TRACE_MERGE_MISMATCH_WANDB_LIMIT: + template_mismatch, retoken_mismatch, others_mismatch = _classify_trace_merge_mismatch( + self.tokenizer, current_context, prompt_ids + ) merge_mismatch_rows.append( { "global_steps": global_steps, "rollout_id": rollout.rollout_id, "data_id": rollout.data_id, "turn_index": turn_index, - # Token-prefix failures are classified as other mismatches. - "template_mismatch": False, - "retoken_mismatch": False, - "others_mismatch": True, + "template_mismatch": template_mismatch, + "retoken_mismatch": retoken_mismatch, + "others_mismatch": others_mismatch, "prompt_length": len(prompt_ids), "response_length": len(response_ids), "previous_trace_length": len(current_context), diff --git a/docs/20-trainer-configuration.md b/docs/20-trainer-configuration.md index 63af662b4..0123753e3 100644 --- a/docs/20-trainer-configuration.md +++ b/docs/20-trainer-configuration.md @@ -139,6 +139,14 @@ agentlightning: The aggregator automatically merges consecutive calls when the next prompt starts with the exact token sequence of the previous prompt and response. Tokens added between calls, such as tool observations, are retained as context but masked from the policy loss. If exact token-prefix continuity is broken, the aggregator starts a new training row instead of merging incompatible calls. +For up to 100 such failures per batch, `training/trace_merge_mismatches` reports diagnostic flags: + +- `retoken_mismatch`: the decoded next prompt still starts with the decoded previous prompt and response; +- `template_mismatch`: that prefix is recovered only after collapsing whitespace and removing the literal ``, ``, `` and `` delimiters; +- `others_mismatch`: neither check matches, no tokenizer is available, or decoding fails. + +These flags are conservative heuristics, not proof of the underlying cause. Template normalization keeps reasoning/tool-call content and unknown tags; an empty normalized context is not a match. Classification never changes the strict token-prefix split or the resulting training tensors. + In this mode: - `trajectory_max_prompt_length` limits the initial prompt in each merged training row; diff --git a/tests/verl/test_rollout_adapter.py b/tests/verl/test_rollout_adapter.py index 5b4bca405..e0eae5597 100644 --- a/tests/verl/test_rollout_adapter.py +++ b/tests/verl/test_rollout_adapter.py @@ -336,6 +336,115 @@ def test_trajectory_prefix_mismatch_wandb_table_is_capped_at_100( assert table.rows[-1][3] == 100 +@pytest.mark.parametrize( + ("prompt", "response", "next_prompt", "use_tokenizer", "expected"), + [ + ([1], [2], [3, 4], True, (False, True, False)), + ([1, 6], [7, 2], [1, 2, 4], True, (True, False, False)), + ([1, 10], [2, 11], [1, 2, 4], True, (True, False, False)), + ([9], [5], [3, 5, 4], True, (True, False, False)), + ([1], [2], [8, 4], True, (False, False, True)), + ([1, 6], [7, 2], [8, 4], True, (False, False, True)), + ([1, 12], [13, 2], [1, 2, 4], True, (False, False, True)), + ([6], [7], [8, 4], True, (False, False, True)), + ([-1], [2], [3, 4], True, (False, False, True)), + ([1], [2], [3, 4], False, (False, False, True)), + ], + ids=[ + "retoken", + "think-delimiters", + "tool-call-delimiters", + "whitespace", + "content", + "template-and-content", + "unknown-delimiters", + "empty-normalized-context", + "decode-error", + "no-tokenizer", + ], +) +def test_trajectory_mismatch_classification_preserves_training_data( + monkeypatch: pytest.MonkeyPatch, + prompt: list[int], + response: list[int], + next_prompt: list[int], + use_tokenizer: bool, + expected: tuple[bool, bool, bool], +) -> None: + from tokenizers import Tokenizer, models + from transformers import PreTrainedTokenizerFast + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer( + models.WordLevel( + { + "[UNK]": 0, + "hello": 1, + "world": 2, + "hello world": 3, + "obs": 4, + "done": 5, + "": 6, + "": 7, + "changed": 8, + "hello\nworld": 9, + "": 10, + "": 11, + "": 12, + "": 13, + }, + unk_token="[UNK]", + ) + ), + unk_token="[UNK]", + ) + logged: list[tuple[dict[str, FakeTable], int]] = [] + _install_fake_wandb(monkeypatch, logged) + triplets = [_triplet(prompt, response), _triplet(next_prompt, [5])] + for triplet in triplets: + routes = np.arange( + (len(triplet.prompt["token_ids"]) + len(triplet.response["token_ids"])) * 4, + dtype=np.uint8, + ).reshape(-1, 2, 2) + triplet.response["routed_experts"] = _encode_routes(routes) + rollout = CompletedRollout( + rollout_id="classification", + data_id="data-1", + step=0, + sample_idx_in_step=0, + enqueue_time=0.0, + final_reward=0.75, + triplets=triplets, + ) + original = rollout.model_copy(deep=True) + adapter = RolloutAdapter( + max_prompt_length=12, + max_response_length=12, + device=torch.device("cpu"), + pad_token_id=0, + trace_aggregator_level="trajectory", + require_routed_experts=True, + ) + baseline, baseline_metrics = adapter.get_train_data_batch([rollout]) + logged.clear() + adapter.tokenizer = tokenizer if use_tokenizer else None + batch, metrics = adapter.get_train_data_batch([rollout]) + + table, _ = _logged_table(logged, "training/trace_merge_mismatches") + row = dict(zip(table.columns, table.rows[0], strict=True)) + assert (row["template_mismatch"], row["retoken_mismatch"], row["others_mismatch"]) == expected + assert len(table.rows) == 1 + assert len(batch) == 2 # Diagnostic classification never relaxes the exact token-prefix split. + assert rollout == original + assert metrics == baseline_metrics + assert batch.batch is not None and baseline.batch is not None + assert set(batch.batch.keys()) == set(baseline.batch.keys()) + for key, value in baseline.batch.items(): + torch.testing.assert_close(batch.batch[key], value) + for key, value in baseline.non_tensor_batch.items(): + np.testing.assert_array_equal(batch.non_tensor_batch[key], value) + + def test_training_step_uploads_24_compact_rollout_trajectories_to_wandb_zip( monkeypatch: pytest.MonkeyPatch, ) -> None: From e5eeb2d672123a4f2ffa90ec2b4526230805a82e Mon Sep 17 00:00:00 2001 From: Daniel Peng Date: Tue, 20 Oct 2026 11:23:43 +0800 Subject: [PATCH 2/2] fix(verl): bound retokenization mismatch diagnostics --- agentlightning/verl/rollout_adapter.py | 20 ++--- docs/20-trainer-configuration.md | 8 +- tests/verl/test_rollout_adapter.py | 115 +++++++++++++++++++++++-- 3 files changed, 121 insertions(+), 22 deletions(-) diff --git a/agentlightning/verl/rollout_adapter.py b/agentlightning/verl/rollout_adapter.py index 68649ae33..36b5b4340 100644 --- a/agentlightning/verl/rollout_adapter.py +++ b/agentlightning/verl/rollout_adapter.py @@ -7,7 +7,6 @@ import base64 # [multimodal-patch] import io import json -import re import zipfile from typing import Any, cast @@ -20,6 +19,7 @@ _TRACE_MERGE_MISMATCH_WANDB_LIMIT = 100 _TRACE_MERGE_MISMATCH_TEXT_LIMIT = 4000 +_TRACE_MERGE_MISMATCH_TOKEN_LIMIT = 128 _ROLLOUT_TRAJECTORY_WANDB_LIMIT = 24 _TRACE_MERGE_MISMATCH_COLUMNS = [ "global_steps", @@ -66,10 +66,16 @@ def _classify_trace_merge_mismatch( ) -> tuple[bool, bool, bool]: """Return diagnostic (template, retokenization, other) flags, never merge decisions. - Only decoded-prefix equality, whitespace and the explicit think/tool-call - delimiters are recognized. These are heuristics, not proof of the cause. + Compare complete short inputs only; long or unrecognized cases stay other. + Template causes are not inferred, and no extra full-history decoding is done. """ - if tokenizer is None: + if ( + tokenizer is None + or not previous_ids + or not prompt_ids + or len(previous_ids) > _TRACE_MERGE_MISMATCH_TOKEN_LIMIT + or len(prompt_ids) > _TRACE_MERGE_MISMATCH_TOKEN_LIMIT + ): return False, False, True try: try: @@ -85,12 +91,6 @@ def _classify_trace_merge_mismatch( if previous_text and prompt_text.startswith(previous_text): return False, True, False - # Keep all content, including reasoning/tool-call bodies and unknown tags. - # An empty normalized context cannot establish a meaningful prefix match. - previous_normalized = " ".join(re.sub(r"", "", previous_text).split()) - prompt_normalized = " ".join(re.sub(r"", "", prompt_text).split()) - if previous_normalized and prompt_normalized.startswith(previous_normalized): - return True, False, False return False, False, True diff --git a/docs/20-trainer-configuration.md b/docs/20-trainer-configuration.md index 0123753e3..cc7aa8757 100644 --- a/docs/20-trainer-configuration.md +++ b/docs/20-trainer-configuration.md @@ -139,13 +139,9 @@ agentlightning: The aggregator automatically merges consecutive calls when the next prompt starts with the exact token sequence of the previous prompt and response. Tokens added between calls, such as tool observations, are retained as context but masked from the policy loss. If exact token-prefix continuity is broken, the aggregator starts a new training row instead of merging incompatible calls. -For up to 100 such failures per batch, `training/trace_merge_mismatches` reports diagnostic flags: +For up to 100 such failures per batch, `training/trace_merge_mismatches` reports `retoken_mismatch` when the decoded next prompt still starts with the decoded previous prompt and response. This check only runs when both complete token sequences contain at most 128 tokens, limiting the extra decoding to short inputs. Longer inputs, absent tokenizers, decoding failures and other unrecognized cases remain `others_mismatch`. -- `retoken_mismatch`: the decoded next prompt still starts with the decoded previous prompt and response; -- `template_mismatch`: that prefix is recovered only after collapsing whitespace and removing the literal ``, ``, `` and `` delimiters; -- `others_mismatch`: neither check matches, no tokenizer is available, or decoding fails. - -These flags are conservative heuristics, not proof of the underlying cause. Template normalization keeps reasoning/tool-call content and unknown tags; an empty normalized context is not a match. Classification never changes the strict token-prefix split or the resulting training tensors. +`template_mismatch` remains false: no template normalization or root-cause inference is attempted. These flags are diagnostic hints, not proof of a cause. Classification never changes the strict token-prefix split or the resulting training tensors; long trajectories may therefore have retokenization mismatches that are still reported as `other`. In this mode: diff --git a/tests/verl/test_rollout_adapter.py b/tests/verl/test_rollout_adapter.py index e0eae5597..f76d57b36 100644 --- a/tests/verl/test_rollout_adapter.py +++ b/tests/verl/test_rollout_adapter.py @@ -9,7 +9,7 @@ import weakref import zipfile from types import SimpleNamespace -from typing import ClassVar +from typing import Any, ClassVar import pytest @@ -340,15 +340,19 @@ def test_trajectory_prefix_mismatch_wandb_table_is_capped_at_100( ("prompt", "response", "next_prompt", "use_tokenizer", "expected"), [ ([1], [2], [3, 4], True, (False, True, False)), - ([1, 6], [7, 2], [1, 2, 4], True, (True, False, False)), - ([1, 10], [2, 11], [1, 2, 4], True, (True, False, False)), - ([9], [5], [3, 5, 4], True, (True, False, False)), + ([1, 6], [7, 2], [1, 2, 4], True, (False, False, True)), + ([1, 10], [2, 11], [1, 2, 4], True, (False, False, True)), + ([9], [5], [3, 5, 4], True, (False, False, True)), ([1], [2], [8, 4], True, (False, False, True)), ([1, 6], [7, 2], [8, 4], True, (False, False, True)), ([1, 12], [13, 2], [1, 2, 4], True, (False, False, True)), ([6], [7], [8, 4], True, (False, False, True)), ([-1], [2], [3, 4], True, (False, False, True)), ([1], [2], [3, 4], False, (False, False, True)), + ([1], [2], [1], True, (False, False, True)), + ([1], [2] + [4] * 126, [3] + [4] * 126, True, (False, True, False)), + ([1], [2] + [4] * 127, [3] + [4] * 127, True, (False, False, True)), + ([1], [2], [3] + [4] * 128, True, (False, False, True)), ], ids=[ "retoken", @@ -358,9 +362,13 @@ def test_trajectory_prefix_mismatch_wandb_table_is_capped_at_100( "content", "template-and-content", "unknown-delimiters", - "empty-normalized-context", + "empty-text", "decode-error", "no-tokenizer", + "response-cannot-complete-prefix", + "previous-128", + "previous-129", + "prompt-129", ], ) def test_trajectory_mismatch_classification_preserves_training_data( @@ -400,7 +408,7 @@ def test_trajectory_mismatch_classification_preserves_training_data( ) logged: list[tuple[dict[str, FakeTable], int]] = [] _install_fake_wandb(monkeypatch, logged) - triplets = [_triplet(prompt, response), _triplet(next_prompt, [5])] + triplets = [_triplet(prompt, response), _triplet(next_prompt, [2])] for triplet in triplets: routes = np.arange( (len(triplet.prompt["token_ids"]) + len(triplet.response["token_ids"])) * 4, @@ -445,6 +453,101 @@ def test_trajectory_mismatch_classification_preserves_training_data( np.testing.assert_array_equal(batch.non_tensor_batch[key], value) +@pytest.mark.parametrize( + ("previous_size", "prompt_size", "expected_calls"), + [(128, 128, [128, 128]), (129, 128, []), (128, 129, []), (4096, 4096, [])], +) +def test_mismatch_classification_bounds_decoding( + previous_size: int, prompt_size: int, expected_calls: list[int] +) -> None: + from tokenizers import Tokenizer, models + from transformers import PreTrainedTokenizerFast + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer(models.WordLevel({"[UNK]": 0, "hello": 1}, unk_token="[UNK]")), + unk_token="[UNK]", + ) + calls: list[int] = [] + + class CountedTokenizer: + def decode(self, ids: list[int], **kwargs: Any) -> str: + calls.append(len(ids)) + text = tokenizer.decode(ids, **kwargs) + assert isinstance(text, str) + return text + + result = rollout_adapter_module._classify_trace_merge_mismatch( + CountedTokenizer(), [1] * previous_size, [1] * prompt_size + ) + assert calls == expected_calls + assert result == ((False, True, False) if expected_calls else (False, False, True)) + + +def test_mismatch_classification_does_not_use_response_or_numeric_fallback() -> None: + assert rollout_adapter_module._classify_trace_merge_mismatch(None, [1], [1]) == (False, False, True) + # The next response could complete a decoded prefix; only the prompt is eligible. + assert rollout_adapter_module._classify_trace_merge_mismatch(FakeTokenizer(), [1, 2], [1]) == ( + False, + False, + True, + ) + + +def test_mismatch_classification_keeps_legacy_decoder_and_rejects_invalid_text() -> None: + class LegacyTokenizer: + def decode(self, ids: list[int]) -> str: + return " ".join(str(token) for token in ids) + + assert rollout_adapter_module._classify_trace_merge_mismatch(LegacyTokenizer(), [1], [1, 2]) == ( + False, + True, + False, + ) + invalid = SimpleNamespace(decode=lambda ids, **kwargs: None) + assert rollout_adapter_module._classify_trace_merge_mismatch(invalid, [1], [1, 2]) == (False, False, True) + + +@pytest.mark.parametrize(("context_size", "expected_decode_calls"), [(2, 402), (129, 202)]) +def test_mismatch_decoding_stops_at_table_limit( + monkeypatch: pytest.MonkeyPatch, context_size: int, expected_decode_calls: int +) -> None: + from tokenizers import Tokenizer, models + from transformers import PreTrainedTokenizerFast + + tokenizer = PreTrainedTokenizerFast( + tokenizer_object=Tokenizer(models.WordLevel({"[UNK]": 0, "hello": 1, "world": 2}, unk_token="[UNK]")), + unk_token="[UNK]", + ) + calls: list[int] = [] + + class CountedTokenizer: + def decode(self, ids: list[int], **kwargs: Any) -> str: + calls.append(len(ids)) + text = tokenizer.decode(ids, **kwargs) + assert isinstance(text, str) + return text + + logged: list[tuple[dict[str, FakeTable], int]] = [] + _install_fake_wandb(monkeypatch, logged) + rollout = CompletedRollout( + rollout_id="bounded", + data_id="data-1", + step=0, + sample_idx_in_step=0, + enqueue_time=0.0, + final_reward=0.5, + triplets=[_triplet([1 + turn % 2] * (context_size - 1), [1]) for turn in range(102)], + ) + adapter = _adapter() + adapter.tokenizer = CountedTokenizer() + adapter.get_train_data_batch([rollout]) + table, _ = _logged_table(logged, "training/trace_merge_mismatches") + assert len(table.rows) == 100 + # Each row decodes two traces; the trajectory summary decodes its final prompt/response. + # Only short rows add two classifier decodes, and none are added beyond the 100-row cap. + assert len(calls) == expected_decode_calls + + def test_training_step_uploads_24_compact_rollout_trajectories_to_wandb_zip( monkeypatch: pytest.MonkeyPatch, ) -> None: