Fix special token id persistence and id range validation

Y0oshi · Y0oshi · commit f7299e4e2647 · 2026-07-14T11:59:01.000+03:00
diff --git a/configs/base.json b/configs/base.json
@@ -1,6 +1,6 @@
 {
   "model": {
-    "vocab_size": 8192,
+    "vocab_size": 8256,
     "context_length": 512,
     "n_embed": 768,
     "n_head": 12,
diff --git a/configs/small.json b/configs/small.json
@@ -1,6 +1,6 @@
 {
   "model": {
-    "vocab_size": 8192,
+    "vocab_size": 8256,
     "context_length": 256,
     "n_embed": 384,
     "n_head": 6,
diff --git a/tests/test_alignment.py b/tests/test_alignment.py
@@ -28,7 +28,10 @@ def test_sft_with_no_trainable_examples_raises_clearly(tmp_path):
     tok.add_special(SPECIAL_TOKENS)
     tok.save(tmp_path / "tok.json")
 
-    cfg = Config(_model_cfg(), TrainConfig(device="cpu"))
+    from textllm.config import ModelConfig
+
+    model_cfg = ModelConfig(vocab_size=512, context_length=32, n_embed=64, n_head=4, n_kv_head=2, n_blocks=2)
+    cfg = Config(model_cfg, TrainConfig(device="cpu"))
     save_checkpoint(tmp_path / "base.pt", Transformer(cfg.model), None, 0, cfg)
 
     # user-only conversations: prompt masking leaves nothing to supervise
diff --git a/tests/test_cli.py b/tests/test_cli.py
@@ -89,3 +89,15 @@ def test_pretrain_rejects_token_file_too_short_for_a_window(tmp_path, monkeypatc
 
     with pytest.raises(ValueError, match="context_length"):
         main(["pretrain", *TINY_OVERRIDES, "--set", "train.data_dir=data", "--set", "train.out_dir=ckpt"])
+
+
+def test_pretrain_rejects_token_ids_beyond_model_vocab(tmp_path, monkeypatch):
+    import numpy as np
+    import pytest
+
+    monkeypatch.chdir(tmp_path)
+    (tmp_path / "data").mkdir()
+    np.full(200, 600, dtype=np.uint16).tofile(tmp_path / "data" / "train.bin")  # vocab is 512
+
+    with pytest.raises(ValueError, match="vocab_size"):
+        main(["pretrain", *TINY_OVERRIDES, "--set", "train.data_dir=data", "--set", "train.out_dir=ckpt"])
diff --git a/tests/test_data.py b/tests/test_data.py
@@ -116,3 +116,19 @@ def test_preference_pair_dropped_when_prompt_fills_context():
     long_prompt = [Message("user", "hello world how are you doing today friend " * 30)]
     ds = PreferenceDataset([(long_prompt, "fine", "bad")], tok, block_size=16)
     assert len(ds) == 0  # truncation left no supervised response tokens on either side
+
+
+def test_build_bin_bounds_the_id_range_not_the_count(tmp_path):
+    import pytest
+
+    from textllm.data import build_bin
+
+    class SparseIds:
+        vocab_size = 300          # small count
+        max_token_id = 70_000     # but an id beyond uint16
+
+        def encode(self, text):
+            return [0]
+
+    with pytest.raises(ValueError, match="overflow"):
+        build_bin(SparseIds(), ["hello"], tmp_path / "t.bin")
diff --git a/tests/test_tokenizer.py b/tests/test_tokenizer.py
@@ -32,3 +32,90 @@ def test_save_load_preserves_encoding(tmp_path):
     s = "the quick café <|end|>"
     assert reloaded.encode(s) == tok.encode(s)
     assert reloaded.decode(reloaded.encode(s)) == s
+
+
+def test_load_restores_saved_special_ids_exactly(tmp_path):
+    import json
+
+    tok = _trained(tmp_path)
+    path = tmp_path / "tok.json"
+    tok.save(path)
+
+    # perturb the stored ids so they are no longer the contiguous re-derivable block
+    data = json.loads(path.read_text())
+    data["special"] = {t: i + 100 for t, i in data["special"].items()}
+    path.write_text(json.dumps(data))
+
+    reloaded = Tokenizer.load(path)
+    assert reloaded.special == data["special"]
+    for text, i in data["special"].items():
+        assert reloaded.encode(text) == [i]
+
+
+def test_add_special_across_calls_never_collides(tmp_path):
+    tok = _trained(tmp_path)
+    tok.add_special(["<|tool|>"])
+    tok.add_special(["<|result|>"])
+    ids = list(tok.special.values())
+    assert len(ids) == len(set(ids))
+
+
+def test_add_special_after_load_does_not_collide(tmp_path):
+    tok = _trained(tmp_path)
+    path = tmp_path / "tok.json"
+    tok.save(path)
+
+    reloaded = Tokenizer.load(path)
+    reloaded.add_special(["<|tool|>"])
+    ids = list(reloaded.special.values())
+    assert len(ids) == len(set(ids))
+    assert reloaded.special["<|tool|>"] > max(tok.special.values())
+
+
+def test_train_after_add_special_raises(tmp_path):
+    import pytest
+
+    tok = Tokenizer()
+    tok.add_special(["<|end|>"])
+    with pytest.raises(ValueError, match="after training"):
+        tok.train("some text " * 20, vocab_size=300)
+
+
+def test_load_rejects_colliding_special_ids(tmp_path):
+    import json
+
+    import pytest
+
+    tok = _trained(tmp_path)
+    path = tmp_path / "tok.json"
+    tok.save(path)
+
+    data = json.loads(path.read_text())
+    first = next(iter(data["special"]))
+    data["special"][first] = 260  # a merge id
+    path.write_text(json.dumps(data))
+    with pytest.raises(ValueError, match="special token ids"):
+        Tokenizer.load(path)
+
+
+def test_longer_special_token_wins_over_its_prefix(tmp_path):
+    tok = _trained(tmp_path)
+    tok.add_special(["<|t|>", "<|t|>x"])
+    assert tok.encode("<|t|>x") == [tok.special["<|t|>x"]]
+
+
+def test_special_id_and_vocab_fit_helpers(tmp_path):
+    import pytest
+
+    from textllm.tokenizer import check_vocab_fit, special_id
+
+    tok = _trained(tmp_path)
+    assert special_id(tok, "<|pad|>") == tok.special["<|pad|>"]
+    check_vocab_fit(tok, tok.max_token_id + 1)
+    with pytest.raises(ValueError, match="vocabulary"):
+        check_vocab_fit(tok, tok.max_token_id)
+
+    bare = Tokenizer()
+    bare.train("some text here " * 20, vocab_size=300)
+    with pytest.raises(ValueError, match="special"):
+        special_id(bare, "<|pad|>")
diff --git a/textllm/config.py b/textllm/config.py
@@ -16,7 +16,7 @@
 
 @dataclass
 class ModelConfig:
-    vocab_size: int = 8192
+    vocab_size: int = 8256  # headroom above a vocab-8192 tokenizer plus its special ids
     context_length: int = 512
     n_embed: int = 384
     n_head: int = 6
diff --git a/textllm/data/shards.py b/textllm/data/shards.py
@@ -21,12 +21,16 @@ def build_bin(tokenizer, texts: Iterable[str], out_path: str | Path, eot: str |
     Returns the number of tokens written. Documents are flushed as they're encoded so a
     large corpus never has to sit in memory all at once.
     """
-    max_id = np.iinfo(DTYPE).max
-    vocab_size = getattr(tokenizer, "vocab_size", None)
-    if vocab_size is not None and vocab_size > max_id + 1:
+    limit = np.iinfo(DTYPE).max
+    # vocab_size counts entries; special ids can sit far above it, so bound the id range
+    top = getattr(tokenizer, "max_token_id", None)
+    if top is None:
+        vocab_size = getattr(tokenizer, "vocab_size", None)
+        top = vocab_size - 1 if vocab_size else None
+    if top is not None and top > limit:
         raise ValueError(
-            f"tokenizer vocab ({vocab_size}) exceeds what {np.dtype(DTYPE).name} can store "
-            f"({max_id + 1} ids) — token ids would silently overflow"
+            f"tokenizer ids reach {top}, beyond what {np.dtype(DTYPE).name} can store "
+            f"({limit}) — token ids would silently overflow"
         )
 
     out_path = Path(out_path)
diff --git a/textllm/eval/reasoning.py b/textllm/eval/reasoning.py
@@ -22,7 +22,8 @@ def accuracy(ckpt_path: str, tokenizer_spec: str, qa_path: str, max_new_tokens:
     device = pick_device()
     model, cfg = load_model(ckpt_path, device)
     tokenizer = get_tokenizer(tokenizer_spec)
-    stops = set(tokenizer.encode(EOT))
+    eot_ids = tokenizer.encode(EOT)
+    stops = set(eot_ids) if len(eot_ids) == 1 else set()
 
     rows = [json.loads(line) for line in Path(qa_path).read_text().splitlines() if line.strip()]
     correct = 0
diff --git a/textllm/infer/chat.py b/textllm/infer/chat.py
@@ -12,13 +12,9 @@
 
 
 def _stop_ids(tokenizer) -> set[int]:
-    stops = set()
-    for marker in (EOT,):
-        try:
-            stops.update(tokenizer.encode(marker))
-        except Exception:
-            pass
-    return stops
+    # a multi-token EOT would stop generation on ordinary text, so require a single id
+    ids = tokenizer.encode(EOT)
+    return set(ids) if len(ids) == 1 else set()
 
 
 def chat_repl(ckpt_path: str, tokenizer_spec: str, system: str | None = None, **sampling) -> None:
diff --git a/textllm/tokenizer/__init__.py b/textllm/tokenizer/__init__.py
@@ -1,6 +1,13 @@
 from textllm.tokenizer.bpe import SPLIT_PATTERN, Tokenizer, TiktokenTokenizer
 
-__all__ = ["Tokenizer", "TiktokenTokenizer", "SPLIT_PATTERN", "get_tokenizer"]
+__all__ = [
+    "Tokenizer",
+    "TiktokenTokenizer",
+    "SPLIT_PATTERN",
+    "get_tokenizer",
+    "special_id",
+    "check_vocab_fit",
+]
 
 
 def get_tokenizer(spec: str):
@@ -12,3 +19,24 @@ def get_tokenizer(spec: str):
     if spec.startswith("tiktoken:"):
         return TiktokenTokenizer(spec.split(":", 1)[1])
     return Tokenizer.load(spec)
+
+
+def special_id(tokenizer, token: str) -> int:
+    """Id of a registered special token; raises if the tokenizer lacks it."""
+    ids = tokenizer.encode(token)
+    if len(ids) != 1:
+        raise ValueError(
+            f"{token!r} is not a special token in this tokenizer; chat training and "
+            f"inference need a tokenizer trained with the textllm chat specials"
+        )
+    return ids[0]
+
+
+def check_vocab_fit(tokenizer, vocab_size: int) -> None:
+    """Reject a tokenizer whose ids the model's embedding cannot index."""
+    top = getattr(tokenizer, "max_token_id", None)
+    if top is not None and top >= vocab_size:
+        raise ValueError(
+            f"tokenizer ids reach {top} but the model vocabulary has {vocab_size} "
+            f"entries; raise model.vocab_size or use a smaller tokenizer"
+        )
diff --git a/textllm/tokenizer/bpe.py b/textllm/tokenizer/bpe.py
@@ -50,6 +50,9 @@ def __init__(self, pattern: str = SPLIT_PATTERN) -> None:
     def train(self, text: str, vocab_size: int, verbose: bool = False) -> None:
         if vocab_size < 256:
             raise ValueError("vocab_size must be at least 256 (the byte alphabet)")
+        if self.special:
+            # merge ids are assigned from 256 upward and would collide with them
+            raise ValueError("add special tokens after training, not before")
         n_merges = vocab_size - 256
 
         # Count how often each pre-split chunk appears, then only ever work on the unique
@@ -78,19 +81,33 @@ def train(self, text: str, vocab_size: int, verbose: bool = False) -> None:
     def add_special(self, tokens: list[str]) -> None:
         """Register special tokens (chat markers, end-of-text) above the learned vocab."""
         next_id = max(self.vocab) + 1 if self.vocab else 256
+        if self.special:
+            # existing specials live above the vocab too; never hand out their ids again
+            next_id = max(next_id, max(self.special.values()) + 1)
         for tok in tokens:
             if tok in self.special:
                 continue
             self.special[tok] = next_id
             next_id += 1
+        self._compile_special_pattern()
+
+    def _compile_special_pattern(self) -> None:
         if self.special:
-            joined = "|".join(re.escape(t) for t in self.special)
+            # longest first, so one special token can't shadow a longer one in the regex
+            joined = "|".join(re.escape(t) for t in sorted(self.special, key=len, reverse=True))
             self._special_re = re.compile(f"({joined})")
 
     @property
     def vocab_size(self) -> int:
         return len(self.vocab) + len(self.special)
 
+    @property
+    def max_token_id(self) -> int:
+        top = max(self.vocab)
+        if self.special:
+            top = max(top, max(self.special.values()))
+        return top
+
     # -- encode / decode ----------------------------------------------------------
 
     def _encode_chunk(self, ids: list[int]) -> list[int]:
@@ -154,7 +171,15 @@ def load(cls, path: str | Path) -> "Tokenizer":
             a, b = (int(x) for x in key.split(","))
             tok.merges[(a, b)] = new_id
             tok.vocab[new_id] = tok.vocab[a] + tok.vocab[b]
-        tok.add_special(list(data.get("special", {})))
+        # a trained model depends on these exact ids, so restore them verbatim
+        tok.special = {t: int(i) for t, i in data.get("special", {}).items()}
+        ids = list(tok.special.values())
+        if len(set(ids)) != len(ids) or any(i < 0 for i in ids) or set(ids) & set(tok.vocab):
+            raise ValueError(
+                f"invalid special token ids in {path}: they must be unique, non-negative, "
+                f"and outside the merge vocabulary"
+            )
+        tok._compile_special_pattern()
         return tok
 
     # -- alternative: reuse a tiktoken encoding -----------------------------------
@@ -177,12 +202,17 @@ def __init__(self, name: str = "gpt2") -> None:
     def vocab_size(self) -> int:
         return self._enc.n_vocab
 
+    @property
+    def max_token_id(self) -> int:
+        return self._enc.n_vocab - 1
+
     def encode_ordinary(self, text: str) -> list[int]:
         return self._enc.encode_ordinary(text)
 
     def encode(self, text: str, allowed_special: bool = True) -> list[int]:
-        special = "all" if allowed_special else set()
-        return self._enc.encode(text, allowed_special=special)
+        if not allowed_special:
+            return self._enc.encode_ordinary(text)  # special strings become plain text
+        return self._enc.encode(text, allowed_special="all")
 
     def decode(self, ids: list[int]) -> str:
         return self._enc.decode(ids)
diff --git a/textllm/train/dpo.py b/textllm/train/dpo.py
@@ -20,7 +20,7 @@
 from textllm.data import PreferenceDataset, collate_preference
 from textllm.device import describe, pick_device
 from textllm.model import Transformer
-from textllm.tokenizer import get_tokenizer
+from textllm.tokenizer import check_vocab_fit, get_tokenizer, special_id
 from textllm.train.loop import build_optimizer, infinite, load_checkpoint, model_state, train
 from textllm.train.reward import load_preferences
 
@@ -67,7 +67,8 @@ def run_dpo(cfg: Config, base_ckpt: str, tokenizer_spec: str, pref_path: str, be
         p.requires_grad_(False)
 
     tokenizer = get_tokenizer(tokenizer_spec)
-    pad_id = tokenizer.encode(PAD)[0]
+    check_vocab_fit(tokenizer, model_cfg.vocab_size)
+    pad_id = special_id(tokenizer, PAD)
     dataset = PreferenceDataset(load_preferences(pref_path), tokenizer, model_cfg.context_length)
     if len(dataset) == 0:
         raise ValueError(f"no preference pairs in {pref_path}")
diff --git a/textllm/train/grpo.py b/textllm/train/grpo.py
@@ -18,7 +18,7 @@
 from textllm.device import describe, pick_device
 from textllm.infer.generate import generate
 from textllm.model import Transformer
-from textllm.tokenizer import get_tokenizer
+from textllm.tokenizer import check_vocab_fit, get_tokenizer
 from textllm.train.loop import build_optimizer, load_checkpoint, model_state, save_checkpoint
 from textllm.train.rewards import RewardFn
 
@@ -58,6 +58,7 @@ def run_grpo(
         p.requires_grad_(False)
 
     tokenizer = get_tokenizer(tokenizer_spec)
+    check_vocab_fit(tokenizer, model_cfg.vocab_size)
     pad_id = tokenizer.encode(" ")[0]  # any real token id; padded positions are masked out
     optimizer = build_optimizer(policy, cfg.train)
 
diff --git a/textllm/train/pretrain.py b/textllm/train/pretrain.py
@@ -45,6 +45,12 @@ def run_pretrain(cfg: Config, resume: str | None = None) -> None:
             f"form one window of context_length + 1 = {block + 1}; add more data or lower "
             f"model.context_length"
         )
+    peak = int(train_set.data.max())
+    if peak >= cfg.model.vocab_size:
+        raise ValueError(
+            f"train.bin contains token id {peak} but model.vocab_size is "
+            f"{cfg.model.vocab_size}; the embedding cannot index it"
+        )
     train_loader = DataLoader(
         train_set,
         batch_size=min(cfg.train.batch_size, len(train_set)),
diff --git a/textllm/train/reward.py b/textllm/train/reward.py
diff --git a/textllm/train/sft.py b/textllm/train/sft.py

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,6 @@`
`1`	`1`	`{`
`2`	`2`	`"model": {`
`3`		`- "vocab_size": 8192,`
	`3`	`+ "vocab_size": 8256,`
`4`	`4`	`"context_length": 512,`
`5`	`5`	`"n_embed": 768,`
`6`	`6`	`"n_head": 12,`