Skip to content

Commit fdc651c

Browse files
committed
Merge branch 'main' into codex/issue5434-gtp-stage1
Signed-off-by: yaoyu-33 <yaoyu.094@gmail.com>
2 parents 152fced + 46c4255 commit fdc651c

5 files changed

Lines changed: 731 additions & 30 deletions

File tree

‎examples/models/qwen/qwen35_vl/README.md‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,15 @@ convergence.
102102

103103
See the [slurm_sft.sh](slurm_sft.sh) script for full parameter fine-tuning with configurable model sizes.
104104

105+
Text-only chat rows may provide official Qwen template controls under `chat_template_kwargs`. For example,
106+
`{"enable_thinking": true, "truncate_history_thinking": false}` preserves historical reasoning on Qwen3.6 by
107+
translating the canonical history option to the template-native `preserve_thinking=true`. The Qwen3.5 template does
108+
not expose that control and continues to omit reasoning before the latest user query.
109+
110+
Tool definitions remain in the top-level `tools` field. OpenAI JSON-string `function.arguments` are normalized to
111+
objects before rendering. Assistant headers, system/user turns, and tool responses are masked from loss; assistant
112+
thinking boundaries, retained reasoning, tool calls, answers, and `<|im_end|>` are supervised.
113+
105114
### Parameter-Efficient Fine-Tuning (PEFT) with LoRA
106115

107116
See the [slurm_peft.sh](slurm_peft.sh) script for LoRA fine-tuning with configurable model sizes.

‎src/megatron/bridge/data/conversation_processing.py‎

Lines changed: 116 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,9 @@ class AssistantMaskBoundaryConfig:
7777
``role_end_token_variants`` provides ordered fallback end sequences for
7878
formats whose final turn may omit a trailing delimiter such as a
7979
newline. The primary ``role_end_tokens`` sequence is always preferred.
80+
``rendered_turn_end_roles`` covers formats such as GLM-5.2 whose
81+
assistant turn has no explicit end token. For those roles, the
82+
conversation-prefix render determines the exclusive content boundary.
8083
8184
Example::
8285
@@ -105,6 +108,7 @@ class AssistantMaskBoundaryConfig:
105108
trim_leading_token_ids: Sequence[int] = ()
106109
trim_leading_token_sequences: Sequence[Sequence[int]] = ()
107110
role_end_token_variants: Mapping[str, Sequence[Sequence[int]]] = field(default_factory=dict)
111+
rendered_turn_end_roles: Sequence[str] = ()
108112

109113

