fix(data): scope GLM assistant boundary scan to the current turn

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)
 

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)