sft_prepare_dataset: fix double-BOS detection reading one character - #959
sft_prepare_dataset: fix double-BOS detection reading one character#959BardiaKoopah wants to merge 2 commits into
Conversation
unslothai#106 (@tamewild, Mar 2025) identified this bug and the same one-line fix. That PR no longer applies to current main (its diff is against line 547) and has no tests. This carries the fix forward with a regression suite, CI wiring, and the reachability analysis below. Credit for the diagnosis is theirs. `sft_prepare_dataset` passes `add_special_tokens=False` when a dataset's text already carries the BOS token. That detection was broken for plain-string text fields, and for Llama-3 it was the only working detector - so every row was tokenized with a doubled BOS. Verified on unsloth/Llama-3.2-1B-Instruct: unpatched, `sft_prepare_dataset` tokenizes with `add_special_tokens=[True]` on text that already begins with `<|begin_of_text|>`. Mechanism. For the plain-text-field branch, `test_text` was built as `next(iter(dataset))[dataset_text_field][0]`. That indexes a string, yielding its first character, so `startswith(bos_token)` could never be true for any multi-character BOS. The `[0]` is correct one branch up, where `formatting_func` returns a list of strings; it was copied into a branch where the value is a string. Reachability. Detection has two arms: (test_text is not None and test_text.startswith(bos_token)) or bos_token in chat_template The second arm only fires when the template holds the BOS token as a literal string. The broken first arm is therefore the only detector whenever a chat template emits BOS via the Jinja variable `{{ bos_token }}` instead. Llama-3 is confirmed by execution: on unsloth/Llama-3.2-1B-Instruct `'<|begin_of_text|>' in chat_template` is False while `'bos_token' in chat_template` is True. Other templates using the variable form are likely affected, but I have not tested them. Change. Use the field value directly; unwrap list/tuple-valued text fields (element 0, None if empty) so that shape keeps working. Tested. New tests/test_sft_prepare_dataset_double_bos.py (5 tests, stub tokenizer, no weights, no network) fails on unpatched source with `add_special_tokens=[True]` and passes after. Guards cover: no-BOS text, literal-BOS template, list-valued field, tokenizer without BOS. Wired into the behavioral-gates CI step so it is executed, not just collected (unslothai#669/unslothai#739 precedent). Full suite run twice per side: no test regressed. Refs unslothai#106
![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: b3b57e770e
ℹ️ 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".
| class Args: | ||
| def __init__(self, dataset_text_field="text"): | ||
| self.max_length = 64 | ||
| self.dataset_text_field = dataset_text_field | ||
| self.remove_unused_columns = True |
There was a problem hiding this comment.
Because this stub Args does not set dataset_num_proc, the production helper auto-selects num_proc on Linux when the multiprocessing start method is fork, which is exactly the Ubuntu CI hard-gate environment. Dataset.map then calls _tokenize in child processes, so the child copies mutate RecordingTokenizer.add_special_tokens_seen while the parent list remains empty, making these newly-gated tests fail with “tokenizer was never invoked”; set dataset_num_proc = 1 or assert from the returned dataset instead.
Useful? React with 👍 / 👎.
The test added by the previous commit passed on macOS and failed on Linux CI,
taking the behavioral-gates HARD GATE red with all 5 cases reporting
"tokenizer was never invoked".
`sft_prepare_dataset` hands `num_proc` to `Dataset.map` based on the
multiprocessing start method:
if _mp.get_start_method() != 'fork': dataset_num_proc = None
else: dataset_num_proc = min(max((psutil.cpu_count() or 1)+4, 2), 64)
macOS defaults to spawn, so `num_proc` is None, the map body runs in-process,
and a stub tokenizer recording `add_special_tokens` is visible to the assertion.
Linux defaults to fork, so the body runs in child processes and the parent's
recorder stays empty. The old test asserted on that recorder, so it could only
ever pass under spawn - the platform I verified on could not exercise the branch
that CI takes.
Assert on the returned dataset instead. The stub now encodes the decision in its
output: one leading BOS id for text that already carries the BOS string, one more
when `add_special_tokens` is set. A correct caller therefore yields exactly one
and the bug yields two. That reads back through the dataset, which crosses the
process boundary, and doubled BOS ids are the user-visible harm anyway.
Verified in all four combinations: fails unpatched under both spawn and fork
(2 leading BOS ids, [128000, 128000, 10, 11, 12]), passes patched under both.
Full suite twice per side: no regression, collected total 2977 -> 2982 for the
5 tests added. The CI gate list was also run under forced fork; its two
test_hf_xet_fallback failures are pre-existing on baseline here and green on
Linux CI (macOS case-insensitive filesystem).
This also corrects the previous commit's "full suite twice per side, no test
regressed" - true on macOS, and structurally unable to catch this.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b84e96d378
ℹ️ 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".
| test_text = next(iter(dataset))[dataset_text_field] | ||
| if isinstance(test_text, (list, tuple)): | ||
| test_text = test_text[0] if len(test_text) != 0 else None |
There was a problem hiding this comment.
When the first plain-text example starts with the BOS token but later examples do not, this probe now disables add_special_tokens for the entire Dataset.map call, so those later rows are tokenized without any BOS at all. This can happen with concatenated or partially preformatted SFT datasets; the decision needs to be made per row (or after normalizing the whole column), rather than using the first row as a dataset-wide invariant.
Useful? React with 👍 / 👎.
1 participant
Impact
sft_prepare_datasetpassesadd_special_tokens = Falsewhen a dataset's text already carries the BOS token. That detection was broken for plain-string text fields, and for Llama-3 it was the only working detector — so every row was tokenized with a doubled BOS.Verified on
unsloth/Llama-3.2-1B-Instruct: unpatched,sft_prepare_datasettokenizes withadd_special_tokens = [True]on text that already begins with<|begin_of_text|>.Mechanism
Detection is two-armed:
For the plain-text-field branch,
test_textwas built asnext(iter(dataset))[dataset_text_field][0](dataset_utils.py:1000). That indexes a string, yielding its first character, sostartswith(bos_token)could never be true for any multi-character BOS. The[0]is correct one branch up at line 981, whereformatting_funcreturns a list of strings; it was copied into a branch where the value is a plain string.Reachability
The second arm only fires when the chat template holds the BOS token as a literal. Llama-3 templates open with
{{- bos_token }}, a Jinja variable, so the literal is absent — verified:'<|begin_of_text|>' in chat_templateisFalsewhile'bos_token' in chat_templateisTrue.So the broken arm is the only detector whenever a template emits BOS via
{{ bos_token }}rather than as a literal string. Llama-3 is confirmed by execution; other templates using the variable form are likely affected but untested here.Change
Use the field value directly. List/tuple-valued text fields are unwrapped (element 0,
Noneif empty) so that shape keeps working.Tested
New
tests/test_sft_prepare_dataset_double_bos.py— 5 tests, stub tokenizer, no weights, no network. Fails on unpatched source (add_special_tokens = [True]), passes after. Guards cover: text without BOS, template with a literal BOS, list-valued field, and a tokenizer with no BOS token.Wired into the behavioral-gates CI step so the file is executed rather than merely collected (#669/#739 precedent: an unexecuted test was born broken in #669 and stayed invisible until #739).
Full zoo suite run twice per side with
-p no:randomly: no test failing in both patched runs that passed in both baseline runs. Collected totals reconcile — 2977 → 2982, the 5 tests added.Note on the test harness:
sft_prepare_datasetpassesnum_proctoDataset.maponly when the multiprocessing start method isfork. macOS defaults tospawn, so the map body runs in-process; Linux CI defaults tofork, so it runs in child processes. An earlier version of these tests asserted on a recorder in the parent process, which could only ever pass underspawn. The tests now assert on doubled BOS ids in the returned dataset — the user-visible harm rather than a proxy for it — which reads back correctly across the process boundary. Verified passing under both start methods, and failing under both on unpatched source.Inferred, not tested: that no caller depends on receiving a single character from this expression. Every use is a BOS prefix check, so reading the whole string is the intended behaviour.
Refs #106