110114
@dataclass(frozen=True)
@@ -491,7 +495,9 @@ def resolve_chat_template_kwargs(template_owner: Any, template_kwargs: Mapping[s
491495
492496
MBridge exposes ``truncate_history_thinking`` consistently. Templates that
493497
implement the inverse ``preserve_thinking`` control receive the equivalent
494-
negated value only at the rendering boundary.
498+
negated value only at the rendering boundary, while GLM-style
499+
``clear_thinking`` receives the same value. Custom tokenizers such as
500+
Kimi-K3 use ``thinking`` instead of ``enable_thinking``.
495501
496502
Args:
497503
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
501507
A copied mapping adapted to the selected template.
502508
"""
503509
kwargs = dict(template_kwargs)
504-
if "truncate_history_thinking" not in kwargs:
505-
return kwargs
506510
tokenizer = get_processor_tokenizer(template_owner)
507511
template = _get_chat_template(template_owner, tokenizer)
508-
if template is not None and "preserve_thinking" in template and "truncate_history_thinking" not in template:
509-
truncate_history = kwargs.pop("truncate_history_thinking")
510-
kwargs["preserve_thinking"] = not truncate_history
512+
513+
if "truncate_history_thinking" in kwargs and (template is None or "truncate_history_thinking" not in template):
514+
truncate_history = kwargs["truncate_history_thinking"]
515+
if template is not None and "preserve_thinking" in template:
516+
kwargs.pop("truncate_history_thinking")
517+
kwargs["preserve_thinking"] = not truncate_history
518+
elif template is not None and "clear_thinking" in template:
519+
kwargs.pop("truncate_history_thinking")
520+
kwargs["clear_thinking"] = truncate_history
521+
522+
if "enable_thinking" in kwargs and (template is None or "enable_thinking" not in template):
523+
parameters: set[str] = set()
524+
for owner in (template_owner, tokenizer):
525+
apply_chat_template = _get_explicit_attribute(owner, "apply_chat_template")
526+
try:
527+
parameters.update(inspect.signature(apply_chat_template).parameters if apply_chat_template else {})
528+
except (TypeError, ValueError):
529+
continue
530+
if "thinking" in parameters and "enable_thinking" not in parameters:
531+
enable_thinking = kwargs.pop("enable_thinking")
532+
kwargs.setdefault("thinking", enable_thinking)
511533
return kwargs
512534

513535

@@ -556,7 +578,8 @@ def normalize_chat_conversation(
556578
"""Normalize supported text-chat schemas to OpenAI-style messages.
557579
558580
Rows may use ``messages``, singular ``conversation``, or legacy plural
559-
``conversations`` with ``from``/``value`` keys.
581+
``conversations`` with ``from``/``value`` keys. Null OpenAI message content
582+
is converted to an empty string before chat-template rendering.
560583
"""
561584
if isinstance(example_or_conversation, Mapping):
562585
source = example_or_conversation
@@ -581,6 +604,8 @@ def normalize_chat_conversation(
581604
raise ValueError("Chat conversation turns must be dictionaries.")
582605
turn = dict(raw_turn)
583606
if "role" in turn and "content" in turn:
607+
if turn["content"] is None:
608+
turn["content"] = ""
584609
_normalize_assistant_tool_call_arguments(turn)
585610
normalized.append(turn)
586611
continue
@@ -853,7 +878,12 @@ def tokenize_chat_example(
853878
)
854879

855880
chat_example: dict[str, Any] = {"conversation": conversation}
856-
chat_example.update(template_kwargs)
881+
mask_template_kwargs = dict(template_kwargs)
882+
mask_tools = mask_template_kwargs.pop("tools", None)
883+
if mask_template_kwargs:
884+
chat_example["chat_template_kwargs"] = mask_template_kwargs
885+
if mask_tools is not None:
886+
chat_example["tools"] = mask_tools
857887
effective_boundary_config = boundary_config or infer_assistant_mask_boundary_config(processor)
858888
if returned_mask is not None and effective_boundary_config is None:
859889
assistant_mask = returned_mask
@@ -1038,9 +1068,42 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
10381068
"""Infer assistant role boundaries from known HF chat-template formats."""
10391069
tokenizer = get_processor_tokenizer(processor)
10401070
template = _get_chat_template(tokenizer, processor)
1071+
10411072
if template is None:
1073+
# Kimi-K3 owns rendering in a custom tokenizer rather than a Jinja
1074+
# ``chat_template``. Detect its XTML structural vocabulary instead of a
1075+
# Python class name so local/remote-code wrappers retain the same behavior.
1076+
apply_chat_template = _get_explicit_attribute(tokenizer, "apply_chat_template")
1077+
try:
1078+
template_parameters = inspect.signature(apply_chat_template).parameters if apply_chat_template else {}
1079+
except (TypeError, ValueError):
1080+
template_parameters = {}
1081+
if "thinking" in template_parameters:
1082+
get_vocab = _get_explicit_attribute(tokenizer, "get_vocab")
1083+
try:
1084+
vocab = get_vocab() if callable(get_vocab) else {}
1085+
except (AttributeError, TypeError, ValueError):
1086+
vocab = {}
1087+
kimi_markers = ("<|open|>", "<|close|>", "<|sep|>", "<|end_of_msg|>")
1088+
if isinstance(vocab, Mapping) and all(marker in vocab for marker in kimi_markers):
1089+
return assistant_mask_boundary_config_from_markers(
1090+
processor,
1091+
assistant_start='<|open|>message role="assistant"<|sep|>',
1092+
assistant_end="<|close|>message<|sep|><|end_of_msg|>",
1093+
)
10421094
return None
10431095

1096+
# GLM-5.2 does not terminate completed assistant turns with a structural
1097+
# token. Prefix rendering is therefore the only loss boundary that remains
1098+
# exact for both terminal and historical turns.
1099+
if all(marker in template for marker in ("<|assistant|>", "clear_thinking", "<think>")):
1100+
return assistant_mask_boundary_config_from_markers(
1101+
processor,
1102+
assistant_start="<|assistant|>",
1103+
assistant_end=None,
1104+
rendered_turn_end_roles=(CHATML_ASSISTANT_ROLE,),
1105+
)
1106+
10441107
candidates: tuple[tuple[tuple[str, ...], str, str, tuple[str, ...], Mapping[str, str]], ...] = (
10451108
(
10461109
("<|im_start|>assistant", "<|im_end|>"),
@@ -1087,6 +1150,18 @@ def infer_assistant_mask_boundary_config(processor: Any) -> AssistantMaskBoundar
10871150
("<|end|>", "<|call|>"),
10881151
{role: f"<|start|>{role}" for role in ("system", "developer", "user", "tool")},
10891152
),
1153+
(
1154+
("]~!b[", "]~b]", "[e~[", "<mm:think>"),
1155+
"]~b]ai\n",
1156+
"[e~[\n",
1157+
(),
1158+
{
1159+
"system": "]~!b[]~b]system\n",
1160+
"developer": "]~b]developer\n",
1161+
"user": "]~b]user\n",
1162+
"tool": "]~b]tool",
1163+
},
1164+
),
10901165
)
10911166
for required_markers, assistant_start, assistant_end, assistant_end_fallbacks, other_role_starts in candidates:
10921167
if not all(marker in template for marker in required_markers):
@@ -1126,25 +1201,26 @@ def assistant_mask_boundary_config_from_markers(
11261201
processor: Any,
11271202
*,
11281203
assistant_start: str,
1129-
assistant_end: str,
1204+
assistant_end: str | None,
11301205
assistant_end_fallbacks: Sequence[str] = (),
11311206
role_start_markers: Mapping[str, str] | None = None,
11321207
loss_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,),
11331208
include_start_tokens_for_roles: Sequence[str] = (),
11341209
include_end_tokens_for_roles: Sequence[str] = (CHATML_ASSISTANT_ROLE,),
11351210
trim_leading_token_ids: Sequence[int] = (),
11361211
trim_leading_token_sequences: Sequence[Sequence[int]] = (),
1212+
rendered_turn_end_roles: Sequence[str] = (),
11371213
) -> AssistantMaskBoundaryConfig:
11381214
"""Build explicit assistant boundary config from model-declared marker strings."""
11391215
tokenizer = get_processor_tokenizer(processor)
11401216
start_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_start)
1141-
end_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_end)
1217+
end_tokens = tokenize_text_without_special_tokens(tokenizer, assistant_end) if assistant_end is not None else []
11421218
end_token_variants = [
11431219
token_ids
11441220
for fallback in assistant_end_fallbacks
11451221
if (token_ids := tokenize_text_without_special_tokens(tokenizer, fallback))
11461222
]
1147-
if not start_tokens or not end_tokens:
1223+
if not start_tokens or (not end_tokens and CHATML_ASSISTANT_ROLE not in rendered_turn_end_roles):
11481224
raise ValueError(
11491225
"Unable to tokenize assistant loss-mask boundary markers. "
11501226
"Check the model collator's assistant_start and assistant_end markers."
@@ -1157,13 +1233,14 @@ def assistant_mask_boundary_config_from_markers(
11571233
role_start_tokens[role] = marker_tokens
11581234
return AssistantMaskBoundaryConfig(
11591235
role_start_tokens=role_start_tokens,
1160-
role_end_tokens={role: end_tokens for role in role_start_tokens},
1236+
role_end_tokens={role: end_tokens for role in role_start_tokens if role not in rendered_turn_end_roles},
11611237
loss_roles=loss_roles,
11621238
include_start_tokens_for_roles=include_start_tokens_for_roles,
11631239
include_end_tokens_for_roles=include_end_tokens_for_roles,
11641240
trim_leading_token_ids=trim_leading_token_ids,
11651241
trim_leading_token_sequences=trim_leading_token_sequences,
11661242
role_end_token_variants={role: end_token_variants for role in role_start_tokens},
1243+
rendered_turn_end_roles=rendered_turn_end_roles,
11671244
)
11681245

11691246

@@ -1316,13 +1393,14 @@ def _assistant_mask_from_conversation_turns(
13161393
}
13171394
include_start_roles = set(boundary_config.include_start_tokens_for_roles)
13181395
include_end_roles = set(boundary_config.include_end_tokens_for_roles)
1396+
rendered_turn_end_roles = set(boundary_config.rendered_turn_end_roles)
13191397
rendered_mask = torch.zeros(len(rendered_ids), dtype=torch.float32)
13201398

13211399
for turn_index, turn in loss_turns:
13221400
role = turn["role"]
13231401
start_tokens = role_start_tokens.get(role)
13241402
end_patterns = [role_end_tokens.get(role, []), *role_end_token_variants.get(role, [])]
1325-
if not start_tokens or not any(end_patterns):
1403+
if not start_tokens or (not any(end_patterns) and role not in rendered_turn_end_roles):
13261404
return None
13271405

13281406
try:
@@ -1360,20 +1438,23 @@ def _assistant_mask_from_conversation_turns(
13601438
if content_start > turn_limit:
13611439
return None
13621440

1363-
candidate_ends: list[tuple[int, int, int, list[int]]] = []
1364-
for priority, pattern in enumerate(end_patterns):
1365-
if not pattern:
1366-
continue
1367-
search_start = content_start
1368-
while search_start < turn_limit:
1369-
end_start, after_end = find_token_span(rendered_ids, pattern, search_start)
1370-
if end_start < 0 or after_end > turn_limit:
1371-
break
1372-
candidate_ends.append((end_start, -priority, after_end, pattern))
1373-
search_start = end_start + 1
1374-
if not candidate_ends:
1375-
return None
1376-
content_end, _, after_end, _ = max(candidate_ends)
1441+
if role in rendered_turn_end_roles:
1442+
content_end = after_end = turn_limit
1443+
else:
1444+
candidate_ends: list[tuple[int, int, int, list[int]]] = []
1445+
for priority, pattern in enumerate(end_patterns):
1446+
if not pattern:
1447+
continue
1448+
search_start = content_start
1449+
while search_start < turn_limit:
1450+
end_start, after_end = find_token_span(rendered_ids, pattern, search_start)
1451+
if end_start < 0 or after_end > turn_limit:
1452+
break
1453+
candidate_ends.append((end_start, -priority, after_end, pattern))
1454+
search_start = end_start + 1
1455+
if not candidate_ends:
1456+
return None
1457+
content_end, _, after_end, _ = max(candidate_ends)
13771458

13781459
loss_start = _trim_leading_token_sequences(
13791460
rendered_ids,
@@ -1391,7 +1472,7 @@ def _assistant_mask_from_conversation_turns(
13911472
rendered_mask[loss_start:content_end] = 1.0
13921473
if role in include_start_roles:
13931474
rendered_mask[role_start : role_start + matched_start_length] = 1.0
1394-
if role in include_end_roles:
1475+
if role in include_end_roles and role not in rendered_turn_end_roles:
13951476
rendered_mask[content_end:after_end] = 1.0
13961477

13971478
aligned_mask = _align_assistant_mask_to_padded_ids(rendered_ids, rendered_mask.tolist(), ids, tokenizer)
@@ -1551,6 +1632,12 @@ def _assistant_mask_from_boundary_config(
15511632
for role, raw_variants in boundary_config.role_end_token_variants.items()
15521633
}
15531634
loss_roles = set(boundary_config.loss_roles)
1635+
rendered_turn_end_roles = set(boundary_config.rendered_turn_end_roles)
1636+
if rendered_turn_end_roles:
1637+
# A raw token scan cannot distinguish the end of an unterminated turn
1638+
# from following non-loss roles. Only the conversation-prefix renderer
1639+
# can establish this boundary safely.
1640+
return None
15541641
if not role_start_tokens or not role_end_tokens or not loss_roles:
15551642
raise ValueError(
15561643
"AssistantMaskBoundaryConfig must provide non-empty role_start_tokens, role_end_tokens, and loss_roles."
@@ -1657,7 +1744,7 @@ def build_assistant_loss_mask(
16571744

16581745
mask = _assistant_mask_from_hf_chat_template(example_or_conversation, ids, processor, tokenizer)
16591746
if boundary_config is not None:
1660-
boundary_scan_is_safe = (
1747+
boundary_scan_is_safe = not boundary_config.rendered_turn_end_roles and (
16611748
not conversation
16621749
or _conversation_contains_boundary_tokens(example_or_conversation, tokenizer, boundary_config) is False
16631750
)

‎tests/unit_tests/data/test_conversation_processing.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1439,7 +1439,7 @@ def test_shared_chat_preprocessing_normalizes_sharegpt_roles_before_templating()
14391439
assert tokenized.assistant_mask.any()
14401440

14411441

1442-
def test_normalize_chat_conversation_parses_openai_tool_call_arguments_without_mutating_input():
1442+
def test_normalize_chat_conversation_normalizes_openai_tool_calls_without_mutating_input():
14431443
row = {
14441444
"messages": [
14451445
{"role": "user", "content": "Weather?"},
@@ -1459,7 +1459,9 @@ def test_normalize_chat_conversation_parses_openai_tool_call_arguments_without_m
14591459

14601460
normalized = normalize_chat_conversation(row)
14611461

1462+
assert normalized[1]["content"] == ""
14621463
assert normalized[1]["tool_calls"][0]["function"]["arguments"] == {"city": "Seattle"}
1464+
assert row["messages"][1]["content"] is None
14631465
assert row["messages"][1]["tool_calls"][0]["function"]["arguments"] == '{"city":"Seattle"}'
14641466

14651467

0 commit comments

Comments
 (0)