[data] fix: preserve GLM historical answer loss

Signed-off-by: Yu Yao <yaoyu.094@gmail.com>

Author: yaoyu-33

src/megatron/bridge/data/conversation_processing.py

@@ -1101,6 +1101,11 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
             processor,
             assistant_start="<|assistant|>",
             assistant_end=None,
+            role_start_markers={
+                "system": "<|system|>",
+                "user": "<|user|>",
+                "tool": "<|observation|>",
+            },
             rendered_turn_end_roles=(CHATML_ASSISTANT_ROLE,),
         )
 
@@ -1440,6 +1445,18 @@ 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):
+                boundary_tokens_in_payload = _conversation_contains_boundary_tokens(
+                    example_or_conversation, tokenizer, boundary_config
+                )
+                if boundary_tokens_in_payload is not False:
+                    return None
+                next_role = conversation[turn_index + 1].get("role")
+                next_role_start_tokens = role_start_tokens.get(next_role, [])
+                next_role_start, _ = find_token_span(rendered_ids, next_role_start_tokens, content_start)
+                if next_role_start < 0:
+                    return None
+                content_end = after_end = next_role_start
         else:
             candidate_ends: list[tuple[int, int, int, list[int]]] = []
             for priority, pattern in enumerate(end_patterns):

tests/unit_tests/models/test_glm52_kimi_k3_minimax_m3_chat.py

@@ -69,6 +69,9 @@ def encode(self, text: str, add_special_tokens: bool = False) -> list[int]:
         assert add_special_tokens is False
         marker_ids = {
             "<|assistant|>": [ASSISTANT_START],
+            "<|system|>": [40],
+            "<|user|>": [USER_START],
+            "<|observation|>": [TOOL_RESULT],
             '<|open|>message role="assistant"<|sep|>': [ASSISTANT_START],
             "<|close|>message<|sep|><|end_of_msg|>": [TURN_END],
             "]~b]ai\n": [ASSISTANT_START],
@@ -268,6 +271,29 @@ def test_model_chat_four_way_rendering_and_loss_boundaries(model: str, tools: bo
         assert all("enable_thinking" not in entry for entry in semantic_kwargs)
 
 
+def test_glm_history_thinking_truncation_preserves_historical_answer_loss() -> None:
+    tokenizer = _ModelChatTokenizer("glm")
+    example = {
+        "messages": _messages(tools=False, thinking=True),
+        "chat_template_kwargs": {
+            "enable_thinking": True,
+            "truncate_history_thinking": True,
+        },
+    }
+
+    tokenized = tokenize_chat_example(example, tokenizer, warn_on_all_masked=False)
+
+    assert tokenized.input_ids[tokenized.assistant_mask].tolist() == [
+        THINK_OPEN,
+        THINK_CLOSE,
+        TEXT_IDS["answer one"],
+        THINK_OPEN,
+        TEXT_IDS["reason two"],
+        THINK_CLOSE,
+        TEXT_IDS["answer two"],
+    ]
+
+
 @pytest.mark.parametrize("model", ["glm", "kimi", "minimax"])
 def test_model_chat_right_truncation_and_padding_preserve_loss_alignment(model: str) -> None:
     tokenizer = _ModelChatTokenizer(model)