bpe: do not slice merge tokens at a byte offset that may not be a boundary - #2398
apollo-2006 wants to merge 1 commit into
Conversation
…ndary
`BpeBuilder::build` strips the continuing subword prefix off the right hand
side of every merge by byte offset, without checking that the token actually
carries the prefix:
let b_len = b.len() - prefix_len;
let merge_len = a.len() + b_len;
buffer[a.len()..merge_len].copy_from_slice(&b.as_bytes()[prefix_len..]);
// SAFETY: buffer contains a concatenation of two valid UTF-8 strings, so
// it is itself valid UTF-8, even considering prefix_len
let new_token = unsafe { from_utf8_unchecked(&buffer[..merge_len]) };
The safety comment does not hold. Both `continuing_subword_prefix` and the
merge list come from the tokenizer JSON, and nothing relates them, so three
things go wrong on input that is merely malformed rather than malicious.
1. The offset can land inside a multi-byte character. For prefix "ab" and
merge ("x", "\u{65e5}"), which is E6 97 A5 and does not start with "ab",
`&b.as_bytes()[2..]` is the lone continuation byte A5. The bytes handed to
`from_utf8_unchecked` are then [78, A5], which `str::from_utf8` rejects
with "invalid utf-8 sequence of 1 bytes from index 1". Building a `str`
from those bytes violates its validity invariant.
2. `b.len() - prefix_len` underflows when the prefix is longer than the
token. A 4-byte prefix against a 1-byte token panics with "attempt to
subtract with overflow" in debug, and wraps in release.
3. `buffer` is sized to the longest vocabulary entry, but the merged token
need not be in the vocabulary at all, so `merge_len` can exceed it:
"range end index 4 out of range for slice of length 3".
Strip the prefix with `str::strip_prefix`, which only removes it when it is
actually there and therefore always cuts on a character boundary, and build
the token in a reused `String` so growth is handled and no unchecked
conversion is needed. The `unsafe` block goes away entirely.
Well formed vocabularies are unaffected: when `b` carries the prefix the
result is byte for byte what it was. `test_bpe_with_continuing_subword_prefix`
covers that and is unchanged.
Adds three regression tests, one per case above. Each fails without this
change with its own distinct error.
|
@ArthurZucker hows your morning been? 🙂 The PR removes the shared It still applies: What I could not work out from the outside is whether main is the right target anymore. #2178 landed on 2026-09-15 against So: do you still want BPE bugfixes on main while that rewrite is in flight, or should I retarget this at |
BpeBuilder::buildstrips the continuing subword prefix off the right hand side of every merge by byte offset, without checking that the token actually carries the prefix:The safety comment does not hold.
continuing_subword_prefixand the merge list are both deserialized from the tokenizer JSON and nothing relates them, so a file that is merely malformed rather than malicious reaches three separate failures.1. The offset can land inside a character
With prefix
"ab"and the merge("x", "日"), the right hand side isE6 97 A5and does not start with"ab", but two bytes are removed anyway. The remainder is the lone continuation byteA5, so the bytes handed tofrom_utf8_uncheckedare:Constructing a
strfrom those bytes violates its validity invariant. Note that Miri does not check UTF-8 validity of&str, so this does not show up undercargo miri test; the output above is from runningstr::from_utf8on the exact slice the unsafe call receives.2. The subtraction underflows
b.len() - prefix_lenhas no lower bound check. A 4-byte prefix against a 1-byte token:Debug panics, release wraps.
3. The scratch buffer can be too small
bufferis sized to the longest vocabulary entry, but the merged token need not be in the vocabulary, somerge_lencan exceed it:Fix
Use
str::strip_prefix, which removes the prefix only when it is actually present and therefore always cuts on a character boundary, and build the token in a reusedStringso growth is handled automatically. That addresses all three, and theunsafeblock is no longer needed at all.The allocation behaviour is the same as before: one buffer reused across every merge,
String::with_capacity(max_len)in place ofvec![0; max_len].Behaviour
Well formed vocabularies are unaffected. When
bcarries the prefix,strip_prefixremoves exactly the same bytes the old slice did, so the merged token is byte for byte identical. The existingtest_bpe_with_continuing_subword_prefixcovers that path and is unchanged.Where behaviour does change is the malformed cases above, which previously produced a corrupted token, a panic, or undefined behaviour, and now produce the ordinary
MergeTokenOutOfVocabularyerror.Tests
Three regression tests, one per case. Against the unfixed code each fails with its own distinct error:
test_bpe_prefix_not_present_on_mergetest_bpe_prefix_longer_than_tokenattempt to subtract with overflowtest_bpe_merged_token_longer_than_longest_vocab_entryrange end index 4 out of range for slice of length 3Verified with the full test data downloaded (
make testresources, including the GPT-2 vocab and merges, so a real 50k-entry byte-level BPE goes through the changed path): 258 passed, 0 failed.cargo fmt --checkclean andcargo clippy --all-targets --all-features -- -D warningsclean.