feat(dataset): Kimi-K3 conversation renderer - #164
Conversation
Kimi-K3 serializes conversations as XTML, a nested structure a ChatTemplate cannot express, and the only correct serializer is the tokenizer shipped with the checkpoint. K3Renderer therefore delegates all prompt construction to apply_chat_template and limits itself to the two things the tokenizer does not provide: deriving TorchSpec's assistant loss mask from the resulting token stream, and applying a truncation that keeps the mask aligned with it. Assistant spans are located structurally rather than by string matching. The outer assistant header is derived once at construction time by probing the tokenizer, and supervision covers each assistant body between its <|sep|> and <|end_of_msg|>, excluding both. The same header/end pair is exposed through get_assistant_token_ids so a mask recomputed at training time against the inference engine's token IDs reproduces the offline mask exactly, including after multimodal placeholder expansion shifts every offset. Two properties are enforced rather than assumed, since both fail silently: - The constructor verifies that a user-supplied "<|open|>" or "<|sep|>" literal is encoded as text and not as a control token. Without this, message content could fabricate an assistant span and pull unsupervised text into the loss. - generation_config is validated as a unit. thinking_effort must be one of the three levels the chat template implements (low/high/max) when thinking is on and null when it is off, and reasoning_effort must agree with it. A mismatched or unsupported level would otherwise render a prompt that disagrees with the effort the corpus was generated at. Renderer-agnostic behavior (registry lookup, tools and generation_config forwarding, multimodal mask deferral) is already covered by tests/test_renderer_registry.py, so these tests assert only K3 specifics against a stub tokenizer implementing the remote-code interface. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cd98dcd10c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
The renderer forced thinking_effort="max" whenever a row carried no generation_config, which is the default dataset path. The tokenizer treats any non-None effort as an explicit rendering option and injects an internal thinking-effort system message for it, so default training prompts carried a system message that the checkpoint does not render at serving time for a request that named no effort -- a silent train/serve prompt mismatch on the most common path. Default to None instead, in both _resolve_generation_config and the _apply_chat_template signature, so the tokenizer applies its own default. The constructor's probes go through the same default, so the assistant header is now derived against the same rendering the common path produces. An explicit generation_config is unchanged and still validated: thinking_effort remains required when thinking is on, so a row that means to pin an effort cannot lose it to a typo. Signed-off-by: torchspec-bot <262938024+torchspec-bot@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dbd6350a79
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| body_start += 1 | ||
| if body_start == len(input_ids): | ||
| raise ValueError("Kimi-K3 assistant header is missing its <|sep|> terminator") | ||
| body_start += 1 |
There was a problem hiding this comment.
With the real K3 tokenizer, the outer assistant message separator is immediately followed by a nested channel opener such as <|open|>think<|sep|> (or <|open|>response<|sep|> when thinking is disabled), and serving's generation prompt already includes that opener before the model starts producing tokens. Starting the loss span right after the outer <|sep|> therefore supervises prompt-prefix XTML wrapper tokens on every assistant turn instead of only the generated reasoning/response stream, which corrupts K3 training labels; the span should advance past the same inner channel prefix that add_generation_prompt provides.
Useful? React with 👍 / 👎.
| add_generation_prompt=False, | ||
| thinking=thinking, | ||
| preserve_thinking=True, | ||
| thinking_effort=thinking_effort, |
There was a problem hiding this comment.
For rows whose generation_config includes K3 rendering controls such as tool_choice, response_format, or response_schema, render() validates the thinking fields but this tokenizer call forwards only thinking_effort, so those controls are silently dropped. K3's renderer turns those options into internal XTML system messages, so training prompts for required/no-tools or JSON-constrained samples no longer match the serving request that produced the row; preserve and forward the remaining rendering kwargs after validating the effort fields.
Useful? React with 👍 / 👎.
1 participant
Summary
Adds
K3Renderer, the first concreteConversationRendereron top of the framework merged in #163. The registry has shipped empty until now; this registerskimi-k3.Kimi-K3 serializes conversations as XTML — a nested structure with
<|open|>/<|sep|>/<|close|>/<|end_of_msg|>control tokens — which aChatTemplatecannot express. The only correct serializer is the tokenizer shipped with the checkpoint, so the renderer delegates all prompt construction toapply_chat_templateand owns only the two things the tokenizer does not provide:<|sep|>and<|end_of_msg|>, excluding both.get_assistant_token_ids()returns that same header/end pair, so a mask recomputed at training time against the IDs the inference engine actually produced reproduces the offline mask exactly — including after multimodal placeholder expansion shifts every offset. K3's derived header deliberately stops before the outer<|sep|>, so the returned header appends it to keep that structural token out of the loss.Two failure modes are enforced rather than assumed
Both are silent if left unchecked, which is why they're in the constructor / validation path rather than in review notes:
<|open|><|close|><|sep|><|end_of_msg|>literal and fails if any of them encode as control tokens rather than text. Without this, message content could fabricate an assistant span and pull unsupervised text into the loss.generation_configconsistency.thinking_effortmust be one of the three levels the chat template actually implements (low/high/max) whenthinkingis on, andnullwhen it is off;reasoning_effortmust agree with it. Notably"medium"is rejected: it is a common level in generic four-level effort scales but K3's template does not implement it, so accepting it would render a prompt that disagrees with the effort the corpus was generated at. Thinking-disabled rendering also refuses to run if any assistant message still carries non-empty reasoning, rather than silently discarding it.Reviewability caveat
K3Rendererdepends on the checkpoint'strust_remote_codetokenizer exposing aspecial_tokensdict and acceptingthinking/preserve_thinking/thinking_effortkwargs. That is not a publictransformerscontract, so the class cannot be exercised in CI against a real checkpoint — the tests drive it through a stub tokenizer implementing that interface. Flagging this explicitly since it bounds what the test suite can prove.Test plan
tests/test_kimi_k3_renderer.py— 21 tests, all passing.tests/test_renderer_registry.pystill passes (17 tests); nothing there assumed an empty registry.test_renderer_registry.py,test_pretokenized_dataset.py,test_conversation_normalization.py,test_kimi_k3_renderer.py= 60 passed.ruff checkandruff format --checkclean on all four files.Tests cover the K3-specific surface only — XTML special-token contract, header derivation from the remote tokenizer, structural span detection, last-turn selection before truncation, the
thinking/thinking_effortvalidation matrix, and the control-token injection defense. Renderer-agnostic behavior (registry lookup,tools/generation_configforwarding, multimodal mask deferral) is already covered bytests/test_renderer_registry.pyfrom #163 and is not re-tested here through a K3 stub.