fix(data): align GLM Kimi and MiniMax chat masks

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

Author: yaoyu-33

src/megatron/bridge/data/conversation_processing.py

@@ -77,6 +77,9 @@ class AssistantMaskBoundaryConfig:
         ``role_end_token_variants`` provides ordered fallback end sequences for
         formats whose final turn may omit a trailing delimiter such as a
         newline. The primary ``role_end_tokens`` sequence is always preferred.
+        ``rendered_turn_end_roles`` covers formats such as GLM-5.2 whose
+        assistant turn has no explicit end token. For those roles, the
+        conversation-prefix render determines the exclusive content boundary.
 
         Example::
 
@@ -105,6 +108,7 @@ class AssistantMaskBoundaryConfig:
     trim_leading_token_ids: Sequence[int] = ()
     trim_leading_token_sequences: Sequence[Sequence[int]] = ()
     role_end_token_variants: Mapping[str, Sequence[Sequence[int]]] = field(default_factory=dict)
+    rendered_turn_end_roles: Sequence[str] = ()
 
 
 @dataclass(frozen=True)
@@ -491,7 +495,9 @@ def resolve_chat_template_kwargs(template_owner: Any, template_kwargs: Mapping[s
 
     MBridge exposes ``truncate_history_thinking`` consistently. Templates that
     implement the inverse ``preserve_thinking`` control receive the equivalent
-    negated value only at the rendering boundary.
+    negated value only at the rendering boundary, while GLM-style
+    ``clear_thinking`` receives the same value. Custom tokenizers such as
+    Kimi-K3 use ``thinking`` instead of ``enable_thinking``.
 
     Args:
         template_owner: Processor or tokenizer that owns the chat template.
@@ -501,13 +507,29 @@ def resolve_chat_template_kwargs(template_owner: Any, template_kwargs: Mapping[s
         A copied mapping adapted to the selected template.
     """
     kwargs = dict(template_kwargs)
-    if "truncate_history_thinking" not in kwargs:
-        return kwargs
     tokenizer = get_processor_tokenizer(template_owner)
     template = _get_chat_template(template_owner, tokenizer)
-    if template is not None and "preserve_thinking" in template and "truncate_history_thinking" not in template:
-        truncate_history = kwargs.pop("truncate_history_thinking")
-        kwargs["preserve_thinking"] = not truncate_history
+
+    if "truncate_history_thinking" in kwargs and (template is None or "truncate_history_thinking" not in template):
+        truncate_history = kwargs["truncate_history_thinking"]
+        if template is not None and "preserve_thinking" in template:
+            kwargs.pop("truncate_history_thinking")
+            kwargs["preserve_thinking"] = not truncate_history
+        elif template is not None and "clear_thinking" in template:
+            kwargs.pop("truncate_history_thinking")
+            kwargs["clear_thinking"] = truncate_history
+
+    if "enable_thinking" in kwargs and (template is None or "enable_thinking" not in template):
+        parameters: set[str] = set()
+        for owner in (template_owner, tokenizer):
+            apply_chat_template = _get_explicit_attribute(owner, "apply_chat_template")
+            try:
+                parameters.update(inspect.signature(apply_chat_template).parameters if apply_chat_template else {})
+            except (TypeError, ValueError):
+                continue
+        if "thinking" in parameters and "enable_thinking" not in parameters:
+            enable_thinking = kwargs.pop("enable_thinking")
+            kwargs.setdefault("thinking", enable_thinking)
     return kwargs
 
 
@@ -852,7 +874,12 @@ def tokenize_chat_example(
         )
 
     chat_example: dict[str, Any] = {"conversation": conversation}
-    chat_example.update(template_kwargs)
+    mask_template_kwargs = dict(template_kwargs)
+    mask_tools = mask_template_kwargs.pop("tools", None)
+    if mask_template_kwargs:
+        chat_example["chat_template_kwargs"] = mask_template_kwargs
+    if mask_tools is not None:
+        chat_example["tools"] = mask_tools
     effective_boundary_config = boundary_config or infer_assistant_mask_boundary_config(processor)
     if returned_mask is not None and effective_boundary_config is None:
         assistant_mask = returned_mask
@@ -1037,9 +1064,42 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
     """Infer assistant role boundaries from known HF chat-template formats."""
     tokenizer = get_processor_tokenizer(processor)
     template = _get_chat_template(tokenizer, processor)
+
     if template is None:
+        # Kimi-K3 owns rendering in a custom tokenizer rather than a Jinja
+        # ``chat_template``. Detect its XTML structural vocabulary instead of a
+        # Python class name so local/remote-code wrappers retain the same behavior.
+        apply_chat_template = _get_explicit_attribute(tokenizer, "apply_chat_template")
+        try:
+            template_parameters = inspect.signature(apply_chat_template).parameters if apply_chat_template else {}
+        except (TypeError, ValueError):
+            template_parameters = {}
+        if "thinking" in template_parameters:
+            get_vocab = _get_explicit_attribute(tokenizer, "get_vocab")
+            try:
+                vocab = get_vocab() if callable(get_vocab) else {}
+            except (AttributeError, TypeError, ValueError):
+                vocab = {}
+            kimi_markers = ("<|open|>", "<|close|>", "<|sep|>", "<|end_of_msg|>")
+            if isinstance(vocab, Mapping) and all(marker in vocab for marker in kimi_markers):
+                return assistant_mask_boundary_config_from_markers(
+                    processor,
+                    assistant_start='<|open|>message role="assistant"<|sep|>',
+                    assistant_end="<|close|>message<|sep|><|end_of_msg|>",
+                )
         return None
 
+    # GLM-5.2 does not terminate completed assistant turns with a structural
+    # token. Prefix rendering is therefore the only loss boundary that remains
+    # exact for both terminal and historical turns.
+    if all(marker in template for marker in ("<|assistant|>", "clear_thinking", "<think>")):
+        return assistant_mask_boundary_config_from_markers(
+            processor,
+            assistant_start="<|assistant|>",
+            assistant_end=None,
+            rendered_turn_end_roles=(CHATML_ASSISTANT_ROLE,),
+        )
+
     candidates: tuple[tuple[tuple[str, ...], str, str, tuple[str, ...], Mapping[str, str]], ...] = (
         (
             ("<|im_start|>assistant", "<|im_end|>"),
@@ -1086,6 +1146,18 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
             ("<|end|>", "<|call|>"),
             {role: f"<|start|>{role}" for role in ("system", "developer", "user", "tool")},
         ),
+        (
+            ("]~!b[", "]~b]", "[e~[", "<mm:think>"),
+            "]~b]ai\n",
+            "[e~[\n",
+            (),
+            {
+                "system": "]~!b[]~b]system\n",
+                "developer": "]~b]developer\n",
+                "user": "]~b]user\n",
+                "tool": "]~b]tool",
+            },
+        ),
     )
     for required_markers, assistant_start, assistant_end, assistant_end_fallbacks, other_role_starts in candidates:
         if not all(marker in template for marker in required_markers):
@@ -1118,25 +1190,26 @@ def assistant_mask_boundary_config_from_markers(
     processor: Any,
     *,
     assistant_start: str,
-    assistant_end: str,
+    assistant_end: str | None,
     assistant_end_fallbacks: Sequence[str] = (),
     role_start_markers: Mapping[str, str] | None = None,
     loss_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,),
     include_start_tokens_for_roles: Sequence[str] = (),
     include_end_tokens_for_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,),
     trim_leading_token_ids: Sequence[int] = (),
     trim_leading_token_sequences: Sequence[Sequence[int]] = (),
+    rendered_turn_end_roles: Sequence[str] = (),
 ) -> AssistantMaskBoundaryConfig:
     """Build explicit assistant boundary config from model-declared marker strings."""
     tokenizer = get_processor_tokenizer(processor)
     start_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_start)
-    end_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_end)
+    end_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_end) if assistant_end is not None else []
     end_token_variants = [
         token_ids
         for fallback in assistant_end_fallbacks
         if (token_ids := tokenize_text_without_special_tokens(tokenizer, fallback))
     ]
-    if not start_tokens or not end_tokens:
+    if not start_tokens or (not end_tokens and CHATML_ASSISTANT_ROLE not in rendered_turn_end_roles):
         raise ValueError(
             "Unable to tokenize assistant loss-mask boundary markers. "
             "Check the model collator's assistant_start and assistant_end markers."
@@ -1149,13 +1222,14 @@ def assistant_mask_boundary_config_from_markers(
         role_start_tokens[role] = marker_tokens
     return AssistantMaskBoundaryConfig(
         role_start_tokens=role_start_tokens,
-        role_end_tokens={role: end_tokens for role in role_start_tokens},
+        role_end_tokens={role: end_tokens for role in role_start_tokens if role not in rendered_turn_end_roles},
         loss_roles=loss_roles,
         include_start_tokens_for_roles=include_start_tokens_for_roles,
         include_end_tokens_for_roles=include_end_tokens_for_roles,
         trim_leading_token_ids=trim_leading_token_ids,
         trim_leading_token_sequences=trim_leading_token_sequences,
         role_end_token_variants={role: end_token_variants for role in role_start_tokens},
+        rendered_turn_end_roles=rendered_turn_end_roles,
     )
 
 
@@ -1308,13 +1382,14 @@ def _assistant_mask_from_conversation_turns(
     }
     include_start_roles = set(boundary_config.include_start_tokens_for_roles)
     include_end_roles = set(boundary_config.include_end_tokens_for_roles)
+    rendered_turn_end_roles = set(boundary_config.rendered_turn_end_roles)
     rendered_mask = torch.zeros(len(rendered_ids), dtype=torch.float32)
 
     for turn_index, turn in loss_turns:
         role = turn["role"]
         start_tokens = role_start_tokens.get(role)
         end_patterns = [role_end_tokens.get(role, []), *role_end_token_variants.get(role, [])]
-        if not start_tokens or not any(end_patterns):
+        if not start_tokens or (not any(end_patterns) and role not in rendered_turn_end_roles):
             return None
 
         try:
@@ -1344,20 +1419,23 @@ def _assistant_mask_from_conversation_turns(
         if role_start < 0 or content_start > turn_limit:
             return None
 
-        candidate_ends: list[tuple[int, int, int, list[int]]] = []
-        for priority, pattern in enumerate(end_patterns):
-            if not pattern:
-                continue
-            search_start = content_start
-            while search_start < turn_limit:
-                end_start, after_end = find_token_span(rendered_ids, pattern, search_start)
-                if end_start < 0 or after_end > turn_limit:
-                    break
-                candidate_ends.append((end_start, -priority, after_end, pattern))
-                search_start = end_start + 1
-        if not candidate_ends:
-            return None
-        content_end, _, after_end, _ = max(candidate_ends)
+        if role in rendered_turn_end_roles:
+            content_end = after_end = turn_limit
+        else:
+            candidate_ends: list[tuple[int, int, int, list[int]]] = []
+            for priority, pattern in enumerate(end_patterns):
+                if not pattern:
+                    continue
+                search_start = content_start
+                while search_start < turn_limit:
+                    end_start, after_end = find_token_span(rendered_ids, pattern, search_start)
+                    if end_start < 0 or after_end > turn_limit:
+                        break
+                    candidate_ends.append((end_start, -priority, after_end, pattern))
+                    search_start = end_start + 1
+            if not candidate_ends:
+                return None
+            content_end, _, after_end, _ = max(candidate_ends)
 
         loss_start = _trim_leading_token_sequences(
             rendered_ids,
@@ -1375,7 +1453,7 @@ def _assistant_mask_from_conversation_turns(
             rendered_mask[loss_start:content_end] = 1.0
         if role in include_start_roles:
             rendered_mask[role_start : role_start + len(start_tokens)] = 1.0
-        if role in include_end_roles:
+        if role in include_end_roles and role not in rendered_turn_end_roles:
             rendered_mask[content_end:after_end] = 1.0
 
     aligned_mask = _align_assistant_mask_to_padded_ids(rendered_ids, rendered_mask.tolist(), ids, tokenizer)
@@ -1535,6 +1613,12 @@ def _assistant_mask_from_boundary_config(
         for role, raw_variants in boundary_config.role_end_token_variants.items()
     }
     loss_roles = set(boundary_config.loss_roles)
+    rendered_turn_end_roles = set(boundary_config.rendered_turn_end_roles)
+    if rendered_turn_end_roles:
+        # A raw token scan cannot distinguish the end of an unterminated turn
+        # from following non-loss roles. Only the conversation-prefix renderer
+        # can establish this boundary safely.
+        return None
     if not role_start_tokens or not role_end_tokens or not loss_roles:
         raise ValueError(
             "AssistantMaskBoundaryConfig must provide non-empty role_start_tokens, role_end_tokens, and loss_roles."
@@ -1641,7 +1725,7 @@ def build_assistant_loss_mask(
 
     mask = _assistant_mask_from_hf_chat_template(example_or_conversation, ids, processor, tokenizer)
     if boundary_config is not None:
-        boundary_scan_is_safe = (
+        boundary_scan_is_safe = not boundary_config.rendered_turn_end_roles and (
             not conversation
             or _conversation_contains_boundary_tokens(example_or_conversation, tokenizer, boundary_config) is False
         )