diff --git a/src/megatron/bridge/data/conversation_processing.py b/src/megatron/bridge/data/conversation_processing.py index f80d54637b..cc4b124a8c 100644 --- a/src/megatron/bridge/data/conversation_processing.py +++ b/src/megatron/bridge/data/conversation_processing.py @@ -1446,8 +1446,10 @@ def _assistant_mask_from_conversation_turns( if role in rendered_turn_end_roles: content_end = after_end = turn_limit if turn_index + 1 < len(conversation): + # Only this turn's payload renders between ``content_start`` and the next role + # header, so a marker quoted in any other turn cannot shift that boundary. boundary_tokens_in_payload = _conversation_contains_boundary_tokens( - example_or_conversation, tokenizer, boundary_config + example_or_conversation, tokenizer, boundary_config, turns=[turn] ) if boundary_tokens_in_payload is not False: return None @@ -1567,8 +1569,14 @@ def _conversation_contains_boundary_tokens( example_or_conversation: Mapping[str, Any] | Sequence[Mapping[str, Any]], tokenizer: Any, boundary_config: AssistantMaskBoundaryConfig, + *, + turns: Sequence[Mapping[str, Any]] | None = None, ) -> bool | None: - """Return whether rendered conversation payloads contain structural token sequences.""" + """Return whether rendered conversation payloads contain structural token sequences. + + ``turns`` restricts the scan to those turns instead of the whole conversation. Template + kwargs are always scanned. + """ token_sequences = [ *(_token_map_from_boundary_config(boundary_config.role_start_tokens).values()), *(_token_map_from_boundary_config(boundary_config.role_end_tokens).values()), @@ -1579,10 +1587,8 @@ def _conversation_contains_boundary_tokens( if (token_ids := _as_token_id_list(raw_token_ids)) ), ] - payloads: list[Any] = [ - {key: value for key, value in turn.items() if key != "role"} - for turn in _conversation_from_example(example_or_conversation) - ] + scanned_turns = _conversation_from_example(example_or_conversation) if turns is None else turns + payloads: list[Any] = [{key: value for key, value in turn.items() if key != "role"} for turn in scanned_turns] payloads.append(chat_template_kwargs_from_example(example_or_conversation)) return _value_contains_token_sequence(payloads, tokenizer, token_sequences) diff --git a/tests/unit_tests/models/test_glm52_kimi_k3_minimax_m3_chat.py b/tests/unit_tests/models/test_glm52_kimi_k3_minimax_m3_chat.py index 6c5ba39f3b..f3aaf40614 100644 --- a/tests/unit_tests/models/test_glm52_kimi_k3_minimax_m3_chat.py +++ b/tests/unit_tests/models/test_glm52_kimi_k3_minimax_m3_chat.py @@ -35,6 +35,12 @@ "Seattle is 12 C.": 38, } +# Texts that quote a role marker and therefore tokenize to more than one id. +QUOTED_MARKER_TEXTS = { + "question two about <|observation|>": [TEXT_IDS["question two"], TOOL_RESULT], + "answer one about <|observation|>": [TEXT_IDS["answer one"], TOOL_RESULT], +} + class _ModelChatTokenizer: """Small semantic renderer for three official chat-template boundary formats.""" @@ -81,6 +87,8 @@ def encode(self, text: str, add_special_tokens: bool = False) -> list[int]: "]~b]user\n": [USER_START], "]~b]tool": [TOOL_RESULT], } + if text in QUOTED_MARKER_TEXTS: + return list(QUOTED_MARKER_TEXTS[text]) return marker_ids.get(text, [TEXT_IDS.get(text, 99)]) def __call__(self, text: str, add_special_tokens: bool = False) -> dict[str, list[int]]: @@ -107,7 +115,7 @@ def _assistant_tokens( result.append(THINK_CLOSE) if isinstance(content, str) and content: - result.append(TEXT_IDS[content]) + result.extend(self.encode(content)) if self.model == "kimi": result.append(RESPONSE_CLOSE) for tool_call in message.get("tool_calls", []): @@ -147,7 +155,7 @@ def apply_chat_template( elif role == "tool": input_ids.extend([TOOL_RESULT, TEXT_IDS[message["content"]], TURN_END]) else: - input_ids.extend([USER_START, TEXT_IDS[message["content"]], TURN_END]) + input_ids.extend([USER_START, *self.encode(message["content"]), TURN_END]) if truncation and max_length is not None: input_ids = input_ids[:max_length] return {"input_ids": input_ids} if return_dict else input_ids @@ -294,6 +302,44 @@ def test_glm_history_thinking_truncation_preserves_historical_answer_loss() -> N ] +def _glm_history_truncation_example(messages: list[dict[str, Any]]) -> dict[str, Any]: + return { + "messages": messages, + "chat_template_kwargs": {"enable_thinking": True, "truncate_history_thinking": True}, + } + + +def test_glm_role_marker_quoted_in_a_user_turn_still_masks_the_historical_answer() -> None: + """A role marker quoted outside the assistant turn does not veto its loss boundary.""" + tokenizer = _ModelChatTokenizer("glm") + messages = _messages(tools=False, thinking=True) + messages[2] = {**messages[2], "content": "question two about <|observation|>"} + + tokenized = tokenize_chat_example(_glm_history_truncation_example(messages), tokenizer, warn_on_all_masked=False) + + supervised_ids = tokenized.input_ids[tokenized.assistant_mask].tolist() + assert supervised_ids == [ + THINK_OPEN, + THINK_CLOSE, + TEXT_IDS["answer one"], + THINK_OPEN, + TEXT_IDS["reason two"], + THINK_CLOSE, + TEXT_IDS["answer two"], + ] + assert TOOL_RESULT not in supervised_ids + + +def test_glm_role_marker_quoted_in_an_assistant_turn_still_fails_closed() -> None: + """A role marker quoted inside the assistant payload keeps refusing an ambiguous boundary.""" + tokenizer = _ModelChatTokenizer("glm") + messages = _messages(tools=False, thinking=True) + messages[1] = {**messages[1], "content": "answer one about <|observation|>"} + + with pytest.raises(ValueError, match="did not match any loss-contributing spans"): + tokenize_chat_example(_glm_history_truncation_example(messages), tokenizer, warn_on_all_masked=False) + + @pytest.mark.parametrize("model", ["glm", "kimi", "minimax"]) def test_model_chat_right_truncation_and_padding_preserve_loss_alignment(model: str) -> None: tokenizer = _ModelChatTokenizer(model)