From 525d45f5dfc66e8eace67f139d9b138be2862fd5 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 10 Sep 2026 06:39:10 +0000 Subject: [PATCH 1/2] Preserve captured causal context across split histories --- src/art/trajectories/_history.py | 49 +++----- tests/unit/trajectories/test_history.py | 21 ++-- tests/unit/trajectories/test_tokenize.py | 143 +++++++++++++++++++---- 3 files changed, 151 insertions(+), 62 deletions(-) diff --git a/src/art/trajectories/_history.py b/src/art/trajectories/_history.py index bab127e0c..f7f950f2d 100644 --- a/src/art/trajectories/_history.py +++ b/src/art/trajectories/_history.py @@ -343,34 +343,21 @@ def _contains_tokens(tokens: Sequence[int], sampled: Sequence[int]) -> bool: return False -def _retains_output_suffix( +def _retains_captured_output( prompt: Sequence[int] | None, output: Sequence[int] | None, later_prompt: Sequence[int] | None, ) -> bool: - if ( - prompt is None - or output is None - or later_prompt is None - or not _is_prefix(prompt, later_prompt) - ): - return False - continuation = later_prompt[len(prompt) :] - if not output or not continuation: - return False - - prefix_lengths = _token_prefix_lengths(continuation) - matched = 0 - for index, token in enumerate(output): - while matched and token != continuation[matched]: - matched = prefix_lengths[matched - 1] - if token == continuation[matched]: - matched += 1 - if matched == len(continuation): - if index + 1 == len(output): - return True - matched = prefix_lengths[matched - 1] - return matched > 0 + # A suffix sampled under a longer prefix has different conditional + # probabilities. Keep it as request conditioning in the later branch; + # the existing split branch preserves the complete original generation. + return ( + prompt is not None + and output is not None + and bool(output) + and later_prompt is not None + and _is_prefix([*prompt, *output], later_prompt) + ) def _chat_generation_tokens( @@ -464,7 +451,7 @@ def _chat_retains_sampled_reasoning( return True -def _chat_retains_sampled_suffix( +def _chat_retains_captured_output( branch: _Branch[Message, ChatCompletionsMessageSource, _ChatContext], prompt_ids: Sequence[int] | None, prompt_length: int, @@ -487,7 +474,7 @@ def _chat_retains_sampled_suffix( prior_prompt, prior_output = _chat_generation_tokens( prior_source.exchange, prior_source.choice_index, cache ) - return _retains_output_suffix(prior_prompt, prior_output, prompt_ids) + return _retains_captured_output(prior_prompt, prior_output, prompt_ids) def _chat_structured_generation_hit_limit( @@ -578,7 +565,7 @@ def _anthropic_generation_tokens( return cache[key] -def _anthropic_retains_sampled_suffix( +def _anthropic_retains_captured_output( branch: _Branch[AnthropicMessageParam, AnthropicMessageSource, _AnthropicContext], prompt_ids: Sequence[int] | None, prompt_length: int, @@ -597,7 +584,7 @@ def _anthropic_retains_sampled_suffix( prior_prompt, prior_output = _anthropic_generation_tokens( prior_source.exchange, cache ) - return _retains_output_suffix(prior_prompt, prior_output, prompt_ids) + return _retains_captured_output(prior_prompt, prior_output, prompt_ids) def _chat_message_key(message: Message, *, visible_only: bool = False) -> str: @@ -719,7 +706,7 @@ def chat_completions_histories( source_continuation = lambda branch: ( reconcile or exact_continuation(branch) - or _chat_retains_sampled_suffix( + or _chat_retains_captured_output( branch, prompt_ids, len(prompt), token_cache ) ) @@ -826,7 +813,7 @@ def anthropic_messages_histories( continuation = lambda branch: reconcile or exact_continuation(branch) source_continuation = lambda branch: ( continuation(branch) - or _anthropic_retains_sampled_suffix( + or _anthropic_retains_captured_output( branch, prompt_ids, len(prompt), token_cache ) ) @@ -1321,7 +1308,7 @@ def _responses_split_prompt_source( if not 0 <= source.generation_index < len(generations): raise ValueError("Responses generation source index is out of bounds") generation = generations[source.generation_index] - if _retains_output_suffix( + if _retains_captured_output( generation.prompt_token_ids, generation.output_token_ids, current_prompt_ids, diff --git a/tests/unit/trajectories/test_history.py b/tests/unit/trajectories/test_history.py index 089cdbd1f..b71fa7965 100644 --- a/tests/unit/trajectories/test_history.py +++ b/tests/unit/trajectories/test_history.py @@ -334,14 +334,15 @@ def test_contains_tokens(tokens: list[int], sampled: list[int], expected: bool) [ ([0], [], [0, 1], False), ([0], [1], [0], False), - ([0], [7, 1, 2], [0, 1, 2], True), - ([0], [1, 1, 1, 2], [0, 1, 1, 2], True), + ([0], [7, 1, 2], [0, 1, 2], False), + ([0], [1, 1, 1, 2], [0, 1, 1, 2], False), ([0], [1, 2, 3], [0, 1, 2], False), - ([0], [9, 1], [0, 1, 2], True), + ([0], [9, 1], [0, 1, 2], False), ([0], [1], [9, 1], False), + ([0], [1, 2], [0, 1, 2, 3], True), ], ) -def test_retains_output_suffix( +def test_retains_captured_output( prompt: list[int], output: list[int], later_prompt: list[int], @@ -350,7 +351,8 @@ def test_retains_output_suffix( history_module = importlib.import_module("art.trajectories._history") assert ( - history_module._retains_output_suffix(prompt, output, later_prompt) is expected + history_module._retains_captured_output(prompt, output, later_prompt) + is expected ) @@ -383,7 +385,7 @@ def __iter__(self): output = CountingTokens([1] * 10_000 + [2]) later_prompt = CountingTokens([0, *([1] * 1_000), 2]) - assert history_module._retains_output_suffix( + assert not history_module._retains_captured_output( [0], cast(Any, output), cast(Any, later_prompt) ) assert output.accesses + later_prompt.accesses < 10 * ( @@ -1051,7 +1053,7 @@ def test_cross_exchange_responses_reasoning_stripping_splits_histories() -> None first_answer_source = histories[1].input_sources[1] assert first_answer_source is not None assert first_answer_source.exchange is first - assert first_answer_source.generation_index == 0 + assert first_answer_source.generation_index is None @pytest.mark.parametrize( @@ -1630,8 +1632,9 @@ def test_chat_template_stripped_reasoning_splits_exact_histories() -> None: assert len(histories) == 2 assert [len(history.messages) for history in histories] == [2, 4] assert histories[1].message_sources[1] is not None - assert histories[1].message_sources[1].exchange is first - assert histories[1].message_sources[1].choice_index == 0 + assert histories[1].message_sources[1].exchange is second + assert histories[1].message_sources[1].request_index == 1 + assert histories[1].message_sources[1].choice_index is None with pytest.raises(ValueError, match="exactly one history"): trajectory.tokenize() diff --git a/tests/unit/trajectories/test_tokenize.py b/tests/unit/trajectories/test_tokenize.py index 8b0711546..5cc53885d 100644 --- a/tests/unit/trajectories/test_tokenize.py +++ b/tests/unit/trajectories/test_tokenize.py @@ -3434,13 +3434,16 @@ def apply_chat_template( tokenized = history.tokenize(tokenizer=Tokenizer()) assert tokenized.tokens == [10, 101, 102, 11, 91, 201] - assert tokenized.logprobs[1:3] == pytest.approx([-10.1, -10.2]) + if top_level_only: + assert all(math.isnan(value) for value in tokenized.logprobs[1:3]) + else: + assert tokenized.logprobs[1:3] == pytest.approx([-10.1, -10.2]) assert tokenized.logprobs[-2] == pytest.approx(-10.0) assert tokenized.logprobs[-1] == pytest.approx(-20.1) - assert tokenized.flags[1:3] == [ - _SAMPLED_ASSISTANT_OUTPUT, - _SAMPLED_ASSISTANT_OUTPUT, - ] + assert ( + tokenized.flags[1:3] + == [tr.TokenFlag.EXACT if top_level_only else _SAMPLED_ASSISTANT_OUTPUT] * 2 + ) assert tokenized.flags[-2:] == [ _SAMPLED_ASSISTANT_OUTPUT, _SAMPLED_ASSISTANT_OUTPUT, @@ -4330,12 +4333,12 @@ def test_cross_exchange_responses_reasoning_split_uses_later_prompt_backbone() - [1, 3, 4, 5], ] assert math.isnan(tokenized.histories[1].logprobs[0]) - assert tokenized.histories[1].logprobs[1] == -0.3 + assert math.isnan(tokenized.histories[1].logprobs[1]) assert math.isnan(tokenized.histories[1].logprobs[2]) assert tokenized.histories[1].logprobs[3] == -0.1 assert tokenized.histories[1].flags == [ tr.TokenFlag.EXACT, - _SAMPLED_ASSISTANT_OUTPUT, + tr.TokenFlag.EXACT, tr.TokenFlag.EXACT, _SAMPLED_ASSISTANT_OUTPUT, ] @@ -6366,13 +6369,8 @@ def apply_chat_template( [1, 2, 101, 102, 9], [1, 101, 102, 9, 4, 5, 6, 9], ] - assert tokenized.histories[1].flags[1] & tr.TokenFlag.SAMPLED - assert tokenized.histories[1].flags[1] & tr.TokenFlag.EXACT - assert tokenized.histories[1].logprobs[1:3] == [-10.1, -10.2] - assert tokenized.histories[1].flags[3] == ( - _SAMPLED_ASSISTANT_OUTPUT | tr.TokenFlag.STOP - ) - assert tokenized.histories[1].logprobs[3] == -0.9 + assert tokenized.histories[1].flags[1:4] == [tr.TokenFlag.EXACT] * 3 + assert all(math.isnan(value) for value in tokenized.histories[1].logprobs[1:4]) assert 2 not in tokenized.histories[1].tokens assert 500 not in tokenized.histories[1].tokens @@ -6534,11 +6532,8 @@ def apply_chat_template( ) second_history = tokenized.histories[1] assert second_history.tokens == [1, 7, 8, 4, 5] - assert second_history.logprobs[1:3] == [-0.7, -0.8] - assert second_history.flags[1:3] == [ - _SAMPLED_ASSISTANT_OUTPUT, - _SAMPLED_ASSISTANT_OUTPUT, - ] + assert all(math.isnan(value) for value in second_history.logprobs[1:3]) + assert second_history.flags[1:3] == [tr.TokenFlag.EXACT] * 2 preprocessing = list( tokenize_trajectory_groups( @@ -7496,9 +7491,17 @@ def test_reasoning_split_trajectory_reuses_prevalidated_projections( exchanges: list[ChatCompletionsExchange] = [] request_messages: list[ChatCompletionMessageParam] = [] prompt: list[int] = [] + expected: dict[tuple[int, ...], float] = {} for index in range(40): request_messages.append({"role": "user", "content": f"u{index}"}) prompt.append(3000 + index) + output = [1000 + index, 2000 + index] + expected.update( + { + tuple([*prompt, *output[: position + 1]]): -token / 10 + for position, token in enumerate(output) + } + ) exchange = _chat_exchange( list(prompt), [1000 + index, 2000 + index], offset=index ) @@ -7526,10 +7529,10 @@ def test_reasoning_split_trajectory_reuses_prevalidated_projections( any( source is not None and source.choice_index == 0 - and source.exchange is exchanges[0] + and source.exchange is exchanges[index] for source in history.message_sources ) - for history in projected + for index, history in enumerate(projected) ) from art.trajectories import _tokenize @@ -7549,8 +7552,19 @@ def counted(history: tr.History) -> bool: first_key = next(key for key in traces[0].source_keys if key is not None) assert len(tokenized.histories) == 40 - assert all(first_key in trace.sources for trace in traces) + assert all(first_key not in trace.sources for trace in traces[1:]) assert calls == 0 + selected = { + tuple(history.tokens[: index + 1]): history.logprobs[index] + for history, mask in zip( + tokenized.histories, + tokenized.tensorize().first_occurrence_masks(where=tr.TokenFlag.SAMPLED), + strict=True, + ) + for index, chosen in enumerate(mask.tolist()) + if chosen + } + assert selected == expected def test_explicit_template_override_rerenders_exact_exchange_scaffold() -> None: @@ -8128,3 +8142,88 @@ def apply_chat_template( == _SAMPLED_ASSISTANT_OUTPUT for i in positions ) + + +@pytest.mark.parametrize("protocol", ["chat", "messages", "responses"]) +@pytest.mark.parametrize("later_prompt", [[1, 3, 4], [1, 2, 3, 4], [9, 3, 4]]) +def test_divergent_captured_context_preserves_all_original_sampled_prefixes( + protocol: str, later_prompt: list[int] +) -> None: + captures = [([1], [2, 3]), (later_prompt, [5])] + if protocol == "chat": + exchanges = TrajectoryExchanges( + chat_completions=[ + _chat_exchange(prompt, output, offset=index) + for index, (prompt, output) in enumerate(captures) + ] + ) + elif protocol == "messages": + exchanges = TrajectoryExchanges( + messages=[ + _message_exchange( + MessagesRequest( + model="test/model", + max_tokens=16, + messages=( + [{"role": "user", "content": "turn 0"}] + if index == 0 + else [ + {"role": "user", "content": "turn 0"}, + { + "role": "assistant", + "content": [{"type": "text", "text": "answer"}], + }, + {"role": "user", "content": "turn 1"}, + ] + ), + ), + identifier=f"message-{index}", + offset=index, + prompt_token_ids=prompt, + token_ids=output, + logprobs=[-token / 10 for token in output], + ) + for index, (prompt, output) in enumerate(captures) + ] + ) + else: + responses = [] + for index, (prompt, output) in enumerate(captures): + exchange = _response_exchange( + f"response-{index}", + output[-1], + offset=index, + previous_response_id="response-0" if index else None, + prompt_token_ids=prompt, + ) + data = exchange.response.model_dump(mode="python") + data["token_generations"][0]["output_tokens"] = [ + {"token_id": token, "logprob": -token / 10} for token in output + ] + exchange.response = Response.model_validate(data) + responses.append(exchange) + exchanges = TrajectoryExchanges(responses=responses) + trajectory = art.Trajectory(exchanges=exchanges) + tokenized = trajectory.tokenize(multi_history=True) + expected = { + tuple([*prompt, *output[: index + 1]]): -token / 10 + for prompt, output in captures + for index, token in enumerate(output) + } + actual = {} + masks = tokenized.tensorize().first_occurrence_masks(where=tr.TokenFlag.SAMPLED) + for history, mask in zip(tokenized.histories, masks, strict=True): + for index, selected in enumerate(mask.tolist()): + if selected: + prefix = tuple(history.tokens[: index + 1]) + assert prefix not in actual + actual[prefix] = history.logprobs[index] + assert actual == expected # Complete coverage as well as no changed-context sample. + if later_prompt == [1, 2, 3, 4]: + assert len(tokenized.histories) == 1 + assert tokenized.histories[0].tokens == [1, 2, 3, 4, 5] + else: + assert [history.tokens for history in tokenized.histories] == [ + [1, 2, 3], + [*later_prompt, 5], + ] From da4f165a0a98c5d9539c74d854cce1d15e6810f7 Mon Sep 17 00:00:00 2001 From: Brad Hilton Date: Thu, 10 Sep 2026 07:45:16 +0000 Subject: [PATCH 2/2] Align preprocessing controls with captured causal lineage --- .../test_exchange_training_model_selection.py | 114 ++++++++++++++---- 1 file changed, 92 insertions(+), 22 deletions(-) diff --git a/tests/unit/test_exchange_training_model_selection.py b/tests/unit/test_exchange_training_model_selection.py index 9d9935663..e7b6264ce 100644 --- a/tests/unit/test_exchange_training_model_selection.py +++ b/tests/unit/test_exchange_training_model_selection.py @@ -18,7 +18,7 @@ from art.dev.model import InternalModelConfig from art.local import LocalBackend from art.openai import ART_MOE_ROUTING_METADATA_KEY -from art.preprocessing.moe_routing import MoeRouteArray +from art.preprocessing.moe_routing import MoeRouteArray, MoeRouteSegments from art.preprocessing.tokenize import ( TokenizedResult, _chat_choice_trace, @@ -149,7 +149,7 @@ def _routed_exchange( return exchange -def _reasoning_stripped_group() -> art.TrajectoryGroup: +def _reasoning_stripped_group(*, complete_prefix: bool = False) -> art.TrajectoryGroup: def set_choice( exchange: ChatCompletionsExchange, token_ids: list[int], @@ -187,7 +187,7 @@ def set_choice( ) set_choice( first, - [2, 101, 102, 103, 104, 9], + [9] if complete_prefix else [2, 101, 102, 103, 104, 9], content="first", reasoning="long reasoning", ) @@ -207,10 +207,29 @@ def set_choice( content="second", reasoning="short reasoning", ) + exchanges = [first, second] + if complete_prefix: + # Token 9 has the same captured conditional prefix in both branches. + # The earlier overlength branch must not claim it from the fitting one. + long = _routed_exchange( + prompt_token_ids=[1], + output_token=9, + messages=[{"role": "user", "content": "one"}], + content="long answer", + ) + set_choice( + long, [9, 10, 11, 12, 13, 14], content="long answer", reasoning="long" + ) + # Preserve the full first message so first+second form one history; + # no separate short history can claim token 9 before the fitting one. + second.request["messages"][1] = first.response.choices[0].message.model_dump( + mode="python", exclude_none=True + ) + exchanges.insert(0, long) return art.TrajectoryGroup( [ art.Trajectory( - exchanges=tr.TrajectoryExchanges(chat_completions=[first, second]), + exchanges=tr.TrajectoryExchanges(chat_completions=exchanges), reward=reward, ) for reward in (1.0, 0.0) @@ -218,6 +237,27 @@ def set_choice( ) +def _assert_captured_training_prefixes( + results: list[TokenizedResult], group: art.TrajectoryGroup +) -> None: + # Derive eligible conditional logprobs directly from captured exchanges, + # independently of history lineage, token flags, and preprocessing masks. + captured = {} + for exchange in group.trajectories[0].exchanges.chat_completions: + for choice in exchange.response.choices: + extra = choice.model_extra + assert extra is not None and choice.logprobs is not None + prompt, output = extra["prompt_token_ids"], extra["token_ids"] + for index, logprob in enumerate(choice.logprobs.content or []): + captured[tuple(prompt + output[: index + 1])] = logprob.logprob + for result in results: + for index, selected in enumerate(result.assistant_mask): + if selected: + prefix = tuple(result.token_ids[: index + 1]) + assert prefix in captured, f"Uncaptured training prefix: {prefix}" + assert result.logprobs[index] == captured[prefix] + + def _group() -> art.TrajectoryGroup: trajectories = [ art.Trajectory( @@ -361,11 +401,15 @@ def counted_public( assert public_calls == len(group.trajectories) -def test_overlength_history_does_not_claim_sources_from_fitting_history() -> None: +@pytest.mark.parametrize("complete_prefix", [False, True], ids=["shifted", "captured"]) +def test_overlength_history_does_not_claim_sources_from_fitting_history( + complete_prefix: bool, +) -> None: + group = _reasoning_stripped_group(complete_prefix=complete_prefix) results = list( tokenize_trajectory_groups( cast(PreTrainedTokenizerBase, _Tokenizer()), - [_reasoning_stripped_group()], + [group], allow_training_without_logprobs=False, scale_rewards=False, shuffle_group_trajectories=False, @@ -374,18 +418,28 @@ def test_overlength_history_does_not_claim_sources_from_fitting_history() -> Non _max_sequence_length=5, ) ) + _assert_captured_training_prefixes(results, group) long = [result for result in results if len(result.token_ids) > 5] fitting = [result for result in results if len(result.token_ids) <= 5] assert len(long) == len(fitting) == 2 assert all(result.assistant_mask == [0] * 7 for result in long) assert all(result.token_ids == [1, 9, 4, 5, 6] for result in fitting) - assert all(result.assistant_mask == [0, 1, 0, 1, 1] for result in fitting) - assert all(result.weight == pytest.approx(1 / 3) for result in results) + # Token 9 is eligible only when sampled under this exact prefix, [1]. + assert all( + result.assistant_mask == [0, int(complete_prefix), 0, 1, 1] + for result in fitting + ) + assert all( + result.weight == pytest.approx(1 / (2 + int(complete_prefix))) + for result in results + ) +@pytest.mark.parametrize("complete_prefix", [False, True], ids=["shifted", "captured"]) def test_local_backend_trains_retained_source_after_overlength_history( tmp_path: Path, + complete_prefix: bool, ) -> None: backend = LocalBackend(path=str(tmp_path)) model = TrainableModel( @@ -413,7 +467,7 @@ def test_local_backend_trains_retained_source_after_overlength_history( ): packed = backend._get_packed_tensors( model, - [_reasoning_stripped_group()], + [_reasoning_stripped_group(complete_prefix=complete_prefix)], advantage_balance=0.0, allow_training_without_logprobs=False, scale_rewards=False, @@ -424,7 +478,10 @@ def test_local_backend_trains_retained_source_after_overlength_history( assert packed is not None assert packed["tokens"].tolist() == [[1, 9, 4, 5, 6]] * 2 - assert packed["assistant_mask"].tolist() == [[False, True, False, True, True]] * 2 + assert ( + packed["assistant_mask"].tolist() + == [[False, complete_prefix, False, True, True]] * 2 + ) def test_training_rejects_multiple_concrete_policy_versions() -> None: @@ -744,7 +801,9 @@ def test_preprocessing_preserves_moe_routes_for_reasoning_stripped_suffix() -> N "completion_token_ids": [5, 6], "num_experts": 2048, "routed_experts": np.asarray( - [[[10]], [[1010]], [[1020]], [[90]], [[40]], [[50]], [[60]]], + # Same token IDs, different prefixes: use distinct prompt routes + # to detect an invalid overlay from the earlier generation. + [[[10]], [[1110]], [[1120]], [[190]], [[40]], [[50]], [[60]]], dtype=np.uint16, ), } @@ -797,22 +856,33 @@ def apply_chat_template( initial = [result for result in results if result.token_ids[1] == 2] stripped = [result for result in results if result.token_ids[1] == 101] + _assert_captured_training_prefixes(results, group) assert len(initial) == 2 assert len(stripped) == 2 assert all(result.choice_offsets == [1] for result in initial) - # The retained response has a different complete visible prefix after its - # reasoning is stripped, so it is independently eligible in this history. - assert all(result.choice_offsets == [1, 5] for result in stripped) + # The suffix is conditioning only: its captured logprobs belong to the + # complete original prefix. The original generation remains trainable. + assert all(result.choice_offsets == [5] for result in stripped) assert all(result.assistant_mask == [0, 1, 1, 1, 1] for result in initial) - assert all(result.assistant_mask == [0, 1, 1, 1, 0, 1, 1] for result in stripped) - assert all(result.weight == pytest.approx(1 / 9) for result in results) - expected_routes = np.asarray( - [[[10]], [[1010]], [[1020]], [[90]], [[40]], [[50]], [[60]]], + assert all(result.assistant_mask == [0, 0, 0, 0, 0, 1, 1] for result in stripped) + assert all(result.weight == pytest.approx(1 / 6) for result in results) + stripped_routes = np.asarray( + [[[10]], [[1110]], [[1120]], [[190]], [[40]], [[50]], [[60]]], dtype=np.uint16, ) - for result in stripped: - assert isinstance(result.moe_routed_experts, MoeRouteArray) - assert np.array_equal(result.moe_routed_experts, expected_routes) + for histories, expected_routes in ( + (initial, first_extra[ART_MOE_ROUTING_METADATA_KEY]["routed_experts"]), + (stripped, stripped_routes), + ): + for result in histories: + routes = result.moe_routed_experts + assert isinstance(routes, (MoeRouteArray, MoeRouteSegments)) + assert routes.num_experts == 2048 + segments = ( + routes.segments if isinstance(routes, MoeRouteSegments) else (routes,) + ) + assert all(not segment.flags.writeable for segment in segments) + assert np.array_equal(np.concatenate(segments), expected_routes) datums = trajectory_groups_to_datums( [group], @@ -824,7 +894,7 @@ def apply_chat_template( ) masks = [datum.loss_fn_inputs["mask"].to_torch().tolist() for datum in datums] assert masks.count([1, 1, 1, 1]) == 2 - assert masks.count([1, 1, 1, 0, 1, 1]) == 2 + assert masks.count([0, 0, 0, 0, 1, 1]) == 2 def test_ambiguous_non_moe_suffix_falls_back_to_sampled_spans() -> None: