From 1968beaf4b0cba90a42dbd98af4efd6d65cf6a0a Mon Sep 17 00:00:00 2001 From: Yu Yao Date: Tue, 18 Aug 2026 00:40:45 -0700 Subject: [PATCH] [data] fix: preserve GLM historical answer loss Signed-off-by: Yu Yao --- .../bridge/data/conversation_processing.py | 17 ++++++++++++ .../test_glm52_kimi_k3_minimax_m3_chat.py | 26 +++++++++++++++++++ 2 files changed, 43 insertions(+) diff --git a/src/megatron/bridge/data/conversation_processing.py b/src/megatron/bridge/data/conversation_processing.py index 6d4dc7f200..f80d54637b 100644 --- a/src/megatron/bridge/data/conversation_processing.py +++ b/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): 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 967bae58b2..6c5ba39f3b 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 @@ -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)