Studio: fix the manual response-template markers that never match their rendered templates - #7062
Conversation
Six template families in TEMPLATE_TO_RESPONSES_MAPPER shipped markers that
never match what their chat templates actually render, so the manual
train_on_completions path masked every assistant token and the run died on
the all-labels-masked safety net:
- mistral, llama: '[INST] ' / ' [/INST]' - the surrounding spaces fold into
the neighbouring tokens ('[INST]'/'[/INST]' are single special tokens in
Mistral v0.3, SentencePiece pieces in Llama-2), so the padded strings
never match. Now '[INST]' / '[/INST]'.
- starling: trailing space after 'GPT4 Correct Assistant:' folds into the
next content token. Now no trailing space.
- glm: '[gMASK]<sop>' renders once at text start, never before later user
turns, and '<think>' is generation scaffolding rendered as a lone
'</think>' on non-final turns. Now '<|user|>' / '<|assistant|>'.
- qwen3-thinking: '<think>' is stripped from non-final assistant turns
(Qwen3-Thinking-2507) and never rendered by QwQ. Now the bare assistant
header, matching the other qwen entries.
- zephyr: role tags are plain text and SentencePiece tokenizes them
differently at text start than after '</s>' + newline mid-conversation;
the markers need the leading newline anchor. Now '\n<|user|>\n' /
'\n<|assistant|>\n'.
Validated token-level on each family's representative tokenizer with a
two-turn fixture plus system message: user and system content fully masked,
every assistant turn trained, and the final EOS label never -100. The
fixed mistral, llama, starling and glm markers produce labels identical to
zoo auto-detection; qwen3-thinking differs only in one turn-separator
newline token. All 22 unchanged entries produce byte-identical labels to
before this change.
Adds tests/test_response_template_markers.py pinning the fixed and key
unchanged marker literals (dependency-free) plus token-level masking checks
that skip when tokenizers or unsloth_zoo are unavailable offline.
for more information, see https://pre-commit.ci
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request corrects several template markers in TEMPLATE_TO_RESPONSES_MAPPER (for qwen3-thinking, mistral, llama, zephyr, starling, and glm) to ensure they align with actual tokenizer outputs and prevent unintended masking of assistant tokens. It also adds a new test suite (test_response_template_markers.py) to verify these markers literally and at the token level. The feedback recommends using a context manager with explicit UTF-8 encoding when opening the tokenizer configuration file in the test helper to prevent unclosed file handles and platform-dependent encoding issues.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| return AutoTokenizer.from_pretrained(repo) | ||
| except OSError as e: | ||
| pytest.skip(f"tokenizer {repo} unavailable (offline?): {e}") | ||
| except Exception: |
There was a problem hiding this comment.
The file tokenizer_config.json is opened using open() without a with statement or explicitly closing it. This can lead to unclosed file handles, especially if an exception occurs during JSON parsing. Additionally, it is recommended to specify encoding="utf-8" when opening text files to prevent platform-dependent encoding issues (e.g., on Windows).
Using a with statement ensures the file is properly closed after reading.
| except Exception: | |
| with open(hf_hub_download(repo, "tokenizer_config.json"), encoding="utf-8") as f: | |
| cfg = _json.load(f) |
There was a problem hiding this comment.
Fixed in 51925d2: the config read now uses a with block and encoding utf-8. The encoding half is the real risk since chat templates in tokenizer_config.json are rarely ASCII-only and the Windows default codec would fail the GLM fallback loader.
Chat templates in tokenizer_config.json are rarely ASCII-only, so the default locale codec could fail the GLM fallback loader on Windows.
for more information, see https://pre-commit.ci
|
@codex review |
1 similar comment
|
@codex review |
|
Codex Review: Didn't find any major issues. Nice work! Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
|
Checked the six families against the real 1. "llama": {
- "instruction": "[INST]",
+ "instruction": "<s>[INST]",
"response": "[/INST]",
},If that holds up on your side, the description/comment probably shouldn't list llama as an "all-masked" case (it trains, just leaks). 2. The token test may not be running. On current transformers - labels = fn({"input_ids": [list(ids)]})["labels"][0]
+ if hasattr(ids, "keys"):
+ ids = ids["input_ids"]
+ labels = fn({"input_ids": [list(ids)]})["labels"][0]With that, llama would still fail its own 3. Optional: "unsloth": {
- "instruction": ">>> User: ",
- "response": ">>> Assistant: ",
+ "instruction": ">>> User:",
+ "response": ">>> Assistant:",
},
"vicuna": {
- "instruction": "USER: ",
- "response": "ASSISTANT: ",
+ "instruction": "USER:",
+ "response": "ASSISTANT:",
},Gemini's |
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Spot-checked the mistral space-folding claim against the actual tokenizer rather than taking it on faith:
tok = AutoTokenizer.from_pretrained('unsloth/mistral-7b-instruct-v0.3')
tok.encode('[INST] hello', add_special_tokens=False)
# -> [3, 7080, 29477], decoded per-token: ['[INST]', 'hell', 'o']Token 3 is the single [INST] special token, and there's no separate token for the space before "hello" -- it gets folded into the following text the same way the PR describes, so the old marker's trailing space wasn't recoverable as its own token boundary. That matches the stated root cause for mistral/llama.
Didn't re-derive all six template diffs at the token level (would need the GLM-4.7 and Qwen3-Thinking-2507 checkpoints too, and the PR's own new test file already does exactly that with graceful skip when a tokenizer isn't reachable), but the diff matches the table in the description exactly, and the new test_response_template_markers.py pins both the fixed entries and the unchanged ones so a future refactor can't silently regress either. Looks correct.
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Went back and verified both of @oobabooga's points at the token level, since my original approval explicitly hadn't re-derived all six template diffs individually, only spot-checked mistral. Both hold up, and the current head commit hasn't addressed either yet.
1. The llama entry is still "instruction": "[INST]", unchanged. Checked against the real unsloth/llama-2-7b-chat tokenizer:
encode('[INST]') standalone -> [518, 25580, 29962] -> ['▁[', 'INST', ']']
encode('<s>[INST] hi [/INST]') -> ['<s>', '[', 'INST', ']', '▁hi', '▁[', '/', 'INST', ']']
[INST] tokenized on its own gets the space-prefixed ▁[ (518), but at the real turn boundary, right after <s>, it tokenizes as bare [ (29961) instead. A marker built from the standalone encoding never matches the actual boundary, so this entry is a no-op exactly as described, it doesn't fix the leak the PR claims to fix for this model family. Anchoring on <s>[INST] instead of [INST] is the right call.
2. The chat-template test path is broken independent of the above. Reproduced directly:
ids = tok.apply_chat_template(msgs, tokenize=True, add_generation_prompt=False)
type(ids) # BatchEncoding, not a list
list(ids) # ['input_ids', 'attention_mask'] -- dict keys, not token idsSo fn({"input_ids": [list(ids)]}) is being handed the string keys ['input_ids', 'attention_mask'] instead of actual token ids, and the later ids[i] indexing doesn't return an int either. This means the chat-template-based half of the test isn't validating what it looks like it's validating, which also means it wouldn't have caught issue 1.
Given the PR's whole purpose is fixing markers that don't match their rendered templates, and llama is one of the six families in scope, I'm walking back my approval until this entry and the test path are both fixed. Everything else I checked (mistral, and the general structure of the fix) still looks correct.
On transformers 5.x llama-2 tokenizes [INST] after <s> as a bare left bracket while the standalone encoding gives the space-prefixed piece, so the unanchored marker missed every turn boundary and later user turns leaked into training; 4.57 masked this. Anchoring on <s>[INST] matches both tokenizations, verified token-level under 4.57.6 and 5.5.0. The test now unwraps the BatchEncoding that apply_chat_template returns on 5.x before indexing, and the latent trailing spaces in the unreachable unsloth and vicuna entries are dropped for table consistency.
|
All three landed in 5b87037, with one refinement on the diagnosis: the boundary tokenization is transformers-version dependent. On 4.57.6 llama-2 renders the turn boundary as |
|
Both points are fixed in 5b87037. Reproducing your bare-bracket encode surfaced the missing variable: it is the transformers major version. 4.57.6 tokenizes the boundary as |
|
@codex review |
|
Codex Review: Didn't find any major issues. Bravo. Reviewed commit: ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
If Codex has suggestions, it will comment; otherwise it will react with 👍. Codex can also answer questions or update the PR. Try commenting "@codex address that feedback". |
ErenAta16
left a comment
ErenAta16
left a comment
There was a problem hiding this comment.
Verified the two claims I could check directly against real tokenizers rather than just reading the table (unsloth/mistral-7b-instruct-v0.3, unsloth/llama-2-7b-chat, both loadable without gating).
Mistral: rendered apply_chat_template() output has [INST] as a single token (['<s>', '[INST]', ...]). The old marker "[INST] " tokenizes to ['[INST]', '▁'] (the trailing space becomes its own token) and has zero occurrences in the real sequence. The new marker "[INST]" matches at both turn boundaries. Confirms the bug and the fix.
Llama: same idea but sharper. Old instruction marker "[INST] " tokenizes to ['▁[', 'INST', ']', '▁'], zero occurrences in the real rendered sequence. New marker "<s>[INST]" tokenizes to ['<s>', '[', 'INST', ']'] and matches at exactly the two turn starts (positions 0 and 23 in my two-turn fixture) — [INST] really does tokenize as a bare [ right after <s>, vs. ▁[ in any other context, so the anchor is necessary, not just defensive.
One thing I noticed while checking llama's response marker specifically: " [/INST]" and "[/INST]" tokenize identically here (['▁[', '/', 'INST', ']'] either way — the leading Python-string space doesn't change how SentencePiece represents it), and both match at the same positions in the real sequence. So for llama, the response_part edit in the table is cosmetic; the instruction_part <s> anchor is what actually fixes the bug. Doesn't affect correctness, just worth knowing if the table gets edited later.
Which brings up the one real discrepancy I found: the PR description's "Marker changes" table lists llama the same as mistral — '[INST] ' → '[INST]', no <s> mentioned. But the actual diff and EXPECTED_FIXED in the test file both have "llama": {"instruction": "<s>[INST]"}. The code and tests agree with each other and with what I verified above; it's just the summary table that's out of sync. Worth a one-line fix to the table so it doesn't mislead someone skimming the description later.
Didn't have time to independently verify starling/glm/qwen3-thinking/zephyr against their tokenizers the same way, but the reasoning holds up for the two I did check, and the token-level test (test_fixed_markers_token_level) applies the same real methodology — masks checked against actual rendered+tokenized fixtures, not just marker literals — to all six, plus the "final token never masked" regression check for the infinite-generation failure mode. That's the right test to have here.
Approving based on the two families I verified directly plus the soundness of the test methodology for the rest.
3 participants
Summary
Complements #7054 (auto-first completion masking): that PR makes
train_on_completionsauto-detect the instruction/response markers from the chat template at runtime, with Studio's manualTEMPLATE_TO_RESPONSES_MAPPERtable as the fallback. This PR makes that fallback table correct. Six template families shipped manual markers that do not match what their chat templates actually render: five masked every assistant token (the run then died on the all-labels-masked safety net), and qwen3-thinking on Thinking-2507 checkpoints matched at most the final turn, silently skipping every earlier assistant turn.Marker changes (old -> new)
'[INST] '->'[INST]'' [/INST]'->'[/INST]'[INST]/[/INST]are single special tokens in v0.3; the surrounding spaces fold into neighbouring text tokens'[INST] '->'[INST]'' [/INST]'->'[/INST]''GPT4 Correct User: '->'GPT4 Correct User:''GPT4 Correct Assistant: '->'GPT4 Correct Assistant:'▁Hello)->'<->'\n<All markers were derived empirically from what each family's representative tokenizer renders for a two-turn conversation, and cross-checked against
unsloth_zoo.get_chat_template_partswhere it succeeds (the zephyr auto-detect fix is unslothai/unsloth-zoo#899).Validation
Token-level, on each representative tokenizer (
unsloth/mistral-7b-instruct-v0.3,unsloth/llama-2-7b-chat,unsloth/Starling-LM-7B-beta,unsloth/GLM-4.7-Flash,unsloth/Qwen3-4B-Thinking-2507,Qwen/QwQ-32B,unsloth/zephyr-sft), two-turn fixture plus system message, released unsloth_zoo:Consumers audited:
core/training/trainer.pyandcore/training/worker.pyonly pass the strings through totrain_on_responses_only; no other code depends on the old literals (the[INST]incore/inference/inference.pyis an independent fallback formatter).Tests
studio/backend/tests/test_response_template_markers.py: pins the fixed and key unchanged marker literals (dependency-free, runs everywhere) and runs the token-level masking checks above, skipping when tokenizers or unsloth_zoo are unavailable offline.