diff --git a/.claude/skills/add-chat-template/SKILL.md b/.claude/skills/add-chat-template/SKILL.md new file mode 100644 index 000000000..11e086a2a --- /dev/null +++ b/.claude/skills/add-chat-template/SKILL.md @@ -0,0 +1,271 @@ +--- +name: add-chat-template +description: Add or audit an XTuner chat-template implementation and its loss mask from an official Hugging Face model repository or local model directory. Use when supporting a new model under xtuner/v1/data_proto/messages, checking an existing implementation against tokenizer or processor apply_chat_template, deciding which assistant output and stop/EOS tokens receive loss, adding regression tests, or validating that exported training behavior matches vLLM and SGLang inference. +--- + +# Add Chat Template + +## Goal + +Treat the official Hugging Face tokenizer, or processor for a multimodal model, +as the rendering oracle. Produce the same conversation token sequence in +XTuner, supervise model-generated output by default, and supervise the +context-appropriate token that ends each assistant generation. Prove the +behavior through XTuner's public tokenize API and, when available, the target +vLLM and SGLang versions. + +Keep the implementation small and model-local. Do not infer a template or stop +contract from another model in the family. + +## 1. Establish the reference + +1. Record the official HF repo or local directory, immutable revision, tokenizer + class, Transformers version, and whether `trust_remote_code=True` is required. + Follow the repository's requested Python environment; otherwise use + `conda activate pt29_glm1`. +2. Inspect the actual artifacts used by inference: + `tokenizer_config.json`, `chat_template.jinja` or `chat_template`, + `special_tokens_map.json`, `generation_config.json`, and `config.json`. + Inspect `processor_config.json` and remote processor/template code when the + official multimodal path uses `AutoProcessor.apply_chat_template`. +3. Run the bundled inventory against the same revision: + + ```bash + python .claude/skills/add-chat-template/scripts/audit_hf_chat_template.py \ + --trust-remote-code + ``` + + Omit `--trust-remote-code` unless the official model requires it. For + model-specific branches such as tools, thinking, or multimodal inputs, pass a + JSON case file with `--cases `; see the script's `--help` output. + The script inventories the tokenizer side; for a processor-owned multimodal + template, run the corresponding `AutoProcessor` renderer separately. +4. Render every supported branch with the same `tokenizer.apply_chat_template` + or `processor.apply_chat_template` entry point used by the official example. + For multimodal models, do not assume tokenizer-only rendering reproduces + media placeholders or processor expansion. Record all required template + kwargs and their defaults, including + `add_generation_prompt`, `continue_final_message`, tool definitions, + thinking/reasoning flags, named templates, and multimodal options. Do not + copy rendered strings from a model card when the executable tokenizer is + available. +5. When the selected official template contains Jinja `{% generation %}` + regions, call its `apply_chat_template` with `tokenize=True`, + `return_dict=True`, and + `return_assistant_tokens_mask=True`. Use its assistant mask as an additional + oracle, then separately audit stop boundaries: a next-role stop token can sit + outside the official generation region while still belonging to the previous + assistant for SFT. +6. Pin the reference revision in tests or CI configuration. A moving HF branch + is not a stable oracle. + +## 2. Audit before implementing + +Search these integration points first: + +- `xtuner/v1/data_proto/templates/__init__.py` +- `xtuner/v1/data_proto/messages/` +- `xtuner/v1/data_proto/messages/__init__.py` +- `xtuner/v1/datasets/sft_tokenize_fn/openai.py` +- multimodal tokenize functions when applicable +- `tests/chat_template/` and `tests/datasets/` + +If the model already exists, run its public tokenize path before changing it: + +```python +tokenize_fn = OpenaiTokenizeFunctionConfig(chat_template="").build(tokenizer) +result = tokenize_fn({"messages": messages, "tools": tools}) +``` + +Compare the result with the official renderer and inspect existing tests. Find a +concrete failing message sequence before fixing a mismatch. Do not rewrite a +working implementation because its structure looks unusual. + +Use `HybridChatTemplate` plus `ChatMessages` only when fixed role wrappers fully +express the official template. Add a dedicated +`xtuner/v1/data_proto/messages/_chat.py` when rendering depends on +message history, tools, reasoning, content types, or the next role. Follow +`glm52_chat.py` and `qwen35_chat.py` as integration examples, not as rendering +specifications for another model. + +## 3. Write the stop contract first + +Inventory three distinct sources; do not collapse them into one `eos_token`: + +1. **Configured generation stops**: `tokenizer.eos_token_id`, `config.json`, + `generation_config.json`, including integer lists and stop strings. +2. **Template boundaries**: the exact token or token sequence following an + assistant when the next item is user, tool/observation, system/developer, or + the end of the sample. +3. **Engine stops**: the IDs and strings actually installed by the target vLLM + and SGLang versions after loading the exported checkpoint. + +Create an explicit table before writing the mask: + +| Transition | Official serialized boundary | Token IDs | Engine mechanism | Loss when assistant loss is true | +|---|---|---|---|---| +| assistant → user | exact role/EOT sequence | exact IDs | EOS ID, stop ID/string, or parser | yes for the generated stop target | +| assistant tool call → tool | exact observation boundary | exact IDs | EOS ID, stop ID/string, or parser | yes for the generated stop target | +| final assistant → end | exact EOT/EOS, or documented training-only EOS | exact IDs | EOS ID or stop ID/string | yes | + +Verify each token by both encoding the exact boundary and converting its ID back +to a token. Reject unknown-token conversions. Handle a multi-token stop as a +sequence rather than assuming every stop is one special token. + +### Role tokens can belong to the preceding assistant + +Some official templates have no dedicated per-turn EOT. The next role token is +then the token the model must generate to stop the previous assistant. In that +case, keep the official serialized order and assign loss only on that boundary +token or sequence to the preceding assistant: + +```text +assistant answer user text + ^ assistant loss ^ masked +``` + +For a GLM-like protocol this can mean supervising the user boundary after a +normal answer, the observation boundary after a tool call, and the configured +terminal EOS only when no following role supplies a boundary. + +Do not append a generic EOS after every assistant when the official inference +template places a role/observation boundary there. That creates training history +that vLLM and SGLang do not render. Conversely, when the official full-message +renderer omits the final token that the model must generate to terminate, +append exactly that one terminal token for SFT and document this single +training-only suffix. + +Configured EOS IDs are alternatives, not a sequence to append together. +Supervise the actual context-specific stop token present in each training path. + +## 4. Implement rendering and loss together + +Prefer one cohesive renderer that returns rendered text plus a character- or +token-level loss mask. Keep context-dependent boundary ownership beside the +rendering branch that emits the boundary; avoid a collection of shallow helper +functions. + +Apply these defaults: + +- Treat `assistant` messages with no `loss` field as `loss=True`. +- Mask system/developer/user/tool inputs and role scaffolding unless the stop + contract proves that a boundary is generated by the preceding assistant. +- Supervise all model-generated assistant components: visible content, + reasoning that remains in the rendered sample, reasoning closing syntax, + tool calls, and the appropriate stop target. +- Treat `loss=False` as masking the entire assistant output and its stop target. +- Mask `add_generation_prompt` completely; it is prompt scaffolding, not a + demonstrated model output. +- Preserve explicit model behavior such as clearing historical reasoning only + when the official template does so for the chosen kwargs. + +If character offsets are used, verify special-token and Unicode boundaries with +the real tokenizer. A token overlapping a loss span needs a deliberate, +tested rule. Keep a slow, independent token-level oracle for regression tests; +do not call the production renderer from both sides of the same parity test. + +Maintain exact rendering parity before any documented training suffix: + +```mermaid +flowchart LR + A["Official HF messages"] --> B["tokenizer.apply_chat_template"] + A --> C["XTuner public tokenize API"] + B --> D["Exact text and token IDs"] + C --> D + D --> E["Optional one terminal SFT stop"] + E --> F["Loss-mask contract"] +``` + +Do not use decoded-text equality as the sole check: decoding can hide a token-ID +mismatch. + +## 5. Register the smallest integration + +For a dedicated message implementation, update only the required seams: + +1. Export the message class from `xtuner/v1/data_proto/messages/__init__.py`. +2. Add the public name and serving stop metadata to `CHAT_TEMPLATE_MAP`. +3. Dispatch that name in `OpenaiTokenizeFunction`; update multimodal dispatch + only if the model uses it. +4. Update typed CLI/config literals only where they restrict the new name. + +Keep `CHAT_TEMPLATE_MAP.stop_words` consistent with the discovered contract, +but do not treat it as proof of engine behavior. vLLM and SGLang load their +runtime stops from the exported HF artifacts and request/engine configuration. + +## 6. Add good regression tests + +Use the real official tokenizer/processor and XTuner's public tokenize +function. Mock only unavailable external services. Cover every supported +template branch with the smallest useful matrix: + +- user → assistant, with `loss` omitted; +- assistant → user in a multi-turn conversation; +- assistant tool call → tool/observation → assistant; +- final assistant; +- `loss=False` followed by another role and at sample end; +- system/developer input; +- `add_generation_prompt=True`; +- reasoning enabled/disabled and historical reasoning behavior; +- tools and multimodal content when supported; +- repeated assistant text, Unicode, and adjacent special tokens to expose + faulty substring or offset alignment. + +Assert public behavior: + +1. Official tokenizer/processor rendered text equals XTuner rendered text, + except for the one explicitly documented final SFT suffix if required. +2. Official token IDs equal XTuner `input_ids` before that suffix. +3. `len(input_ids) == len(labels)` and every label is either the matching input + ID or `IGNORE_INDEX`. +4. Assistant output receives loss by default; `loss=False` removes all of it. +5. Each transition's exact stop token IDs receive or do not receive loss as the + stop table specifies. +6. User/tool content and generation prompts remain masked. +7. The fast implementation matches an independent slow oracle. + +Test token positions and IDs directly. A substring assertion can supplement but +must not replace the loss assertion. Avoid snapshots that pass after both the +renderer and expected string are changed to the same incorrect value. + +## 7. Validate vLLM and SGLang + +Resolve exact target versions. Prefer the user's production versions; otherwise +inspect the current installed versions and latest stable official source. For +each engine: + +1. Load the same exported tokenizer, `config.json`, and + `generation_config.json`. +2. Compare the engine's chat-render/tokenize output with the official + tokenizer/processor renderer for every stop-table transition. +3. Inspect or log the final stop token IDs/strings installed in the generation + request. Distinguish tokenizer EOS, model generation-config EOS, request stop + strings, and model-specific parsers. +4. Run generation smoke tests for a normal final answer and a tool call. Use a + sufficiently high output limit and assert the engine reports a stop rather + than a length limit, with no role leakage. + +Obtain the repository GPU lock before a local GPU server or generation test. If +the exact runtime cannot be executed, audit the exact tagged engine source and +report the result as a static source audit, not runtime compatibility. + +Switching from Chat Completions to Responses does not remove this requirement: +for ordinary HF models both interfaces ultimately render model messages through +the chat template. + +## 8. Completion checklist + +Finish only when the report includes: + +- official HF repo, revision, and Transformers version; +- implementation path and whether it is generic or dedicated; +- rendering-parity matrix; +- stop-contract table with token strings, IDs, and loss ownership; +- tests proving default assistant loss and `loss=False`; +- vLLM/SGLang versions and runtime or exact-tag source validation; +- every intentional difference from the official renderer, normally at most a + final training stop suffix. + +If auditing an existing implementation, state which requirements already had +tests, which were missing, and the concrete mismatches found. Do not claim the +implementation is correct merely because a test file exists. diff --git a/.claude/skills/add-chat-template/agents/openai.yaml b/.claude/skills/add-chat-template/agents/openai.yaml new file mode 100644 index 000000000..4873cfae0 --- /dev/null +++ b/.claude/skills/add-chat-template/agents/openai.yaml @@ -0,0 +1,4 @@ +interface: + display_name: "Add Chat Template" + short_description: "Add and verify XTuner chat templates" + default_prompt: "Use $add-chat-template to add or audit XTuner chat-template rendering, loss masks, and stop-token behavior for this official Hugging Face model." diff --git a/.claude/skills/add-chat-template/scripts/audit_hf_chat_template.py b/.claude/skills/add-chat-template/scripts/audit_hf_chat_template.py new file mode 100644 index 000000000..779ffa405 --- /dev/null +++ b/.claude/skills/add-chat-template/scripts/audit_hf_chat_template.py @@ -0,0 +1,233 @@ +#!/usr/bin/env python3 +"""Inspect an official HF chat template, rendered cases, and stop-token sources.""" + +from __future__ import annotations + +import argparse +import json +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import transformers +from transformers import AutoConfig, AutoTokenizer, GenerationConfig + + +DEFAULT_CASES = [ + { + "name": "final_assistant", + "messages": [ + {"role": "user", "content": "Question one"}, + {"role": "assistant", "content": "Answer one"}, + ], + "kwargs": {"add_generation_prompt": False}, + }, + { + "name": "assistant_to_user", + "messages": [ + {"role": "user", "content": "Question one"}, + {"role": "assistant", "content": "Answer one"}, + {"role": "user", "content": "Question two"}, + ], + "kwargs": {"add_generation_prompt": False}, + }, + { + "name": "generation_prompt", + "messages": [{"role": "user", "content": "Question one"}], + "kwargs": {"add_generation_prompt": True}, + }, +] + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser( + description=( + "Load an official Hugging Face tokenizer and print a JSON audit of " + "its selected chat template, EOS sources, special tokens, and rendered cases." + ) + ) + parser.add_argument("model", help="Official HF repo id or local model directory") + parser.add_argument("--revision", help="Immutable HF revision or commit") + parser.add_argument("--trust-remote-code", action="store_true") + parser.add_argument("--template", help="Named template or explicit template passed to apply_chat_template") + parser.add_argument( + "--cases", + type=Path, + help=( + "JSON file containing a list of cases. Each case has name, messages, " + "optional tools, and optional kwargs for apply_chat_template." + ), + ) + return parser.parse_args() + + +def load_cases(path: Path | None) -> list[dict[str, Any]]: + if path is None: + return DEFAULT_CASES + cases = json.loads(path.read_text()) + if not isinstance(cases, list): + raise ValueError("--cases must contain a JSON list") + return cases + + +def normalize_ids(value: Any) -> list[int]: + if value is None: + return [] + if isinstance(value, int): + return [value] + return [int(item) for item in value] + + +def id_details(tokenizer, value: Any) -> list[dict[str, Any]]: + return [ + { + "id": token_id, + "token": tokenizer.convert_ids_to_tokens(token_id), + "decoded": tokenizer.decode([token_id], skip_special_tokens=False), + } + for token_id in normalize_ids(value) + ] + + +def load_generation_config(model: str, load_kwargs: dict[str, Any]) -> tuple[Any, str | None]: + try: + return GenerationConfig.from_pretrained(model, **load_kwargs), None + except Exception as exc: + return None, f"{type(exc).__name__}: {exc}" + + +def render_case(tokenizer, case: dict[str, Any], template: str | None) -> dict[str, Any]: + kwargs = dict(case.get("kwargs", {})) + if "tools" in case: + kwargs["tools"] = case["tools"] + if template is not None: + kwargs["chat_template"] = template + + try: + text = tokenizer.apply_chat_template(case["messages"], tokenize=False, **kwargs) + applied_ids = tokenizer.apply_chat_template(case["messages"], tokenize=True, **kwargs) + if isinstance(applied_ids, Mapping): + applied_ids = applied_ids["input_ids"] + encoded_ids = tokenizer.encode(text, add_special_tokens=False) + + assistant_mask = None + assistant_mask_error = None + try: + masked = tokenizer.apply_chat_template( + case["messages"], + tokenize=True, + return_dict=True, + return_assistant_tokens_mask=True, + **kwargs, + ) + assistant_mask = masked.get("assistant_masks") + if hasattr(assistant_mask, "tolist"): + assistant_mask = assistant_mask.tolist() + except Exception as exc: + assistant_mask_error = f"{type(exc).__name__}: {exc}" + + return { + "name": case.get("name", "unnamed"), + "kwargs": kwargs, + "rendered": text, + "apply_chat_template_ids": applied_ids, + "encode_rendered_ids": encoded_ids, + "ids_match": applied_ids == encoded_ids, + "official_assistant_mask": assistant_mask, + "official_assistant_mask_error": assistant_mask_error, + "tokens": tokenizer.convert_ids_to_tokens(applied_ids), + "decoded": tokenizer.decode(applied_ids, skip_special_tokens=False), + } + except Exception as exc: + return { + "name": case.get("name", "unnamed"), + "kwargs": kwargs, + "error": f"{type(exc).__name__}: {exc}", + } + + +def main() -> None: + args = parse_args() + load_kwargs = { + "revision": args.revision, + "trust_remote_code": args.trust_remote_code, + } + load_kwargs = {key: value for key, value in load_kwargs.items() if value is not None} + + tokenizer = AutoTokenizer.from_pretrained(args.model, **load_kwargs) + config = AutoConfig.from_pretrained(args.model, **load_kwargs) + generation_config, generation_config_error = load_generation_config(args.model, load_kwargs) + selected_template = tokenizer.get_chat_template(chat_template=args.template) + + raw_template = getattr(tokenizer, "chat_template", None) + template_sources = [selected_template] + if isinstance(raw_template, dict): + template_sources.extend(str(value) for value in raw_template.values()) + elif raw_template is not None: + template_sources.append(str(raw_template)) + template_text = "\n".join(template_sources) + + eos_sources = { + "tokenizer": { + "value": tokenizer.eos_token, + "ids": id_details(tokenizer, tokenizer.eos_token_id), + }, + "config": { + "raw": getattr(config, "eos_token_id", None), + "ids": id_details(tokenizer, getattr(config, "eos_token_id", None)), + }, + "generation_config": None, + } + if generation_config is not None: + eos_sources["generation_config"] = { + "raw": generation_config.eos_token_id, + "ids": id_details(tokenizer, generation_config.eos_token_id), + "stop_strings": getattr(generation_config, "stop_strings", None), + } + + special_tokens = [] + for token in dict.fromkeys(tokenizer.all_special_tokens): + token_id = tokenizer.convert_tokens_to_ids(token) + special_tokens.append( + { + "token": token, + "id": token_id, + "appears_in_selected_or_raw_template": token in template_text, + "encoded_ids": tokenizer.encode(token, add_special_tokens=False), + } + ) + + added_tokens = [] + for token_id, token in sorted(tokenizer.added_tokens_decoder.items()): + token_text = str(token) + added_tokens.append( + { + "token": token_text, + "id": token_id, + "special": bool(getattr(token, "special", False)), + "appears_in_selected_or_raw_template": token_text in template_text, + "encoded_ids": tokenizer.encode(token_text, add_special_tokens=False), + } + ) + + report = { + "model": args.model, + "requested_revision": args.revision, + "resolved_revision": getattr(config, "_commit_hash", None) or tokenizer.init_kwargs.get("_commit_hash"), + "transformers_version": transformers.__version__, + "tokenizer_class": type(tokenizer).__name__, + "model_type": getattr(config, "model_type", None), + "selected_template": selected_template, + "raw_chat_template": raw_template, + "special_tokens_map": tokenizer.special_tokens_map, + "eos_sources": eos_sources, + "generation_config_error": generation_config_error, + "special_tokens": special_tokens, + "added_tokens": added_tokens, + "cases": [render_case(tokenizer, case, args.template) for case in load_cases(args.cases)], + } + print(json.dumps(report, ensure_ascii=False, indent=2, default=str)) + + +if __name__ == "__main__": + main() diff --git a/tests/datasets/test_glm52_openai_tokenize_fn.py b/tests/datasets/test_glm52_openai_tokenize_fn.py index f62aa751b..f62a4be92 100644 --- a/tests/datasets/test_glm52_openai_tokenize_fn.py +++ b/tests/datasets/test_glm52_openai_tokenize_fn.py @@ -1,7 +1,8 @@ """GLM-5.2 OpenAI 对话分词行为测试。 TestGlm52Rendering - test_plain_text_matches_hf_and_golden_labels: 普通对话与 HF 模板及慢速 golden 对齐。 + test_all_generation_eos_are_supervised_at_assistant_boundaries: 三类停止 token 均按边界参与训练。 + test_plain_text_matches_hf_plus_final_eos_and_golden_labels: 普通对话仅比 HF 推理模板多最终 EOS。 test_multiturn_reasoning_defaults_to_preserved_and_can_be_cleared: 默认保留历史推理且支持显式清除。 test_tools_and_loss_switch_follow_template_masking: 工具对话与 loss 开关生成正确标签。 TestGlm52MessageOptions @@ -54,8 +55,49 @@ def _label_flags_for_span(tokenizer, text, labels, substring): class TestGlm52Rendering: - def test_plain_text_matches_hf_and_golden_labels(self, tokenizer, tokenize_fn): - # 验证普通对话的 token 与标签同时对齐 HF 模板和独立慢速实现。 + def test_all_generation_eos_are_supervised_at_assistant_boundaries(self, tokenizer, tokenize_fn): + # user/observation 作为轮间停止目标,只有无后继角色的 assistant 才补 endoftext。 + messages = [ + {"role": "user", "content": "First question"}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Call a tool"}, + { + "role": "assistant", + "content": "Calling now", + "tool_calls": [{"function": {"name": "lookup", "arguments": {"key": "value"}}}], + }, + {"role": "tool", "content": "result"}, + {"role": "assistant", "content": "Unsupervised answer", "loss": False}, + {"role": "user", "content": "Final question"}, + {"role": "assistant", "content": "Final answer"}, + ] + + tokenized = tokenize_fn({"messages": messages}) + rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False) + hf_rendered = _render_from_hf(tokenizer, messages, add_generation_prompt=False) + stop_ids = { + token: tokenizer.convert_tokens_to_ids(token) for token in ("<|endoftext|>", "<|user|>", "<|observation|>") + } + stop_labels = { + token: [ + tokenized["labels"][index] + for index, token_id in enumerate(tokenized["input_ids"]) + if token_id == stop_id + ] + for token, stop_id in stop_ids.items() + } + + assert tokenizer.eos_token == "<|endoftext|>" + assert rendered == hf_rendered + tokenizer.eos_token + assert "First answer<|user|>" in rendered + assert "<|observation|>" in rendered + assert rendered.endswith("Final answer<|endoftext|>") + assert stop_labels["<|user|>"] == [-100, stop_ids["<|user|>"], -100] + assert stop_labels["<|observation|>"] == [stop_ids["<|observation|>"]] + assert stop_labels["<|endoftext|>"] == [stop_ids["<|endoftext|>"]] + + def test_plain_text_matches_hf_plus_final_eos_and_golden_labels(self, tokenizer, tokenize_fn): + # HF 推理模板不带最终 EOS;XTuner 在最后一个 assistant 末尾显式补齐。 messages = [ {"role": "user", "content": "Hello!"}, {"role": "assistant", "content": "Hi there."}, @@ -64,8 +106,9 @@ def test_plain_text_matches_hf_and_golden_labels(self, tokenizer, tokenize_fn): tokenized = tokenize_fn({"messages": messages}) slow_input_ids, slow_labels = glm52_tokenize_fn_slowspeed(tokenizer, messages) - rendered = _render_from_hf(tokenizer, messages, add_generation_prompt=False) - assert tokenized["input_ids"] == tokenizer.encode(rendered, add_special_tokens=False) + hf_rendered = _render_from_hf(tokenizer, messages, add_generation_prompt=False) + rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False) + assert rendered == hf_rendered + tokenizer.eos_token assert tokenized["input_ids"] == slow_input_ids assert tokenized["labels"] == slow_labels assert ( @@ -73,7 +116,7 @@ def test_plain_text_matches_hf_and_golden_labels(self, tokenizer, tokenize_fn): [label for label in tokenized["labels"] if label != -100], skip_special_tokens=False, ) - == "Hi there." + == "Hi there.<|endoftext|>" ) def test_multiturn_reasoning_defaults_to_preserved_and_can_be_cleared(self, tokenizer, tokenize_fn): @@ -86,11 +129,13 @@ def test_multiturn_reasoning_defaults_to_preserved_and_can_be_cleared(self, toke ] tokenized = tokenize_fn({"messages": messages}) - rendered = _render_from_hf(tokenizer, messages, add_generation_prompt=False) + hf_rendered = _render_from_hf(tokenizer, messages, add_generation_prompt=False) + rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False) slow_input_ids, slow_labels = glm52_tokenize_fn_slowspeed(tokenizer, messages) assert "old trace" in rendered - assert tokenized["input_ids"] == tokenizer.encode(rendered, add_special_tokens=False) + assert rendered == hf_rendered + tokenizer.eos_token + assert rendered.count(tokenizer.eos_token) == 1 assert tokenized["input_ids"] == slow_input_ids assert tokenized["labels"] == slow_labels assert all(_label_flags_for_span(tokenizer, rendered, tokenized["labels"], "old trace")) @@ -99,12 +144,13 @@ def test_multiturn_reasoning_defaults_to_preserved_and_can_be_cleared(self, toke assert all(_label_flags_for_span(tokenizer, rendered, tokenized["labels"], "Final answer.")) cleared = Glm52ChatMessages(messages=messages).tokenize(tokenizer, clear_thinking=True) - cleared_rendered = _render_from_hf( + cleared_hf_rendered = _render_from_hf( tokenizer, messages, add_generation_prompt=False, clear_thinking=True, ) + cleared_rendered = tokenizer.decode(cleared["input_ids"], skip_special_tokens=False) cleared_slow_ids, cleared_slow_labels = glm52_tokenize_fn_slowspeed( tokenizer, messages, @@ -112,7 +158,8 @@ def test_multiturn_reasoning_defaults_to_preserved_and_can_be_cleared(self, toke ) assert "old trace" not in cleared_rendered - assert cleared["input_ids"] == tokenizer.encode(cleared_rendered, add_special_tokens=False) + assert cleared_rendered == cleared_hf_rendered + tokenizer.eos_token + assert cleared_rendered.count(tokenizer.eos_token) == 1 assert cleared["input_ids"] == cleared_slow_ids assert cleared["labels"] == cleared_slow_labels assert not any(_label_flags_for_span(tokenizer, cleared_rendered, cleared["labels"], "Old answer.")) @@ -152,17 +199,22 @@ def test_tools_and_loss_switch_follow_template_masking(self, tokenizer, tokenize ] tokenized = tokenize_fn({"messages": messages, "tools": tools}) - rendered = _render_from_hf( + hf_rendered = _render_from_hf( tokenizer, messages, tools=tools, add_generation_prompt=False, ) + rendered = tokenizer.decode(tokenized["input_ids"], skip_special_tokens=False) slow_input_ids, slow_labels = glm52_tokenize_fn_slowspeed(tokenizer, messages, tools=tools) - assert tokenized["input_ids"] == tokenizer.encode(rendered, add_special_tokens=False) + assert rendered == hf_rendered + tokenizer.eos_token + assert rendered.count(tokenizer.eos_token) == 1 assert tokenized["input_ids"] == slow_input_ids assert tokenized["labels"] == slow_labels + observation_id = tokenizer.convert_tokens_to_ids("<|observation|>") + assert tokenized["labels"][tokenized["input_ids"].index(observation_id)] == observation_id + assert tokenized["labels"][-1] == -100 assert not any( _label_flags_for_span(tokenizer, rendered, tokenized["labels"], '"description": "Gets the weather."') ) @@ -214,15 +266,15 @@ def test_default_system_is_inserted_or_replaced(self, tokenizer): {"role": "system", "content": "Default system instruction."}, *inserted_messages, ] - rendered = _render_from_hf( + hf_rendered = _render_from_hf( tokenizer, expected_messages, add_generation_prompt=False, ) + rendered = tokenizer.decode(inserted["input_ids"], skip_special_tokens=False) - expected_ids = tokenizer.encode(rendered, add_special_tokens=False) - assert inserted["input_ids"] == expected_ids - assert replaced["input_ids"] == expected_ids + assert rendered == hf_rendered + tokenizer.eos_token + assert replaced["input_ids"] == inserted["input_ids"] assert not any( _label_flags_for_span( tokenizer, diff --git a/xtuner/v1/data_proto/messages/glm52_chat.py b/xtuner/v1/data_proto/messages/glm52_chat.py index 1c7c34966..f642c8a4f 100644 --- a/xtuner/v1/data_proto/messages/glm52_chat.py +++ b/xtuner/v1/data_proto/messages/glm52_chat.py @@ -12,6 +12,8 @@ _MEDIA_TYPES = {"image", "image_url", "video", "video_url", "audio", "audio_url", "input_audio"} +_END_OF_TEXT = "<|endoftext|>" +_NEXT_ROLE_STOP_TOKENS = {"user": "<|user|>", "tool": "<|observation|>"} def _visible_text(content: Any) -> str: @@ -167,10 +169,14 @@ def append(value: str, loss: bool) -> None: if message.get("role") == "user": last_user_index = index + # user/observation 角色 token 同时是生成停止目标,其 loss 归属前一个 assistant。 + previous_assistant_loss = False for index, message in enumerate(messages): role = message.get("role") if role == "user": - append(f"<|user|>{_visible_text(message.get('content', ''))}", False) + boundary_loss = index > 0 and messages[index - 1].get("role") == "assistant" and previous_assistant_loss + append(_NEXT_ROLE_STOP_TOKENS[role], boundary_loss) + append(_visible_text(message.get("content", "")), False) elif role == "system": append(f"<|system|>{_visible_text(message.get('content', ''))}", False) elif role == "assistant": @@ -191,9 +197,17 @@ def append(value: str, loss: bool) -> None: append(content.strip(), loss) if message.get("tool_calls"): append(_render_tool_calls(message["tool_calls"]), loss) + previous_assistant_loss = loss + next_role = messages[index + 1].get("role") if index + 1 < len(messages) else None + if next_role not in _NEXT_ROLE_STOP_TOKENS: + # 没有 user/observation 角色边界时,用标准 EOS 结束 assistant。 + append(_END_OF_TEXT, loss) elif role == "tool": if index == 0 or messages[index - 1].get("role") != "tool": - append("<|observation|>", False) + boundary_loss = ( + index > 0 and messages[index - 1].get("role") == "assistant" and previous_assistant_loss + ) + append(_NEXT_ROLE_STOP_TOKENS[role], boundary_loss) append(_render_tool_result(message.get("content", ""), tools), False) if add_generation_prompt: @@ -255,24 +269,21 @@ def glm52_tokenize_fn_slowspeed( ) -> tuple[list[int], list[int]]: """慢速 golden 参考实现:基于 token 级别前缀 diff 对齐 labels。 - 这份逻辑刻意保持和旧 `golden_tokenize_fn.py` 一致,用来校验 fast path: + 渲染复用 GLM 多停止边界的 SFT 语义,标签仍通过独立的 token 前缀 diff 计算: 1. 先渲染完整对话,得到唯一的 total_ids 参考序列。 2. 对每条需要 loss 的 assistant 消息,渲染其历史前缀并加 generation prompt。 3. 再渲染“历史 + 当前 assistant”,用 token 前缀差得到当前 assistant 应监督的 suffix。 4. 从上次命中位置开始在 total_ids 里顺序查找 suffix,找到后复制到 labels。 """ - # 显式传递确定值,避免 Jinja 将已定义的 None 解释成 False。 - hf_kwargs: dict[str, Any] = dict( - tokenize=False, - add_generation_prompt=add_generation_prompt, + # SFT 渲染保持官方轮间格式,并在没有后继角色边界时补标准 EOS。 + full_text, _ = render_glm52_chat( + messages, tools=tools, + add_generation_prompt=add_generation_prompt, enable_thinking=enable_thinking, reasoning_effort=reasoning_effort, clear_thinking=clear_thinking, ) - hf_kwargs.update(kwargs) - - full_text = tokenizer.apply_chat_template(messages, **hf_kwargs) total_ids = tokenizer.encode(full_text, add_special_tokens=False) labels = [IGNORE_INDEX] * len(total_ids) @@ -283,17 +294,25 @@ def glm52_tokenize_fn_slowspeed( continue # 历史前缀以 generation prompt 结束;prefix 之后的 token 才是当前 assistant 生成区间。 - prompt_kwargs = dict(hf_kwargs) - prompt_kwargs["add_generation_prompt"] = True - prompt_kwargs["tools"] = tools if index == 0 else None - prefix_text = tokenizer.apply_chat_template(messages[:index], **prompt_kwargs) + prefix_text, _ = render_glm52_chat( + messages[:index], + tools=tools if index == 0 else None, + add_generation_prompt=True, + enable_thinking=enable_thinking, + reasoning_effort=reasoning_effort, + clear_thinking=clear_thinking, + ) # 当前截断渲染可能和 full render 不同,例如历史 thinking 会被模板清掉;是否能在 full # render 中匹配上 suffix,正是 golden 语义的一部分。 - message_kwargs = dict(hf_kwargs) - message_kwargs["add_generation_prompt"] = False - message_kwargs["tools"] = tools if index == 0 else None - message_text = tokenizer.apply_chat_template([m.copy() for m in messages[: index + 1]], **message_kwargs) + message_text, _ = render_glm52_chat( + [m.copy() for m in messages[: index + 1]], + tools=tools if index == 0 else None, + add_generation_prompt=False, + enable_thinking=enable_thinking, + reasoning_effort=reasoning_effort, + clear_thinking=clear_thinking, + ) prefix_ids = tokenizer.encode(prefix_text, add_special_tokens=False) message_ids = tokenizer.encode(message_text, add_special_tokens=False) @@ -301,6 +320,11 @@ def glm52_tokenize_fn_slowspeed( if not content_ids: continue + next_role = messages[index + 1].get("role") if index + 1 < len(messages) else None + if next_role in _NEXT_ROLE_STOP_TOKENS: + # 截断渲染以 endoftext 结尾;完整对话改用实际的下一角色停止 token。 + content_ids[-1] = tokenizer.convert_tokens_to_ids(_NEXT_ROLE_STOP_TOKENS[next_role]) + # 在完整 token 序列中做 token 级绝对对齐,避免字符 offset 对特殊 token 的边界解释差异。 for start in range(curr_ptr, len(total_ids) - len(content_ids) + 1): if total_ids[start : start + len(content_ids)] == content_ids: